diff --git a/CMakeLists.txt b/CMakeLists.txt index 01199cade..34aaa817d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1636,6 +1636,14 @@ audiocpp_add_model(firered_audio engine::models::firered_audio::make_firered_audio_loader ) +# VibeASR VAE encoder. No loader yet: the LM half runs on I2_S weights whose +# kernel is not ported, so there is no session to register a family for. +audiocpp_add_model(vibeasr + SOURCES + src/community_models/vibeasr/assets.cpp + src/community_models/vibeasr/vae_encoder.cpp +) + set(AUDIOCPP_ENABLED_MODELS "") if (AUDIOCPP_MODEL_SET STREQUAL "full") set(AUDIOCPP_ENABLED_MODELS ${AUDIOCPP_MODEL_TARGETS}) @@ -2279,6 +2287,20 @@ if (ENGINE_BUILD_TESTS) COMMAND conv_lowering_matrix_test ) + add_engine_unittest(i8_s_fused_ops_test tests/unittests/test_i8_s_fused_ops.cpp) + + add_test( + NAME i8_s_fused_ops_test + COMMAND i8_s_fused_ops_test + ) + + add_engine_unittest(i2_s_mul_mat_test tests/unittests/test_i2_s_mul_mat.cpp) + + add_test( + NAME i2_s_mul_mat_test + COMMAND i2_s_mul_mat_test + ) + add_engine_unittest(gguf_tensor_source_test tests/unittests/test_gguf_tensor_source.cpp) target_include_directories(gguf_tensor_source_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tests/unittests) @@ -2418,6 +2440,34 @@ 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/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 + ) + 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 7c95c6024..6620222ee 100644 --- a/README.md +++ b/README.md @@ -167,6 +167,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 | [@WalkingCat](https://github.com/WalkingCat) | [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 (encoder only) | n/a | GGUF I8_S | [@XsquirrelC](https://github.com/XsquirrelC) | [VibeASR VAE encoder](docs/community_models/vibeasr.md) INT8-weight and INT8-activation port of the VibeVoice acoustic/semantic tokenizers from [VibeASR.cpp](https://github.com/microsoft/VibeASR.cpp); encoder only, the ternary I2_S decoder is not ported yet | | **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..5ea156426 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. +An INT8-activation port of the same encoder, from Microsoft's +[VibeASR.cpp](https://github.com/microsoft/VibeASR.cpp), lives under community +models: see [VibeASR VAE encoder](community_models/vibeasr.md). It is not a +separate model, and it has no CLI family yet. + | Field | Value | |---|---| | Family | `vibevoice_asr` | diff --git a/docs/community_models/models.md b/docs/community_models/models.md index 994556b52..8298b2439 100644 --- a/docs/community_models/models.md +++ b/docs/community_models/models.md @@ -34,3 +34,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 (encoder only) | n/a (waveform in, LM features out) | [@XsquirrelC](https://github.com/XsquirrelC) | [VibeASR VAE encoder](vibeasr.md) INT8-weight *and* INT8-activation port of the VibeVoice acoustic/semantic tokenizers from [VibeASR.cpp](https://github.com/microsoft/VibeASR.cpp), on the fused `GGML_TYPE_I8_S` ops. Encoder only: the ternary I2_S decoder is not ported, so there is no CLI family yet | diff --git a/docs/community_models/vibeasr.md b/docs/community_models/vibeasr.md new file mode 100644 index 000000000..89df1b4eb --- /dev/null +++ b/docs/community_models/vibeasr.md @@ -0,0 +1,207 @@ +# VibeASR VAE encoder 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 page covers the first half of +that port — the VAE encoder — which is what this entry currently adds. + +## 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 and supports streaming. + +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 | Q8_0 Qwen2 | ternary `GGML_TYPE_I2_S` (kernel in tree, graph not yet) | +| Backends | CPU, CUDA, Metal | CPU only — the I8_S ops have no GPU kernels | +| Streaming | yes | no | + +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 much smaller decoder once the +I2_S kernel lands). + +It stays a separate community entry rather than becoming a weight path inside +`vibevoice_asr`, because the two share no graph code: every activation here 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 — is data and conventions, and this entry follows +`vibevoice_asr` on all of it. + +## Architecture + +Both branches are identical in shape and differ only in latent width: + +- **Input**: mono 16 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** — 5 frames per second. 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. + +## Usage + +VibeASR.cpp already ships the encoder 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 42/43. 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_vae.py \ + --input vibeasr-vae-encoder-i8_s.gguf --list + +# produce the audio.cpp package (~703 MB, both branches) +python3 tools/community_models/convert_vibeasr_vae.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_vae.py \ + --input models/vibeasr/vae_encoder-i8_s.gguf --check +``` + +There is no CLI family yet (see [Status](#status)); the encoder is reached +through `VibeASRVaeEncoderRuntime` or the parity probe: + +```bash +cmake -B build -DCMAKE_BUILD_TYPE=Release -DENGINE_BUILD_MODEL_TESTS=ON +cmake --build build -j --target test_vibeasr_vae_encoder + +./build/bin/test_vibeasr_vae_encoder \ + --model models/vibeasr/vae_encoder-i8_s.gguf \ + --audio assets/asr_validation/librispeech/librispeech_test_clean_6930-75918-0000.wav \ + --threads 8 +``` + +Without `--reference-*` the probe checks shape, finiteness, and frame count only. +It exits 125 (SKIP) when the model or audio is missing, so it is safe in ctest. + +## Parity + +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/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, 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` (under ctest) 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-core AMD EPYC 7V13, 3.505 s clip, both branches +(the encoder is run twice — once per branch — because that is what the decoder +consumes): + +| Threads | acoustic | semantic | both | RTF | +|---|---|---|---|---| +| 8 | 276 ms | 270 ms | 546 ms | 0.156 | +| 1 | 1657 ms | 1600 ms | 3257 ms | 0.929 | + +Peak RSS is 1.31 GB against 703 MB of weights: `BackendWeightStore` stages each +tensor before upload, so weight loading briefly holds roughly two copies. The +graph arena itself is 64 MB by default. + +## Status + +Ported: + +- I8_S VAE encoder graph, both branches, CPU backend. +- GGUF type remapping tool and geometry-from-tensors asset loader. +- Parity probe against upstream, plus op-level unit tests. + +Not ported yet: + +- **The decoder.** The ternary `GGML_TYPE_I2_S` matmul is in tree (plain + `ggml_mul_mat` with an I2_S weight; see `i2_s_mul_mat_test`), but the Qwen2 + graph that uses it is not, so there is no loader, no session, and no + `--family vibeasr` — nothing can transcribe through this path today. The + encoder output is the LM input, so the halves are independently reviewable but + only useful together. + +Known limitations: + +- **CPU only.** The fused I8_S ops have no CUDA or Metal kernels; the probe pins + the backend to CPU. +- **Offline only.** No streaming; upstream's encoder is causal, so streaming is + implementable, but the state machine is not ported. +- Bit-exact parity with upstream is out of reach by design; see + [Parity](#parity). + +## Upstream + +- Model port: (`src/vae.cpp`) +- Base model: VibeVoice ASR, also in tree as [`vibevoice_asr`](../asr.md#vibevoice-asr) diff --git a/external/ggml/include/ggml.h b/external/ggml/include/ggml.h index 0de79aed5..1626f0105 100644 --- a/external/ggml/include/ggml.h +++ b/external/ggml/include/ggml.h @@ -429,7 +429,15 @@ extern "C" { GGML_TYPE_MXFP4 = 39, // MXFP4 (1 block) GGML_TYPE_NVFP4 = 40, // NVFP4 (4 blocks, E4M3 scale) GGML_TYPE_Q1_0 = 41, - GGML_TYPE_COUNT = 42, + // INT8 / ternary with a single per-tensor scale, stored as one F32 + // immediately after the payload rather than interleaved per block. + // Used by the VibeASR CPU pipeline; see ggml_mul_mat_add(). + // NOTE: 36/37 are deliberately avoided even though VibeASR's own fork + // uses them -- they are retired IQ4_NL_4_4/4_8 slots and reusing an ID + // would silently misread GGUF files that still carry the old type. + GGML_TYPE_I8_S = 42, + GGML_TYPE_I2_S = 43, + GGML_TYPE_COUNT = 44, }; // precision @@ -588,6 +596,14 @@ extern "C" { GGML_OP_GLU, GGML_OP_CONVROT_LINEAR, + + // VibeASR CPU INT8 pipeline. Appended at the tail so every existing + // op keeps its value -- GGML_OP_NAME and GGML_OP_SYMBOL are positional. + GGML_OP_ADD_SCALED, + GGML_OP_RMS_NORM_SCALED, + GGML_OP_MUL_MAT_ADD, + GGML_OP_MUL_MAT_ADD_RELU, + GGML_OP_IM2COL_ASYM, GGML_OP_COUNT, }; @@ -2472,6 +2488,77 @@ extern "C" { struct ggml_tensor * bias, int group_size); + // VibeASR CPU INT8 pipeline (GGML_TYPE_I8_S / GGML_TYPE_I2_S). + // + // Ported from https://github.com/microsoft/VibeASR.cpp, an end-to-end INT8 + // ASR stack (INT8 VAE encoder, ternary-weight language model) built for CPU + // inference on edge devices. + // + // These are CPU-only, mirroring how ggml_convrot_linear and + // ggml_sage_attn2_i8 above are CUDA-only. + // + // Unlike those two, the scale is NOT a separate F32 src tensor. An I8_S + // activation's scale is recomputed from that activation at run time, and a + // ggml node has exactly one output, so the scale has to travel with the + // data: every I8_S/I2_S tensor stores one F32 immediately after its int8 + // payload (see ggml_type_extra_bytes in ggml.c). Weight scales could have + // used the separate-tensor convention, but sharing one representation with + // activations keeps a single kernel per op instead of two. + + // y = a*scale + b, fusing a ConvNeXt LayerScale into its residual add. + // a and b are I8_S and same-shape, scale is F32 per-channel broadcast on + // ne[0]. Output is I8_S and carries a freshly computed per-tensor scale. + GGML_API struct ggml_tensor * ggml_add_scaled( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * scale); + + // y = rms_norm(a) * scale, fused. a is I8_S, scale is F32 per-channel, + // output is I8_S. Equivalent to ggml_mul(ggml_rms_norm(a), scale) but + // avoids materializing the F32 intermediate. + GGML_API struct ggml_tensor * ggml_rms_norm_scaled( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * scale, + float eps); + + // y = a*b + bias, with a I8_S weights and b I8_S activations. bias is F32 + // and broadcasts on ne[0]. Output is I8_S. The ternary I2_S weights of the + // language model go through plain ggml_mul_mat, which has no bias to fuse. + // + // a with ne[1] == 1 and ne[2] > 1 selects a depthwise contraction: one + // length-ne[0] filter per channel, output indexed channel-major. + GGML_API struct ggml_tensor * ggml_mul_mat_add( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * bias); + + // As ggml_mul_mat_add, with ReLU folded into the epilogue. + GGML_API struct ggml_tensor * ggml_mul_mat_add_relu( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * bias); + + // im2col with independent left/right padding on the width axis. ggml_im2col + // only takes a single symmetric p0, so causal 1D convolutions otherwise need + // a separate ggml_pad_ext node and a full copy of the activation. + GGML_API struct ggml_tensor * ggml_im2col_asym( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int s1, + int lp0, + int rp0, + int p1, + int d0, + int d1, + bool is_2D, + enum ggml_type dst_type); + // MINITTS_FLASH_BIAS_WRAPPER: // Helper for models that already assemble a dense additive attention bias // (for example relative-position scores). The helper expands an optional diff --git a/external/ggml/src/ggml-cpu/ggml-cpu.c b/external/ggml/src/ggml-cpu/ggml-cpu.c index d9ec09939..05639b05a 100644 --- a/external/ggml/src/ggml-cpu/ggml-cpu.c +++ b/external/ggml/src/ggml-cpu/ggml-cpu.c @@ -1255,6 +1255,17 @@ void ggml_compute_forward_mul_mat( return; } + // Ternary weights own their activation quantization, so they cannot use the + // vec_dot_type path below: that quantizes src1 one row at a time into a + // fixed row_size and hands the kernel nothing but two row pointers, while + // I2_S needs the per-row activation scale and int8 row sum to survive into + // the epilogue. Branching here rather than adding an op keeps the language + // model graph on plain ggml_mul_mat. + if (src0->type == GGML_TYPE_I2_S) { + ggml_compute_forward_mul_mat_i2_s(params, dst); + return; + } + GGML_TENSOR_BINARY_OP_LOCALS const int ith = params->ith; @@ -1825,7 +1836,7 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm ggml_compute_forward_l2_norm(params, tensor); } break; case GGML_OP_MUL_MAT: - case GGML_OP_MUL_MAT_PACK4: + case GGML_OP_MUL_MAT_PACK4: { ggml_compute_forward_mul_mat(params, tensor); } break; @@ -1901,18 +1912,18 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_conv_transpose_1d(params, tensor); } break; - case GGML_OP_IM2COL: - { - ggml_compute_forward_im2col(params, tensor); - } break; - case GGML_OP_IM2COL_FAST_1D: - { - ggml_compute_forward_im2col_fast_1d(params, tensor); - } break; - case GGML_OP_IM2COL_BACK: - { - ggml_compute_forward_im2col_back_f32(params, tensor); - } break; + case GGML_OP_IM2COL: + { + ggml_compute_forward_im2col(params, tensor); + } break; + case GGML_OP_IM2COL_FAST_1D: + { + ggml_compute_forward_im2col_fast_1d(params, tensor); + } break; + case GGML_OP_IM2COL_BACK: + { + ggml_compute_forward_im2col_back_f32(params, tensor); + } break; case GGML_OP_IM2COL_3D: { ggml_compute_forward_im2col_3d(params, tensor); @@ -2092,6 +2103,26 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm ggml_compute_forward_opt_step_sgd(params, tensor); } break; + case GGML_OP_ADD_SCALED: + { + ggml_compute_forward_add_scaled(params, tensor); + } break; + case GGML_OP_RMS_NORM_SCALED: + { + ggml_compute_forward_rms_norm_scaled(params, tensor); + } break; + case GGML_OP_MUL_MAT_ADD: + { + ggml_compute_forward_mul_mat_add(params, tensor); + } break; + case GGML_OP_MUL_MAT_ADD_RELU: + { + ggml_compute_forward_mul_mat_add_relu(params, tensor); + } break; + case GGML_OP_IM2COL_ASYM: + { + ggml_compute_forward_im2col_asym(params, tensor); + } break; case GGML_OP_NONE: { // nop @@ -2301,7 +2332,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_GROUP_NORM: case GGML_OP_CONCAT: case GGML_OP_MUL_MAT: - case GGML_OP_MUL_MAT_PACK4: + case GGML_OP_MUL_MAT_PACK4: case GGML_OP_MUL_MAT_ID: case GGML_OP_OUT_PROD: { @@ -2343,15 +2374,20 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { { n_tasks = MIN(n_threads, ggml_nrows(node->src[0])); } break; - case GGML_OP_IM2COL: - case GGML_OP_IM2COL_FAST_1D: - case GGML_OP_IM2COL_BACK: - case GGML_OP_IM2COL_3D: - case GGML_OP_CONV_2D: + case GGML_OP_IM2COL: + case GGML_OP_IM2COL_FAST_1D: + case GGML_OP_IM2COL_ASYM: + case GGML_OP_IM2COL_BACK: + case GGML_OP_IM2COL_3D: + case GGML_OP_CONV_2D: case GGML_OP_CONV_3D: case GGML_OP_CONV_2D_DW: case GGML_OP_CONV_TRANSPOSE_1D: case GGML_OP_CONV_TRANSPOSE_2D: + case GGML_OP_ADD_SCALED: + case GGML_OP_RMS_NORM_SCALED: + case GGML_OP_MUL_MAT_ADD: + case GGML_OP_MUL_MAT_ADD_RELU: { n_tasks = n_threads; } break; @@ -2819,8 +2855,19 @@ struct ggml_cplan ggml_graph_plan( cur = ggml_type_size(node->type)*n_tasks; } break; case GGML_OP_MUL_MAT: - case GGML_OP_MUL_MAT_PACK4: + case GGML_OP_MUL_MAT_PACK4: { + if (node->src[0]->type == GGML_TYPE_I2_S) { + // The int8 activation rows, then one scale and one + // int8 row sum each. I2_S has no vec_dot_type entry + // -- see ggml_compute_forward_mul_mat_i2_s. + const int64_t nrows_y = ggml_nrows(node->src[1]); + + cur = GGML_PAD((size_t) ggml_nelements(node->src[1]), sizeof(float)); + cur += nrows_y*(sizeof(float) + sizeof(int32_t)); + break; + } + const enum ggml_type vec_dot_type = type_traits_cpu[node->src[0]->type].vec_dot_type; if (node->src[1]->type != vec_dot_type) { @@ -2948,6 +2995,18 @@ struct ggml_cplan ggml_graph_plan( { cur = ggml_type_size(node->type)*(n_tasks + node->src[0]->ne[0]*n_tasks); } break; + case GGML_OP_ADD_SCALED: + case GGML_OP_RMS_NORM_SCALED: + case GGML_OP_MUL_MAT_ADD: + case GGML_OP_MUL_MAT_ADD_RELU: + { + // F32 staging for the whole output, plus one absmax per + // thread. The output cannot be quantized in place: its + // scale is only known once every element has been + // computed, so all of them have to be held somewhere + // first. + cur = sizeof(float)*ggml_nelements(node) + sizeof(float)*n_tasks; + } break; case GGML_OP_GATED_DELTA_NET: { const int64_t S_v = node->src[2]->ne[0]; diff --git a/external/ggml/src/ggml-cpu/ops.cpp b/external/ggml/src/ggml-cpu/ops.cpp index 0f0f57399..92d976aaf 100644 --- a/external/ggml/src/ggml-cpu/ops.cpp +++ b/external/ggml/src/ggml-cpu/ops.cpp @@ -7,6 +7,7 @@ #include "ggml.h" #include "unary-ops.h" #include "vec.h" +#include "ggml-quants.h" #include #include @@ -531,6 +532,16 @@ void ggml_compute_forward_dup( if (src0->type == dst->type) { ggml_compute_forward_dup_bytes(params, dst); + // I8_S keeps one scale for the whole tensor, stored past the last + // element: a byte copy moves the payload but leaves that float + // uninitialized, so a cont/cpy of an I8_S view has to carry it over. + if (dst->type == GGML_TYPE_I8_S && params->ith == 0) { + // Reading through to the parent keeps this working for the permuted + // views that ggml_cont() is usually handed, which have no scale of + // their own. + const ggml_tensor * scale_src = src0->view_src ? src0->view_src : src0; + *ggml_inband_scale(dst) = *ggml_inband_scale_const(scale_src); + } return; } @@ -1908,36 +1919,36 @@ static void ggml_compute_forward_concat_any( GGML_TENSOR_BINARY_OP_LOCALS - const int32_t dim = ggml_get_op_params_i32(dst, 0); - - GGML_ASSERT(dim >= 0 && dim < 4); - - // MINITTS_CONCAT_FASTPATH: when concatenating along dim 1 with fully contiguous tensors, - // copy each (ne0 x ne1) plane as two bulk memcpys instead of element-wise indexing. - if (dim == 1 && ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst)) { - GGML_ASSERT(ne00 == ne10); - GGML_ASSERT(ne02 == ne12 && ne03 == ne13); - const int64_t planes = ne02 * ne03; - const int64_t plane_begin = planes * ith / nth; - const int64_t plane_end = planes * (ith + 1) / nth; - const size_t src0_bytes = (size_t) ne00 * ne01 * len; - const size_t src1_bytes = (size_t) ne10 * ne11 * len; - const size_t dst_plane_bytes = (size_t) ne0 * ne1 * len; - const size_t src0_plane_bytes = (size_t) ne00 * ne01 * len; - const size_t src1_plane_bytes = (size_t) ne10 * ne11 * len; - - for (int64_t plane = plane_begin; plane < plane_end; ++plane) { - char * dst_plane = (char *) dst->data + (size_t) plane * dst_plane_bytes; - const char * src0_plane = (const char *) src0->data + (size_t) plane * src0_plane_bytes; - const char * src1_plane = (const char *) src1->data + (size_t) plane * src1_plane_bytes; - memcpy(dst_plane, src0_plane, src0_bytes); - memcpy(dst_plane + src0_bytes, src1_plane, src1_bytes); - } - return; - } - - int64_t o[4] = {0, 0, 0, 0}; - o[dim] = src0->ne[dim]; + const int32_t dim = ggml_get_op_params_i32(dst, 0); + + GGML_ASSERT(dim >= 0 && dim < 4); + + // MINITTS_CONCAT_FASTPATH: when concatenating along dim 1 with fully contiguous tensors, + // copy each (ne0 x ne1) plane as two bulk memcpys instead of element-wise indexing. + if (dim == 1 && ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst)) { + GGML_ASSERT(ne00 == ne10); + GGML_ASSERT(ne02 == ne12 && ne03 == ne13); + const int64_t planes = ne02 * ne03; + const int64_t plane_begin = planes * ith / nth; + const int64_t plane_end = planes * (ith + 1) / nth; + const size_t src0_bytes = (size_t) ne00 * ne01 * len; + const size_t src1_bytes = (size_t) ne10 * ne11 * len; + const size_t dst_plane_bytes = (size_t) ne0 * ne1 * len; + const size_t src0_plane_bytes = (size_t) ne00 * ne01 * len; + const size_t src1_plane_bytes = (size_t) ne10 * ne11 * len; + + for (int64_t plane = plane_begin; plane < plane_end; ++plane) { + char * dst_plane = (char *) dst->data + (size_t) plane * dst_plane_bytes; + const char * src0_plane = (const char *) src0->data + (size_t) plane * src0_plane_bytes; + const char * src1_plane = (const char *) src1->data + (size_t) plane * src1_plane_bytes; + memcpy(dst_plane, src0_plane, src0_bytes); + memcpy(dst_plane + src0_bytes, src1_plane, src1_bytes); + } + return; + } + + int64_t o[4] = {0, 0, 0, 0}; + o[dim] = src0->ne[dim]; const char * x; @@ -2054,88 +2065,88 @@ static void ggml_compute_forward_concat_f32( const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; - const size_t len = ggml_type_size(src0->type); - GGML_ASSERT(len == sizeof(float)); + const size_t len = ggml_type_size(src0->type); + GGML_ASSERT(len == sizeof(float)); const int ith = params->ith; const int nth = params->nth; GGML_TENSOR_BINARY_OP_LOCALS - const int32_t dim = ggml_get_op_params_i32(dst, 0); - - GGML_ASSERT(dim >= 0 && dim < 4); - - // MINITTS_CONCAT_FASTPATH: when concatenating along dim 0 with fully contiguous tensors, - // copy each logical row as two bulk memcpys instead of scalar element dispatch. - if (dim == 0 && ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst)) { - GGML_ASSERT(ne01 == ne11 && ne02 == ne12 && ne03 == ne13); - const int64_t rows = ne1 * ne2 * ne3; - const int64_t row_begin = rows * ith / nth; - const int64_t row_end = rows * (ith + 1) / nth; - const size_t src0_bytes = (size_t) ne00 * len; - const size_t src1_bytes = (size_t) ne10 * len; - const size_t dst_row_bytes = (size_t) ne0 * len; - const size_t src0_row_bytes = (size_t) ne00 * len; - const size_t src1_row_bytes = (size_t) ne10 * len; - - for (int64_t row = row_begin; row < row_end; ++row) { - char * dst_row = (char *) dst->data + (size_t) row * dst_row_bytes; - const char * src0_row = (const char *) src0->data + (size_t) row * src0_row_bytes; - const char * src1_row = (const char *) src1->data + (size_t) row * src1_row_bytes; - memcpy(dst_row, src0_row, src0_bytes); - memcpy(dst_row + src0_bytes, src1_row, src1_bytes); - } - return; - } - - // MINITTS_CONCAT_FASTPATH: the f32 specialization gets the same plane-copy shortcut - // so common contiguous dim-1 concat patterns avoid the scalar fallback below. - if (dim == 1 && ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst)) { - GGML_ASSERT(ne00 == ne10); - GGML_ASSERT(ne02 == ne12 && ne03 == ne13); - const int64_t planes = ne02 * ne03; - const int64_t plane_begin = planes * ith / nth; - const int64_t plane_end = planes * (ith + 1) / nth; - const size_t src0_bytes = (size_t) ne00 * ne01 * len; - const size_t src1_bytes = (size_t) ne10 * ne11 * len; - const size_t dst_plane_bytes = (size_t) ne0 * ne1 * len; - const size_t src0_plane_bytes = (size_t) ne00 * ne01 * len; - const size_t src1_plane_bytes = (size_t) ne10 * ne11 * len; - - for (int64_t plane = plane_begin; plane < plane_end; ++plane) { - char * dst_plane = (char *) dst->data + (size_t) plane * dst_plane_bytes; - const char * src0_plane = (const char *) src0->data + (size_t) plane * src0_plane_bytes; - const char * src1_plane = (const char *) src1->data + (size_t) plane * src1_plane_bytes; - memcpy(dst_plane, src0_plane, src0_bytes); - memcpy(dst_plane + src0_bytes, src1_plane, src1_bytes); - } - return; - } - - int64_t o[4] = {0, 0, 0, 0}; - o[dim] = src0->ne[dim]; - - const float * x; - - // TODO: smarter multi-theading - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = ith; i2 < ne2; i2 += nth) { - for (int i1 = 0; i1 < ne1; i1++) { - for (int i0 = 0; i0 < ne0; i0++) { - if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { - x = (const float *) ((const char *)src0->data + (i0 )*nb00 + (i1 )*nb01 + (i2 )*nb02 + (i3 )*nb03); - } else { - x = (const float *) ((const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13); - } - - float * y = (float *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); - - *y = *x; - } - } - } - } + const int32_t dim = ggml_get_op_params_i32(dst, 0); + + GGML_ASSERT(dim >= 0 && dim < 4); + + // MINITTS_CONCAT_FASTPATH: when concatenating along dim 0 with fully contiguous tensors, + // copy each logical row as two bulk memcpys instead of scalar element dispatch. + if (dim == 0 && ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst)) { + GGML_ASSERT(ne01 == ne11 && ne02 == ne12 && ne03 == ne13); + const int64_t rows = ne1 * ne2 * ne3; + const int64_t row_begin = rows * ith / nth; + const int64_t row_end = rows * (ith + 1) / nth; + const size_t src0_bytes = (size_t) ne00 * len; + const size_t src1_bytes = (size_t) ne10 * len; + const size_t dst_row_bytes = (size_t) ne0 * len; + const size_t src0_row_bytes = (size_t) ne00 * len; + const size_t src1_row_bytes = (size_t) ne10 * len; + + for (int64_t row = row_begin; row < row_end; ++row) { + char * dst_row = (char *) dst->data + (size_t) row * dst_row_bytes; + const char * src0_row = (const char *) src0->data + (size_t) row * src0_row_bytes; + const char * src1_row = (const char *) src1->data + (size_t) row * src1_row_bytes; + memcpy(dst_row, src0_row, src0_bytes); + memcpy(dst_row + src0_bytes, src1_row, src1_bytes); + } + return; + } + + // MINITTS_CONCAT_FASTPATH: the f32 specialization gets the same plane-copy shortcut + // so common contiguous dim-1 concat patterns avoid the scalar fallback below. + if (dim == 1 && ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst)) { + GGML_ASSERT(ne00 == ne10); + GGML_ASSERT(ne02 == ne12 && ne03 == ne13); + const int64_t planes = ne02 * ne03; + const int64_t plane_begin = planes * ith / nth; + const int64_t plane_end = planes * (ith + 1) / nth; + const size_t src0_bytes = (size_t) ne00 * ne01 * len; + const size_t src1_bytes = (size_t) ne10 * ne11 * len; + const size_t dst_plane_bytes = (size_t) ne0 * ne1 * len; + const size_t src0_plane_bytes = (size_t) ne00 * ne01 * len; + const size_t src1_plane_bytes = (size_t) ne10 * ne11 * len; + + for (int64_t plane = plane_begin; plane < plane_end; ++plane) { + char * dst_plane = (char *) dst->data + (size_t) plane * dst_plane_bytes; + const char * src0_plane = (const char *) src0->data + (size_t) plane * src0_plane_bytes; + const char * src1_plane = (const char *) src1->data + (size_t) plane * src1_plane_bytes; + memcpy(dst_plane, src0_plane, src0_bytes); + memcpy(dst_plane + src0_bytes, src1_plane, src1_bytes); + } + return; + } + + int64_t o[4] = {0, 0, 0, 0}; + o[dim] = src0->ne[dim]; + + const float * x; + + // TODO: smarter multi-theading + for (int i3 = 0; i3 < ne3; i3++) { + for (int i2 = ith; i2 < ne2; i2 += nth) { + for (int i1 = 0; i1 < ne1; i1++) { + for (int i0 = 0; i0 < ne0; i0++) { + if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { + x = (const float *) ((const char *)src0->data + (i0 )*nb00 + (i1 )*nb01 + (i2 )*nb02 + (i3 )*nb03); + } else { + x = (const float *) ((const char *)src1->data + (i0 - o[0])*nb10 + (i1 - o[1])*nb11 + (i2 - o[2])*nb12 + (i3 - o[3])*nb13); + } + + float * y = (float *)((char *)dst->data + i0*nb0 + i1*nb1 + i2*nb2 + i3*nb3); + + *y = *x; + } + } + } + } } void ggml_compute_forward_concat( @@ -4501,52 +4512,52 @@ static void ggml_compute_forward_scale_f32( const ggml_compute_params * params, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - - GGML_ASSERT(ggml_is_contiguous(src0)); - GGML_ASSERT(ggml_is_contiguous(dst)); - GGML_ASSERT(ggml_can_repeat(src0, dst)); - - GGML_TENSOR_UNARY_OP_LOCALS - - float s; // scale factor - float b; // bias - + const ggml_tensor * src0 = dst->src[0]; + + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(ggml_can_repeat(src0, dst)); + + GGML_TENSOR_UNARY_OP_LOCALS + + float s; // scale factor + float b; // bias + memcpy(&s, (float *) dst->op_params + 0, sizeof(float)); memcpy(&b, (float *) dst->op_params + 1, sizeof(float)); const int ith = params->ith; const int nth = params->nth; - const int64_t nc = ne0; - const int64_t nr = ggml_nrows(dst); - - // rows per thread - const int64_t dr = (nr + nth - 1)/nth; - - // row range for this thread - const int64_t ir0 = dr*ith; - const int64_t ir1 = MIN(ir0 + dr, nr); - - const bool is_src0_full_shape = ggml_are_same_shape(src0, dst); - - if (!is_src0_full_shape) { - for (int64_t ir = ir0; ir < ir1; ++ir) { - const int64_t i3 = ir/(ne2*ne1); - const int64_t i2 = (ir - i3*ne2*ne1)/ne1; - const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1); - float * dst_row = (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1); - const char * src0_row = (const char *) src0->data + - (i3 % ne03)*nb03 + (i2 % ne02)*nb02 + (i1 % ne01)*nb01; - for (int64_t i0 = 0; i0 < nc; ++i0) { - const float value = *(const float *) (src0_row + (i0 % ne00)*nb00); - dst_row[i0] = value * s + b; - } - } - return; - } - - if (b == 0.0f) { + const int64_t nc = ne0; + const int64_t nr = ggml_nrows(dst); + + // rows per thread + const int64_t dr = (nr + nth - 1)/nth; + + // row range for this thread + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); + + const bool is_src0_full_shape = ggml_are_same_shape(src0, dst); + + if (!is_src0_full_shape) { + for (int64_t ir = ir0; ir < ir1; ++ir) { + const int64_t i3 = ir/(ne2*ne1); + const int64_t i2 = (ir - i3*ne2*ne1)/ne1; + const int64_t i1 = (ir - i3*ne2*ne1 - i2*ne1); + float * dst_row = (float *) ((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1); + const char * src0_row = (const char *) src0->data + + (i3 % ne03)*nb03 + (i2 % ne02)*nb02 + (i1 % ne01)*nb01; + for (int64_t i0 = 0; i0 < nc; ++i0) { + const float value = *(const float *) (src0_row + (i0 % ne00)*nb00); + dst_row[i0] = value * s + b; + } + } + return; + } + + if (b == 0.0f) { for (int i1 = ir0; i1 < ir1; i1++) { if (dst->data != src0->data) { // src0 is same shape as dst => same indices @@ -4557,10 +4568,10 @@ static void ggml_compute_forward_scale_f32( } } else { for (int i1 = ir0; i1 < ir1; i1++) { - ggml_vec_mad1_f32(nc, - (float *) ((char *) dst->data + i1*nb1), - (float *) ((char *) src0->data + i1*nb01), - s, b); + ggml_vec_mad1_f32(nc, + (float *) ((char *) dst->data + i1*nb1), + (float *) ((char *) src0->data + i1*nb01), + s, b); } } } @@ -6420,9 +6431,9 @@ static void ggml_compute_forward_im2col_f16( } } -void ggml_compute_forward_im2col( - const ggml_compute_params * params, - ggml_tensor * dst) { +void ggml_compute_forward_im2col( + const ggml_compute_params * params, + ggml_tensor * dst) { switch (dst->type) { case GGML_TYPE_F16: { @@ -6436,188 +6447,188 @@ void ggml_compute_forward_im2col( { GGML_ABORT("fatal error"); } - } -} - -static inline int64_t ggml_div_ceil_nonneg(int64_t num, int64_t den) { - GGML_ASSERT(num >= 0); - GGML_ASSERT(den > 0); - return (num + den - 1) / den; -} - -static inline int64_t ggml_im2col_1d_valid_start(int64_t base, int64_t step) { - if (base >= 0) { - return 0; - } - return ggml_div_ceil_nonneg(-base, step); -} - -static inline int64_t ggml_im2col_1d_valid_end(int64_t base, int64_t step, int64_t width, int64_t kw) { - if (width <= 0 || base > width - 1) { - return 0; - } - const int64_t max_ikw = (width - 1 - base) / step; - return std::min(kw, max_ikw + 1); -} - -static void ggml_compute_forward_im2col_f32_1d( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - - GGML_TENSOR_BINARY_OP_LOCALS; - - const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; - const int32_t p0 = ((const int32_t *)(dst->op_params))[2]; - const int32_t d0 = ((const int32_t *)(dst->op_params))[4]; - - const int ith = params->ith; - const int nth = params->nth; - - const int64_t N = ne12; - const int64_t IC = ne11; - const int64_t IW = ne10; - const int64_t KW = ne00; - const int64_t OW = ne1; - - const int ofs0 = nb12; - const int ofs1 = nb11; - - GGML_ASSERT(nb10 == sizeof(float)); - - float * const wdata = (float *) dst->data; - - for (int64_t in = 0; in < N; ++in) { - for (int64_t iow = 0; iow < OW; ++iow) { - const int64_t base = iow*s0 - p0; - const int64_t ikw0 = ggml_im2col_1d_valid_start(base, d0); - const int64_t ikw1 = ggml_im2col_1d_valid_end(base, d0, IW, KW); - - for (int64_t iic = ith; iic < IC; iic += nth) { - float * const dst_row = wdata + (in*OW + iow)*(IC*KW) + iic*KW; - const float * const src_row = (const float *)((const char *) src1->data + in*ofs0 + iic*ofs1); - - if (ikw0 > 0) { - memset(dst_row, 0, ikw0*sizeof(float)); - } - if (ikw1 > ikw0) { - if (d0 == 1) { - memcpy(dst_row + ikw0, src_row + base + ikw0, (ikw1 - ikw0)*sizeof(float)); - } else { - int64_t iiw = base + ikw0*d0; - for (int64_t ikw = ikw0; ikw < ikw1; ++ikw, iiw += d0) { - dst_row[ikw] = src_row[iiw]; - } - } - } - if (ikw1 < KW) { - memset(dst_row + ikw1, 0, (KW - ikw1)*sizeof(float)); - } - } - } - } -} - -static void ggml_compute_forward_im2col_f16_1d( - const ggml_compute_params * params, - ggml_tensor * dst) { - - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - - GGML_ASSERT(src0->type == GGML_TYPE_F16); - GGML_ASSERT(src1->type == GGML_TYPE_F16 || src1->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F16); - - GGML_TENSOR_BINARY_OP_LOCALS; - - const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; - const int32_t p0 = ((const int32_t *)(dst->op_params))[2]; - const int32_t d0 = ((const int32_t *)(dst->op_params))[4]; - - const int ith = params->ith; - const int nth = params->nth; - - const int64_t N = ne12; - const int64_t IC = ne11; - const int64_t IW = ne10; - const int64_t KW = ne00; - const int64_t OW = ne1; - - const int ofs0 = nb12; - const int ofs1 = nb11; - - GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); - GGML_ASSERT(nb10 == ggml_type_size(src1->type)); - - ggml_fp16_t * const wdata = (ggml_fp16_t *) dst->data; - - for (int64_t in = 0; in < N; ++in) { - for (int64_t iow = 0; iow < OW; ++iow) { - const int64_t base = iow*s0 - p0; - const int64_t ikw0 = ggml_im2col_1d_valid_start(base, d0); - const int64_t ikw1 = ggml_im2col_1d_valid_end(base, d0, IW, KW); - - for (int64_t iic = ith; iic < IC; iic += nth) { - ggml_fp16_t * const dst_row = wdata + (in*OW + iow)*(IC*KW) + iic*KW; - const float * const src_row_f32 = src1->type == GGML_TYPE_F32 - ? (const float *)((const char *) src1->data + in*ofs0 + iic*ofs1) - : nullptr; - const ggml_fp16_t * const src_row_f16 = src1->type == GGML_TYPE_F16 - ? (const ggml_fp16_t *)((const char *) src1->data + in*ofs0 + iic*ofs1) - : nullptr; - - if (ikw0 > 0) { - memset(dst_row, 0, ikw0*sizeof(ggml_fp16_t)); - } - if (ikw1 > ikw0) { - if (src_row_f16 != nullptr && d0 == 1) { - memcpy(dst_row + ikw0, src_row_f16 + base + ikw0, (ikw1 - ikw0)*sizeof(ggml_fp16_t)); - } else if (src_row_f16 != nullptr) { - int64_t iiw = base + ikw0*d0; - for (int64_t ikw = ikw0; ikw < ikw1; ++ikw, iiw += d0) { - dst_row[ikw] = src_row_f16[iiw]; - } - } else { - int64_t iiw = base + ikw0*d0; - for (int64_t ikw = ikw0; ikw < ikw1; ++ikw, iiw += d0) { - dst_row[ikw] = GGML_CPU_FP32_TO_FP16(src_row_f32[iiw]); - } - } - } - if (ikw1 < KW) { - memset(dst_row + ikw1, 0, (KW - ikw1)*sizeof(ggml_fp16_t)); - } - } - } - } -} - -void ggml_compute_forward_im2col_fast_1d( - const ggml_compute_params * params, - ggml_tensor * dst) { - switch (dst->type) { - case GGML_TYPE_F16: - { - ggml_compute_forward_im2col_f16_1d(params, dst); - } break; - case GGML_TYPE_F32: - { - ggml_compute_forward_im2col_f32_1d(params, dst); - } break; - default: - { - GGML_ABORT("fatal error"); - } - } -} - -// ggml_compute_forward_im2col_back_f32 + } +} + +static inline int64_t ggml_div_ceil_nonneg(int64_t num, int64_t den) { + GGML_ASSERT(num >= 0); + GGML_ASSERT(den > 0); + return (num + den - 1) / den; +} + +static inline int64_t ggml_im2col_1d_valid_start(int64_t base, int64_t step) { + if (base >= 0) { + return 0; + } + return ggml_div_ceil_nonneg(-base, step); +} + +static inline int64_t ggml_im2col_1d_valid_end(int64_t base, int64_t step, int64_t width, int64_t kw) { + if (width <= 0 || base > width - 1) { + return 0; + } + const int64_t max_ikw = (width - 1 - base) / step; + return std::min(kw, max_ikw + 1); +} + +static void ggml_compute_forward_im2col_f32_1d( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + GGML_TENSOR_BINARY_OP_LOCALS; + + const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; + const int32_t p0 = ((const int32_t *)(dst->op_params))[2]; + const int32_t d0 = ((const int32_t *)(dst->op_params))[4]; + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t N = ne12; + const int64_t IC = ne11; + const int64_t IW = ne10; + const int64_t KW = ne00; + const int64_t OW = ne1; + + const int ofs0 = nb12; + const int ofs1 = nb11; + + GGML_ASSERT(nb10 == sizeof(float)); + + float * const wdata = (float *) dst->data; + + for (int64_t in = 0; in < N; ++in) { + for (int64_t iow = 0; iow < OW; ++iow) { + const int64_t base = iow*s0 - p0; + const int64_t ikw0 = ggml_im2col_1d_valid_start(base, d0); + const int64_t ikw1 = ggml_im2col_1d_valid_end(base, d0, IW, KW); + + for (int64_t iic = ith; iic < IC; iic += nth) { + float * const dst_row = wdata + (in*OW + iow)*(IC*KW) + iic*KW; + const float * const src_row = (const float *)((const char *) src1->data + in*ofs0 + iic*ofs1); + + if (ikw0 > 0) { + memset(dst_row, 0, ikw0*sizeof(float)); + } + if (ikw1 > ikw0) { + if (d0 == 1) { + memcpy(dst_row + ikw0, src_row + base + ikw0, (ikw1 - ikw0)*sizeof(float)); + } else { + int64_t iiw = base + ikw0*d0; + for (int64_t ikw = ikw0; ikw < ikw1; ++ikw, iiw += d0) { + dst_row[ikw] = src_row[iiw]; + } + } + } + if (ikw1 < KW) { + memset(dst_row + ikw1, 0, (KW - ikw1)*sizeof(float)); + } + } + } + } +} + +static void ggml_compute_forward_im2col_f16_1d( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + GGML_ASSERT(src0->type == GGML_TYPE_F16); + GGML_ASSERT(src1->type == GGML_TYPE_F16 || src1->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F16); + + GGML_TENSOR_BINARY_OP_LOCALS; + + const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; + const int32_t p0 = ((const int32_t *)(dst->op_params))[2]; + const int32_t d0 = ((const int32_t *)(dst->op_params))[4]; + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t N = ne12; + const int64_t IC = ne11; + const int64_t IW = ne10; + const int64_t KW = ne00; + const int64_t OW = ne1; + + const int ofs0 = nb12; + const int ofs1 = nb11; + + GGML_ASSERT(nb00 == sizeof(ggml_fp16_t)); + GGML_ASSERT(nb10 == ggml_type_size(src1->type)); + + ggml_fp16_t * const wdata = (ggml_fp16_t *) dst->data; + + for (int64_t in = 0; in < N; ++in) { + for (int64_t iow = 0; iow < OW; ++iow) { + const int64_t base = iow*s0 - p0; + const int64_t ikw0 = ggml_im2col_1d_valid_start(base, d0); + const int64_t ikw1 = ggml_im2col_1d_valid_end(base, d0, IW, KW); + + for (int64_t iic = ith; iic < IC; iic += nth) { + ggml_fp16_t * const dst_row = wdata + (in*OW + iow)*(IC*KW) + iic*KW; + const float * const src_row_f32 = src1->type == GGML_TYPE_F32 + ? (const float *)((const char *) src1->data + in*ofs0 + iic*ofs1) + : nullptr; + const ggml_fp16_t * const src_row_f16 = src1->type == GGML_TYPE_F16 + ? (const ggml_fp16_t *)((const char *) src1->data + in*ofs0 + iic*ofs1) + : nullptr; + + if (ikw0 > 0) { + memset(dst_row, 0, ikw0*sizeof(ggml_fp16_t)); + } + if (ikw1 > ikw0) { + if (src_row_f16 != nullptr && d0 == 1) { + memcpy(dst_row + ikw0, src_row_f16 + base + ikw0, (ikw1 - ikw0)*sizeof(ggml_fp16_t)); + } else if (src_row_f16 != nullptr) { + int64_t iiw = base + ikw0*d0; + for (int64_t ikw = ikw0; ikw < ikw1; ++ikw, iiw += d0) { + dst_row[ikw] = src_row_f16[iiw]; + } + } else { + int64_t iiw = base + ikw0*d0; + for (int64_t ikw = ikw0; ikw < ikw1; ++ikw, iiw += d0) { + dst_row[ikw] = GGML_CPU_FP32_TO_FP16(src_row_f32[iiw]); + } + } + } + if (ikw1 < KW) { + memset(dst_row + ikw1, 0, (KW - ikw1)*sizeof(ggml_fp16_t)); + } + } + } + } +} + +void ggml_compute_forward_im2col_fast_1d( + const ggml_compute_params * params, + ggml_tensor * dst) { + switch (dst->type) { + case GGML_TYPE_F16: + { + ggml_compute_forward_im2col_f16_1d(params, dst); + } break; + case GGML_TYPE_F32: + { + ggml_compute_forward_im2col_f32_1d(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + +// ggml_compute_forward_im2col_back_f32 void ggml_compute_forward_im2col_back_f32( const ggml_compute_params * params, @@ -11637,3 +11648,663 @@ void ggml_compute_forward_fwht(const ggml_compute_params * params, ggml_tensor * } } } + +// VibeASR CPU INT8 pipeline +// +// Scalar reference implementations of the five fused I8_S ops. They are written +// to be obviously correct rather than fast: the vectorized kernels live with the +// other SIMD code under arch/, and dispatch to them replaces the inner loops +// here without changing the surrounding structure. +// +// Four of the five produce I8_S, and all four share one shape: +// +// 1. compute the result in F32 into params->wdata, tracking this thread's absmax +// 2. barrier, reduce the per-thread absmaxes to a tensor-wide absmax +// 3. quantize this thread's slice into dst; thread 0 writes the output scale +// +// Step 2 is why these are fused ops at all. The output scale cannot be known +// before the whole output has been computed, so an unfused chain would have to +// materialize each intermediate in F32, rescan it, and write it again. +// +// wdata layout, shared by all four: +// +// [0, n_stage) F32 staging for the result +// [n_stage, n_stage + nth) per-thread absmax +// [n_stage + nth, ...) op-specific scratch (int32 accumulators) + +// Steps 2 and 3. The thread's slice of the output is the rectangle +// {row*row_stride + col : row < n_rows, col0 <= col < col1}; contiguous callers +// pass n_rows = 1, row_stride = 0. Both `stage` and dst use the same indexing. +// +// `relu` clamps after rounding rather than before the absmax scan, so the +// negatives that are about to be zeroed still widen the scale: an output +// spanning [-10, +5] scales by 10, leaving the surviving positives to reach 63 +// of the available 127. That is what the reference implementation does, and +// matching it is what makes a layer-by-layer comparison meaningful, so it is +// reproduced here deliberately rather than improved in passing. +// Scale a contiguous run of staged F32 into I8_S. +// +// relu is folded into the low clamp: clamping to 0 before rounding and rounding +// before clamping to 0 agree on every negative input, and the float clamp is +// free here. The absmax that produced id was taken before the clamp, so a +// channel about to be zeroed still counts towards the scale. +// +// roundf() would be a PLT call per element. rintf() inlines to one instruction +// and rounds ties to even, which is what _mm256_cvtps_epi32 does as well, so the +// vector body and the scalar tail agree - and it matches ggml's own +// nearest_int(). VibeASR rounds ties away from zero in its scalar tail but to +// even in its vector body, so no tie convention reproduces it exactly. +static inline void ggml_i8_s_quantize_range( + int8_t * GGML_RESTRICT dst, + const float * GGML_RESTRICT src, + int64_t n, + float id, + bool relu) { + + int64_t i = 0; + + const float lo = relu ? 0.0f : -127.0f; + +#if defined(__AVX2__) + const __m256 v_id = _mm256_set1_ps(id); + const __m256 v_lo = _mm256_set1_ps(lo); + const __m256 v_hi = _mm256_set1_ps(127.0f); + + for (; i + 8 <= n; i += 8) { + __m256 vf = _mm256_mul_ps(_mm256_loadu_ps(src + i), v_id); + vf = _mm256_min_ps(_mm256_max_ps(vf, v_lo), v_hi); + + const __m256i vi32 = _mm256_cvtps_epi32(vf); + + __m256i vi16 = _mm256_permute4x64_epi64(_mm256_packs_epi32(vi32, vi32), 0xD8); + __m256i vi8 = _mm256_permute4x64_epi64(_mm256_packs_epi16(vi16, vi16), 0xD8); + + _mm_storel_epi64((__m128i *)(dst + i), _mm256_castsi256_si128(vi8)); + } +#elif defined(__ARM_NEON) && defined(__aarch64__) + const float32x4_t v_id = vdupq_n_f32(id); + const float32x4_t v_lo = vdupq_n_f32(lo); + const float32x4_t v_hi = vdupq_n_f32(127.0f); + + for (; i + 8 <= n; i += 8) { + float32x4_t f0 = vmulq_f32(vld1q_f32(src + i ), v_id); + float32x4_t f1 = vmulq_f32(vld1q_f32(src + i + 4), v_id); + + f0 = vminq_f32(vmaxq_f32(f0, v_lo), v_hi); + f1 = vminq_f32(vmaxq_f32(f1, v_lo), v_hi); + + const int16x8_t vi16 = vcombine_s16(vqmovn_s32(vcvtnq_s32_f32(f0)), + vqmovn_s32(vcvtnq_s32_f32(f1))); + + vst1_s8(dst + i, vqmovn_s16(vi16)); + } +#endif + + for (; i < n; ++i) { + float v = src[i] * id; + v = v < lo ? lo : (v > 127.0f ? 127.0f : v); + dst[i] = (int8_t) rintf(v); + } +} + +static void ggml_i8_s_requantize( + const ggml_compute_params * params, + ggml_tensor * dst, + const float * stage, + float * thread_max, + int64_t row_stride, + int64_t n_rows, + int64_t col0, + int64_t col1, + float local_absmax, + bool relu) { + + const int ith = params->ith; + const int nth = params->nth; + + thread_max[ith] = local_absmax; + + ggml_barrier(params->threadpool); + + if (ith == 0) { + float global_max = 0.0f; + for (int t = 0; t < nth; t++) { + if (thread_max[t] > global_max) global_max = thread_max[t]; + } + thread_max[0] = global_max; + } + + ggml_barrier(params->threadpool); + + const float global_max = thread_max[0]; + const float id = global_max != 0.0f ? 127.0f / global_max : 0.0f; + + int8_t * dst_data = (int8_t *) dst->data; + + // Each row's [col0, col1) slice is contiguous, so the vectorized helper runs + // once per row regardless of how the rectangle is strided. + for (int64_t row = 0; row < n_rows; row++) { + const int64_t i = row*row_stride + col0; + + ggml_i8_s_quantize_range(dst_data + i, stage + i, col1 - col0, id, relu); + } + + if (ith == 0) { + // Multiplier, so that dequantizing is q * scale like every other ggml + // quantized type -- see ggml_i8_s_to_float. + *ggml_inband_scale(dst) = global_max != 0.0f ? global_max / 127.0f : 0.0f; + } +} + +// ggml_compute_forward_add_scaled + +void ggml_compute_forward_add_scaled( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; // a (I8_S) + const ggml_tensor * src1 = dst->src[1]; // b (I8_S) + const ggml_tensor * src2 = dst->src[2]; // scale (F32, per channel) + + GGML_ASSERT(src0->type == GGML_TYPE_I8_S); + GGML_ASSERT(src1->type == GGML_TYPE_I8_S); + GGML_ASSERT(dst->type == GGML_TYPE_I8_S); + GGML_ASSERT(ggml_are_same_shape(src0, src1)); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(src1)); + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t n = ggml_nelements(src0); + const int64_t ne0 = src0->ne[0]; + + const int8_t * a = (const int8_t *) src0->data; + const int8_t * b = (const int8_t *) src1->data; + + const float a_d = *ggml_inband_scale_const(src0); + const float b_d = *ggml_inband_scale_const(src1); + + // Channel index is i % ne0: the graph is channel-major, so ne[0] is the + // channel axis and one coefficient covers every position of that channel. + const float * gamma = (const float *) src2->data; + + float * stage = (float *) params->wdata; + float * thread_max = stage + n; + + const int64_t dr = (n + nth - 1)/nth; + const int64_t i0 = dr*ith; + const int64_t i1 = MIN(i0 + dr, n); + + float local_absmax = 0.0f; + + for (int64_t i = i0; i < i1; i++) { + const float v = (float) a[i] * a_d * gamma[i % ne0] + (float) b[i] * b_d; + + stage[i] = v; + + const float av = fabsf(v); + if (av > local_absmax) local_absmax = av; + } + + ggml_i8_s_requantize(params, dst, stage, thread_max, 0, 1, i0, i1, local_absmax, false); +} + +// ggml_compute_forward_rms_norm_scaled + +void ggml_compute_forward_rms_norm_scaled( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; // a (I8_S) + const ggml_tensor * src1 = dst->src[1]; // scale (F32, per channel) + + GGML_ASSERT(src0->type == GGML_TYPE_I8_S); + GGML_ASSERT(dst->type == GGML_TYPE_I8_S); + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(ggml_is_contiguous(src0)); + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t ne00 = src0->ne[0]; + const int64_t n = ggml_nelements(src0); + const int64_t nr = n / ne00; + + float eps; + memcpy(&eps, dst->op_params, sizeof(float)); + + const int8_t * x = (const int8_t *) src0->data; + const float * gamma = (const float *) src1->data; + + const float x_d = *ggml_inband_scale_const(src0); + + // The normalization is scale-invariant, so the input scale cancels between + // numerator and denominator and never has to be applied to the int8 values: + // + // real[i] = q[i] * x_d + // out[i] = real[i] / sqrt(mean(real^2) + eps) * gamma[i] + // = q[i] / sqrt(mean(q^2) + eps/x_d^2) * gamma[i] + // + // which keeps the sum of squares in exact int64 arithmetic. Only eps has to + // move into the int8 domain. + const float eps_q = (float) ((double) eps / ((double) x_d * (double) x_d)); + + float * stage = (float *) params->wdata; + float * thread_max = stage + n; + + const int64_t dr = (nr + nth - 1)/nth; + const int64_t ir0 = dr*ith; + const int64_t ir1 = MIN(ir0 + dr, nr); + + float local_absmax = 0.0f; + + for (int64_t ir = ir0; ir < ir1; ir++) { + const int8_t * x_row = x + ir*ne00; + float * stage_row = stage + ir*ne00; + + int64_t sum_sq = 0; + for (int64_t i = 0; i < ne00; i++) { + sum_sq += (int32_t) x_row[i] * (int32_t) x_row[i]; + } + + const float rms_inv = 1.0f/sqrtf((float) sum_sq/(float) ne00 + eps_q); + + for (int64_t i = 0; i < ne00; i++) { + const float v = (float) x_row[i] * rms_inv * gamma[i]; + + stage_row[i] = v; + + const float av = fabsf(v); + if (av > local_absmax) local_absmax = av; + } + } + + ggml_i8_s_requantize(params, dst, stage, thread_max, 0, 1, ir0*ne00, ir1*ne00, local_absmax, false); +} + +// ggml_compute_forward_mul_mat_add + +#define GGML_MUL_MAT_ADD_OC_CHUNK 64 + +// stage[j] = acc[j]*d + bias[j], returning max(|stage[j]|) folded into amax. +// +// This is the second pass over every output, so leaving it scalar costs about as +// much as the dot products do at the small contraction lengths the conv layers +// use. Vectorizing it is exact: max over floats is order-independent, and the +// multiply and add are kept separate so the tail cannot disagree with the body. +static inline float ggml_i8_s_scale_bias_absmax( + float * GGML_RESTRICT stage, + const int32_t * GGML_RESTRICT acc, + const float * GGML_RESTRICT bias, + int64_t n, + float d, + float amax) { + + int64_t i = 0; + +#if defined(__AVX2__) + const __m256 v_d = _mm256_set1_ps(d); + const __m256 v_abs = _mm256_castsi256_ps(_mm256_set1_epi32(0x7fffffff)); + + __m256 v_amax = _mm256_setzero_ps(); + + for (; i + 8 <= n; i += 8) { + const __m256 v = _mm256_add_ps( + _mm256_mul_ps(_mm256_cvtepi32_ps(_mm256_loadu_si256((const __m256i *)(acc + i))), v_d), + _mm256_loadu_ps(bias + i)); + + _mm256_storeu_ps(stage + i, v); + + v_amax = _mm256_max_ps(v_amax, _mm256_and_ps(v, v_abs)); + } + + if (i > 0) { + __m128 r = _mm_max_ps(_mm256_castps256_ps128(v_amax), _mm256_extractf128_ps(v_amax, 1)); + r = _mm_max_ps(r, _mm_movehl_ps(r, r)); + r = _mm_max_ss(r, _mm_shuffle_ps(r, r, 1)); + + const float v = _mm_cvtss_f32(r); + if (v > amax) amax = v; + } +#elif defined(__ARM_NEON) && defined(__aarch64__) + const float32x4_t v_d = vdupq_n_f32(d); + + float32x4_t v_amax = vdupq_n_f32(0.0f); + + for (; i + 4 <= n; i += 4) { + const float32x4_t v = vaddq_f32( + vmulq_f32(vcvtq_f32_s32(vld1q_s32(acc + i)), v_d), + vld1q_f32(bias + i)); + + vst1q_f32(stage + i, v); + + v_amax = vmaxq_f32(v_amax, vabsq_f32(v)); + } + + if (i > 0) { + const float v = vmaxvq_f32(v_amax); + if (v > amax) amax = v; + } +#endif + + for (; i < n; ++i) { + const float v = (float) acc[i]*d + bias[i]; + + stage[i] = v; + + const float av = fabsf(v); + if (av > amax) amax = av; + } + + return amax; +} + +// Shared by GGML_OP_MUL_MAT_ADD and GGML_OP_MUL_MAT_ADD_RELU. +// +// Two shapes are handled. When src0 is [K, 1, C] the op is depthwise: C +// independent length-K dot products per position, with src1 laid out +// [K, N, C] and the output [C, N] channel-major. Otherwise src0 is an ordinary +// [IC, OC] weight matrix, src1 is [IC, N], and the output is [OC, N]. +static void ggml_compute_forward_mul_mat_add_impl( + const ggml_compute_params * params, + ggml_tensor * dst, + bool relu) { + + const ggml_tensor * src0 = dst->src[0]; // weight (I8_S) + const ggml_tensor * src1 = dst->src[1]; // input (I8_S) + const ggml_tensor * src2 = dst->src[2]; // bias (F32) + + GGML_ASSERT(src0->type == GGML_TYPE_I8_S); + GGML_ASSERT(src1->type == GGML_TYPE_I8_S); + GGML_ASSERT(dst->type == GGML_TYPE_I8_S); + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(ggml_is_contiguous(src1)); + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t ne00 = src0->ne[0]; // IC, or K in the depthwise case + const int64_t ne01 = src0->ne[1]; // OC, or 1 + const int64_t ne02 = src0->ne[2]; // 1, or C + const int64_t ne11 = src1->ne[1]; // N + + const float * bias = (const float *) src2->data; + + // Both scales are multipliers, so they combine by multiplication and the + // int32 accumulator only has to be scaled once. + const float combined_d = *ggml_inband_scale_const(src0) * *ggml_inband_scale_const(src1); + + const int8_t * w = (const int8_t *) src0->data; + const int8_t * x = (const int8_t *) src1->data; + + const bool depthwise = ne01 == 1 && ne02 > 1; + + // Output channel count, and the number of outputs per column. + const int64_t OC = depthwise ? ne02 : ne01; + + // The loops below walk a single batch. Batched matmul has no caller in the + // VibeASR graphs, so refuse it here rather than silently leaving + // dst->ne[2]/ne[3] beyond the first as uninitialized garbage. In the + // depthwise case ne[2] is the channel axis, which is handled. + GGML_ASSERT(dst->ne[3] == 1); + GGML_ASSERT(depthwise || dst->ne[2] == 1); + GGML_ASSERT(ne11*OC == ggml_nelements(dst)); + + float * stage = (float *) params->wdata; + float * thread_max = stage + ne11*OC; + + // Split over columns: every thread owns whole columns of the output, so the + // int32 accumulators need no cross-thread coordination. + const int64_t dr = (ne11 + nth - 1)/nth; + const int64_t col0 = dr*ith; + const int64_t col1 = MIN(col0 + dr, ne11); + + float local_absmax = 0.0f; + + if (depthwise) { + // stage/dst index is ch*ne11 + col, so each channel is a row of length + // ne11 and this thread owns the [col0, col1) band of every row. + // + // The batched roles are inverted here relative to the matmul path: the + // filter is the single row and the positions are the nrc rows, since it is + // the positions that are strided by ne00 and the filter that is reused. + // That also makes the results contiguous, so the scale/bias pass vectorizes. + int32_t acc[GGML_MUL_MAT_ADD_OC_CHUNK]; + float bias_v[GGML_MUL_MAT_ADD_OC_CHUNK]; + + for (int64_t ch = 0; ch < ne02; ch++) { + const int8_t * w_ch = w + ch*ne00; + + // Broadcast once per channel, not per position. + for (int64_t j = 0; j < GGML_MUL_MAT_ADD_OC_CHUNK; j++) { + bias_v[j] = bias[ch]; + } + + for (int64_t col = col0; col < col1; col += GGML_MUL_MAT_ADD_OC_CHUNK) { + const int64_t ncol = MIN((int64_t) GGML_MUL_MAT_ADD_OC_CHUNK, col1 - col); + + // src1 is [K, N, C]: the channel stride is ne11*ne00. + const int8_t * x_col = x + ch*ne11*ne00 + col*ne00; + + ggml_vec_dot_i8_i8(ne00, acc, 1, x_col, ne00, w_ch, ncol); + + local_absmax = ggml_i8_s_scale_bias_absmax( + stage + ch*ne11 + col, acc, bias_v, ncol, combined_d, local_absmax); + } + } + + ggml_i8_s_requantize(params, dst, stage, thread_max, ne11, ne02, col0, col1, local_absmax, relu); + } else { + // stage/dst index is col*ne01 + oc, so this thread owns the contiguous + // block [col0*ne01, col1*ne01). + const size_t nb11 = src1->nb[1]; + + // Output channels are handled a chunk at a time so the int32 + // accumulators fit a fixed stack buffer and params->wdata does not have + // to grow. Within a chunk the weight rows are contiguous and x_col stays + // hot, which is the whole reason for batching them into one call. + int32_t acc[GGML_MUL_MAT_ADD_OC_CHUNK]; + + for (int64_t col = col0; col < col1; col++) { + const int8_t * x_col = (const int8_t *) ((const char *) src1->data + col*nb11); + + for (int64_t oc0 = 0; oc0 < ne01; oc0 += GGML_MUL_MAT_ADD_OC_CHUNK) { + const int64_t noc = MIN((int64_t) GGML_MUL_MAT_ADD_OC_CHUNK, ne01 - oc0); + + ggml_vec_dot_i8_i8(ne00, acc, 1, w + oc0*ne00, ne00, x_col, noc); + + local_absmax = ggml_i8_s_scale_bias_absmax( + stage + col*ne01 + oc0, acc, bias + oc0, noc, combined_d, local_absmax); + } + } + + ggml_i8_s_requantize(params, dst, stage, thread_max, 0, 1, col0*ne01, col1*ne01, local_absmax, relu); + } +} + +void ggml_compute_forward_mul_mat_add( + const ggml_compute_params * params, + ggml_tensor * dst) { + ggml_compute_forward_mul_mat_add_impl(params, dst, false); +} + +void ggml_compute_forward_mul_mat_add_relu( + const ggml_compute_params * params, + ggml_tensor * dst) { + ggml_compute_forward_mul_mat_add_impl(params, dst, true); +} + +// ggml_compute_forward_mul_mat_i2_s + +// Output features per int32 accumulator batch. Batching them means the +// activation row is read once for 64 weight rows instead of once per row, and 64 +// int32s is a small enough stack buffer that params->wdata does not have to +// carry it. +#define GGML_MUL_MAT_I2_S_OC_CHUNK 64 + +void ggml_compute_forward_mul_mat_i2_s( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; // weights (I2_S, ternary) + const ggml_tensor * src1 = dst->src[1]; // activations (F32) + + GGML_ASSERT(src0->type == GGML_TYPE_I2_S); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(src0)); + + GGML_TENSOR_BINARY_OP_LOCALS + + GGML_ASSERT(ne0 == ne01); + GGML_ASSERT(ne1 == ne11); + GGML_ASSERT(ne2 == ne12); + GGML_ASSERT(ne3 == ne13); + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(ne00 % 128 == 0); + + // One scale covers the whole weight tensor, so broadcasting src0 over a + // batch would work, but the language model never has more than a 2D weight + // and silently supporting an untested shape is worse than refusing it. + GGML_ASSERT(ne02 == 1 && ne03 == 1); + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t nrows_y = ne11*ne12*ne13; + + // wdata holds the quantized activation rows, then a sidecar with one scale + // and one int8 row sum each. The sidecar cannot be folded into the rows the + // way I8_S folds its scale in: both values are needed in the epilogue, after + // the kernel has already consumed the row. + int8_t * const qy = (int8_t *) params->wdata; + const size_t qy_bytes = GGML_PAD((size_t) ne10*nrows_y, sizeof(float)); + float * const act_scale = (float *) ((char *) qy + qy_bytes); + int32_t * const act_sum = (int32_t *) (act_scale + nrows_y); + + GGML_ASSERT(params->wsize >= qy_bytes + nrows_y*(sizeof(float) + sizeof(int32_t))); + + // Striped by whole rows: the scale is a row-wide absmax, so a row cannot be + // split across threads. + for (int64_t ir = ith; ir < nrows_y; ir += nth) { + const int64_t i11 = ir % ne11; + const int64_t i12 = (ir / ne11) % ne12; + const int64_t i13 = ir / (ne11*ne12); + + const float * y_row = (const float *) ((const char *) src1->data + i11*nb11 + i12*nb12 + i13*nb13); + + ggml_i8_s_quantize_act(y_row, qy + ir*ne10, ne10, act_scale + ir, act_sum + ir); + } + + ggml_barrier(params->threadpool); + + const float w_scale = *ggml_inband_scale_const(src0); + const size_t w_row = nb01; // bytes per packed weight row + + // Split over output features and sweep the whole batch inside, so a weight + // row is loaded once and reused across every column. Splitting over columns + // instead would re-stream the weights per thread, and the weights are what + // this op is bandwidth-bound on. + const int64_t dr = (ne01 + nth - 1)/nth; + const int64_t oc0 = dr*ith; + const int64_t oc1 = MIN(oc0 + dr, ne01); + + if (oc0 >= oc1) { + return; + } + + int32_t acc[GGML_MUL_MAT_I2_S_OC_CHUNK]; + + for (int64_t ir = 0; ir < nrows_y; ++ir) { + const int64_t i11 = ir % ne11; + const int64_t i12 = (ir / ne11) % ne12; + const int64_t i13 = ir / (ne11*ne12); + + const int8_t * y_row = qy + ir*ne10; + + // The kernel returns sum(code*q) with codes {0,1,2}; subtracting the row + // sum turns that into sum(w*q) with w in {-1,0,+1}. Both scales are + // multipliers, so they combine once at the end. + const float d = w_scale * act_scale[ir]; + const int32_t bias = act_sum[ir]; + + float * dst_row = (float *) ((char *) dst->data + i11*nb1 + i12*nb2 + i13*nb3); + + for (int64_t oc = oc0; oc < oc1; oc += GGML_MUL_MAT_I2_S_OC_CHUNK) { + const int64_t noc = MIN((int64_t) GGML_MUL_MAT_I2_S_OC_CHUNK, oc1 - oc); + + ggml_vec_dot_i2_i8(ne00, acc, 1, + (const uint8_t *) src0->data + oc*w_row, w_row, + y_row, noc); + + for (int64_t j = 0; j < noc; ++j) { + dst_row[oc + j] = (float) (acc[j] - bias) * d; + } + } + } +} + +// ggml_compute_forward_im2col_asym + +void ggml_compute_forward_im2col_asym( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; // kernel, for its shape only + const ggml_tensor * src1 = dst->src[1]; // input (I8_S) + + GGML_ASSERT(src1->type == GGML_TYPE_I8_S); + GGML_ASSERT(dst->type == GGML_TYPE_I8_S); + + const int32_t s0 = ggml_get_op_params_i32(dst, 0); + const int32_t lp0 = ggml_get_op_params_i32(dst, 2); + const int32_t d0 = ggml_get_op_params_i32(dst, 4); + const bool is_2D = ggml_get_op_params_i32(dst, 6) == 1; + + // 1D only. The 2D case has no caller, and guessing at its indexing would + // mean shipping an untested path. + GGML_ASSERT(!is_2D && "ggml_im2col_asym: only the 1D case is implemented"); + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t IC = src0->ne[1]; + const int64_t KW = src0->ne[0]; + const int64_t IW = src1->ne[0]; + const int64_t N = src1->ne[2]; + const int64_t OW = dst->ne[1]; + + const int64_t total = N*OW; + const int64_t per_thread = (total + nth - 1)/nth; + const int64_t start = per_thread*ith; + const int64_t end = MIN(start + per_thread, total); + + int8_t * dst_data = (int8_t *) dst->data; + + for (int64_t idx = start; idx < end; idx++) { + const int64_t in = idx/OW; + const int64_t iow = idx%OW; + + int8_t * dst_col = dst_data + idx*(IC*KW); + + for (int64_t iic = 0; iic < IC; iic++) { + const int8_t * src_ch = (const int8_t *) ((const char *) src1->data + + iic*src1->nb[1] + in*src1->nb[2]); + + for (int64_t ikw = 0; ikw < KW; ikw++) { + const int64_t iiw = iow*s0 + ikw*d0 - lp0; + + // Zero is exact in both directions, so padding needs no + // special handling when the scale is applied later. + dst_col[iic*KW + ikw] = (iiw < 0 || iiw >= IW) ? 0 : src_ch[iiw]; + } + } + } + + if (ith == 0) { + // Pure rearrangement: no value changes, so the scale carries over. + *ggml_inband_scale(dst) = *ggml_inband_scale_const(src1); + } +} diff --git a/external/ggml/src/ggml-cpu/ops.h b/external/ggml/src/ggml-cpu/ops.h index e0b3a7bd2..0e3324c97 100644 --- a/external/ggml/src/ggml-cpu/ops.h +++ b/external/ggml/src/ggml-cpu/ops.h @@ -63,11 +63,11 @@ void ggml_compute_forward_soft_max(const struct ggml_compute_params * params, st void ggml_compute_forward_soft_max_ext_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_rope(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_rope_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_clamp(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_conv_transpose_1d(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_im2col(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_im2col_fast_1d(const struct ggml_compute_params * params, struct ggml_tensor * dst); -void ggml_compute_forward_im2col_back_f32(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_clamp(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_conv_transpose_1d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_im2col(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_im2col_fast_1d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_im2col_back_f32(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_im2col_3d(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_conv_2d(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_conv_3d(const struct ggml_compute_params * params, struct ggml_tensor * dst); @@ -115,6 +115,14 @@ void ggml_compute_forward_opt_step_adamw(const struct ggml_compute_params * para void ggml_compute_forward_mul_mat(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_fwht(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_opt_step_sgd(const struct ggml_compute_params * params, struct ggml_tensor * dst); + +// VibeASR CPU INT8 pipeline +void ggml_compute_forward_add_scaled(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_rms_norm_scaled(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_mul_mat_add(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_mul_mat_add_relu(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_mul_mat_i2_s(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_im2col_asym(const struct ggml_compute_params * params, struct ggml_tensor * dst); #ifdef __cplusplus } #endif diff --git a/external/ggml/src/ggml-cpu/vec.cpp b/external/ggml/src/ggml-cpu/vec.cpp index 7e8135617..0d42f4ede 100644 --- a/external/ggml/src/ggml-cpu/vec.cpp +++ b/external/ggml/src/ggml-cpu/vec.cpp @@ -393,6 +393,238 @@ void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * G *s = sumf; } +#if defined(__AVX2__) +static inline int ggml_i8_hsum_i32_4(const __m128i a) { + const __m128i hi64 = _mm_unpackhi_epi64(a, a); + const __m128i sum64 = _mm_add_epi32(hi64, a); + const __m128i hi32 = _mm_shuffle_epi32(sum64, _MM_SHUFFLE(2, 3, 0, 1)); + return _mm_cvtsi128_si32(_mm_add_epi32(sum64, hi32)); +} +#endif + +void ggml_vec_dot_i8_i8(int n, int32_t * GGML_RESTRICT s, size_t bs, const int8_t * GGML_RESTRICT x, size_t bx, const int8_t * GGML_RESTRICT y, int nrc) { + for (int row = 0; row < nrc; ++row) { + const int8_t * xr = x + (size_t)row*bx; + + int i = 0; + int32_t sumi = 0; + +#if defined(__AVX2__) + // The sign trick: maddubs wants the left operand unsigned, so move the + // sign of x onto y and take |x|. Each maddubs lane holds a sum of two + // products, at most 2*128*127 = 32512, which still fits int16 - but a + // second int16 accumulation could not, so widen to int32 every block + // instead of batching in int16. + // + // Everything reduces into one 128-bit accumulator so that short rows + // (the depthwise filters are 4 or 8 long) neither run the 256-bit loop + // nor pay for reducing a register that stayed zero. + __m128i acc = _mm_setzero_si128(); + + if (i + 32 <= n) { + const __m256i one16 = _mm256_set1_epi16(1); + + __m256i acc256 = _mm256_setzero_si256(); + for (; i + 32 <= n; i += 32) { + const __m256i xq = _mm256_loadu_si256((const __m256i *)(xr + i)); + const __m256i yq = _mm256_loadu_si256((const __m256i *)(y + i)); + + const __m256i ax = _mm256_sign_epi8(xq, xq); + const __m256i sy = _mm256_sign_epi8(yq, xq); + const __m256i dot = _mm256_maddubs_epi16(ax, sy); + + acc256 = _mm256_add_epi32(acc256, _mm256_madd_epi16(dot, one16)); + } + acc = _mm_add_epi32(_mm256_castsi256_si128(acc256), + _mm256_extracti128_si256(acc256, 1)); + } + + if (i + 8 <= n) { + const __m128i one16 = _mm_set1_epi16(1); + + for (; i + 16 <= n; i += 16) { + const __m128i xq = _mm_loadu_si128((const __m128i *)(xr + i)); + const __m128i yq = _mm_loadu_si128((const __m128i *)(y + i)); + + const __m128i ax = _mm_sign_epi8(xq, xq); + const __m128i sy = _mm_sign_epi8(yq, xq); + const __m128i dot = _mm_maddubs_epi16(ax, sy); + + acc = _mm_add_epi32(acc, _mm_madd_epi16(dot, one16)); + } + for (; i + 8 <= n; i += 8) { + const __m128i xq = _mm_loadl_epi64((const __m128i *)(xr + i)); + const __m128i yq = _mm_loadl_epi64((const __m128i *)(y + i)); + + const __m128i ax = _mm_sign_epi8(xq, xq); + const __m128i sy = _mm_sign_epi8(yq, xq); + const __m128i dot = _mm_maddubs_epi16(ax, sy); + + acc = _mm_add_epi32(acc, _mm_madd_epi16(dot, one16)); + } + } + + // i > 0 exactly when some vector block ran. Reducing a register that + // stayed zero costs four instructions per row, which is not free when the + // row is a 4-tap depthwise filter and there is one row per output. + if (i > 0) { + sumi = ggml_i8_hsum_i32_4(acc); + } +#elif defined(__ARM_NEON) && defined(__aarch64__) + int32x4_t acc = vdupq_n_s32(0); + for (; i + 16 <= n; i += 16) { + const int8x16_t xv = vld1q_s8(xr + i); + const int8x16_t yv = vld1q_s8(y + i); + #if defined(__ARM_FEATURE_DOTPROD) + acc = vdotq_s32(acc, xv, yv); + #else + // vmull_s8 tops out at 128*128 = 16384, so the int16 products are + // safe; vpadalq_s16 folds them pairwise straight into int32. + acc = vpadalq_s16(acc, vmull_s8(vget_low_s8 (xv), vget_low_s8 (yv))); + acc = vpadalq_s16(acc, vmull_s8(vget_high_s8(xv), vget_high_s8(yv))); + #endif + } + for (; i + 8 <= n; i += 8) { + const int8x8_t xv = vld1_s8(xr + i); + const int8x8_t yv = vld1_s8(y + i); + acc = vpadalq_s16(acc, vmull_s8(xv, yv)); + } + if (i > 0) { + sumi = vaddvq_s32(acc); + } +#endif + + for (; i < n; ++i) { + sumi += (int32_t)xr[i] * (int32_t)y[i]; + } + + s[(size_t)row*bs] = sumi; + } +} + +void ggml_vec_dot_i2_i8(int n, int32_t * GGML_RESTRICT s, size_t bs, const uint8_t * GGML_RESTRICT x, size_t bx, const int8_t * GGML_RESTRICT y, int nrc) { + // 128 values per 32-byte group, and every I2_S row in the language model is + // a multiple of that, so there is no partial-group path to get wrong. + assert(n % 128 == 0); + + const int nb = n / 128; + + for (int row = 0; row < nrc; ++row) { + const uint8_t * xr = x + (size_t)row*bx; + + int b = 0; + int32_t sumi = 0; + +#if defined(__AVX2__) + // Shifting by 16-bit lanes pulls bits down from the neighbouring byte, + // but the mask discards them: after >>6 the low two bits of each byte + // are exactly that byte's top code. + const __m256i mask = _mm256_set1_epi8(0x03); + const __m256i one16 = _mm256_set1_epi16(1); + + __m256i acc = _mm256_setzero_si256(); + + // Each maddubs lane holds a sum of two code*int8 products, at most + // 2*2*127 = 508 in magnitude, and eight groups contribute 32 of them: + // 32*508 = 16256, still inside int16. Widening once per eight groups + // rather than once per group keeps the madd out of the inner loop. + while (b < nb) { + const int bend = b + 8 < nb ? b + 8 : nb; + + __m256i acc16 = _mm256_setzero_si256(); + + for (; b < bend; ++b) { + const __m256i xq = _mm256_loadu_si256((const __m256i *)(xr + (size_t)b*32)); + + const __m256i c0 = _mm256_and_si256(_mm256_srli_epi16(xq, 6), mask); + const __m256i c1 = _mm256_and_si256(_mm256_srli_epi16(xq, 4), mask); + const __m256i c2 = _mm256_and_si256(_mm256_srli_epi16(xq, 2), mask); + const __m256i c3 = _mm256_and_si256(xq, mask); + + const int8_t * py = y + (size_t)b*128; + + const __m256i y0 = _mm256_loadu_si256((const __m256i *)(py + 0)); + const __m256i y1 = _mm256_loadu_si256((const __m256i *)(py + 32)); + const __m256i y2 = _mm256_loadu_si256((const __m256i *)(py + 64)); + const __m256i y3 = _mm256_loadu_si256((const __m256i *)(py + 96)); + + acc16 = _mm256_add_epi16(acc16, _mm256_add_epi16(_mm256_maddubs_epi16(c0, y0), + _mm256_maddubs_epi16(c1, y1))); + acc16 = _mm256_add_epi16(acc16, _mm256_add_epi16(_mm256_maddubs_epi16(c2, y2), + _mm256_maddubs_epi16(c3, y3))); + } + + acc = _mm256_add_epi32(acc, _mm256_madd_epi16(acc16, one16)); + } + + sumi = ggml_i8_hsum_i32_4(_mm_add_epi32(_mm256_castsi256_si128(acc), + _mm256_extracti128_si256(acc, 1))); +#elif defined(__ARM_NEON) && defined(__aarch64__) + // Half a group per iteration, since a NEON register holds 16 of the 32 + // bytes. The four code lanes still map to the four 32-value slices of y, + // offset by which half of the group is loaded. + const uint8x16_t mask = vdupq_n_u8(3); + + int32x4_t acc = vdupq_n_s32(0); + + for (; b < nb; ++b) { + for (int h = 0; h < 2; ++h) { + const uint8x16_t xq = vld1q_u8(xr + (size_t)b*32 + h*16); + + const int8x16_t c0 = vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(xq, 6), mask)); + const int8x16_t c1 = vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(xq, 4), mask)); + const int8x16_t c2 = vreinterpretq_s8_u8(vandq_u8(vshrq_n_u8(xq, 2), mask)); + const int8x16_t c3 = vreinterpretq_s8_u8(vandq_u8(xq, mask)); + + const int8_t * py = y + (size_t)b*128 + h*16; + + const int8x16_t y0 = vld1q_s8(py + 0); + const int8x16_t y1 = vld1q_s8(py + 32); + const int8x16_t y2 = vld1q_s8(py + 64); + const int8x16_t y3 = vld1q_s8(py + 96); + + #if defined(__ARM_FEATURE_DOTPROD) + acc = vdotq_s32(acc, c0, y0); + acc = vdotq_s32(acc, c1, y1); + acc = vdotq_s32(acc, c2, y2); + acc = vdotq_s32(acc, c3, y3); + #else + // Products top out at 2*127 = 254, so the int16 halves are safe + // and vpadalq_s16 folds them pairwise straight into int32. + acc = vpadalq_s16(acc, vmull_s8(vget_low_s8 (c0), vget_low_s8 (y0))); + acc = vpadalq_s16(acc, vmull_s8(vget_high_s8(c0), vget_high_s8(y0))); + acc = vpadalq_s16(acc, vmull_s8(vget_low_s8 (c1), vget_low_s8 (y1))); + acc = vpadalq_s16(acc, vmull_s8(vget_high_s8(c1), vget_high_s8(y1))); + acc = vpadalq_s16(acc, vmull_s8(vget_low_s8 (c2), vget_low_s8 (y2))); + acc = vpadalq_s16(acc, vmull_s8(vget_high_s8(c2), vget_high_s8(y2))); + acc = vpadalq_s16(acc, vmull_s8(vget_low_s8 (c3), vget_low_s8 (y3))); + acc = vpadalq_s16(acc, vmull_s8(vget_high_s8(c3), vget_high_s8(y3))); + #endif + } + } + + sumi = vaddvq_s32(acc); +#endif + + // Portable path, and the only path on targets without either ISA. + for (; b < nb; ++b) { + const uint8_t * px = xr + (size_t)b*32; + const int8_t * py = y + (size_t)b*128; + + for (int gp = 0; gp < 32; ++gp) { + const uint8_t v = px[gp]; + + sumi += (int32_t)((v >> 6) & 3) * (int32_t)py[ gp]; + sumi += (int32_t)((v >> 4) & 3) * (int32_t)py[32 + gp]; + sumi += (int32_t)((v >> 2) & 3) * (int32_t)py[64 + gp]; + sumi += (int32_t)( v & 3) * (int32_t)py[96 + gp]; + } + } + + s[(size_t)row*bs] = sumi; + } +} + void ggml_vec_silu_f32(const int n, float * y, const float * x) { int i = 0; #if defined(__AVX512F__) && defined(__AVX512DQ__) diff --git a/external/ggml/src/ggml-cpu/vec.h b/external/ggml/src/ggml-cpu/vec.h index 741902479..8943e86c6 100644 --- a/external/ggml/src/ggml-cpu/vec.h +++ b/external/ggml/src/ggml-cpu/vec.h @@ -43,6 +43,30 @@ void ggml_vec_dot_f32(int n, float * GGML_RESTRICT s, size_t bs, const float * G void ggml_vec_dot_bf16(int n, float * GGML_RESTRICT s, size_t bs, ggml_bf16_t * GGML_RESTRICT x, size_t bx, ggml_bf16_t * GGML_RESTRICT y, size_t by, int nrc); void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * GGML_RESTRICT x, size_t bx, ggml_fp16_t * GGML_RESTRICT y, size_t by, int nrc); +// int8 x int8 dot products accumulating in int32, for the GGML_TYPE_I8_S ops. +// The scale is per-tensor and cancels out of the contraction, so these stay +// integral and the caller applies it once. nrc rows of x, each bx bytes apart, +// are contracted against the single row y; result row stride is bs int32s. +// n is arbitrary: whatever the vector width does not cover is done scalar. +void ggml_vec_dot_i8_i8(int n, int32_t * GGML_RESTRICT s, size_t bs, const int8_t * GGML_RESTRICT x, size_t bx, const int8_t * GGML_RESTRICT y, int nrc); + +// Packed ternary x int8 dot products for GGML_TYPE_I2_S, accumulating in int32. +// +// x holds the 2-bit codes as they sit on disk: 128 values per 32-byte group, +// byte gp of a group carrying the values at group-relative positions gp, 32+gp, +// 64+gp and 96+gp in bit pairs 6, 4, 2, 0. y is plain sequential int8, so the +// four code lanes of a group line up with four consecutive 32-value slices of y +// and no shuffling is needed on either side. +// +// The result is sum(code*y), NOT sum(w*y): the codes are {0,1,2} where the +// weights are {-1,0,+1}, so the caller subtracts sum(y) to recover the real +// contraction. Keeping the bias out of the kernel is what lets the unsigned +// multiply-add instructions be used directly. +// +// nrc rows of x, each bx bytes apart, are contracted against the single row y; +// result row stride is bs int32s. n must be a multiple of 128. +void ggml_vec_dot_i2_i8(int n, int32_t * GGML_RESTRICT s, size_t bs, const uint8_t * GGML_RESTRICT x, size_t bx, const int8_t * GGML_RESTRICT y, int nrc); + void ggml_vec_silu_f32(const int n, float * y, const float * x); ggml_float ggml_vec_cvar_f32(const int n, float * y, const float * x, const float mean); //it will also center y ( y = y - mean ) ggml_float ggml_vec_soft_max_f32(const int n, float * y, const float * x, float max); diff --git a/external/ggml/src/ggml-impl.h b/external/ggml/src/ggml-impl.h index d4cb6c9b3..2f50da73a 100644 --- a/external/ggml/src/ggml-impl.h +++ b/external/ggml/src/ggml-impl.h @@ -354,6 +354,26 @@ struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph, int i0, int i1); // ggml-alloc.c: true if the operation can reuse memory from its sources GGML_API bool ggml_op_can_inplace(enum ggml_op op); +// Bytes a type needs past its payload. Non-zero only for GGML_TYPE_I8_S and +// GGML_TYPE_I2_S, which append one F32 per-tensor scale; see ggml.c. +GGML_API size_t ggml_type_extra_bytes(enum ggml_type type); + +// The in-band scale of an I8_S or I2_S tensor, which sits immediately after the +// payload. ggml_nbytes already counts the padded scale, so subtracting it back +// out gives the payload end without duplicating either type's row arithmetic. +// +// The scale is a multiplier in both cases: dequantizing is q * scale. +static inline float * ggml_inband_scale(struct ggml_tensor * tensor) { + const size_t extra = ggml_type_extra_bytes(tensor->type); + GGML_ASSERT(extra > 0 && "type has no in-band scale"); + GGML_ASSERT(ggml_is_contiguous(tensor)); + return (float *) ((char *) tensor->data + ggml_nbytes(tensor) - extra); +} + +static inline const float * ggml_inband_scale_const(const struct ggml_tensor * tensor) { + return ggml_inband_scale((struct ggml_tensor *) tensor); +} + // Memory allocation diff --git a/external/ggml/src/ggml-quants.c b/external/ggml/src/ggml-quants.c index 51f08fa3b..facdc2473 100644 --- a/external/ggml/src/ggml-quants.c +++ b/external/ggml/src/ggml-quants.c @@ -2407,6 +2407,163 @@ void dequantize_row_tq2_0(const block_tq2_0 * GGML_RESTRICT x, float * GGML_REST } } +// ====================== GGML_TYPE_I8_S / GGML_TYPE_I2_S +// +// Unlike every block-quantized type above, these two carry a single F32 scale +// for the whole tensor rather than one per block, stored immediately after the +// payload (ggml_type_extra_bytes reserves the room). The four functions below +// are therefore whole-tensor: `n` is ggml_nelements(), and the scale lives at +// byte offset ggml_row_size(type, n) rounded down to the payload end. +// +// The scale is a multiplier in both directions -- real = q * scale -- matching +// the `d` of every other ggml quant type. VibeASR's own fork uses a multiplier +// for weights but stores the reciprocal for op-produced activations; the two +// are reconciled by a division deep inside each kernel. Only the multiplier +// convention is used here. Activation scales never reach a file, so this +// changes nothing about how an existing GGUF is read. + +#define I2_S_GROUP 128 // ternary values per packed group +#define I2_S_BYTES 32 // bytes those 128 values occupy + +// Scale slot: the F32 sits directly after the payload. Kept as helpers so the +// offset arithmetic is not repeated in four places. The payload size is the +// only difference between the two types -- I8_S is one byte per value, I2_S +// packs 128 values into 32 bytes. +static inline size_t i8_s_payload_bytes(int64_t n) { return (size_t)n; } +static inline size_t i2_s_payload_bytes(int64_t n) { return (size_t)(n / I2_S_GROUP) * I2_S_BYTES; } + +static inline float * i8_s_scale_ptr (void * x, int64_t n) { return (float *)((char *)x + i8_s_payload_bytes(n)); } +static inline const float * i8_s_scale_ptr_const(const void * x, int64_t n) { return (const float *)((const char *)x + i8_s_payload_bytes(n)); } +static inline float * i2_s_scale_ptr (void * x, int64_t n) { return (float *)((char *)x + i2_s_payload_bytes(n)); } +static inline const float * i2_s_scale_ptr_const(const void * x, int64_t n) { return (const float *)((const char *)x + i2_s_payload_bytes(n)); } + +void ggml_i8_s_to_float(const void * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t n) { + const int8_t * q = (const int8_t *)x; + const float d = *i8_s_scale_ptr_const(x, n); + + for (int64_t i = 0; i < n; ++i) { + y[i] = (float)q[i] * d; + } +} + +size_t ggml_i8_s_from_float(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t n) { + int8_t * q = (int8_t *)y; + + float amax = 0.0f; + for (int64_t i = 0; i < n; ++i) { + const float a = fabsf(x[i]); + if (a > amax) amax = a; + } + + // Symmetric range: -127..127 rather than -128..127, so that negating a + // tensor is exactly representable and the kernels can widen to int16 + // without a special case for the one asymmetric value. + const float d = amax / 127.0f; + const float id = d != 0.0f ? 1.0f / d : 0.0f; + + for (int64_t i = 0; i < n; ++i) { + int v = nearest_int(x[i] * id); + if (v > 127) v = 127; + if (v < -127) v = -127; + q[i] = (int8_t)v; + } + + *i8_s_scale_ptr(y, n) = d; + + return i8_s_payload_bytes(n) + ggml_type_extra_bytes(GGML_TYPE_I8_S); +} + +// Bit pair -> ternary value. Codes 1 and 3 both mean zero; the packer only ever +// emits 1, but 3 is accepted so a stray pattern dequantizes to 0 instead of +// reading out of bounds. +static const float i2_s_map[4] = { -1.0f, 0.0f, +1.0f, 0.0f }; + +void ggml_i2_s_to_float(const void * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t n) { + assert(n % I2_S_GROUP == 0); + + const uint8_t * q = (const uint8_t *)x; + const float d = *i2_s_scale_ptr_const(x, n); + + for (int64_t base = 0; base < n; base += I2_S_GROUP) { + // Byte gp holds the values at base+gp, base+32+gp, base+64+gp and + // base+96+gp, in bit pairs 6, 4, 2, 0 -- strided, not consecutive, so + // that a SIMD load of 32 bytes yields 4 aligned lanes of 32 values. + for (int gp = 0; gp < I2_S_BYTES; ++gp) { + const uint8_t b = q[gp]; + + y[base + 0 + gp] = d * i2_s_map[(b >> 6) & 3]; + y[base + 32 + gp] = d * i2_s_map[(b >> 4) & 3]; + y[base + 64 + gp] = d * i2_s_map[(b >> 2) & 3]; + y[base + 96 + gp] = d * i2_s_map[(b >> 0) & 3]; + } + + q += I2_S_BYTES; + } +} + +size_t ggml_i2_s_from_float(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t n) { + assert(n % I2_S_GROUP == 0); + + uint8_t * q = (uint8_t *)y; + + float amax = 0.0f; + for (int64_t i = 0; i < n; ++i) { + const float a = fabsf(x[i]); + if (a > amax) amax = a; + } + + // The type stores {-1, 0, +1} * d, so d is the absmax and a ternary input + // round-trips exactly. Non-ternary input is rounded to the nearest of the + // three representable values, which is all this type can express. + const float d = amax; + const float id = d != 0.0f ? 1.0f / d : 0.0f; + + memset(q, 0, i2_s_payload_bytes(n)); + + for (int64_t base = 0; base < n; base += I2_S_GROUP) { + for (int j = 0; j < I2_S_GROUP; ++j) { + int v = nearest_int(x[base + j] * id); + if (v > 1) v = 1; + if (v < -1) v = -1; + + const uint8_t code = (uint8_t)(v + 1); // -1,0,1 -> 0,1,2 + + q[j % I2_S_BYTES] |= (uint8_t)(code << (6 - 2*(j / I2_S_BYTES))); + } + + q += I2_S_BYTES; + } + + *i2_s_scale_ptr(y, n) = d; + + return i2_s_payload_bytes(n) + ggml_type_extra_bytes(GGML_TYPE_I2_S); +} + +void ggml_i8_s_quantize_act(const float * GGML_RESTRICT x, int8_t * GGML_RESTRICT q, int64_t n, + float * GGML_RESTRICT scale, int32_t * GGML_RESTRICT sum) { + float amax = 0.0f; + for (int64_t i = 0; i < n; ++i) { + const float a = fabsf(x[i]); + if (a > amax) amax = a; + } + + // -127..127 rather than -128..127, as in ggml_i8_s_from_float. + const float d = amax / 127.0f; + const float id = d != 0.0f ? 1.0f / d : 0.0f; + + int32_t s = 0; + for (int64_t i = 0; i < n; ++i) { + int v = nearest_int(x[i] * id); + if (v > 127) v = 127; + if (v < -127) v = -127; + q[i] = (int8_t)v; + s += v; + } + + *scale = d; + *sum = s; +} + // ====================== "True" 2-bit (de)-quantization void dequantize_row_iq2_xxs(const block_iq2_xxs * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { diff --git a/external/ggml/src/ggml-quants.h b/external/ggml/src/ggml-quants.h index 3fb6140b7..cda127695 100644 --- a/external/ggml/src/ggml-quants.h +++ b/external/ggml/src/ggml-quants.h @@ -102,6 +102,36 @@ GGML_API size_t quantize_q8_0(const float * GGML_RESTRICT src, void * GGML_RESTR GGML_API size_t quantize_mxfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); GGML_API size_t quantize_nvfp4(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +// GGML_TYPE_I8_S / GGML_TYPE_I2_S conversions. +// +// These are whole-tensor, not per-row: both types carry a single F32 scale for +// the entire tensor, stored immediately after the payload (see +// ggml_type_extra_bytes), so `n` is ggml_nelements() and a row pointer alone +// cannot locate the scale. That is also why the two traits entries leave +// .to_float / .from_float_ref NULL instead of pointing here, and why +// ggml_quantize_chunk does not list these types -- its +// `result == nrows * row_size` invariant cannot hold for a per-tensor scale. +// +// The from_float direction returns the number of bytes written, payload plus +// the padded scale. +GGML_API void ggml_i8_s_to_float (const void * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t n); +GGML_API size_t ggml_i8_s_from_float(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t n); +GGML_API void ggml_i2_s_to_float (const void * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t n); +GGML_API size_t ggml_i2_s_from_float(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t n); + +// Activation quantizer for the I2_S matmul: one row in, int8 payload out, with +// the scale and the row's int8 sum handed back out of band. +// +// Out of band because neither fits the in-band convention above. The scale is +// per row here, not per tensor, since a language model activation row is a +// single token and its dynamic range has nothing to do with its neighbours'. +// The sum is needed because I2_S stores the ternary values as the codes +// {0, 1, 2} rather than {-1, 0, +1}: an integer dot against the codes gives +// sum(w*q) + sum(q), so the row sum has to be subtracted back out. Computing it +// here costs nothing -- the values are already in registers. +GGML_API void ggml_i8_s_quantize_act(const float * GGML_RESTRICT x, int8_t * GGML_RESTRICT q, int64_t n, + float * GGML_RESTRICT scale, int32_t * GGML_RESTRICT sum); + GGML_API void iq2xs_init_impl(enum ggml_type type); GGML_API void iq2xs_free_impl(enum ggml_type type); GGML_API void iq3xs_init_impl(int grid_size); diff --git a/external/ggml/src/ggml.c b/external/ggml/src/ggml.c index 0ba560368..869645bdc 100644 --- a/external/ggml/src/ggml.c +++ b/external/ggml/src/ggml.c @@ -927,8 +927,49 @@ static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = { .type_size = 0, .is_quantized = false, }, + // .to_float / .from_float_ref are deliberately left NULL for the two types + // below. Both carry a single F32 scale for the whole tensor rather than one + // per block, so a row-scoped callback has no way to find it: the scale is + // past the end of the last row, not the row it is handed. Use + // ggml_i8_s_to_float / ggml_i2_s_to_float (ggml-quants.h), which take an + // element count of ggml_nelements(). For the same reason neither type + // appears in ggml_quantize_chunk -- its `result == nrows * row_size` + // invariant cannot hold when a per-tensor scale is in play. + [GGML_TYPE_I8_S] = { + .type_name = "i8_s", + .blck_size = 1, + .type_size = sizeof(int8_t), + .is_quantized = true, + }, + [GGML_TYPE_I2_S] = { + .type_name = "i2_s", + // 128 ternary values packed into 32 bytes: byte gp holds the values at + // positions gp, 32+gp, 64+gp and 96+gp in bit pairs 6,4,2,0. A partial + // group still consumes a full 32 bytes, so the row size is + // ceil(ne0/128)*32 -- declaring the block as 128/32 makes ggml_row_size + // compute exactly that and assert the ne0 % 128 == 0 the packing needs. + .blck_size = 128, + .type_size = 32, + .is_quantized = true, + }, }; +// I8_S and I2_S append one F32 per-tensor scale after the payload. Padded to 32 +// bytes so the next tensor in an arena stays aligned. +// +// This is the only place the layout is encoded; ggml_nbytes and +// ggml_new_tensor_impl both go through it. Returns 0 for every other type, so +// no existing type changes size. +size_t ggml_type_extra_bytes(enum ggml_type type) { + switch (type) { + case GGML_TYPE_I8_S: + case GGML_TYPE_I2_S: + return 32; + default: + return 0; + } +} + const struct ggml_type_traits * ggml_get_type_traits(enum ggml_type type) { assert(type >= 0); assert(type < GGML_TYPE_COUNT); @@ -1084,9 +1125,15 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", "CONVROT_LINEAR", + + "ADD_SCALED", + "RMS_NORM_SCALED", + "MUL_MAT_ADD", + "MUL_MAT_ADD_RELU", + "IM2COL_ASYM", }; -static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); +static_assert(GGML_OP_COUNT == 107, "GGML_OP_COUNT != 107"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1200,9 +1247,15 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", "convrot_linear(weight_i8, input, weight_scale, bias)", + + "a*scale+b", + "rms_norm(x)*scale", + "a*b+bias", + "relu(a*b+bias)", + "im2col_asym(x)", }; -static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102"); +static_assert(GGML_OP_COUNT == 107, "GGML_OP_COUNT != 107"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -1303,7 +1356,7 @@ size_t ggml_nbytes(const struct ggml_tensor * tensor) { } } - return nbytes; + return nbytes + ggml_type_extra_bytes(tensor->type); } size_t ggml_nbytes_pad(const struct ggml_tensor * tensor) { @@ -1757,6 +1810,7 @@ static struct ggml_tensor * ggml_new_tensor_impl( for (int i = 1; i < n_dims; i++) { data_size *= ne[i]; } + data_size += ggml_type_extra_bytes(type); GGML_ASSERT(view_src == NULL || data_size == 0 || data_size + view_offs <= ggml_nbytes(view_src)); @@ -5594,6 +5648,161 @@ struct ggml_tensor * ggml_convrot_linear( return result; } +// VibeASR CPU INT8 pipeline +// +// Five fused ops that keep an activation chain entirely in GGML_TYPE_I8_S: each +// one consumes int8 with a per-tensor scale, does its arithmetic in int32/F32, +// and re-quantizes its own output to int8 with a freshly measured scale. Fusing +// matters because the requantization needs the output absmax, so an unfused +// chain would have to write, read back and rescan every intermediate. +// +// That per-tensor output scale is why these need GGML_TYPE_I8_S rather than the +// GGML_TYPE_I8-plus-separate-scale-tensor pair used by ggml_convrot_linear +// above: a scale computed from the op's own output has nowhere to live but with +// the data, since a node has exactly one output tensor. + +struct ggml_tensor * ggml_add_scaled( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * scale) { + GGML_ASSERT(ggml_are_same_shape(a, b)); + GGML_ASSERT(scale->type == GGML_TYPE_F32); + // Broadcast along ne[0], the channel axis: one LayerScale coefficient per + // channel, applied to every position. + GGML_ASSERT(scale->ne[0] == a->ne[0]); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_ADD_SCALED; + result->src[0] = a; + result->src[1] = b; + result->src[2] = scale; + + return result; +} + +struct ggml_tensor * ggml_rms_norm_scaled( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * scale, + float eps) { + GGML_ASSERT(a->type == GGML_TYPE_I8_S); + GGML_ASSERT(scale->type == GGML_TYPE_F32); + GGML_ASSERT(scale->ne[0] == a->ne[0]); + GGML_ASSERT(eps > 0.0f); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + ggml_set_op_params(result, &eps, sizeof(eps)); + + result->op = GGML_OP_RMS_NORM_SCALED; + result->src[0] = a; + result->src[1] = scale; + + return result; +} + +// Shared by ggml_mul_mat_add and ggml_mul_mat_add_relu: identical shape rules, +// only the epilogue differs. +static struct ggml_tensor * ggml_mul_mat_add_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * bias, + enum ggml_op op) { + GGML_ASSERT(ggml_can_mul_mat(a, b)); + // I8_S weights only. The ternary I2_S weights of the language model go + // through plain ggml_mul_mat, which has no bias to fuse. + GGML_ASSERT(a->type == GGML_TYPE_I8_S); + GGML_ASSERT(b->type == GGML_TYPE_I8_S); + GGML_ASSERT(bias->type == GGML_TYPE_F32); + + // Bias is per output channel, and which ne holds the output channels + // depends on the shape of a: a [IC, OC] is an ordinary matmul, whereas + // a [K, 1, C] is the depthwise case where the channels sit in ne[2]. + GGML_ASSERT(bias->ne[0] == (a->ne[1] == 1 && a->ne[2] > 1 ? a->ne[2] : a->ne[1])); + + const int64_t ne[4] = { a->ne[1], b->ne[1], b->ne[2], b->ne[3] }; + + // Output is int8 with its own scale, so the chain continues without a + // detour through F32. + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_I8_S, 4, ne); + + result->op = op; + result->src[0] = a; + result->src[1] = b; + result->src[2] = bias; + + return result; +} + +struct ggml_tensor * ggml_mul_mat_add( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * bias) { + return ggml_mul_mat_add_impl(ctx, a, b, bias, GGML_OP_MUL_MAT_ADD); +} + +struct ggml_tensor * ggml_mul_mat_add_relu( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * bias) { + return ggml_mul_mat_add_impl(ctx, a, b, bias, GGML_OP_MUL_MAT_ADD_RELU); +} + +struct ggml_tensor * ggml_im2col_asym( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int s1, + int lp0, + int rp0, + int p1, + int d0, + int d1, + bool is_2D, + enum ggml_type dst_type) { + if (is_2D) { + GGML_ASSERT(a->ne[2] == b->ne[2]); + } else { + GGML_ASSERT(a->ne[1] == b->ne[1]); + GGML_ASSERT(b->ne[3] == 1); + } + + const int64_t OH = is_2D ? ggml_calc_conv_output_size(b->ne[1], a->ne[1], s1, p1, d1) : 0; + // Same formula as ggml_calc_conv_output_size but with the two width pads + // counted separately instead of as 2*p0. + const int64_t OW = (b->ne[0] + lp0 + rp0 - d0 * (a->ne[0] - 1) - 1) / s0 + 1; + + GGML_ASSERT((!is_2D || OH > 0) && "b too small compared to a"); + GGML_ASSERT((OW > 0) && "b too small compared to a"); + + const int64_t ne[4] = { + is_2D ? (a->ne[2] * a->ne[1] * a->ne[0]) : a->ne[1] * a->ne[0], + OW, + is_2D ? OH : b->ne[2], + is_2D ? b->ne[3] : 1, + }; + + struct ggml_tensor * result = ggml_new_tensor(ctx, dst_type, 4, ne); + + // Note the order: lp0 and rp0 occupy slots 2 and 3, where GGML_OP_IM2COL + // keeps p0 and p1. The two ops read their own params and never share a + // forward, so the layouts do not have to agree. + int32_t params[] = { s0, s1, lp0, rp0, d0, d1, (is_2D ? 1 : 0), p1 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_IM2COL_ASYM; + result->src[0] = a; + result->src[1] = b; + + return result; +} + struct ggml_tensor * ggml_flash_attn_ext_with_bias_mask( struct ggml_context * ctx, struct ggml_tensor * q, diff --git a/include/engine/community_models/vibeasr/assets.h b/include/engine/community_models/vibeasr/assets.h new file mode 100644 index 000000000..744eb70e5 --- /dev/null +++ b/include/engine/community_models/vibeasr/assets.h @@ -0,0 +1,74 @@ +#pragma once + +// VibeASR VAE encoder assets. +// +// Ported from https://github.com/microsoft/VibeASR.cpp (src/vae.cpp). + +#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); + +} // 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..23ee9b164 --- /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 16 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/src/community_models/vibeasr/assets.cpp b/src/community_models/vibeasr/assets.cpp new file mode 100644 index 000000000..2bc8d5830 --- /dev/null +++ b/src/community_models/vibeasr/assets.cpp @@ -0,0 +1,148 @@ +#include "engine/community_models/vibeasr/assets.h" + +#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; +} + +} // 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) { + auto assets = std::make_shared(); + assets->source = engine::assets::open_tensor_source(model_path); + assets->config = derive_vae_config(*assets->source); + return assets; +} + +} // 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/unittests/test_i2_s_mul_mat.cpp b/tests/unittests/test_i2_s_mul_mat.cpp new file mode 100644 index 000000000..e20928fb0 --- /dev/null +++ b/tests/unittests/test_i2_s_mul_mat.cpp @@ -0,0 +1,373 @@ +// Numeric checks for the ternary GGML_TYPE_I2_S matmul, which is what the +// VibeASR language model runs on. +// +// The op is reachable through plain ggml_mul_mat with an I2_S src0 and an F32 +// src1, and it quantizes src1 to int8 itself. Its arithmetic is entirely +// integral up to a single float multiply at the end, so the references below are +// exact rather than approximate: the reference computes sum(w*q) as int32 with +// w in {-1,0,+1} and multiplies by the same combined scale, which is bit for bit +// what the kernel does after it subtracts the row sum out of its {0,1,2} codes. +// A tolerance here would hide a wrong packing that happens to be close. + +#include "test_assert.h" + +#include +#include + +#include +#include +#include +#include +#include +#include + +extern "C" { +size_t ggml_type_extra_bytes(enum ggml_type type); +void ggml_i2_s_to_float (const void * x, float * y, int64_t n); +size_t ggml_i2_s_from_float(const float * x, void * y, int64_t n); +void ggml_i8_s_quantize_act(const float * x, int8_t * q, int64_t n, float * scale, int32_t * sum); +} + +using engine::test::require; +using engine::test::require_close; +using engine::test::require_eq; + +namespace { + +constexpr size_t kCtxBytes = 512u * 1024 * 1024; + +const float * inband_scale(const ggml_tensor * t) { + return reinterpret_cast( + static_cast(t->data) + ggml_nbytes(t) - ggml_type_extra_bytes(t->type)); +} + +// Decode the packing by hand instead of going through ggml_i2_s_to_float, so +// that the layout the kernel reads is pinned independently of the dequantizer +// that was written alongside it: 128 values per 32-byte group, byte gp holding +// the values at group-relative positions gp, 32+gp, 64+gp and 96+gp in bit pairs +// 6, 4, 2, 0. +std::vector unpack_i2_s(const void * data, int64_t n) { + const uint8_t * q = static_cast(data); + + std::vector out(static_cast(n)); + for (int64_t base = 0; base < n; base += 128) { + const uint8_t * group = q + base / 4; + + for (int gp = 0; gp < 32; ++gp) { + const uint8_t b = group[gp]; + + out[base + 0 + gp] = static_cast((b >> 6) & 3) - 1; + out[base + 32 + gp] = static_cast((b >> 4) & 3) - 1; + out[base + 64 + gp] = static_cast((b >> 2) & 3) - 1; + out[base + 96 + gp] = static_cast((b >> 0) & 3) - 1; + } + } + return out; +} + +// Ternary weights with a deliberate mix of all three values and a stride that is +// coprime with 32 and 128, so no code lands in the same bit position of every +// byte and a swapped shift shows up immediately. +std::vector ternary_pattern(size_t n, float scale, int phase) { + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + const int k = static_cast((i * 7 + phase) % 3); + v[i] = scale * static_cast(k - 1); + } + return v; +} + +std::vector activation_pattern(size_t n, float phase, float scale) { + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + const float x = static_cast(i); + v[i] = scale * (std::sin(phase + 0.11f * x) + 0.35f * std::cos(0.05f * x - phase)); + } + return v; +} + +void compute(ggml_context * ctx, ggml_tensor * result, int n_threads) { + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, result); + ggml_graph_compute_with_ctx(ctx, gf, n_threads); +} + +// Reference for one batch of ggml_mul_mat(weight_i2_s, activations_f32). +// +// The activation quantization is the op's own, called directly: the point of +// this test is the packing, the code-to-weight bias, and the scale combination, +// not to re-derive an absmax. +std::vector reference(const std::vector & w, // [K*N], ternary + const std::vector & a, // [K*M] + int64_t K, int64_t N, int64_t M, + float w_scale) { + std::vector out(static_cast(N * M)); + + std::vector q(static_cast(K)); + for (int64_t col = 0; col < M; ++col) { + float act_scale = 0.0f; + int32_t act_sum = 0; + ggml_i8_s_quantize_act(a.data() + col * K, q.data(), K, &act_scale, &act_sum); + + const float d = w_scale * act_scale; + + for (int64_t oc = 0; oc < N; ++oc) { + int32_t dot = 0; + for (int64_t k = 0; k < K; ++k) { + dot += w[static_cast(oc * K + k)] * static_cast(q[static_cast(k)]); + } + out[static_cast(col * N + oc)] = static_cast(dot) * d; + } + } + return out; +} + +std::string shape_label(const char * what, int64_t K, int64_t N, int64_t M, int nth) { + return std::string(what) + " K=" + std::to_string(K) + " N=" + std::to_string(N) + + " M=" + std::to_string(M) + " nth=" + std::to_string(nth); +} + +// The packer is whole-tensor, so this also checks that a row-major [K, N] weight +// packs into groups that never straddle a row -- true because every I2_S row +// length in the model is a multiple of 128, and false the moment K is not. +void test_pack_layout() { + ggml_init_params ip = { kCtxBytes, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + const int64_t K = 256; + const int64_t N = 3; + + ggml_tensor * w = ggml_new_tensor_2d(ctx, GGML_TYPE_I2_S, K, N); + + const std::vector values = ternary_pattern(static_cast(K * N), 0.75f, 1); + ggml_i2_s_from_float(values.data(), w->data, K * N); + + require_close(*inband_scale(w), 0.75f, 0.0f, "pack layout scale"); + + // Row size is ceil(K/128)*32 bytes, and nb[1] has to agree or the kernel + // walks into the wrong row. + require_eq(static_cast(w->nb[1]), K / 4, "pack layout row stride"); + + const std::vector codes = unpack_i2_s(w->data, K * N); + for (int64_t i = 0; i < K * N; ++i) { + const float expected = values[static_cast(i)]; + require_close(static_cast(codes[static_cast(i)]) * 0.75f, expected, 0.0f, + "pack layout value " + std::to_string(i)); + } + + // And the shipped dequantizer agrees with the hand decode. + std::vector dequantized(static_cast(K * N)); + ggml_i2_s_to_float(w->data, dequantized.data(), K * N); + for (int64_t i = 0; i < K * N; ++i) { + require_close(dequantized[static_cast(i)], values[static_cast(i)], 0.0f, + "pack layout dequantize " + std::to_string(i)); + } + + ggml_free(ctx); +} + +void test_mul_mat(int n_threads, int64_t K, int64_t N, int64_t M) { + ggml_init_params ip = { kCtxBytes, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + ggml_tensor * w = ggml_new_tensor_2d(ctx, GGML_TYPE_I2_S, K, N); + ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, M); + + const float w_scale = 0.031f; + const std::vector w_values = ternary_pattern(static_cast(K * N), w_scale, static_cast(N)); + ggml_i2_s_from_float(w_values.data(), w->data, K * N); + + const std::vector a_values = activation_pattern(static_cast(K * M), 0.4f, 1.7f); + std::memcpy(a->data, a_values.data(), a_values.size() * sizeof(float)); + + ggml_tensor * result = ggml_mul_mat(ctx, w, a); + require_eq(result->ne[0], N, "mul_mat ne0"); + require_eq(result->ne[1], M, "mul_mat ne1"); + + compute(ctx, result, n_threads); + + const std::vector expected = + reference(unpack_i2_s(w->data, K * N), a_values, K, N, M, *inband_scale(w)); + + const float * got = static_cast(result->data); + for (int64_t i = 0; i < N * M; ++i) { + require_close(got[i], expected[static_cast(i)], 0.0f, + shape_label("mul_mat", K, N, M, n_threads) + " element " + std::to_string(i)); + } + + ggml_free(ctx); +} + +// src1 with a third dimension, so the row indexing has to walk nb12 and the +// output has to walk nb2. The weight is shared across the batch. +void test_mul_mat_batched(int n_threads) { + ggml_init_params ip = { kCtxBytes, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + const int64_t K = 384; + const int64_t N = 70; + const int64_t M = 3; + const int64_t B = 4; + + ggml_tensor * w = ggml_new_tensor_2d(ctx, GGML_TYPE_I2_S, K, N); + ggml_tensor * a = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, K, M, B); + + const std::vector w_values = ternary_pattern(static_cast(K * N), 0.02f, 2); + ggml_i2_s_from_float(w_values.data(), w->data, K * N); + + const std::vector a_values = activation_pattern(static_cast(K * M * B), 1.1f, 0.9f); + std::memcpy(a->data, a_values.data(), a_values.size() * sizeof(float)); + + ggml_tensor * result = ggml_mul_mat(ctx, w, a); + require_eq(result->ne[2], B, "batched ne2"); + + compute(ctx, result, n_threads); + + const std::vector codes = unpack_i2_s(w->data, K * N); + const float * got = static_cast(result->data); + + for (int64_t ib = 0; ib < B; ++ib) { + const std::vector slice(a_values.begin() + static_cast(ib * K * M), + a_values.begin() + static_cast((ib + 1) * K * M)); + const std::vector expected = reference(codes, slice, K, N, M, *inband_scale(w)); + + for (int64_t i = 0; i < N * M; ++i) { + require_close(got[ib * N * M + i], expected[static_cast(i)], 0.0f, + "batched batch " + std::to_string(ib) + " element " + std::to_string(i)); + } + } + + ggml_free(ctx); +} + +// Every code at its maximum (2) against every activation at its maximum (127) is +// the worst case for the int16 accumulation inside the AVX2 body: 32 lanes each +// summing 32 products of 2*127, i.e. 16256, which is why the kernel widens to +// int32 every eight groups rather than at the end of the row. K spans more than +// eight groups so the flush actually has to happen. +void test_accumulator_headroom(int n_threads) { + ggml_init_params ip = { kCtxBytes, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + const int64_t K = 1536; + const int64_t N = 5; + + ggml_tensor * w = ggml_new_tensor_2d(ctx, GGML_TYPE_I2_S, K, N); + ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, 1); + + // All +1: every packed code is 2. + const std::vector w_values(static_cast(K * N), 0.5f); + ggml_i2_s_from_float(w_values.data(), w->data, K * N); + + // Flat, so the absmax is 1 and every quantized activation is exactly +127. + const std::vector a_values(static_cast(K), 1.0f); + std::memcpy(a->data, a_values.data(), a_values.size() * sizeof(float)); + + ggml_tensor * result = ggml_mul_mat(ctx, w, a); + compute(ctx, result, n_threads); + + const float expected = static_cast(K * 127) * (0.5f * (1.0f / 127.0f)); + const float * got = static_cast(result->data); + for (int64_t oc = 0; oc < N; ++oc) { + require_close(got[oc], expected, 0.0f, "headroom output " + std::to_string(oc)); + } + + ggml_free(ctx); +} + +// An all-zero weight tensor has scale 0, and an all-zero activation row has +// scale 0. Neither may produce a NaN: the scales are multipliers here, and a +// reciprocal convention would divide by zero in both cases. +void test_degenerate_scales(int n_threads) { + ggml_init_params ip = { kCtxBytes, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + const int64_t K = 128; + const int64_t N = 8; + const int64_t M = 2; + + ggml_tensor * w = ggml_new_tensor_2d(ctx, GGML_TYPE_I2_S, K, N); + ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, K, M); + + const std::vector zeros_w(static_cast(K * N), 0.0f); + ggml_i2_s_from_float(zeros_w.data(), w->data, K * N); + require_close(*inband_scale(w), 0.0f, 0.0f, "degenerate weight scale"); + + std::vector a_values = activation_pattern(static_cast(K * M), 0.2f, 1.0f); + std::memcpy(a->data, a_values.data(), a_values.size() * sizeof(float)); + + ggml_tensor * zero_w_result = ggml_mul_mat(ctx, w, a); + compute(ctx, zero_w_result, n_threads); + + const float * got = static_cast(zero_w_result->data); + for (int64_t i = 0; i < N * M; ++i) { + require(std::isfinite(got[i]), "degenerate weight output finite " + std::to_string(i)); + require_close(got[i], 0.0f, 0.0f, "degenerate weight output " + std::to_string(i)); + } + + // Now a real weight against a zero activation row. Only the second column is + // zeroed, so the first still has to come out right. + const std::vector w_values = ternary_pattern(static_cast(K * N), 0.25f, 0); + ggml_i2_s_from_float(w_values.data(), w->data, K * N); + + for (int64_t k = 0; k < K; ++k) { + a_values[static_cast(K + k)] = 0.0f; + } + std::memcpy(a->data, a_values.data(), a_values.size() * sizeof(float)); + + ggml_tensor * zero_act_result = ggml_mul_mat(ctx, w, a); + compute(ctx, zero_act_result, n_threads); + + const std::vector expected = + reference(unpack_i2_s(w->data, K * N), a_values, K, N, M, *inband_scale(w)); + + got = static_cast(zero_act_result->data); + for (int64_t i = 0; i < N * M; ++i) { + require(std::isfinite(got[i]), "degenerate activation output finite " + std::to_string(i)); + require_close(got[i], expected[static_cast(i)], 0.0f, + "degenerate activation output " + std::to_string(i)); + } + for (int64_t oc = 0; oc < N; ++oc) { + require_close(got[N + oc], 0.0f, 0.0f, "degenerate activation zero column " + std::to_string(oc)); + } + + ggml_free(ctx); +} + +} // namespace + +int main() { + try { + test_pack_layout(); + + // Single- and multi-threaded: the activation quantization and the output + // pass are separated by a barrier, and threads split the output features, + // so a missing barrier or an overlapping split shows up as a thread-count + // dependent result. + for (int nth : {1, 4}) { + // K=128 is one group. K=1024 is exactly the eight groups the AVX2 + // body accumulates in int16 before widening; K=1152 is that plus one, + // so the second, shorter flush runs too. + // N=1 is a single output row (the decode-time lm_head shape), N=64 is + // exactly one output-channel chunk, N=100 spans two, and N=70 is an + // unaligned span. + test_mul_mat(nth, 128, 1, 1); + test_mul_mat(nth, 128, 64, 1); + test_mul_mat(nth, 128, 100, 5); + test_mul_mat(nth, 1024, 70, 1); + test_mul_mat(nth, 1152, 70, 3); + // Fewer output features than threads: the tail threads get an empty + // range and must not write outside it. + test_mul_mat(nth, 256, 2, 2); + test_mul_mat_batched(nth); + test_accumulator_headroom(nth); + test_degenerate_scales(nth); + } + } catch (const std::exception & e) { + std::cerr << "FAILED: " << e.what() << "\n"; + return 1; + } + + std::cout << "all i2_s mul_mat tests passed\n"; + return 0; +} diff --git a/tests/unittests/test_i8_s_fused_ops.cpp b/tests/unittests/test_i8_s_fused_ops.cpp new file mode 100644 index 000000000..8f58519e6 --- /dev/null +++ b/tests/unittests/test_i8_s_fused_ops.cpp @@ -0,0 +1,550 @@ +// Numeric checks for the VibeASR INT8 pipeline additions to ggml: +// GGML_TYPE_I8_S / GGML_TYPE_I2_S and the five fused ops built on them. +// +// Every op is compared against a plain-loop reference computed from the same +// inputs. The references duplicate the requantization ordering deliberately, +// including the parts that lose range -- ggml_mul_mat_add_relu clamps after +// rounding, so the absmax that sets the output scale still counts the negatives +// that are about to become zero. The point of these tests is to pin the +// arithmetic that the SIMD kernels and the model port have to reproduce, not to +// judge whether that arithmetic is optimal. + +#include "test_assert.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +// The in-band scale layout is internal to ggml, so the tests reach for the same +// declarations the implementation uses rather than re-deriving the offsets. +extern "C" { +size_t ggml_type_extra_bytes(enum ggml_type type); +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); +void ggml_i2_s_to_float (const void * x, float * y, int64_t n); +size_t ggml_i2_s_from_float(const float * x, void * y, int64_t n); +} + +using engine::test::require; +using engine::test::require_close; +using engine::test::require_eq; + +namespace { + +constexpr size_t kCtxBytes = 256u * 1024 * 1024; + +float * inband_scale(ggml_tensor * t) { + return reinterpret_cast( + static_cast(t->data) + ggml_nbytes(t) - ggml_type_extra_bytes(t->type)); +} + +// Quantize into an existing I8_S tensor, returning the scale that was chosen. +float fill_i8_s(ggml_tensor * t, const std::vector & values) { + require_eq(static_cast(values.size()), ggml_nelements(t), "fill_i8_s size"); + ggml_i8_s_from_float(values.data(), t->data, ggml_nelements(t)); + return *inband_scale(t); +} + +std::vector read_i8_s(ggml_tensor * t) { + std::vector out(ggml_nelements(t)); + ggml_i8_s_to_float(t->data, out.data(), ggml_nelements(t)); + return out; +} + +std::vector patterned(size_t n, float phase, float scale) { + std::vector v(n); + for (size_t i = 0; i < n; ++i) { + const float x = static_cast(i); + v[i] = scale * (std::sin(phase + 0.13f * x) + 0.4f * std::cos(0.07f * x - phase)); + } + return v; +} + +// The reference requantizer, mirroring ggml_i8_s_requantize: absmax over the +// whole result -- taken before the relu clamp, so a value about to be zeroed +// still widens the scale -- then clamp, then round ties to even. +// +// std::rint, not std::round: the op rounds ties to even so that its AVX2 and +// NEON bodies agree with their scalar tails. std::round would round ties away +// from zero and make the byte-exact comparisons below fail on any tie. +std::vector requantize_ref(const std::vector & values, bool relu, float * scale_out) { + float amax = 0.0f; + for (float v : values) { + amax = std::max(amax, std::fabs(v)); + } + + const float id = amax != 0.0f ? 127.0f / amax : 0.0f; + const float lo = relu ? 0.0f : -127.0f; + + std::vector q(values.size()); + for (size_t i = 0; i < values.size(); ++i) { + float v = values[i] * id; + v = std::max(lo, std::min(127.0f, v)); + q[i] = static_cast(std::rint(v)); + } + + *scale_out = amax != 0.0f ? amax / 127.0f : 0.0f; + return q; +} + +void compare_i8(ggml_tensor * got, const std::vector & want, float want_scale, + const std::string & label) { + require_eq(static_cast(want.size()), ggml_nelements(got), label + " size"); + + const auto * q = static_cast(got->data); + for (size_t i = 0; i < want.size(); ++i) { + // Exact: both sides do the same rounding on the same floats. A single + // LSB of drift would mean the scale differs, which is what matters. + require_eq(static_cast(q[i]), static_cast(want[i]), + label + " value at " + std::to_string(i)); + } + + require_close(*inband_scale(got), want_scale, 1e-9f, label + " scale"); +} + +// Runs a single-node graph on the CPU backend with the given thread count. The +// four requantizing ops reduce an absmax across threads, so the result has to be +// independent of nth -- these tests run every case at 1 and 4. +void compute(ggml_context * ctx, ggml_tensor * result, int n_threads) { + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, result); + ggml_graph_compute_with_ctx(ctx, gf, n_threads); +} + +// ---------------------------------------------------------------- round trips + +void test_i8_s_round_trip() { + const int64_t n = 4096; + const auto values = patterned(n, 0.3f, 2.5f); + + std::vector buf(n + ggml_type_extra_bytes(GGML_TYPE_I8_S)); + const size_t written = ggml_i8_s_from_float(values.data(), buf.data(), n); + require_eq(written, buf.size(), "i8_s written bytes"); + + std::vector back(n); + ggml_i8_s_to_float(buf.data(), back.data(), n); + + float amax = 0.0f; + for (float v : values) { + amax = std::max(amax, std::fabs(v)); + } + // One step of the quantization grid is amax/127; rounding puts every value + // within half of that. + const float tol = amax / 127.0f * 0.5f + 1e-6f; + + for (int64_t i = 0; i < n; ++i) { + require_close(back[i], values[i], tol, "i8_s round trip at " + std::to_string(i)); + } + + std::cout << "i8_s round trip: n=" << n << " step=" << amax / 127.0f << " OK\n"; +} + +void test_i2_s_round_trip() { + // Ternary input: the type represents {-d, 0, +d} exactly, so this has to + // round trip bit for bit rather than approximately. + const int64_t n = 128 * 7; + const float d = 0.0731f; + + std::vector values(n); + for (int64_t i = 0; i < n; ++i) { + const int code = static_cast(i * 7 % 3) - 1; // -1, 0, +1 cycling + values[i] = static_cast(code) * d; + } + + const size_t payload = static_cast(n / 128) * 32; + std::vector buf(payload + ggml_type_extra_bytes(GGML_TYPE_I2_S)); + const size_t written = ggml_i2_s_from_float(values.data(), buf.data(), n); + require_eq(written, buf.size(), "i2_s written bytes"); + + std::vector back(n); + ggml_i2_s_to_float(buf.data(), back.data(), n); + + for (int64_t i = 0; i < n; ++i) { + require_close(back[i], values[i], 1e-7f, "i2_s round trip at " + std::to_string(i)); + } + + // Check the packing itself, not just the round trip: byte gp of a group + // holds positions gp, 32+gp, 64+gp, 96+gp in bit pairs 6, 4, 2, 0. A + // self-consistent but differently-ordered packing would pass the round trip + // above while being unreadable by the SIMD kernels. + for (int gp = 0; gp < 32; ++gp) { + const uint8_t b = buf[gp]; + for (int sub = 0; sub < 4; ++sub) { + const int code = (b >> (6 - 2 * sub)) & 3; + const int want = static_cast((sub * 32 + gp) * 7 % 3) - 1; + const float value = static_cast(code - 1); + require_close(value, static_cast(want), 1e-7f, + "i2_s bit layout at byte " + std::to_string(gp) + + " pair " + std::to_string(sub)); + } + } + + std::cout << "i2_s round trip: n=" << n << " exact, bit layout OK\n"; +} + +// ---------------------------------------------------------------------- ops + +void test_add_scaled(int n_threads) { + const int64_t C = 96; // channels, ne[0] + const int64_t L = 37; // positions + + ggml_init_params ip = { kCtxBytes, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_I8_S, C, L); + ggml_tensor * b = ggml_new_tensor_2d(ctx, GGML_TYPE_I8_S, C, L); + ggml_tensor * gamma = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, C); + + const auto a_f = patterned(C * L, 0.1f, 1.7f); + const auto b_f = patterned(C * L, 1.9f, 0.6f); + + fill_i8_s(a, a_f); + fill_i8_s(b, b_f); + + const auto g = patterned(C, 0.5f, 0.25f); + std::memcpy(gamma->data, g.data(), g.size() * sizeof(float)); + + ggml_tensor * result = ggml_add_scaled(ctx, a, b, gamma); + compute(ctx, result, n_threads); + + // Reference works from the dequantized inputs, so it sees exactly what the + // op sees -- the input quantization error is shared, not compounded. + const auto a_q = read_i8_s(a); + const auto b_q = read_i8_s(b); + + std::vector want(C * L); + for (int64_t i = 0; i < C * L; ++i) { + want[i] = a_q[i] * g[i % C] + b_q[i]; + } + + float want_scale = 0.0f; + const auto want_q = requantize_ref(want, false, &want_scale); + compare_i8(result, want_q, want_scale, "add_scaled nth=" + std::to_string(n_threads)); + + ggml_free(ctx); + std::cout << "add_scaled: C=" << C << " L=" << L << " nth=" << n_threads << " OK\n"; +} + +void test_rms_norm_scaled(int n_threads) { + const int64_t C = 128; + const int64_t L = 29; + const float eps = 1e-5f; + + ggml_init_params ip = { kCtxBytes, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_I8_S, C, L); + ggml_tensor * gamma = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, C); + + const auto a_f = patterned(C * L, 0.7f, 3.1f); + fill_i8_s(a, a_f); + + const auto g = patterned(C, 1.1f, 0.9f); + std::memcpy(gamma->data, g.data(), g.size() * sizeof(float)); + + ggml_tensor * result = ggml_rms_norm_scaled(ctx, a, gamma, eps); + compute(ctx, result, n_threads); + + // Reference normalizes in the float domain. The op instead cancels the input + // scale and keeps the sum of squares in integers, which is the same value + // computed a different way -- so this also checks that the eps rescaling is + // right, since eps is the one term that does not cancel. + const auto a_q = read_i8_s(a); + + std::vector want(C * L); + for (int64_t row = 0; row < L; ++row) { + double sum_sq = 0.0; + for (int64_t i = 0; i < C; ++i) { + const double v = a_q[row * C + i]; + sum_sq += v * v; + } + const float rms_inv = 1.0f / std::sqrt(static_cast(sum_sq / C) + eps); + for (int64_t i = 0; i < C; ++i) { + want[row * C + i] = a_q[row * C + i] * rms_inv * g[i]; + } + } + + float want_scale = 0.0f; + const auto want_q = requantize_ref(want, false, &want_scale); + + // The two paths reach the same value through different arithmetic, so a + // borderline element can round to either side of a grid step. Compare the + // dequantized results with a tolerance of one step instead of demanding + // identical bytes. + require_close(*inband_scale(result), want_scale, want_scale * 1e-4f, + "rms_norm_scaled scale nth=" + std::to_string(n_threads)); + + const auto * q = static_cast(result->data); + for (int64_t i = 0; i < C * L; ++i) { + require_close(static_cast(q[i]), static_cast(want_q[i]), 1.0f, + "rms_norm_scaled value at " + std::to_string(i)); + } + + ggml_free(ctx); + std::cout << "rms_norm_scaled: C=" << C << " L=" << L << " nth=" << n_threads << " OK\n"; +} + +// IC and OC are varied by the caller rather than fixed: IC decides how much of +// the contraction the SIMD path covers and how much falls to the scalar tail, +// and OC decides whether the output-channel chunking loop wraps around. +void test_mul_mat_add(bool relu, int n_threads, int64_t IC, int64_t OC) { + const int64_t N = 23; + + ggml_init_params ip = { kCtxBytes, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + ggml_tensor * w = ggml_new_tensor_2d(ctx, GGML_TYPE_I8_S, IC, OC); + ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_I8_S, IC, N); + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, OC); + + fill_i8_s(w, patterned(IC * OC, 0.2f, 0.8f)); + fill_i8_s(x, patterned(IC * N, 1.3f, 2.2f)); + + const auto b = patterned(OC, 0.9f, 0.15f); + std::memcpy(bias->data, b.data(), b.size() * sizeof(float)); + + ggml_tensor * result = relu ? ggml_mul_mat_add_relu(ctx, w, x, bias) + : ggml_mul_mat_add(ctx, w, x, bias); + compute(ctx, result, n_threads); + + require_eq(result->ne[0], OC, "mul_mat_add ne0"); + require_eq(result->ne[1], N, "mul_mat_add ne1"); + require_eq(static_cast(result->type), static_cast(GGML_TYPE_I8_S), + "mul_mat_add type"); + + const auto w_q = read_i8_s(w); + const auto x_q = read_i8_s(x); + + std::vector want(OC * N); + for (int64_t col = 0; col < N; ++col) { + for (int64_t oc = 0; oc < OC; ++oc) { + double acc = 0.0; + for (int64_t k = 0; k < IC; ++k) { + acc += static_cast(w_q[oc * IC + k]) * x_q[col * IC + k]; + } + want[col * OC + oc] = static_cast(acc) + b[oc]; + } + } + + float want_scale = 0.0f; + const auto want_q = requantize_ref(want, relu, &want_scale); + + const std::string label = std::string(relu ? "mul_mat_add_relu" : "mul_mat_add") + + " IC=" + std::to_string(IC) + " OC=" + std::to_string(OC) + + " nth=" + std::to_string(n_threads); + + require_close(*inband_scale(result), want_scale, want_scale * 1e-4f, label + " scale"); + + const auto * q = static_cast(result->data); + bool saw_negative_input = false; + for (int64_t i = 0; i < OC * N; ++i) { + require_close(static_cast(q[i]), static_cast(want_q[i]), 1.0f, + label + " value at " + std::to_string(i)); + if (want[i] < 0.0f) saw_negative_input = true; + } + + if (relu) { + // Without negatives in the pre-activation the clamp would be untested. + require(saw_negative_input, label + ": inputs never went negative"); + for (int64_t i = 0; i < OC * N; ++i) { + require(q[i] >= 0, label + " left a negative at " + std::to_string(i)); + } + } + + ggml_free(ctx); + std::cout << label << " N=" << N << " OK\n"; +} + +// K is the contraction length, so it decides how much of the dot product the SIMD +// path covers. N is the number of positions, which is what the depthwise path +// batches -- so it decides whether the chunking loop wraps around. +void test_mul_mat_add_depthwise(int n_threads, int64_t K, int64_t N) { + const int64_t C = 40; // channels + + ggml_init_params ip = { kCtxBytes, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + // ne[1] == 1 with ne[2] > 1 is what selects the depthwise path. + ggml_tensor * w = ggml_new_tensor_3d(ctx, GGML_TYPE_I8_S, K, 1, C); + ggml_tensor * x = ggml_new_tensor_3d(ctx, GGML_TYPE_I8_S, K, N, C); + ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, C); + + fill_i8_s(w, patterned(K * C, 0.4f, 1.1f)); + fill_i8_s(x, patterned(K * N * C, 1.7f, 1.9f)); + + const auto b = patterned(C, 0.2f, 0.3f); + std::memcpy(bias->data, b.data(), b.size() * sizeof(float)); + + ggml_tensor * result = ggml_mul_mat_add(ctx, w, x, bias); + compute(ctx, result, n_threads); + + const auto w_q = read_i8_s(w); + const auto x_q = read_i8_s(x); + + // Output index is ch*N + col: channel-major, one row of N per channel. + std::vector want(C * N); + for (int64_t ch = 0; ch < C; ++ch) { + for (int64_t col = 0; col < N; ++col) { + double acc = 0.0; + for (int64_t k = 0; k < K; ++k) { + acc += static_cast(w_q[ch * K + k]) * x_q[ch * N * K + col * K + k]; + } + want[ch * N + col] = static_cast(acc) + b[ch]; + } + } + + float want_scale = 0.0f; + const auto want_q = requantize_ref(want, false, &want_scale); + + const std::string label = "mul_mat_add depthwise K=" + std::to_string(K) + + " N=" + std::to_string(N) + + " nth=" + std::to_string(n_threads); + require_close(*inband_scale(result), want_scale, want_scale * 1e-4f, label + " scale"); + + const auto * q = static_cast(result->data); + for (int64_t i = 0; i < C * N; ++i) { + require_close(static_cast(q[i]), static_cast(want_q[i]), 1.0f, + label + " value at " + std::to_string(i)); + } + + ggml_free(ctx); + std::cout << label << ": C=" << C << " OK\n"; +} + +void test_im2col_asym(int n_threads) { + const int64_t IW = 33; + const int64_t IC = 5; + const int64_t KW = 4; + const int s0 = 2; + const int lp0 = 3; // asymmetric on purpose: the whole reason this op + const int rp0 = 0; // exists is that ggml_im2col cannot express it + const int d0 = 1; + + ggml_init_params ip = { kCtxBytes, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + ggml_tensor * kernel = ggml_new_tensor_3d(ctx, GGML_TYPE_I8_S, KW, IC, 1); + ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_I8_S, IW, IC, 1); + + fill_i8_s(kernel, patterned(KW * IC, 0.1f, 1.0f)); + const float in_scale = fill_i8_s(input, patterned(IW * IC, 0.8f, 1.4f)); + + ggml_tensor * result = ggml_im2col_asym(ctx, kernel, input, s0, 0, lp0, rp0, 0, + d0, 0, false, GGML_TYPE_I8_S); + + const int64_t OW = (IW + lp0 + rp0 - d0 * (KW - 1) - 1) / s0 + 1; + require_eq(result->ne[0], IC * KW, "im2col_asym ne0"); + require_eq(result->ne[1], OW, "im2col_asym ne1"); + + compute(ctx, result, n_threads); + + const auto * in = static_cast(input->data); + const auto * out = static_cast(result->data); + + for (int64_t iow = 0; iow < OW; ++iow) { + for (int64_t iic = 0; iic < IC; ++iic) { + for (int64_t ikw = 0; ikw < KW; ++ikw) { + const int64_t iiw = iow * s0 + ikw * d0 - lp0; + const int8_t want = (iiw < 0 || iiw >= IW) ? 0 : in[iic * IW + iiw]; + require_eq(static_cast(out[iow * (IC * KW) + iic * KW + ikw]), + static_cast(want), + "im2col_asym at ow=" + std::to_string(iow) + + " ic=" + std::to_string(iic) + " kw=" + std::to_string(ikw)); + } + } + } + + // Rearrangement only, so the scale must pass through untouched. + require_close(*inband_scale(result), in_scale, 0.0f, "im2col_asym scale passthrough"); + + ggml_free(ctx); + std::cout << "im2col_asym: IW=" << IW << " KW=" << KW << " lp0=" << lp0 + << " OW=" << OW << " nth=" << n_threads << " OK\n"; +} + +// The VibeASR encoder flips its activations between channel-major and +// length-major with ggml_cont(ggml_permute(...)), which lands in +// ggml_compute_forward_dup. A byte copy moves the payload but not the in-band +// scale, and the permuted view it is handed has no scale of its own, so the +// scale has to be read through view_src. Getting this wrong leaves the copy +// holding whatever was in the buffer -- values stay right while everything +// downstream is off by an arbitrary factor. +void test_i8_s_cont_permute(int n_threads) { + const int64_t C = 5; + const int64_t L = 33; + + ggml_init_params ip = { kCtxBytes, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + ggml_tensor * src = ggml_new_tensor_2d(ctx, GGML_TYPE_I8_S, C, L); + const float src_scale = fill_i8_s(src, patterned(C * L, 0.6f, 1.9f)); + + ggml_tensor * result = ggml_cont(ctx, ggml_permute(ctx, src, 1, 0, 2, 3)); + require_eq(result->ne[0], L, "cont(permute) ne0"); + require_eq(result->ne[1], C, "cont(permute) ne1"); + + compute(ctx, result, n_threads); + + const auto * in = static_cast(src->data); + const auto * out = static_cast(result->data); + for (int64_t ic = 0; ic < C; ++ic) { + for (int64_t il = 0; il < L; ++il) { + require_eq(static_cast(out[ic * L + il]), static_cast(in[il * C + ic]), + "cont(permute) value at c=" + std::to_string(ic) + " l=" + std::to_string(il)); + } + } + + // Rearrangement only: the same scale has to come out the other side. + require_close(*inband_scale(result), src_scale, 0.0f, "cont(permute) scale carried over"); + + ggml_free(ctx); + std::cout << "cont(permute) i8_s: C=" << C << " L=" << L + << " nth=" << n_threads << " OK\n"; +} + +} // namespace + +int main() { + try { + test_i8_s_round_trip(); + test_i2_s_round_trip(); + + // Every requantizing op reduces an absmax across threads, so each is run + // single-threaded and multi-threaded; a missing barrier shows up as a + // scale that depends on the thread count. + for (int nth : {1, 4}) { + test_add_scaled(nth); + test_rms_norm_scaled(nth); + // IC=64: SIMD only, no tail. IC=67: three elements land in the + // scalar tail, which a truncating block count would silently drop. + // IC=13: shorter than one 32-byte step, so entirely scalar. + // OC=48 stays inside one output-channel chunk, OC=100 spans two. + test_mul_mat_add(false, nth, 64, 48); + test_mul_mat_add(true, nth, 64, 48); + test_mul_mat_add(false, nth, 67, 48); + test_mul_mat_add(false, nth, 13, 48); + test_mul_mat_add(false, nth, 67, 100); + // K=7 is shorter than any vector step, so the dot is all scalar; + // K=36 runs the 32-byte body plus a 4-element tail. N=19 stays inside + // one position chunk, N=150 spans three. + test_mul_mat_add_depthwise(nth, 7, 19); + test_mul_mat_add_depthwise(nth, 36, 19); + test_mul_mat_add_depthwise(nth, 7, 150); + test_im2col_asym(nth); + test_i8_s_cont_permute(nth); + } + } catch (const std::exception & e) { + std::cerr << "FAILED: " << e.what() << "\n"; + return 1; + } + + std::cout << "all i8_s fused op tests passed\n"; + return 0; +} diff --git a/tests/vibeasr/test_vibeasr_vae_encoder.cpp b/tests/vibeasr/test_vibeasr_vae_encoder.cpp new file mode 100644 index 000000000..f42384d45 --- /dev/null +++ b/tests/vibeasr/test_vibeasr_vae_encoder.cpp @@ -0,0 +1,259 @@ +// 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 <16 kHz wav>.\n" + " Convert a VibeASR.cpp checkpoint with tools/community_models/convert_vibeasr_vae.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; + } + if (wav.sample_rate != 16000) { + std::fprintf(stderr, "SKIP: %s is %d Hz, the encoder expects 16 kHz\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) / 16000.0; + + 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_vae.py b/tools/community_models/convert_vibeasr_vae.py new file mode 100644 index 000000000..6a28d9678 --- /dev/null +++ b/tools/community_models/convert_vibeasr_vae.py @@ -0,0 +1,232 @@ +#!/usr/bin/env python3 +"""Convert a VibeASR.cpp VAE encoder GGUF into an audio.cpp GGUF package. + +Upstream: https://github.com/microsoft/VibeASR.cpp + +VibeASR.cpp ships its VAE encoder already quantized to its own ggml fork's +GGML_TYPE_I8_S, 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). + +The on-disk layout is identical either way -- an I8_S tensor is +`nelements` int8 bytes followed by a single padded F32 tensor scale, 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_vae.py \ + --input vibeasr-vae-encoder-i8_s.gguf --list + + # produce the audio.cpp package + python3 tools/community_models/convert_vibeasr_vae.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_vae.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", 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 VAE encoder GGUF") + parser.add_argument("--output", type=Path, help="audio.cpp GGUF to write") + 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() + + 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 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("