From f4d4089a8d5b3a9c75796accbaa885225059e053 Mon Sep 17 00:00:00 2001 From: XSquirrelC Date: Fri, 4 Sep 2026 02:30:13 +0000 Subject: [PATCH 1/7] ggml: add GGML_TYPE_I8_S / I2_S and five fused CPU INT8 ops --- CMakeLists.txt | 7 + external/ggml/include/ggml.h | 89 ++++- external/ggml/src/ggml-cpu/ggml-cpu.c | 37 ++ external/ggml/src/ggml-cpu/ops.cpp | 407 ++++++++++++++++++++ external/ggml/src/ggml-cpu/ops.h | 7 + external/ggml/src/ggml-impl.h | 20 + external/ggml/src/ggml-quants.c | 132 +++++++ external/ggml/src/ggml-quants.h | 17 + external/ggml/src/ggml.c | 215 ++++++++++- tests/unittests/test_i8_s_fused_ops.cpp | 487 ++++++++++++++++++++++++ 10 files changed, 1414 insertions(+), 4 deletions(-) create mode 100644 tests/unittests/test_i8_s_fused_ops.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 313f05b7e..8d1537e5c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2254,6 +2254,13 @@ 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(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) 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..66d2dce64 100644 --- a/external/ggml/src/ggml-cpu/ggml-cpu.c +++ b/external/ggml/src/ggml-cpu/ggml-cpu.c @@ -2092,6 +2092,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 @@ -2345,6 +2365,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { } break; 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: @@ -2352,6 +2373,10 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { 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; @@ -2948,6 +2973,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..fc2ae9177 100644 --- a/external/ggml/src/ggml-cpu/ops.cpp +++ b/external/ggml/src/ggml-cpu/ops.cpp @@ -11637,3 +11637,410 @@ 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. +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; + + for (int64_t row = 0; row < n_rows; row++) { + for (int64_t col = col0; col < col1; col++) { + const int64_t i = row*row_stride + col; + + float v = stage[i] * id; + v = v < -127.0f ? -127.0f : (v > 127.0f ? 127.0f : v); + + int8_t q = (int8_t) roundf(v); + dst_data[i] = relu && q < 0 ? 0 : q; + } + } + + 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 + +// 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. + for (int64_t ch = 0; ch < ne02; ch++) { + const int8_t * w_ch = w + ch*ne00; + + for (int64_t col = col0; col < col1; col++) { + // src1 is [K, N, C]: the channel stride is ne11*ne00. + const int8_t * x_col = x + ch*ne11*ne00 + col*ne00; + + int32_t acc = 0; + for (int64_t k = 0; k < ne00; k++) { + acc += (int32_t) w_ch[k] * (int32_t) x_col[k]; + } + + const float v = (float) acc * combined_d + bias[ch]; + + stage[ch*ne11 + col] = v; + + const float av = fabsf(v); + if (av > local_absmax) local_absmax = av; + } + } + + 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]; + + 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 oc = 0; oc < ne01; oc++) { + const int8_t * w_row = w + oc*ne00; + + int32_t acc = 0; + for (int64_t k = 0; k < ne00; k++) { + acc += (int32_t) w_row[k] * (int32_t) x_col[k]; + } + + const float v = (float) acc * combined_d + bias[oc]; + + stage[col*ne01 + oc] = v; + + const float av = fabsf(v); + if (av > local_absmax) local_absmax = av; + } + } + + 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_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..4783a0134 100644 --- a/external/ggml/src/ggml-cpu/ops.h +++ b/external/ggml/src/ggml-cpu/ops.h @@ -115,6 +115,13 @@ 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_im2col_asym(const struct ggml_compute_params * params, struct ggml_tensor * dst); #ifdef __cplusplus } #endif 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..f717d78f8 100644 --- a/external/ggml/src/ggml-quants.c +++ b/external/ggml/src/ggml-quants.c @@ -2407,6 +2407,138 @@ 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); +} + // ====================== "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..66216d6cd 100644 --- a/external/ggml/src/ggml-quants.h +++ b/external/ggml/src/ggml-quants.h @@ -102,6 +102,23 @@ 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); + 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/tests/unittests/test_i8_s_fused_ops.cpp b/tests/unittests/test_i8_s_fused_ops.cpp new file mode 100644 index 000000000..ad2148c43 --- /dev/null +++ b/tests/unittests/test_i8_s_fused_ops.cpp @@ -0,0 +1,487 @@ +// 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: absmax over the whole result, then round, then +// clamp. Mirrors ggml_i8_s_requantize. +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; + + std::vector q(values.size()); + for (size_t i = 0; i < values.size(); ++i) { + float v = values[i] * id; + v = std::max(-127.0f, std::min(127.0f, v)); + const int8_t r = static_cast(std::round(v)); + q[i] = relu && r < 0 ? 0 : r; + } + + *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"; +} + +void test_mul_mat_add(bool relu, int n_threads) { + const int64_t IC = 64; + const int64_t OC = 48; + 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") + + " 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 << ": IC=" << IC << " OC=" << OC << " N=" << N << " OK\n"; +} + +void test_mul_mat_add_depthwise(int n_threads) { + const int64_t K = 7; // kernel width + const int64_t C = 40; // channels + const int64_t N = 19; // positions + + 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 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 << ": K=" << K << " C=" << C << " N=" << N << " 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"; +} + +} // 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); + test_mul_mat_add(false, nth); + test_mul_mat_add(true, nth); + test_mul_mat_add_depthwise(nth); + test_im2col_asym(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; +} From 42dd4a67c8f0d03a7cb6f1fd2d6af5f18099f9d0 Mon Sep 17 00:00:00 2001 From: XSquirrelC Date: Fri, 4 Sep 2026 02:30:13 +0000 Subject: [PATCH 2/7] ggml-cpu: SIMD kernels for the I8_S fused ops --- external/ggml/src/ggml-cpu/ops.cpp | 205 ++++++++++++++++++++---- external/ggml/src/ggml-cpu/vec.cpp | 109 +++++++++++++ external/ggml/src/ggml-cpu/vec.h | 7 + tests/unittests/test_i8_s_fused_ops.cpp | 56 +++++-- 4 files changed, 329 insertions(+), 48 deletions(-) diff --git a/external/ggml/src/ggml-cpu/ops.cpp b/external/ggml/src/ggml-cpu/ops.cpp index fc2ae9177..735030a78 100644 --- a/external/ggml/src/ggml-cpu/ops.cpp +++ b/external/ggml/src/ggml-cpu/ops.cpp @@ -11671,6 +11671,71 @@ void ggml_compute_forward_fwht(const ggml_compute_params * params, ggml_tensor * // 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, @@ -11705,16 +11770,12 @@ static void ggml_i8_s_requantize( 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++) { - for (int64_t col = col0; col < col1; col++) { - const int64_t i = row*row_stride + col; - - float v = stage[i] * id; - v = v < -127.0f ? -127.0f : (v > 127.0f ? 127.0f : v); + const int64_t i = row*row_stride + col0; - int8_t q = (int8_t) roundf(v); - dst_data[i] = relu && q < 0 ? 0 : q; - } + ggml_i8_s_quantize_range(dst_data + i, stage + i, col1 - col0, id, relu); } if (ith == 0) { @@ -11854,6 +11915,81 @@ void ggml_compute_forward_rms_norm_scaled( // 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 @@ -11919,24 +12055,32 @@ static void ggml_compute_forward_mul_mat_add_impl( 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; - for (int64_t col = col0; col < col1; col++) { - // src1 is [K, N, C]: the channel stride is ne11*ne00. - const int8_t * x_col = x + ch*ne11*ne00 + col*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]; + } - int32_t acc = 0; - for (int64_t k = 0; k < ne00; k++) { - acc += (int32_t) w_ch[k] * (int32_t) x_col[k]; - } + 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); - const float v = (float) acc * combined_d + bias[ch]; + // src1 is [K, N, C]: the channel stride is ne11*ne00. + const int8_t * x_col = x + ch*ne11*ne00 + col*ne00; - stage[ch*ne11 + col] = v; + ggml_vec_dot_i8_i8(ne00, acc, 1, x_col, ne00, w_ch, ncol); - const float av = fabsf(v); - if (av > local_absmax) local_absmax = av; + local_absmax = ggml_i8_s_scale_bias_absmax( + stage + ch*ne11 + col, acc, bias_v, ncol, combined_d, local_absmax); } } @@ -11946,23 +12090,22 @@ static void ggml_compute_forward_mul_mat_add_impl( // 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 oc = 0; oc < ne01; oc++) { - const int8_t * w_row = w + oc*ne00; - - int32_t acc = 0; - for (int64_t k = 0; k < ne00; k++) { - acc += (int32_t) w_row[k] * (int32_t) x_col[k]; - } - - const float v = (float) acc * combined_d + bias[oc]; + 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); - stage[col*ne01 + oc] = v; + ggml_vec_dot_i8_i8(ne00, acc, 1, w + oc0*ne00, ne00, x_col, noc); - const float av = fabsf(v); - if (av > local_absmax) local_absmax = av; + local_absmax = ggml_i8_s_scale_bias_absmax( + stage + col*ne01 + oc0, acc, bias + oc0, noc, combined_d, local_absmax); } } diff --git a/external/ggml/src/ggml-cpu/vec.cpp b/external/ggml/src/ggml-cpu/vec.cpp index 7e8135617..624090e1e 100644 --- a/external/ggml/src/ggml-cpu/vec.cpp +++ b/external/ggml/src/ggml-cpu/vec.cpp @@ -393,6 +393,115 @@ 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_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..bf425d76e 100644 --- a/external/ggml/src/ggml-cpu/vec.h +++ b/external/ggml/src/ggml-cpu/vec.h @@ -43,6 +43,13 @@ 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); + 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/tests/unittests/test_i8_s_fused_ops.cpp b/tests/unittests/test_i8_s_fused_ops.cpp index ad2148c43..25c73ea90 100644 --- a/tests/unittests/test_i8_s_fused_ops.cpp +++ b/tests/unittests/test_i8_s_fused_ops.cpp @@ -67,8 +67,13 @@ std::vector patterned(size_t n, float phase, float scale) { return v; } -// The reference requantizer: absmax over the whole result, then round, then -// clamp. Mirrors ggml_i8_s_requantize. +// 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) { @@ -76,13 +81,13 @@ std::vector requantize_ref(const std::vector & values, bool relu, } 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(-127.0f, std::min(127.0f, v)); - const int8_t r = static_cast(std::round(v)); - q[i] = relu && r < 0 ? 0 : r; + 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; @@ -286,9 +291,10 @@ void test_rms_norm_scaled(int n_threads) { std::cout << "rms_norm_scaled: C=" << C << " L=" << L << " nth=" << n_threads << " OK\n"; } -void test_mul_mat_add(bool relu, int n_threads) { - const int64_t IC = 64; - const int64_t OC = 48; +// 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 }; @@ -331,6 +337,7 @@ void test_mul_mat_add(bool relu, int n_threads) { 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"); @@ -352,13 +359,14 @@ void test_mul_mat_add(bool relu, int n_threads) { } ggml_free(ctx); - std::cout << label << ": IC=" << IC << " OC=" << OC << " N=" << N << " OK\n"; + std::cout << label << " N=" << N << " OK\n"; } -void test_mul_mat_add_depthwise(int n_threads) { - const int64_t K = 7; // kernel width +// 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 - const int64_t N = 19; // positions ggml_init_params ip = { kCtxBytes, nullptr, false }; ggml_context * ctx = ggml_init(ip); @@ -395,7 +403,9 @@ void test_mul_mat_add_depthwise(int n_threads) { float want_scale = 0.0f; const auto want_q = requantize_ref(want, false, &want_scale); - const std::string label = "mul_mat_add depthwise nth=" + std::to_string(n_threads); + 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); @@ -405,7 +415,7 @@ void test_mul_mat_add_depthwise(int n_threads) { } ggml_free(ctx); - std::cout << label << ": K=" << K << " C=" << C << " N=" << N << " OK\n"; + std::cout << label << ": C=" << C << " OK\n"; } void test_im2col_asym(int n_threads) { @@ -472,9 +482,21 @@ int main() { for (int nth : {1, 4}) { test_add_scaled(nth); test_rms_norm_scaled(nth); - test_mul_mat_add(false, nth); - test_mul_mat_add(true, nth); - test_mul_mat_add_depthwise(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); } } catch (const std::exception & e) { From e2dd409e78d1eea8cef6bb057452888b295576d0 Mon Sep 17 00:00:00 2001 From: XSquirrelC Date: Fri, 4 Sep 2026 03:30:05 +0000 Subject: [PATCH 3/7] ggml-cpu: carry the I8_S in-band scale through dup/cont --- external/ggml/src/ggml-cpu/ops.cpp | 10 ++++++ tests/unittests/test_i8_s_fused_ops.cpp | 41 +++++++++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/external/ggml/src/ggml-cpu/ops.cpp b/external/ggml/src/ggml-cpu/ops.cpp index 735030a78..b07bb7823 100644 --- a/external/ggml/src/ggml-cpu/ops.cpp +++ b/external/ggml/src/ggml-cpu/ops.cpp @@ -531,6 +531,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; } diff --git a/tests/unittests/test_i8_s_fused_ops.cpp b/tests/unittests/test_i8_s_fused_ops.cpp index 25c73ea90..8f58519e6 100644 --- a/tests/unittests/test_i8_s_fused_ops.cpp +++ b/tests/unittests/test_i8_s_fused_ops.cpp @@ -469,6 +469,46 @@ void test_im2col_asym(int n_threads) { << " 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() { @@ -498,6 +538,7 @@ int main() { 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"; From a631fab834ee01282452fa1e124163b4202cd1fa Mon Sep 17 00:00:00 2001 From: XSquirrelC Date: Fri, 4 Sep 2026 03:36:25 +0000 Subject: [PATCH 4/7] community_models: VibeASR I8_S VAE encoder --- CMakeLists.txt | 36 ++ README.md | 1 + docs/asr.md | 5 + docs/community_models/models.md | 1 + docs/community_models/vibeasr.md | 205 ++++++++++ .../engine/community_models/vibeasr/assets.h | 74 ++++ .../community_models/vibeasr/vae_encoder.h | 88 ++++ src/community_models/vibeasr/assets.cpp | 148 +++++++ src/community_models/vibeasr/vae_encoder.cpp | 379 ++++++++++++++++++ tests/vibeasr/test_vibeasr_vae_encoder.cpp | 259 ++++++++++++ tools/community_models/convert_vibeasr_vae.py | 232 +++++++++++ 11 files changed, 1428 insertions(+) create mode 100644 docs/community_models/vibeasr.md create mode 100644 include/engine/community_models/vibeasr/assets.h create mode 100644 include/engine/community_models/vibeasr/vae_encoder.h create mode 100644 src/community_models/vibeasr/assets.cpp create mode 100644 src/community_models/vibeasr/vae_encoder.cpp create mode 100644 tests/vibeasr/test_vibeasr_vae_encoder.cpp create mode 100644 tools/community_models/convert_vibeasr_vae.py diff --git a/CMakeLists.txt b/CMakeLists.txt index 8d1537e5c..c9bc62579 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1618,6 +1618,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}) @@ -2400,6 +2408,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 446d157fb..117373cef 100644 --- a/README.md +++ b/README.md @@ -160,6 +160,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 | ## 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 fc6bc07cd..dcad26a3b 100644 --- a/docs/community_models/models.md +++ b/docs/community_models/models.md @@ -33,3 +33,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..4b76de1c5 --- /dev/null +++ b/docs/community_models/vibeasr.md @@ -0,0 +1,205 @@ +# 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` (not ported 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). Converging the two — reusing the `vibevoice_asr` loader, +session, and streaming state machine and treating I8_S as one more weight path — +is the intended direction for the decoder half; see [Status](#status). + +## 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.** VibeASR.cpp runs it on ternary `GGML_TYPE_I2_S` weights whose + kernel is not in tree, 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. +- **Convergence with `vibevoice_asr`.** Once the I2_S kernel lands, the decoder + half should reuse that family's loader, session, and streaming machinery rather + than duplicate it. Maintainer preference on where the I8_S/I2_S path should + live — a separate community family, or a weight path inside `vibevoice_asr` — + is worth settling before that PR. + +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/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/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(" Date: Fri, 4 Sep 2026 05:41:24 +0000 Subject: [PATCH 5/7] docs: state the vibeasr/vibevoice_asr split as a decision --- docs/community_models/vibeasr.md | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/docs/community_models/vibeasr.md b/docs/community_models/vibeasr.md index 4b76de1c5..e5db9fd6f 100644 --- a/docs/community_models/vibeasr.md +++ b/docs/community_models/vibeasr.md @@ -28,9 +28,15 @@ not a different architecture: 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). Converging the two — reusing the `vibevoice_asr` loader, -session, and streaming state machine and treating I8_S as one more weight path — -is the intended direction for the decoder half; see [Status](#status). +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 @@ -180,15 +186,11 @@ Ported: Not ported yet: - **The decoder.** VibeASR.cpp runs it on ternary `GGML_TYPE_I2_S` weights whose - kernel is not in tree, so there is no loader, no session, and no + matmul kernel is not in tree yet, 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. -- **Convergence with `vibevoice_asr`.** Once the I2_S kernel lands, the decoder - half should reuse that family's loader, session, and streaming machinery rather - than duplicate it. Maintainer preference on where the I8_S/I2_S path should - live — a separate community family, or a weight path inside `vibevoice_asr` — - is worth settling before that PR. + only useful together. Two follow-up PRs cover it: the I2_S matmul kernel, then + the decoder graph plus loader and session. Known limitations: From 69de5259dbb9cff8b8cdb2bdb2ccc527eb8f01d1 Mon Sep 17 00:00:00 2001 From: XSquirrelC Date: Fri, 4 Sep 2026 05:51:37 +0000 Subject: [PATCH 6/7] ggml-cpu: ternary I2_S matmul kernel --- CMakeLists.txt | 7 + docs/community_models/vibeasr.md | 10 +- external/ggml/src/ggml-cpu/ggml-cpu.c | 64 ++- external/ggml/src/ggml-cpu/ops.cpp | 781 +++++++++++++++----------- external/ggml/src/ggml-cpu/ops.h | 11 +- external/ggml/src/ggml-cpu/vec.cpp | 123 ++++ external/ggml/src/ggml-cpu/vec.h | 17 + external/ggml/src/ggml-quants.c | 25 + external/ggml/src/ggml-quants.h | 13 + tests/unittests/test_i2_s_mul_mat.cpp | 373 ++++++++++++ 10 files changed, 1058 insertions(+), 366 deletions(-) create mode 100644 tests/unittests/test_i2_s_mul_mat.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 2c0fd403b..34aaa817d 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -2294,6 +2294,13 @@ if (ENGINE_BUILD_TESTS) 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) diff --git a/docs/community_models/vibeasr.md b/docs/community_models/vibeasr.md index e5db9fd6f..89df1b4eb 100644 --- a/docs/community_models/vibeasr.md +++ b/docs/community_models/vibeasr.md @@ -21,7 +21,7 @@ not a different architecture: | 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` (not ported yet) | +| 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 | @@ -185,12 +185,12 @@ Ported: Not ported yet: -- **The decoder.** VibeASR.cpp runs it on ternary `GGML_TYPE_I2_S` weights whose - matmul kernel is not in tree yet, so there is no loader, no session, and no +- **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. Two follow-up PRs cover it: the I2_S matmul kernel, then - the decoder graph plus loader and session. + only useful together. Known limitations: diff --git a/external/ggml/src/ggml-cpu/ggml-cpu.c b/external/ggml/src/ggml-cpu/ggml-cpu.c index 66d2dce64..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); @@ -2321,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: { @@ -2363,12 +2374,12 @@ 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_ASYM: - 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: @@ -2844,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) { diff --git a/external/ggml/src/ggml-cpu/ops.cpp b/external/ggml/src/ggml-cpu/ops.cpp index b07bb7823..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 @@ -1918,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; @@ -2064,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( @@ -4511,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 @@ -4567,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); } } } @@ -6430,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: { @@ -6446,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, @@ -12135,6 +12136,116 @@ void ggml_compute_forward_mul_mat_add_relu( 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( diff --git a/external/ggml/src/ggml-cpu/ops.h b/external/ggml/src/ggml-cpu/ops.h index 4783a0134..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); @@ -121,6 +121,7 @@ void ggml_compute_forward_add_scaled(const struct ggml_compute_params * params, 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 } diff --git a/external/ggml/src/ggml-cpu/vec.cpp b/external/ggml/src/ggml-cpu/vec.cpp index 624090e1e..0d42f4ede 100644 --- a/external/ggml/src/ggml-cpu/vec.cpp +++ b/external/ggml/src/ggml-cpu/vec.cpp @@ -502,6 +502,129 @@ void ggml_vec_dot_i8_i8(int n, int32_t * GGML_RESTRICT s, size_t bs, const int8_ } } +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 bf425d76e..8943e86c6 100644 --- a/external/ggml/src/ggml-cpu/vec.h +++ b/external/ggml/src/ggml-cpu/vec.h @@ -50,6 +50,23 @@ void ggml_vec_dot_f16(int n, float * GGML_RESTRICT s, size_t bs, ggml_fp16_t * G // 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-quants.c b/external/ggml/src/ggml-quants.c index f717d78f8..facdc2473 100644 --- a/external/ggml/src/ggml-quants.c +++ b/external/ggml/src/ggml-quants.c @@ -2539,6 +2539,31 @@ size_t ggml_i2_s_from_float(const float * GGML_RESTRICT x, void * GGML_RESTRICT 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 66216d6cd..cda127695 100644 --- a/external/ggml/src/ggml-quants.h +++ b/external/ggml/src/ggml-quants.h @@ -119,6 +119,19 @@ GGML_API size_t ggml_i8_s_from_float(const float * GGML_RESTRICT x, void * GGML 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/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; +} From 5a2f68845569ae8cc9b4c52660dcb74f0ea3f15d Mon Sep 17 00:00:00 2001 From: XSquirrelC Date: Fri, 4 Sep 2026 07:11:31 +0000 Subject: [PATCH 7/7] vibeasr: add ternary I2_S LM decoder graph, loader, and session --- CMakeLists.txt | 32 +- docs/asr.md | 8 +- docs/community_models/models.md | 2 +- docs/community_models/vibeasr.md | 289 ++++++-- .../engine/community_models/vibeasr/assets.h | 39 +- .../community_models/vibeasr/lm_decoder.h | 62 ++ .../engine/community_models/vibeasr/session.h | 60 ++ .../community_models/vibeasr/vae_encoder.h | 2 +- model_specs/vibeasr.json | 139 ++++ src/community_models/vibeasr/assets.cpp | 129 +++- src/community_models/vibeasr/lm_decoder.cpp | 679 ++++++++++++++++++ src/community_models/vibeasr/session.cpp | 368 ++++++++++ tests/vibeasr/test_vibeasr_asr.cpp | 145 ++++ tests/vibeasr/test_vibeasr_vae_encoder.cpp | 12 +- ...vibeasr_vae.py => convert_vibeasr_gguf.py} | 60 +- 15 files changed, 1920 insertions(+), 106 deletions(-) create mode 100644 include/engine/community_models/vibeasr/lm_decoder.h create mode 100644 include/engine/community_models/vibeasr/session.h create mode 100644 model_specs/vibeasr.json create mode 100644 src/community_models/vibeasr/lm_decoder.cpp create mode 100644 src/community_models/vibeasr/session.cpp create mode 100644 tests/vibeasr/test_vibeasr_asr.cpp rename tools/community_models/{convert_vibeasr_vae.py => convert_vibeasr_gguf.py} (73%) diff --git a/CMakeLists.txt b/CMakeLists.txt index 34aaa817d..14badab01 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1636,12 +1636,16 @@ 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 + src/community_models/vibeasr/lm_decoder.cpp + src/community_models/vibeasr/session.cpp + INCLUDES + engine/community_models/vibeasr/session.h + LOADERS + engine::community_models::vibeasr::make_vibeasr_loader ) set(AUDIOCPP_ENABLED_MODELS "") @@ -2455,7 +2459,7 @@ if (ENGINE_BUILD_TESTS) add_test( NAME test_vibeasr_vae_encoder COMMAND test_vibeasr_vae_encoder - --model ${CMAKE_CURRENT_SOURCE_DIR}/models/vibeasr/vae_encoder-i8_s.gguf + --model ${CMAKE_CURRENT_SOURCE_DIR}/models/vibeasr/vibeasr-vae-encoder-i8_s.gguf --audio ${CMAKE_CURRENT_SOURCE_DIR}/assets/asr_validation/librispeech/librispeech_test_clean_6930-75918-0000.wav ) # Needs the converted 703 MB encoder package, which a normal checkout @@ -2466,6 +2470,28 @@ if (ENGINE_BUILD_TESTS) SKIP_RETURN_CODE 125 TIMEOUT 300 ) + + add_executable(test_vibeasr_asr + tests/vibeasr/test_vibeasr_asr.cpp + ) + target_compile_definitions(test_vibeasr_asr PRIVATE + ENGINE_REPO_ROOT="${CMAKE_CURRENT_SOURCE_DIR}" + ) + target_link_libraries(test_vibeasr_asr PRIVATE engine_runtime ggml) + target_include_directories(test_vibeasr_asr PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}) + if (ENGINE_ENABLE_OPENMP) + target_link_libraries(test_vibeasr_asr PRIVATE OpenMP::OpenMP_CXX) + endif() + add_test( + NAME test_vibeasr_asr + COMMAND test_vibeasr_asr --threads 8 + ) + # Same story as the encoder probe, plus the 993 MB decoder: exits 125 + # (skip) unless both converted GGUFs sit in models/vibeasr/. + set_tests_properties(test_vibeasr_asr PROPERTIES + SKIP_RETURN_CODE 125 + TIMEOUT 600 + ) endif() if (audio8_asr IN_LIST AUDIOCPP_LINKED_MODELS) diff --git a/docs/asr.md b/docs/asr.md index 5ea156426..ffafa7b08 100644 --- a/docs/asr.md +++ b/docs/asr.md @@ -329,10 +329,10 @@ 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. +A fully quantized port of the same model — INT8 activations through the encoder, +ternary BitNet weights in the decoder — lives under community models as +`vibeasr`: see [VibeASR](community_models/vibeasr.md). It is not a separate +model, only a CPU-only alternative numeric pipeline for the same weights. | Field | Value | |---|---| diff --git a/docs/community_models/models.md b/docs/community_models/models.md index 8298b2439..da5c81980 100644 --- a/docs/community_models/models.md +++ b/docs/community_models/models.md @@ -34,4 +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 | +| **vibeasr** | ASR | en | [@XsquirrelC](https://github.com/XsquirrelC) | [VibeASR](vibeasr.md) fully quantized port of [VibeASR.cpp](https://github.com/microsoft/VibeASR.cpp): the VibeVoice acoustic/semantic tokenizers on INT8 weights *and* INT8 activations through the fused `GGML_TYPE_I8_S` ops, feeding a ternary `GGML_TYPE_I2_S` BitNet Qwen2 decoder. Offline, CPU only | diff --git a/docs/community_models/vibeasr.md b/docs/community_models/vibeasr.md index 89df1b4eb..3b675befe 100644 --- a/docs/community_models/vibeasr.md +++ b/docs/community_models/vibeasr.md @@ -1,17 +1,17 @@ -# VibeASR VAE encoder in audio.cpp +# VibeASR in audio.cpp [VibeASR.cpp](https://github.com/microsoft/VibeASR.cpp) is Microsoft's CPU-first port of the VibeVoice ASR stack, quantized end to end for edge inference: the audio VAE encoder runs on INT8 weights *and* INT8 activations, and the Qwen2 -decoder runs on BitNet-style ternary weights. This page covers the first half of -that port — the VAE encoder — which is what this entry currently adds. +decoder runs on BitNet-style ternary weights. This entry ports both halves, so +`--family vibeasr` transcribes end to end on CPU. ## Relation to the existing `vibevoice_asr` family audio.cpp already ships [VibeVoice ASR](../asr.md#vibevoice-asr) in the core model tree, and it is the same model: the same acoustic/semantic causal ConvNeXt tokenizers, the same connectors, the same Qwen2 decoder. That family runs F32 / -Q8_0 weights through the generic ggml ops and supports streaming. +Q8_0 weights through the generic ggml ops. What VibeASR.cpp adds is a different *numeric pipeline* for that architecture, not a different architecture: @@ -21,34 +21,42 @@ not a different architecture: | 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 | +| Decoder weights | Q8_0 Qwen2 | ternary `GGML_TYPE_I2_S`, 993 MB for a 1.5B decoder | +| Backends | CPU, CUDA, Metal | CPU only — the I8_S and I2_S kernels have no GPU variants | +| Decode | greedy, sampling, beam search | greedy | +| Output | text, segments, speaker turns | text | + +Both are offline-only. So this is an alternative execution path for weights that were quantized -upstream, useful where the INT8/ternary package is the point (no F32 -activations anywhere, integer dot products, and a much smaller decoder once the -I2_S kernel lands). +upstream, useful where the INT8/ternary package is the point: no F32 activations +anywhere, integer dot products, and a decoder that fits in under 1 GB. It stays a separate community entry rather than becoming a weight path inside -`vibevoice_asr`, because the two share no 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. +`vibevoice_asr`, because the two share no encoder graph code: every activation +there is I8_S and every node is one of the fused CPU-only ops, so folding it in +would put a second, mutually exclusive graph builder and a second backend policy +behind one family's loader. The reuse that is worth having — tokenizer +vocabulary, prompt layout, feature-injection order, audio normalization — is data +and conventions, and this entry follows `vibevoice_asr` on all of it. The decoder +half needs no new graph code at all: it is +`modules::QwenCausalDecoderModule` unchanged, because every projection goes +through `LinearModule`'s plain `ggml_mul_mat`, which dispatches on the weight +type. ## Architecture +### Encoder + Both branches are identical in shape and differ only in latent width: -- **Input**: mono 16 kHz waveform in `[-1, 1]`, quantized to a single I8_S +- **Input**: mono 24 kHz waveform in `[-1, 1]`, quantized to a single I8_S tensor (one scale for the whole waveform, `amax` floored at 1e-5 to match upstream). - **7 stages**, strides `{1, 2, 2, 4, 5, 5, 8}` (upstream `encoder_ratios` `[8, 5, 5, 4, 2, 2]` reversed, with a stride-1 stem), so **3200 samples per - frame** — 5 frames per second. Channels `32 → 64 → 128 → 256 → 512 → 1024 → - 2048`, depths `3-3-3-3-3-3-8`. + frame** — 7.5 frames per second at 24 kHz. Channels `32 → 64 → 128 → 256 → + 512 → 1024 → 2048`, depths `3-3-3-3-3-3-8`. - Each stage starts with a **strided causal conv** (left pad `K - stride`, right pad 0) and then runs its ConvNeXt-style blocks: RMSNorm → depthwise conv → layer scale → residual → RMSNorm → FC1 → ReLU → FC2 → layer scale → residual. @@ -71,62 +79,200 @@ exist, what shape each weight has), not from GGUF KV metadata — the same approach upstream takes, and it keeps the loader working for any checkpoint with this topology. +### Decoder + +A stock Qwen2 causal decoder, geometry read from the LM GGUF's KV block: 28 +layers, hidden 1536, intermediate 8960, 12 heads over 2 KV heads, head_dim 128, +RMSNorm eps 1e-6, RoPE theta 1e6, context 65536. The checkpoint has no +`qwen2.attention.key_length`, so `head_dim` comes from +`qwen2.rope.dimension_count`, which for this model equals +`embedding_length / head_count`; the loader cross-checks +`head_dim * head_count == embedding_length` and validates the declared geometry +against `token_embd.weight`'s shape. + +Weight types are mixed on purpose, exactly as published: + +| Tensors | Type | +|---|---| +| `blk.N.{attn_q,attn_k,attn_v,attn_output,ffn_gate,ffn_up,ffn_down}.weight` | `I2_S` (ternary) | +| `token_embd.weight` | Q6_K | +| `output.weight` | F16 | +| norms and `blk.N.attn_{q,k,v}.bias` | F32 | + +`I2_S` packs `{-1, 0, +1}` as codes `{0, 1, 2}`, 128 values per 32-byte group, +over the whole flat tensor, with one F32 absmax scale after the payload. The +kernel asserts `ne00 % 128 == 0`; hidden 1536 and intermediate 8960 both satisfy +it, and the weight is always 2-D by the time `LinearModule` calls +`ggml_mul_mat`. + +### Prompt + +Qwen2.5 ChatML, assembled to match `VibeASR.cpp/utils/prompt_builder.h` token for +token: + +``` +<|im_start|>system\nYou are a helpful assistant that transcribes audio input into text output in JSON format.<|im_end|>\n +<|im_start|>user\n<|speech_start|><|speech_pad|>×N<|speech_end|>\nThis is a 3.50 seconds audio, please transcribe it.<|im_end|>\n +``` + +- The special tokens are inserted by numeric id (151643–151648), not through the + tokenizer, because the GGUF vocabulary still carries Qwen2.5's original text + for those slots while the embedding rows are the ones VibeVoice trained. Every + text segment is tokenized with `parse_special = false`. +- There is deliberately **no generation prompt**: the model emits its own + `<|im_start|>assistant\n` header, and the session strips that leading triple + before decoding, as upstream does. +- `N` is the encoder frame count. Upstream builds `ceil(samples / 3200)` pads but + prefills only `min(pads, frames)` of them, so emitting exactly `frames` pads + produces the same sequence. +- The `<|speech_pad|>` rows are replaced in-graph by a `ggml_set_rows` over the + embedding lookup, with the speech features being the **element-wise sum** of + the acoustic and semantic connector outputs — both are 1536 wide, which is what + makes the sum well-defined. +- `output_format=json` swaps the instruction for `please transcribe it with these + keys: Start, End, Speaker, Content`; `context=...` switches to the + `with extra info:` suffix variant. + +Decoding is greedy, stopping at `<|im_end|>` or `<|endoftext|>`. Upstream's +default is temperature 0.7 / top-p 0.9 sampling with `--greedy` as an opt-in; +this port only implements the deterministic path, which is what parity is +measured against. + +### Audio front end + +Mixdown to mono, resample to 24 kHz, RMS-normalize to −25 dBFS with `eps = 1e-6`, +then divide by `max_abs` if it exceeded 1.0. This is audio.cpp's own +`vibevoice_asr` front end, not upstream's: VibeASR.cpp resamples with a naive +linear kernel and omits the clamp. For a clip already at 24 kHz the two agree; +for anything else the resampler differs and so do the encoder features (see +[Parity](#parity)). + ## Usage -VibeASR.cpp already ships the encoder quantized, so there is nothing to +VibeASR.cpp already ships both halves quantized, so there is nothing to re-quantize. The two forks only disagree on the numeric type *ids* — the VibeASR fork put I2_S/I8_S at 36/37, which upstream ggml had already spent on the retired -`IQ4_NL_4_4` / `IQ4_NL_4_8` slots, so audio.cpp registers them at 42/43. The +`IQ4_NL_4_4` / `IQ4_NL_4_8` slots, so audio.cpp registers them at 43/42. The converter rewrites the 4-byte type field in each tensor info and copies everything else through byte for byte: ```bash # inspect first -python3 tools/community_models/convert_vibeasr_vae.py \ +python3 tools/community_models/convert_vibeasr_gguf.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 +# fix both GGUFs in place (703 MB encoder, 993 MB decoder) +python3 tools/community_models/convert_vibeasr_gguf.py \ + --input models/vibeasr/vibeasr-vae-encoder-i8_s.gguf --in-place +python3 tools/community_models/convert_vibeasr_gguf.py \ + --input models/vibeasr/vibeasr-lm-i2_s-embed-q6_k.gguf --in-place # confirm an already-converted package needs no further remapping -python3 tools/community_models/convert_vibeasr_vae.py \ - --input models/vibeasr/vae_encoder-i8_s.gguf --check +python3 tools/community_models/convert_vibeasr_gguf.py \ + --input models/vibeasr/vibeasr-lm-i2_s-embed-q6_k.gguf --check +``` + +Use `--output ` instead of `--in-place` to keep the original. + +The package is two GGUFs plus the tokenizer, so `--model` points at the LM GGUF +and the spec is resolved from the repo — the same invocation shape as +[`minimax_h3`](minimax_h3.md): + +``` +models/vibeasr/ +├── vibeasr-vae-encoder-i8_s.gguf +├── vibeasr-lm-i2_s-embed-q6_k.gguf +├── tokenizer.json +└── tokenizer_config.json +``` + +```bash +cmake -B build -DCMAKE_BUILD_TYPE=Release +cmake --build build -j --target audiocpp_cli + +./build/bin/audiocpp_cli \ + --task asr \ + --family vibeasr \ + --model models/vibeasr/vibeasr-lm-i2_s-embed-q6_k.gguf \ + --model-spec-override model_specs \ + --backend cpu \ + --threads 8 \ + --audio assets/asr_validation/librispeech/librispeech_test_clean_6930-75918-0000.wav \ + --metrics +``` + ``` +text_output=Concord returned to its place amidst the tents. +metrics.wall_ms=1284.65 +metrics.rtf=0.36652 +``` + +Request options: `output_format` (`text` | `json`), `context` (a string folded +into the prompt to bias recognition), `max_new_tokens` (default 1024). Session +options: `vibeasr.encoder_graph_arena_mb` (64), +`vibeasr.prefill_graph_arena_mb` (256), `vibeasr.decode_graph_arena_mb` (256). + +Note that `output_format=json` returns an empty transcript on short +single-speaker clips — the model emits an immediate end-of-turn. VibeASR.cpp +behaves identically on the same input; this port does not paper over it. -There is no CLI family yet (see [Status](#status)); the encoder is reached -through `VibeASRVaeEncoderRuntime` or the parity probe: +## Tests ```bash cmake -B build -DCMAKE_BUILD_TYPE=Release -DENGINE_BUILD_MODEL_TESTS=ON -cmake --build build -j --target test_vibeasr_vae_encoder +cmake --build build -j --target test_vibeasr_asr test_vibeasr_vae_encoder +# end to end: loader, session, prompt, both graphs, greedy decode +./build/bin/test_vibeasr_asr --threads 8 + +# encoder only: shape, finiteness, frame count, and optional upstream parity ./build/bin/test_vibeasr_vae_encoder \ - --model models/vibeasr/vae_encoder-i8_s.gguf \ + --model models/vibeasr/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. +Both exit 125 (SKIP) when the checkpoint is missing, so they are safe in ctest. +`i2_s_mul_mat_test` and `i8_s_fused_ops_test` cover the kernels themselves +against plain-loop references and need no checkpoint. ## Parity +### End to end + +Four LibriSpeech clips, greedy on both sides, against VibeASR.cpp's own +`asr_infer --greedy` on the same two GGUFs: + +| Clip | VibeASR.cpp | this port | +|---|---|---| +| test-clean 6930-75918-0000 | `Concord returned to its place amidst the tents.` | identical | +| test-clean 6930-75918-0001 | `The english forwarded to the french baskets of flowers, of which they had made a plentiful provision to greet the arrival of the young princess. The french, in return, invited the english to a supper, which was to be given the next day.` | identical | +| test-other 7902-96591-0001 | `Don't cry, he said. I was obliged to come.` | identical | +| test-other 7902-96591-0000 | `I'm from the cut or lying off the coast.` | `I'm from the cutter lying off the coast.` | + +Three of four match token for token. The fourth diverges because these clips are +16 kHz and the two resamplers differ — this port uses soxr, upstream uses naive +linear interpolation — which perturbs the encoder features enough to flip one +greedy argmax. (Reference text: `I AM FROM THE CUTTER LYING OFF THE COAST`.) A +clip already at 24 kHz skips resampling entirely and does not have this failure +mode. + +### Encoder + The reference dump is raw F32, `frames * dim`, row-major, produced by calling `vae_encode_acoustic` / `vae_encode_semantic` from VibeASR.cpp's own `vae.h` on the same WAV: ```bash ./build/bin/test_vibeasr_vae_encoder \ - --model models/vibeasr/vae_encoder-i8_s.gguf \ + --model models/vibeasr/vibeasr-vae-encoder-i8_s.gguf \ --audio assets/asr_validation/librispeech/librispeech_test_clean_6930-75918-0000.wav \ --reference-acoustic ref_acoustic.f32 \ --reference-semantic ref_semantic.f32 \ --threads 8 ``` -3.505 s LibriSpeech clip, 17 frames × 1536 per branch: +3.505 s LibriSpeech clip fed at its native 16 kHz, 17 frames × 1536 per branch: | Branch | max abs | mean abs | cosine | |---|---|---|---| @@ -154,54 +300,65 @@ 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. +`i8_s_fused_ops_test` covers the op arithmetic itself against plain-loop +references, including the in-band scale surviving `ggml_cont(ggml_permute(...))` +— the encoder flips activations between channel-major and length-major +constantly, and a byte copy that drops the scale leaves the values right and +everything downstream off by an arbitrary factor. ## Measured performance -Release build, CPU backend, 24-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): +Release build, CPU backend, 24 vCPU AMD EPYC 7V13, 3.505 s clip resampled to +24 kHz (26 speech frames, 72-token prompt, 13 generated tokens): -| Threads | acoustic | semantic | both | RTF | -|---|---|---|---|---| -| 8 | 276 ms | 270 ms | 546 ms | 0.156 | -| 1 | 1657 ms | 1600 ms | 3257 ms | 0.929 | +| Threads | encoder (both branches) | prefill | decode | wall | RTF | +|---|---|---|---|---|---| +| 8 | 758 ms | 214 ms | 239 ms | 1285 ms | 0.367 | +| 1 | 4652 ms | 1415 ms | 941 ms | 7081 ms | 2.020 | -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. +The encoder dominates: it is run twice, once per branch, and it processes raw +samples rather than tokens. Decode is about 18 ms/token at 8 threads. + +The encoder-only probe reports 546 ms for both branches at 8 threads because it +feeds the clip at its native 16 kHz (17 frames); the session resamples to 24 kHz +first (26 frames). + +Peak RSS is 2.20 GB against 1.70 GB of weights: `BackendWeightStore` stages each +tensor before upload, so weight loading briefly holds roughly two copies of the +tensor being uploaded. Graph arenas are 64 MB (encoder) + 256 MB (prefill) + +256 MB (decode) by default. ## Status Ported: - I8_S VAE encoder graph, both branches, CPU backend. +- Ternary I2_S matmul kernel and the Qwen2 decoder graph on top of it, with + prefill + static-cache single-step decode. +- Prompt assembly, speech-feature injection, greedy decode, tokenizer, loader, + session, and `--family vibeasr`. - GGUF type remapping tool and geometry-from-tensors asset loader. -- 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. +- End-to-end and encoder parity probes against upstream, plus op-level unit + tests. 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. +- **CPU only.** The fused I8_S ops and the I2_S matmul have no CUDA or Metal + kernels; the session pins the backend to CPU. +- **Offline only**, like `vibevoice_asr` itself. Upstream's encoder is causal, so + streaming is implementable, but the state machine is not ported. +- **Greedy only.** Upstream's sampling path (temperature, top-p) is not ported, + and neither is `vibevoice_asr`'s beam search. +- **Text only.** No `--segments-out` / `--turns-out` equivalent; `output_format=json` + is a prompt variant, not structured decoding. +- The package is two GGUFs, so it needs `--model ` plus + `--model-spec-override model_specs` rather than a directory path. - Bit-exact parity with upstream is out of reach by design; see [Parity](#parity). ## Upstream -- Model port: (`src/vae.cpp`) +- Model port: (`src/vae.cpp`, + `src/lm.cpp`, `src/asr_server.cpp`, `utils/prompt_builder.h`) - Base model: VibeVoice ASR, also in tree as [`vibevoice_asr`](../asr.md#vibevoice-asr) +- Weights: diff --git a/include/engine/community_models/vibeasr/assets.h b/include/engine/community_models/vibeasr/assets.h index 744eb70e5..d7b77d580 100644 --- a/include/engine/community_models/vibeasr/assets.h +++ b/include/engine/community_models/vibeasr/assets.h @@ -1,9 +1,10 @@ #pragma once -// VibeASR VAE encoder assets. +// VibeASR assets: the I8_S audio VAE encoder and the ternary I2_S Qwen2 decoder. // -// Ported from https://github.com/microsoft/VibeASR.cpp (src/vae.cpp). +// Ported from https://github.com/microsoft/VibeASR.cpp (src/vae.cpp, src/lm.cpp). +#include "engine/framework/assets/resource_bundle.h" #include "engine/framework/assets/tensor_source.h" #include @@ -71,4 +72,38 @@ VibeASRVaeConfig derive_vae_config(const assets::TensorSource & source); std::shared_ptr load_vibeasr_vae_assets(const std::filesystem::path & model_path); +// Same, for a tensor source already opened from a resource bundle. +std::shared_ptr make_vibeasr_vae_assets( + std::shared_ptr source); + +// Decoder geometry. Unlike the encoder, none of this is recoverable from the +// tensor shapes alone -- head_dim, rope_theta and the RMS norm epsilon are not +// implied by any weight -- so it comes from the GGUF qwen2.* metadata block. +struct VibeASRLmConfig { + int64_t vocab_size = 0; + int64_t hidden_size = 0; + int64_t intermediate_size = 0; + int64_t num_hidden_layers = 0; + int64_t num_attention_heads = 0; + int64_t num_key_value_heads = 0; + int64_t head_dim = 0; + int64_t max_position_embeddings = 0; + // 1e-6 for the published checkpoint. Note this is *not* the encoder's + // epsilon: the VAE graph hardcodes 1e-5 (see VibeASRVaeConfig). + float rms_norm_eps = 1e-6f; + float rope_theta = 1e6f; +}; + +// The two GGUF halves plus the tokenizer files, as named by model_specs/vibeasr.json. +struct VibeASRAssets { + assets::ResourceBundle resources; + std::shared_ptr vae; + std::shared_ptr lm_weights; + VibeASRLmConfig lm; +}; + +VibeASRLmConfig derive_lm_config(const assets::TensorSource & source); + +std::shared_ptr load_vibeasr_assets(const std::filesystem::path & model_path); + } // namespace engine::community_models::vibeasr diff --git a/include/engine/community_models/vibeasr/lm_decoder.h b/include/engine/community_models/vibeasr/lm_decoder.h new file mode 100644 index 000000000..7080f4b1d --- /dev/null +++ b/include/engine/community_models/vibeasr/lm_decoder.h @@ -0,0 +1,62 @@ +#pragma once + +// VibeASR language model: the Qwen2 causal decoder whose projections are stored +// as ternary GGML_TYPE_I2_S. Speech features from the VAE encoder replace the +// prompt's <|speech_pad|> placeholders before prefill. +// +// Ported from https://github.com/microsoft/VibeASR.cpp (src/lm.cpp). + +#include "engine/community_models/vibeasr/assets.h" +#include "engine/framework/core/execution_context.h" + +#include +#include +#include +#include + +namespace engine::community_models::vibeasr { + +struct VibeASRLmPrompt { + std::vector input_ids; + // Positions in input_ids occupied by <|speech_pad|>, in order. + std::vector speech_positions; +}; + +// Summed acoustic + semantic connector output, row-major [tokens][hidden_size]. +struct VibeASRSpeechEmbeddings { + int64_t tokens = 0; + int64_t hidden_size = 0; + std::vector values; +}; + +struct VibeASRGenerationOptions { + int64_t max_new_tokens = 1024; + std::vector eos_token_ids; +}; + +class VibeASRLmRuntime { +public: + VibeASRLmRuntime( + std::shared_ptr weights_source, + const VibeASRLmConfig & config, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes); + ~VibeASRLmRuntime(); + + VibeASRLmRuntime(const VibeASRLmRuntime &) = delete; + VibeASRLmRuntime & operator=(const VibeASRLmRuntime &) = delete; + + // Greedy decode. Stops at any eos id or after max_new_tokens. + std::vector generate( + const VibeASRLmPrompt & prompt, + const VibeASRSpeechEmbeddings & speech, + const VibeASRGenerationOptions & options); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::community_models::vibeasr diff --git a/include/engine/community_models/vibeasr/session.h b/include/engine/community_models/vibeasr/session.h new file mode 100644 index 000000000..e35ef8f8a --- /dev/null +++ b/include/engine/community_models/vibeasr/session.h @@ -0,0 +1,60 @@ +#pragma once + +// Offline ASR session for the VibeASR package: I8_S VAE encoder -> ternary I2_S +// Qwen2 decoder, with VibeASR.cpp's ChatML prompt around the speech features. +// +// Ported from https://github.com/microsoft/VibeASR.cpp (src/asr_server.cpp, +// utils/prompt_builder.h). + +#include "engine/community_models/vibeasr/assets.h" +#include "engine/community_models/vibeasr/lm_decoder.h" +#include "engine/community_models/vibeasr/vae_encoder.h" +#include "engine/framework/model_spec/metadata.h" +#include "engine/framework/runtime/model.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/framework/tokenizers/llama_bpe.h" + +#include +#include + +namespace engine::community_models::vibeasr { + +std::shared_ptr make_vibeasr_loader(); + +class VibeASRSession final : public runtime::RuntimeSessionBase, public runtime::IOfflineVoiceTaskSession { +public: + VibeASRSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract); + ~VibeASRSession() override; + + std::string family() const override; + runtime::VoiceTaskKind task_kind() const override; + runtime::RunMode run_mode() const override; + void prepare(const runtime::SessionPreparationRequest & request) override; + runtime::TaskResult run(const runtime::TaskRequest & request) override; + +private: + struct RequestOptions { + std::string output_format = "text"; + std::string context; + int64_t max_new_tokens = 1024; + }; + + RequestOptions parse_request_options(const runtime::TaskRequest & request) const; + runtime::AudioBuffer normalize(const runtime::AudioBuffer & audio) const; + VibeASRSpeechEmbeddings encode_speech(const std::vector & samples); + VibeASRLmPrompt build_prompt(int64_t speech_tokens, float duration_seconds, const RequestOptions & options) const; + std::string decode_tokens(const std::vector & token_ids) const; + + runtime::TaskSpec task_; + std::shared_ptr assets_; + std::shared_ptr contract_; + std::shared_ptr tokenizer_; + VibeASRVaeEncoderRuntime encoder_; + VibeASRLmRuntime lm_; +}; + +} // namespace engine::community_models::vibeasr diff --git a/include/engine/community_models/vibeasr/vae_encoder.h b/include/engine/community_models/vibeasr/vae_encoder.h index 23ee9b164..dd4021dfe 100644 --- a/include/engine/community_models/vibeasr/vae_encoder.h +++ b/include/engine/community_models/vibeasr/vae_encoder.h @@ -65,7 +65,7 @@ class VibeASRVaeEncoderRuntime { 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 + // Both branches consume the same waveform, sampled at 24 kHz and scaled to // [-1, 1], and produce connector_hidden-wide features. VaeEncoderFeatures encode_acoustic(const std::vector & samples); VaeEncoderFeatures encode_semantic(const std::vector & samples); diff --git a/model_specs/vibeasr.json b/model_specs/vibeasr.json new file mode 100644 index 000000000..eb0a46c63 --- /dev/null +++ b/model_specs/vibeasr.json @@ -0,0 +1,139 @@ +{ + "schema_version": 1, + "family": "vibeasr", + "display_name": "VibeVoice-ASR-BitNet", + "description": "VibeASR.cpp's CPU-first VibeVoice ASR port: an INT8 (I8_S) audio VAE encoder feeding a ternary (I2_S) Qwen2 decoder, ported to audio.cpp.", + "category": "asr", + "status": "community", + "tasks": [ + "asr" + ], + "modes": [ + "offline" + ], + "languages": [ + "en", + "zh", + "fr", + "it", + "ko", + "pt", + "vi" + ], + "capabilities": {}, + "options": { + "request": [ + { + "name": "output_format", + "type": "enum", + "description": "Prompt suffix asked of the decoder: plain transcription text, or JSON rows with Start/End/Speaker/Content.", + "values": [ + "text", + "json" + ], + "required": false, + "default": "text" + }, + { + "name": "context", + "type": "string", + "description": "Extra context injected into the prompt (names, jargon) to bias the transcription.", + "required": false, + "default": "" + }, + { + "name": "max_new_tokens", + "type": "int", + "description": "Cap on decoded tokens for one request.", + "required": false, + "min": 1, + "default": 1024 + } + ], + "session": [ + { + "name": "encoder_graph_arena_mb", + "type": "int", + "description": "VAE encoder graph arena size in MB.", + "required": false, + "min": 16, + "default": 64 + }, + { + "name": "prefill_graph_arena_mb", + "type": "int", + "description": "Decoder prefill graph arena size in MB.", + "required": false, + "min": 16, + "default": 256 + }, + { + "name": "decode_graph_arena_mb", + "type": "int", + "description": "Decoder single-step graph arena size in MB.", + "required": false, + "min": 16, + "default": 256 + } + ], + "load": [] + }, + "runtime": { + "tags": [ + "gguf", + "cpu" + ] + }, + "packages": [ + { + "id": "vibeasr_bitnet_i2_s", + "display_name": "VibeVoice-ASR-BitNet I8_S encoder + I2_S decoder", + "description": "Upstream VibeASR.cpp GGUF package. The two GGUFs carry the VibeASR ggml fork's type ids and need one pass of tools/community_models/convert_vibeasr_gguf.py --in-place before audio.cpp can load them.", + "default": true, + "format": "gguf", + "precision": "native", + "target_directory": "VibeVoice-ASR-BitNet", + "files": [ + "vibeasr-vae-encoder-i8_s.gguf", + "vibeasr-lm-i2_s-embed-q6_k.gguf", + "tokenizer.json", + "tokenizer_config.json" + ], + "download": { + "kind": "huggingface_snapshot", + "repo": "XsquirrelC/VibeVoice-ASR-BitNet", + "revision": "main", + "gated": false + } + } + ], + "dependencies": [], + "ui": { + "recommended_package": "vibeasr_bitnet_i2_s", + "tags": [ + "ASR", + "GGUF" + ], + "docs": [ + "docs/community_models/vibeasr.md" + ], + "summary": "INT8 encoder plus ternary Qwen2 decoder transcription on CPU." + }, + "sources": [ + { + "format": "gguf", + "roots": { + "model": "." + }, + "files": { + "tokenizer_json": "model:tokenizer.json", + "tokenizer_config": "model:tokenizer_config.json" + }, + "optional_files": {}, + "tensors": { + "vae_weights": "model:vibeasr-vae-encoder-i8_s.gguf", + "lm_weights": "model:vibeasr-lm-i2_s-embed-q6_k.gguf" + } + } + ] +} diff --git a/src/community_models/vibeasr/assets.cpp b/src/community_models/vibeasr/assets.cpp index 2bc8d5830..8abe937f2 100644 --- a/src/community_models/vibeasr/assets.cpp +++ b/src/community_models/vibeasr/assets.cpp @@ -1,7 +1,13 @@ #include "engine/community_models/vibeasr/assets.h" +#include "engine/framework/model_spec/package.h" + +#include + +#include #include #include +#include #include namespace engine::community_models::vibeasr { @@ -112,6 +118,55 @@ VaeBranchConfig derive_branch(const assets::TensorSource & source, const std::st return branch; } +// assets::TensorSource exposes tensors, not the GGUF KV block, and the decoder +// geometry lives entirely in the KV block. Reading it directly is what the other +// community entries do (see sense_asr/assets.cpp). +class GgufMetadataReader { +public: + explicit GgufMetadataReader(const std::filesystem::path & path) { + gguf_init_params params{}; + params.no_alloc = true; + params.ctx = nullptr; + gguf_context * gguf = gguf_init_from_file(path.string().c_str(), params); + if (gguf == nullptr) { + throw std::runtime_error("Failed to read VibeASR GGUF metadata from " + path.string()); + } + ctx_.reset(gguf); + } + + int64_t require_u32(const char * key) const { + const int64_t id = gguf_find_key(ctx_.get(), key); + if (id < 0) { + throw std::runtime_error(std::string("VibeASR LM GGUF is missing ") + key); + } + return static_cast(gguf_get_val_u32(ctx_.get(), id)); + } + + float require_f32(const char * key) const { + const int64_t id = gguf_find_key(ctx_.get(), key); + if (id < 0) { + throw std::runtime_error(std::string("VibeASR LM GGUF is missing ") + key); + } + return gguf_get_val_f32(ctx_.get(), id); + } + + std::string kv_str(const char * key, std::string fallback) const { + const int64_t id = gguf_find_key(ctx_.get(), key); + return id < 0 ? std::move(fallback) : std::string(gguf_get_val_str(ctx_.get(), id)); + } + +private: + struct GgufDeleter { + void operator()(gguf_context * ctx) const noexcept { + if (ctx != nullptr) { + gguf_free(ctx); + } + } + }; + + std::unique_ptr ctx_; +}; + } // namespace int64_t VaeBranchConfig::frames_for_samples(int64_t num_samples) const { @@ -139,9 +194,79 @@ VibeASRVaeConfig derive_vae_config(const assets::TensorSource & source) { } std::shared_ptr load_vibeasr_vae_assets(const std::filesystem::path & model_path) { + return make_vibeasr_vae_assets(engine::assets::open_tensor_source(model_path)); +} + +std::shared_ptr make_vibeasr_vae_assets( + std::shared_ptr source) { auto assets = std::make_shared(); - assets->source = engine::assets::open_tensor_source(model_path); - assets->config = derive_vae_config(*assets->source); + assets->config = derive_vae_config(*source); + assets->source = std::move(source); + return assets; +} + +VibeASRLmConfig derive_lm_config(const assets::TensorSource & source) { + const GgufMetadataReader reader(source.source_path()); + + const std::string architecture = reader.kv_str("general.architecture", ""); + if (architecture != "qwen2") { + throw std::runtime_error( + "VibeASR LM GGUF declares architecture '" + architecture + "', expected qwen2"); + } + + VibeASRLmConfig config; + config.vocab_size = reader.require_u32("qwen2.vocab_size"); + config.hidden_size = reader.require_u32("qwen2.embedding_length"); + config.intermediate_size = reader.require_u32("qwen2.feed_forward_length"); + config.num_hidden_layers = reader.require_u32("qwen2.block_count"); + config.num_attention_heads = reader.require_u32("qwen2.attention.head_count"); + config.num_key_value_heads = reader.require_u32("qwen2.attention.head_count_kv"); + config.max_position_embeddings = reader.require_u32("qwen2.context_length"); + // The checkpoint has no attention.key_length: Qwen2 stores the per-head width + // only as the RoPE dimension count, which for this model equals + // embedding_length / head_count. + config.head_dim = reader.require_u32("qwen2.rope.dimension_count"); + config.rms_norm_eps = reader.require_f32("qwen2.attention.layer_norm_rms_epsilon"); + config.rope_theta = reader.require_f32("qwen2.rope.freq_base"); + + if (config.head_dim * config.num_attention_heads != config.hidden_size) { + throw std::runtime_error("VibeASR LM head_dim * head_count does not match embedding_length"); + } + if (config.num_key_value_heads <= 0 || config.num_attention_heads % config.num_key_value_heads != 0) { + throw std::runtime_error("VibeASR LM head_count is not a multiple of head_count_kv"); + } + if (config.num_hidden_layers <= 0) { + throw std::runtime_error("VibeASR LM declares no layers"); + } + + // Cross-check the metadata against the one tensor whose shape pins both dims. + const auto embedding = source.require_metadata("token_embd.weight").shape; + if (embedding.size() != 2 || embedding[0] != config.vocab_size || embedding[1] != config.hidden_size) { + throw std::runtime_error("VibeASR LM token_embd.weight does not match the declared geometry"); + } + return config; +} + +std::shared_ptr load_vibeasr_assets(const std::filesystem::path & model_path) { + auto assets = std::make_shared(); + assets->resources = engine::model_spec::load_resource_bundle_for_family(model_path, "vibeasr"); + + // A GGUF still carrying the VibeASR fork's type ids (36/37) fails deep inside + // the reader with an unhelpful message, so name the fix here. + auto open = [&assets](const char * id) { + try { + return assets->resources.open_tensor_source(id); + } catch (const std::exception & error) { + throw std::runtime_error( + std::string("VibeASR could not open the '") + id + "' GGUF (" + error.what() + + "). If it came straight from huggingface.co/XsquirrelC/VibeVoice-ASR-BitNet, run " + "tools/community_models/convert_vibeasr_gguf.py --in-place on it first."); + } + }; + + assets->vae = make_vibeasr_vae_assets(open("vae_weights")); + assets->lm_weights = open("lm_weights"); + assets->lm = derive_lm_config(*assets->lm_weights); return assets; } diff --git a/src/community_models/vibeasr/lm_decoder.cpp b/src/community_models/vibeasr/lm_decoder.cpp new file mode 100644 index 000000000..35bde5ef5 --- /dev/null +++ b/src/community_models/vibeasr/lm_decoder.cpp @@ -0,0 +1,679 @@ +#include "engine/community_models/vibeasr/lm_decoder.h" + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/modules/transformers/qwen_causal_decoder.h" +#include "engine/framework/runtime/errors.h" +#include "engine/framework/runtime/kv_cache.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::community_models::vibeasr { +namespace { + +namespace modules = engine::modules; +using Clock = std::chrono::steady_clock; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct GgmlGallocrDeleter { + void operator()(ggml_gallocr_t alloc) const noexcept { + if (alloc != nullptr) { + ggml_gallocr_free(alloc); + } + } +}; + +struct LmLayerWeights { + core::TensorValue input_norm; + core::TensorValue q_proj; + core::TensorValue q_bias; + core::TensorValue k_proj; + core::TensorValue k_bias; + core::TensorValue v_proj; + core::TensorValue v_bias; + core::TensorValue o_proj; + core::TensorValue post_norm; + core::TensorValue gate_proj; + core::TensorValue up_proj; + core::TensorValue down_proj; +}; + +struct LmWeights { + std::shared_ptr store; + core::TensorValue token_embedding; + std::vector layers; + core::TensorValue norm; + core::TensorValue lm_head; +}; + +struct PrefillOutput { + std::vector logits; + runtime::TransformerKVState kv_state; +}; + +// I2_S is a whole-tensor quantization whose in-band F32 scale sits after the +// packed codes, which is exactly what ggml_nbytes() accounts for, so the GGUF +// payload goes to the backend byte for byte. Same contract as the encoder's +// load_i8_s_tensor(). +core::TensorValue load_i2_s_tensor( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + const std::vector & expected_shape) { + const auto metadata = source.require_metadata(name); + if (metadata.dtype != "i2_s") { + throw std::runtime_error("VibeASR LM tensor " + name + " is " + metadata.dtype + ", expected i2_s"); + } + if (metadata.shape != expected_shape) { + throw std::runtime_error("VibeASR LM tensor " + name + " has an unexpected shape"); + } + + core::TensorShape shape; + shape.rank = expected_shape.size(); + for (size_t i = 0; i < shape.rank; ++i) { + shape.dims[i] = expected_shape[i]; + } + + const auto raw = source.require_tensor_data(name); + return store.make_tensor(shape, GGML_TYPE_I2_S, raw.bytes.data(), raw.bytes.size()); +} + +modules::QwenDecoderLayerWeights to_qwen_layer_weights(const LmLayerWeights & weights) { + modules::QwenDecoderLayerWeights out; + out.input_norm = {weights.input_norm, std::nullopt}; + out.self_attention.q_weight = weights.q_proj; + out.self_attention.q_bias = weights.q_bias; + out.self_attention.k_weight = weights.k_proj; + out.self_attention.k_bias = weights.k_bias; + out.self_attention.v_weight = weights.v_proj; + out.self_attention.v_bias = weights.v_bias; + out.self_attention.out_weight = weights.o_proj; + out.post_norm = {weights.post_norm, std::nullopt}; + out.mlp.gate_proj = {weights.gate_proj, std::nullopt}; + out.mlp.up_proj = {weights.up_proj, std::nullopt}; + out.mlp.down_proj = {weights.down_proj, std::nullopt}; + return out; +} + +// Plain Qwen2: attention biases, no per-head Q/K norms. Nothing here depends on +// the weight type, which is why the framework's decoder runs unmodified on I2_S +// projections -- ggml_mul_mat dispatches on the tensor type. +modules::QwenCausalDecoderConfig make_qwen_decoder_config(const VibeASRLmConfig & config) { + modules::QwenCausalDecoderConfig out; + out.stack.hidden_size = config.hidden_size; + out.stack.num_attention_heads = config.num_attention_heads; + out.stack.num_key_value_heads = config.num_key_value_heads; + out.stack.head_dim = config.head_dim; + out.stack.intermediate_size = config.intermediate_size; + out.stack.layers = config.num_hidden_layers; + out.stack.rms_norm_eps = config.rms_norm_eps; + out.stack.rope_theta = config.rope_theta; + out.stack.use_qk_norm = false; + out.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.logits_size = config.vocab_size; + out.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; + return out; +} + +modules::QwenCausalDecoderWeights make_qwen_decoder_weights(const LmWeights & weights) { + modules::QwenCausalDecoderWeights out; + out.stack.layers.reserve(weights.layers.size()); + for (const auto & layer : weights.layers) { + out.stack.layers.push_back(to_qwen_layer_weights(layer)); + } + out.final_norm = {weights.norm, std::nullopt}; + out.lm_head = {weights.lm_head, std::nullopt}; + return out; +} + +// Token embeddings with the encoder's speech features written over the +// <|speech_pad|> slots. Doing the overwrite in-graph with ggml_set_rows keeps +// the prompt a single I32 upload instead of a host-side embedding matrix. +core::TensorValue prompt_embeddings( + core::ModuleBuildContext & ctx, + const LmWeights & weights, + const VibeASRLmConfig & config, + ggml_tensor * token_ids, + ggml_tensor * speech_embeddings, + ggml_tensor * speech_positions, + int64_t prompt_steps, + int64_t speech_tokens) { + auto ids = core::wrap_tensor(token_ids, core::TensorShape::from_dims({prompt_steps}), GGML_TYPE_I32); + auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size}) + .build(ctx, ids, weights.token_embedding); + if (speech_tokens > 0) { + auto speech = core::wrap_tensor( + speech_embeddings, + core::TensorShape::from_dims({speech_tokens, config.hidden_size}), + GGML_TYPE_F32); + auto positions = core::wrap_tensor( + speech_positions, + core::TensorShape::from_dims({speech_tokens}), + GGML_TYPE_I64); + x = core::wrap_tensor( + ggml_set_rows(ctx.ggml, x.tensor, speech.tensor, positions.tensor), + x.shape, + GGML_TYPE_F32); + } + return core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, prompt_steps, config.hidden_size})); +} + +LmWeights load_weights( + const assets::TensorSource & source, + const VibeASRLmConfig & config, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes) { + LmWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "vibeasr.lm.weights", + weight_context_bytes); + + // The embedding table and the output projection are the two tensors VibeASR + // leaves unternarized -- Q6_K and F16 in the published checkpoint -- so they + // load through the framework's normal path. + weights.token_embedding = weights.store->load_tensor( + source, + "token_embd.weight", + assets::TensorStorageType::Native, + {config.vocab_size, config.hidden_size}); + + const int64_t dim = config.head_dim; + const int64_t q_dim = config.num_attention_heads * dim; + const int64_t kv_dim = config.num_key_value_heads * dim; + weights.layers.reserve(static_cast(config.num_hidden_layers)); + for (int64_t layer = 0; layer < config.num_hidden_layers; ++layer) { + const std::string prefix = "blk." + std::to_string(layer) + "."; + LmLayerWeights w; + w.input_norm = weights.store->load_f32_tensor(source, prefix + "attn_norm.weight", {config.hidden_size}); + w.q_proj = load_i2_s_tensor(*weights.store, source, prefix + "attn_q.weight", {q_dim, config.hidden_size}); + w.q_bias = weights.store->load_f32_tensor(source, prefix + "attn_q.bias", {q_dim}); + w.k_proj = load_i2_s_tensor(*weights.store, source, prefix + "attn_k.weight", {kv_dim, config.hidden_size}); + w.k_bias = weights.store->load_f32_tensor(source, prefix + "attn_k.bias", {kv_dim}); + w.v_proj = load_i2_s_tensor(*weights.store, source, prefix + "attn_v.weight", {kv_dim, config.hidden_size}); + w.v_bias = weights.store->load_f32_tensor(source, prefix + "attn_v.bias", {kv_dim}); + w.o_proj = load_i2_s_tensor(*weights.store, source, prefix + "attn_output.weight", {config.hidden_size, q_dim}); + w.post_norm = weights.store->load_f32_tensor(source, prefix + "ffn_norm.weight", {config.hidden_size}); + w.gate_proj = load_i2_s_tensor( + *weights.store, source, prefix + "ffn_gate.weight", {config.intermediate_size, config.hidden_size}); + w.up_proj = load_i2_s_tensor( + *weights.store, source, prefix + "ffn_up.weight", {config.intermediate_size, config.hidden_size}); + w.down_proj = load_i2_s_tensor( + *weights.store, source, prefix + "ffn_down.weight", {config.hidden_size, config.intermediate_size}); + weights.layers.push_back(std::move(w)); + } + + weights.norm = weights.store->load_f32_tensor(source, "output_norm.weight", {config.hidden_size}); + weights.lm_head = weights.store->load_tensor( + source, + "output.weight", + assets::TensorStorageType::Native, + {config.vocab_size, config.hidden_size}); + weights.store->upload(); + return weights; +} + +int32_t argmax_index(const std::vector & values) { + if (values.empty()) { + throw std::runtime_error("VibeASR LM cannot select from empty logits"); + } + size_t best = 0; + for (size_t i = 1; i < values.size(); ++i) { + if (values[i] > values[best]) { + best = i; + } + } + return static_cast(best); +} + +class LmWeightsRuntime { +public: + LmWeightsRuntime( + std::shared_ptr source, + VibeASRLmConfig config, + core::ExecutionContext & execution, + size_t weight_context_bytes) + : source_(std::move(source)), + config_(std::make_shared(config)), + backend_(execution.backend()), + backend_type_(execution.backend_type()), + threads_(std::max(1, execution.config().threads)), + weights_(std::make_shared(load_weights( + *source_, + *config_, + backend_, + backend_type_, + weight_context_bytes))) {} + + const VibeASRLmConfig & config() const noexcept { return *config_; } + const LmWeights & weights() const noexcept { return *weights_; } + ggml_backend_t backend() const noexcept { return backend_; } + core::BackendType backend_type() const noexcept { return backend_type_; } + int threads() const noexcept { return threads_; } + +private: + std::shared_ptr source_; + std::shared_ptr config_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + std::shared_ptr weights_; +}; + +class PrefillGraph { +public: + PrefillGraph( + std::shared_ptr runtime, + int64_t prompt_steps, + int64_t speech_tokens, + size_t graph_arena_bytes) + : runtime_(std::move(runtime)), + prompt_steps_(prompt_steps), + speech_tokens_(speech_tokens) { + if (prompt_steps_ <= 0) { + throw std::runtime_error("VibeASR LM prefill requires positive prompt length"); + } + if (speech_tokens_ < 0 || speech_tokens_ > prompt_steps_) { + throw std::runtime_error("VibeASR LM prefill speech token count is invalid"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize VibeASR LM prefill graph context"); + } + const auto & config = runtime_->config(); + const auto & weights = runtime_->weights(); + core::ModuleBuildContext ctx{ctx_.get(), "vibeasr.lm.prefill", runtime_->backend_type()}; + token_ids_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_); + speech_embeddings_ = ggml_new_tensor_2d( + ctx_.get(), GGML_TYPE_F32, config.hidden_size, std::max(speech_tokens_, 1)); + speech_positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I64, std::max(speech_tokens_, 1)); + auto x = prompt_embeddings( + ctx, + weights, + config, + token_ids_, + speech_embeddings_, + speech_positions_, + prompt_steps_, + speech_tokens_); + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_); + auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({prompt_steps_}), GGML_TYPE_I32); + + auto decoder_out = modules::QwenCausalDecoderModule(make_qwen_decoder_config(config)) + .build(ctx, x, positions, make_qwen_decoder_weights(weights)); + for (const auto & layer : decoder_out.state.layers) { + if (!layer.key.has_value() || !layer.value.has_value()) { + throw std::runtime_error("VibeASR LM prefill decoder did not return K/V state"); + } + // Copy K/V out of the graph-allocated intermediates and mark them as + // outputs so the allocator cannot recycle them before run() reads + // them back. + auto * key = ggml_cpy(ctx_.get(), layer.key->tensor, ggml_dup_tensor(ctx_.get(), layer.key->tensor)); + auto * value = ggml_cpy(ctx_.get(), layer.value->tensor, ggml_dup_tensor(ctx_.get(), layer.value->tensor)); + ggml_set_output(key); + ggml_set_output(value); + keys_.push_back(key); + values_.push_back(value); + } + logits_ = decoder_out.logits.tensor; + ggml_set_output(logits_); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + ggml_build_forward_expand(graph_, logits_); + for (auto * key : keys_) { + ggml_build_forward_expand(graph_, key); + } + for (auto * value : values_) { + ggml_build_forward_expand(graph_, value); + } + const auto try_alloc = [&]() { + gallocr_.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend()))); + return gallocr_ != nullptr && + ggml_gallocr_reserve(gallocr_.get(), graph_) && + ggml_gallocr_alloc_graph(gallocr_.get(), graph_); + }; + if (!try_alloc() && (engine::core::trim_backend_pools(runtime_->backend()), !try_alloc())) { + throw engine::runtime::CapacityError( + "VibeASR LM prefill graph does not fit in device memory at this size (" + + std::to_string(prompt_steps_) + " prompt steps, of which " + + std::to_string(speech_tokens_) + " are speech tokens)"); + } + position_ids_ = modules::qwen_position_ids(prompt_steps_); + debug::timing_log_scalar("vibeasr.lm.prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("vibeasr.lm.prefill_prompt_steps", prompt_steps_); + } + + ~PrefillGraph() { + engine::core::release_backend_graph_resources(runtime_->backend(), graph_, true); + } + + bool matches(const LmWeightsRuntime & runtime, int64_t prompt_steps, int64_t speech_tokens) const { + return runtime_.get() == &runtime && prompt_steps_ == prompt_steps && speech_tokens_ == speech_tokens; + } + + PrefillOutput run( + const std::vector & token_ids, + const std::vector & speech_embeddings, + const std::vector & speech_positions) { + const auto & config = runtime_->config(); + if (static_cast(token_ids.size()) != prompt_steps_) { + throw std::runtime_error("VibeASR LM prefill token id count mismatch"); + } + if (static_cast(speech_embeddings.size()) != speech_tokens_ * config.hidden_size) { + throw std::runtime_error("VibeASR LM prefill speech embedding size mismatch"); + } + if (static_cast(speech_positions.size()) != speech_tokens_) { + throw std::runtime_error("VibeASR LM prefill speech position count mismatch"); + } + // Re-uploaded on every run: leaves are not pinned by the graph allocator. + ggml_backend_tensor_set(positions_, position_ids_.data(), 0, position_ids_.size() * sizeof(int32_t)); + ggml_backend_tensor_set(token_ids_, token_ids.data(), 0, token_ids.size() * sizeof(int32_t)); + if (speech_tokens_ > 0) { + const std::vector positions(speech_positions.begin(), speech_positions.end()); + ggml_backend_tensor_set( + speech_embeddings_, speech_embeddings.data(), 0, speech_embeddings.size() * sizeof(float)); + ggml_backend_tensor_set(speech_positions_, positions.data(), 0, positions.size() * sizeof(int64_t)); + } + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + const auto compute_start = Clock::now(); + const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); + ggml_backend_synchronize(runtime_->backend()); + debug::timing_log_scalar("vibeasr.lm.prefill.graph.compute_ms", engine::debug::elapsed_ms(compute_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("VibeASR LM prefill graph compute failed"); + } + PrefillOutput out; + out.logits.resize(static_cast(config.vocab_size)); + ggml_backend_tensor_get(logits_, out.logits.data(), 0, out.logits.size() * sizeof(float)); + out.kv_state.current_end = prompt_steps_; + out.kv_state.layers.resize(keys_.size()); + const size_t layer_values = + static_cast(prompt_steps_ * config.num_key_value_heads * config.head_dim); + for (size_t layer = 0; layer < keys_.size(); ++layer) { + auto & state = out.kv_state.layers[layer]; + state.valid_steps = prompt_steps_; + state.key.resize(layer_values); + state.value.resize(layer_values); + ggml_backend_tensor_get(keys_[layer], state.key.data(), 0, state.key.size() * sizeof(float)); + ggml_backend_tensor_get(values_[layer], state.value.data(), 0, state.value.size() * sizeof(float)); + } + return out; + } + +private: + std::shared_ptr runtime_; + int64_t prompt_steps_ = 0; + int64_t speech_tokens_ = 0; + std::unique_ptr ctx_; + ggml_tensor * token_ids_ = nullptr; + ggml_tensor * speech_embeddings_ = nullptr; + ggml_tensor * speech_positions_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * logits_ = nullptr; + std::vector keys_; + std::vector values_; + std::vector position_ids_; + ggml_cgraph * graph_ = nullptr; + std::unique_ptr, GgmlGallocrDeleter> gallocr_; +}; + +class DecodeGraph { +public: + DecodeGraph(std::shared_ptr runtime, int64_t cache_steps, size_t graph_arena_bytes) + : runtime_(std::move(runtime)), + cache_steps_(cache_steps) { + if (cache_steps_ <= 0) { + throw std::runtime_error("VibeASR LM decode requires positive cache length"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize VibeASR LM decode graph context"); + } + const auto & config = runtime_->config(); + const auto & weights = runtime_->weights(); + core::ModuleBuildContext ctx{ctx_.get(), "vibeasr.lm.decode", runtime_->backend_type()}; + token_id_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + auto token_id = core::wrap_tensor(token_id_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size}) + .build(ctx, token_id, weights.token_embedding); + x = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, 1, config.hidden_size})); + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + cache_slot_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + auto cache_slot = core::wrap_tensor(cache_slot_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + attention_mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, cache_steps_, 1, 1, 1); + auto attention_mask = core::wrap_tensor( + attention_mask_, core::TensorShape::from_dims({1, 1, 1, cache_steps_}), GGML_TYPE_F16); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + auto decoder_out = modules::QwenCausalDecoderModule(make_qwen_decoder_config(config)) + .build_static_cache_tail( + ctx, + graph_, + x, + positions, + make_qwen_decoder_weights(weights), + cache_steps_, + attention_mask, + cache_slot); + step_cache_ = std::move(decoder_out.cache); + logits_ = decoder_out.logits.tensor; + ggml_set_output(logits_); + ggml_build_forward_expand(graph_, logits_); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); + if (buffer_ == nullptr) { + engine::core::trim_backend_pools(runtime_->backend()); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); + } + if (buffer_ == nullptr) { + throw engine::runtime::CapacityError( + "VibeASR LM decode graph does not fit in device memory at " + + std::to_string(cache_steps_) + " cache steps"); + } + attention_mask_values_.assign(static_cast(cache_steps_), ggml_fp32_to_fp16(-INFINITY)); + debug::timing_log_scalar("vibeasr.lm.decode.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); + debug::trace_log_scalar("vibeasr.lm.decode_cache_steps", cache_steps_); + } + + ~DecodeGraph() { + engine::core::release_backend_graph_resources(runtime_->backend(), graph_, true); + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + } + + bool can_run(const LmWeightsRuntime & runtime, int64_t required_steps) const { + return runtime_.get() == &runtime && cache_steps_ >= required_steps; + } + + void import_state(const runtime::TransformerKVState & state) { + step_cache_.import_state(state); + } + + std::vector run_step(int32_t token) { + const auto & config = runtime_->config(); + if (step_cache_.valid_steps() >= cache_steps_) { + throw std::runtime_error("VibeASR LM decode cache exhausted"); + } + ggml_backend_tensor_set(token_id_, &token, 0, sizeof(int32_t)); + const int32_t position = static_cast(step_cache_.current_end()); + ggml_backend_tensor_set(positions_, &position, 0, sizeof(int32_t)); + const int32_t cache_slot = static_cast(step_cache_.valid_steps()); + ggml_backend_tensor_set(cache_slot_, &cache_slot, 0, sizeof(int32_t)); + modules::write_qwen_cached_step_mask( + attention_mask_, + attention_mask_values_, + cache_steps_, + step_cache_.valid_steps(), + step_cache_.valid_steps()); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); + ggml_backend_synchronize(runtime_->backend()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("VibeASR LM decode graph compute failed"); + } + logits_buffer_.resize(static_cast(config.vocab_size)); + ggml_backend_tensor_get(logits_, logits_buffer_.data(), 0, logits_buffer_.size() * sizeof(float)); + step_cache_.advance_after_direct_append(1); + // The caller moves out of this buffer before the next step. + return std::move(logits_buffer_); + } + +private: + std::shared_ptr runtime_; + int64_t cache_steps_ = 0; + std::unique_ptr ctx_; + ggml_tensor * token_id_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * cache_slot_ = nullptr; + ggml_tensor * attention_mask_ = nullptr; + ggml_tensor * logits_ = nullptr; + std::vector attention_mask_values_; + std::vector logits_buffer_; + runtime::TransformerKVCache step_cache_; + ggml_cgraph * graph_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; +}; + +} // namespace + +struct VibeASRLmRuntime::Impl { + Impl( + std::shared_ptr weights_source, + const VibeASRLmConfig & config, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes) + : weights(std::make_shared( + std::move(weights_source), + config, + execution, + weight_context_bytes)), + prefill_graph_arena_bytes(prefill_graph_arena_bytes), + decode_graph_arena_bytes(decode_graph_arena_bytes) {} + + void validate_speech(const VibeASRLmPrompt & prompt, const VibeASRSpeechEmbeddings & speech) const { + const auto & config = weights->config(); + if (speech.tokens > 0 && speech.hidden_size != config.hidden_size) { + throw std::runtime_error("VibeASR speech embedding hidden size mismatch"); + } + if (speech.tokens != static_cast(prompt.speech_positions.size())) { + throw std::runtime_error("VibeASR speech embedding count does not match the prompt's speech pads"); + } + if (static_cast(speech.values.size()) != speech.tokens * speech.hidden_size) { + throw std::runtime_error("VibeASR speech embedding value count mismatch"); + } + for (const int32_t position : prompt.speech_positions) { + if (position < 0 || position >= static_cast(prompt.input_ids.size())) { + throw std::runtime_error("VibeASR speech pad position out of range"); + } + } + } + + std::shared_ptr weights; + size_t prefill_graph_arena_bytes = 0; + size_t decode_graph_arena_bytes = 0; + std::unique_ptr prefill_graph; + std::unique_ptr decode_graph; +}; + +VibeASRLmRuntime::VibeASRLmRuntime( + std::shared_ptr weights_source, + const VibeASRLmConfig & config, + core::ExecutionContext & execution, + size_t prefill_graph_arena_bytes, + size_t decode_graph_arena_bytes, + size_t weight_context_bytes) + : impl_(std::make_unique( + std::move(weights_source), + config, + execution, + prefill_graph_arena_bytes, + decode_graph_arena_bytes, + weight_context_bytes)) {} + +VibeASRLmRuntime::~VibeASRLmRuntime() = default; + +std::vector VibeASRLmRuntime::generate( + const VibeASRLmPrompt & prompt, + const VibeASRSpeechEmbeddings & speech, + const VibeASRGenerationOptions & options) { + const auto & config = impl_->weights->config(); + if (prompt.input_ids.empty()) { + throw std::runtime_error("VibeASR LM prompt is empty"); + } + if (options.max_new_tokens <= 0) { + throw std::runtime_error("VibeASR max_new_tokens must be positive"); + } + const int64_t prompt_steps = static_cast(prompt.input_ids.size()); + if (prompt_steps + options.max_new_tokens > config.max_position_embeddings) { + throw std::runtime_error("VibeASR request exceeds the decoder context length"); + } + impl_->validate_speech(prompt, speech); + + if (impl_->prefill_graph == nullptr || + !impl_->prefill_graph->matches(*impl_->weights, prompt_steps, speech.tokens)) { + impl_->prefill_graph.reset(); + impl_->prefill_graph = std::make_unique( + impl_->weights, prompt_steps, speech.tokens, impl_->prefill_graph_arena_bytes); + } + auto prefill = impl_->prefill_graph->run(prompt.input_ids, speech.values, prompt.speech_positions); + + const int64_t required_cache_steps = prompt_steps + options.max_new_tokens; + if (impl_->decode_graph == nullptr || !impl_->decode_graph->can_run(*impl_->weights, required_cache_steps)) { + impl_->decode_graph.reset(); + impl_->decode_graph = + std::make_unique(impl_->weights, required_cache_steps, impl_->decode_graph_arena_bytes); + } + impl_->decode_graph->import_state(prefill.kv_state); + + const auto is_eos = [&options](int32_t token) { + return std::find(options.eos_token_ids.begin(), options.eos_token_ids.end(), token) != + options.eos_token_ids.end(); + }; + + std::vector out; + std::vector logits = std::move(prefill.logits); + const auto decode_start = Clock::now(); + for (int64_t step = 0; step < options.max_new_tokens; ++step) { + const int32_t token = argmax_index(logits); + if (is_eos(token)) { + break; + } + out.push_back(token); + logits = impl_->decode_graph->run_step(token); + } + debug::timing_log_scalar("vibeasr.lm.decode_total_ms", engine::debug::elapsed_ms(decode_start, Clock::now())); + return out; +} + +} // namespace engine::community_models::vibeasr diff --git a/src/community_models/vibeasr/session.cpp b/src/community_models/vibeasr/session.cpp new file mode 100644 index 000000000..b4d4b7989 --- /dev/null +++ b/src/community_models/vibeasr/session.cpp @@ -0,0 +1,368 @@ +#include "engine/community_models/vibeasr/session.h" + +#include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/resampling.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/io/text.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/runtime/spec_backed_model.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::community_models::vibeasr { +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr size_t kWeightContextBytes = 64ull * 1024ull * 1024ull; + +// VibeASR resamples to 24 kHz and RMS-normalizes to -25 dBFS before the encoder; +// both numbers are fixed in the reference implementation, not in the checkpoint. +constexpr int kSampleRate = 24000; +constexpr float kTargetDbFs = -25.0F; +constexpr float kNormalizeEps = 1.0e-6F; + +// Canonical HuggingFace ids for the VibeVoice special tokens. VibeASR inserts +// them numerically rather than through the tokenizer, because the GGUF vocab's +// text for these slots is Qwen2.5's original <|object_ref_start|> family while +// the embedding rows are the ones VibeVoice trained. +constexpr int32_t kEndOfText = 151643; +constexpr int32_t kImStart = 151644; +constexpr int32_t kImEnd = 151645; +constexpr int32_t kSpeechStart = 151646; +constexpr int32_t kSpeechEnd = 151647; +constexpr int32_t kSpeechPad = 151648; + +constexpr const char * kSystemPrompt = + "You are a helpful assistant that transcribes audio input into text output in JSON format."; + +std::shared_ptr require_assets(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("VibeASR session requires assets"); + } + return assets; +} + +const engine::model_spec::ModelContract & require_contract( + const std::shared_ptr & contract) { + if (contract == nullptr) { + throw std::runtime_error("VibeASR session requires a model contract"); + } + return *contract; +} + +runtime::SessionOptions validate_session_setup( + const runtime::TaskSpec & task, + runtime::SessionOptions options, + const engine::model_spec::ModelContract & contract) { + if (task.task != runtime::VoiceTaskKind::Asr) { + throw std::runtime_error("VibeASR only supports VoiceTaskKind::Asr"); + } + if (task.mode != runtime::RunMode::Offline) { + throw std::runtime_error("VibeASR only supports offline sessions"); + } + runtime::validate_spec_backed_session_options(options, contract, "vibeasr", "VibeASR"); + return options; +} + +size_t encoder_graph_arena_bytes(const runtime::SessionOptions & options) { + return runtime::parse_size_mb_option( + options.options, {"vibeasr.encoder_graph_arena_mb"}, 64ull * 1024ull * 1024ull); +} + +size_t prefill_graph_arena_bytes(const runtime::SessionOptions & options) { + return runtime::parse_size_mb_option( + options.options, {"vibeasr.prefill_graph_arena_mb"}, 256ull * 1024ull * 1024ull); +} + +size_t decode_graph_arena_bytes(const runtime::SessionOptions & options) { + return runtime::parse_size_mb_option( + options.options, {"vibeasr.decode_graph_arena_mb"}, 256ull * 1024ull * 1024ull); +} + +std::shared_ptr load_tokenizer(const VibeASRAssets & assets) { + // No merges.txt in the published package, so the tokenizer comes from + // tokenizer.json alone. + return engine::tokenizers::load_llama_bpe_tokenizer(engine::tokenizers::LlamaBpeTokenizerSpec{ + {}, + {}, + assets.resources.require_file("tokenizer_config"), + assets.resources.require_file("tokenizer_json"), + engine::tokenizers::LlamaBpePreTokenizer::Qwen2, + }); +} + +std::string format_duration(float seconds) { + char buffer[64]; + std::snprintf(buffer, sizeof(buffer), "%.2f", static_cast(seconds)); + return std::string(buffer); +} + +} // namespace + +VibeASRSession::VibeASRSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets, + std::shared_ptr contract) + : RuntimeSessionBase(validate_session_setup(task, std::move(options), require_contract(contract))), + task_(std::move(task)), + assets_(require_assets(std::move(assets))), + contract_(std::move(contract)), + tokenizer_(load_tokenizer(*assets_)), + encoder_(assets_->vae, execution_context(), encoder_graph_arena_bytes(RuntimeSessionBase::options())), + lm_(assets_->lm_weights, + assets_->lm, + execution_context(), + prefill_graph_arena_bytes(RuntimeSessionBase::options()), + decode_graph_arena_bytes(RuntimeSessionBase::options()), + kWeightContextBytes) { + // Both weight stores have uploaded by now; drop the resident file blobs. + assets_->vae->source->release_storage(); + assets_->lm_weights->release_storage(); +} + +VibeASRSession::~VibeASRSession() = default; + +std::string VibeASRSession::family() const { + return "vibeasr"; +} + +runtime::VoiceTaskKind VibeASRSession::task_kind() const { + return task_.task; +} + +runtime::RunMode VibeASRSession::run_mode() const { + return task_.mode; +} + +void VibeASRSession::prepare(const runtime::SessionPreparationRequest & request) { + (void)request; + mark_prepared(); +} + +VibeASRSession::RequestOptions VibeASRSession::parse_request_options(const runtime::TaskRequest & request) const { + runtime::validate_spec_backed_request_options(request.options, require_contract(contract_), "VibeASR"); + RequestOptions out; + if (const auto value = runtime::find_option(request.options, {"output_format"}); value.has_value()) { + if (*value != "text" && *value != "json") { + throw std::runtime_error("VibeASR output_format must be text or json"); + } + out.output_format = *value; + } + if (const auto value = runtime::find_option(request.options, {"context"}); value.has_value()) { + out.context = *value; + } + out.max_new_tokens = runtime::parse_positive_i64_option(request.options, {"max_new_tokens"}, out.max_new_tokens); + return out; +} + +runtime::AudioBuffer VibeASRSession::normalize(const runtime::AudioBuffer & audio) const { + if (audio.samples.empty()) { + throw std::runtime_error("VibeASR requires non-empty audio"); + } + auto mono = engine::audio::mixdown_interleaved_to_mono_average(audio.samples, audio.channels); + if (audio.sample_rate != kSampleRate) { + // VibeASR.cpp resamples with a naive linear kernel; audio.cpp's soxr path + // is the better filter, so a non-24 kHz input will not match the + // reference sample for sample. + engine::audio::SoxrResampleOptions options; + options.profile = engine::audio::SoxrResampleProfile::QualityOnly; + options.output_length_policy = engine::audio::SoxrOutputLengthPolicy::ExactExpected; + options.output_padding = 256; + options.reject_empty_output = true; + options.warning_context = "VibeASR audio"; + options.fallback_description = "linear resampling"; + mono = engine::audio::resample_mono_soxr_or_linear(mono, audio.sample_rate, kSampleRate, options); + } + double sum = 0.0; + for (const float sample : mono) { + sum += static_cast(sample) * static_cast(sample); + } + const float rms = std::sqrt(static_cast(sum / std::max(mono.size(), 1))); + if (rms >= kNormalizeEps) { + const float target = std::pow(10.0F, kTargetDbFs / 20.0F); + const float gain = target / (rms + kNormalizeEps); + float max_abs = 0.0F; + for (float & sample : mono) { + sample *= gain; + max_abs = std::max(max_abs, std::abs(sample)); + } + // Not in VibeASR.cpp, which can clip on a loud clip; audio.cpp's own + // vibevoice_asr frontend clamps here and this port follows it. + if (max_abs > 1.0F) { + const float scale = max_abs + kNormalizeEps; + for (float & sample : mono) { + sample /= scale; + } + } + } + return runtime::AudioBuffer{kSampleRate, 1, std::move(mono)}; +} + +VibeASRSpeechEmbeddings VibeASRSession::encode_speech(const std::vector & samples) { + const auto encode_start = Clock::now(); + const auto acoustic = encoder_.encode_acoustic(samples); + const auto semantic = encoder_.encode_semantic(samples); + debug::timing_log_scalar("vibeasr.session.encoder_ms", engine::debug::elapsed_ms(encode_start)); + if (acoustic.frames != semantic.frames || acoustic.dim != semantic.dim) { + throw std::runtime_error("VibeASR encoder branches disagree on the feature shape"); + } + if (acoustic.dim != assets_->lm.hidden_size) { + throw std::runtime_error("VibeASR connector width does not match the decoder hidden size"); + } + + // Both connectors are LM-width, so the reference sums them element-wise. + VibeASRSpeechEmbeddings out; + out.tokens = acoustic.frames; + out.hidden_size = acoustic.dim; + out.values.resize(acoustic.values.size()); + for (size_t i = 0; i < out.values.size(); ++i) { + out.values[i] = acoustic.values[i] + semantic.values[i]; + } + return out; +} + +VibeASRLmPrompt VibeASRSession::build_prompt( + int64_t speech_tokens, + float duration_seconds, + const RequestOptions & options) const { + // Qwen2.5 ChatML, assembled exactly as VibeASR.cpp does it: + // <|im_start|>system\n{SYSTEM}<|im_end|>\n + // <|im_start|>user\n<|speech_start|><|speech_pad|>xN<|speech_end|>{suffix}<|im_end|>\n + // There is deliberately no generation prompt -- the model emits the + // <|im_start|>assistant\n header itself. + const auto encode = [this](const std::string & text) { + return tokenizer_->encode(text, false); + }; + + const std::string instruction = options.output_format == "json" + ? "please transcribe it with these keys: Start, End, Speaker, Content" + : "please transcribe it."; + std::string suffix; + if (options.context.empty()) { + suffix = "\nThis is a " + format_duration(duration_seconds) + " seconds audio, " + instruction; + } else { + suffix = "\nThis is a " + format_duration(duration_seconds) + " seconds audio, with extra info: " + + options.context + "\n\n" + + (options.output_format == "json" + ? "Please transcribe it with these keys: Start, End, Speaker, Content" + : "Please transcribe it."); + } + + const auto system_content = encode(std::string("system\n") + kSystemPrompt); + const auto newline = encode("\n"); + const auto user_prefix = encode("user\n"); + const auto user_suffix = encode(suffix); + + VibeASRLmPrompt prompt; + const auto append = [&prompt](const std::vector & ids) { + prompt.input_ids.insert(prompt.input_ids.end(), ids.begin(), ids.end()); + }; + prompt.input_ids.push_back(kImStart); + append(system_content); + prompt.input_ids.push_back(kImEnd); + append(newline); + prompt.input_ids.push_back(kImStart); + append(user_prefix); + prompt.input_ids.push_back(kSpeechStart); + // The reference builds ceil(samples / 3200) pads but only prefills + // min(pads, frames) of them, so emitting exactly `frames` pads produces the + // same sequence. + for (int64_t i = 0; i < speech_tokens; ++i) { + prompt.speech_positions.push_back(static_cast(prompt.input_ids.size())); + prompt.input_ids.push_back(kSpeechPad); + } + prompt.input_ids.push_back(kSpeechEnd); + append(user_suffix); + prompt.input_ids.push_back(kImEnd); + append(newline); + return prompt; +} + +std::string VibeASRSession::decode_tokens(const std::vector & token_ids) const { + // The prompt carries no generation prompt, so the model emits its own + // "<|im_start|>assistant\n" header; drop it exactly as the reference does. + size_t begin = 0; + const auto piece = [this](int32_t id) { return tokenizer_->decode({id}, true); }; + if (!token_ids.empty() && token_ids[0] == kImStart) { + begin = 1; + if (begin < token_ids.size() && piece(token_ids[begin]) == "assistant") { + ++begin; + if (begin < token_ids.size() && piece(token_ids[begin]) == "\n") { + ++begin; + } + } + } + + std::vector filtered; + filtered.reserve(token_ids.size() - begin); + for (size_t i = begin; i < token_ids.size(); ++i) { + const int32_t id = token_ids[i]; + if (id == kSpeechPad || id == kSpeechStart || id == kSpeechEnd || id == kEndOfText || + tokenizer_->is_control_token_id(id)) { + continue; + } + filtered.push_back(id); + } + if (filtered.empty()) { + return ""; + } + return engine::io::trim_ascii_whitespace(tokenizer_->decode(filtered, true)); +} + +runtime::TaskResult VibeASRSession::run(const runtime::TaskRequest & request) { + require_prepared("VibeASR run()"); + if (!request.audio_input.has_value()) { + throw std::runtime_error("VibeASR run() requires audio_input"); + } + const auto wall_start = Clock::now(); + const auto options = parse_request_options(request); + const auto audio = normalize(*request.audio_input); + const float duration_seconds = + static_cast(audio.samples.size()) / static_cast(kSampleRate); + + auto speech = encode_speech(audio.samples); + if (speech.tokens <= 0) { + throw std::runtime_error("VibeASR audio is too short to produce a single encoder frame"); + } + const auto prompt = build_prompt(speech.tokens, duration_seconds, options); + + VibeASRGenerationOptions generation; + generation.max_new_tokens = options.max_new_tokens; + generation.eos_token_ids = {kImEnd, kEndOfText}; + const auto generated = lm_.generate(prompt, speech, generation); + + runtime::TaskResult result; + result.text_output = runtime::Transcript{decode_tokens(generated), ""}; + debug::trace_log_scalar("vibeasr.session.speech_tokens", speech.tokens); + debug::trace_log_scalar("vibeasr.session.generated_tokens", static_cast(generated.size())); + debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start)); + return result; +} + +std::shared_ptr make_vibeasr_loader() { + runtime::SpecBackedVoiceModelConfig config; + config.family = "vibeasr"; + config.load_assets = [](const std::filesystem::path & model_path) { + return load_vibeasr_assets(model_path); + }; + config.create_session = []( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options, + std::shared_ptr assets, + std::shared_ptr contract) { + return std::make_unique(task, options, std::move(assets), std::move(contract)); + }; + return runtime::make_spec_backed_voice_loader(std::move(config)); +} + +} // namespace engine::community_models::vibeasr diff --git a/tests/vibeasr/test_vibeasr_asr.cpp b/tests/vibeasr/test_vibeasr_asr.cpp new file mode 100644 index 000000000..880babbe4 --- /dev/null +++ b/tests/vibeasr/test_vibeasr_asr.cpp @@ -0,0 +1,145 @@ +// End-to-end probe for the ported VibeASR pipeline: I8_S VAE encoder -> ternary +// I2_S Qwen2 decoder -> transcript. +// +// The package ships two GGUFs, so --model points at the LM GGUF and the spec is +// resolved from the repo (same convention as minimax_h3). Skips with 125 when +// the checkpoint is not installed. +// +// Upstream: https://github.com/microsoft/VibeASR.cpp + +#include "engine/framework/audio/wav_reader.h" +#include "engine/framework/io/filesystem.h" +#include "engine/framework/io/text.h" +#include "engine/framework/runtime/model.h" +#include "engine/framework/runtime/registry.h" +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include +#include +#include +#include + +#ifndef ENGINE_REPO_ROOT +#define ENGINE_REPO_ROOT "." +#endif + +namespace { + +constexpr int kExitPass = 0; +constexpr int kExitFail = 1; +constexpr int kExitSkip = 125; + +// LibriSpeech test-clean 6930-75918-0000, transcribed by VibeASR.cpp's own +// asr_infer --greedy on the same two GGUFs. +const char * kExpectedText = "Concord returned to its place amidst the tents."; + +std::filesystem::path repo_path(const std::string & relative) { + return std::filesystem::path(ENGINE_REPO_ROOT) / relative; +} + +std::string arg_value(int argc, char ** argv, const std::string & name, const std::string & fallback) { + for (int i = 1; i + 1 < argc; ++i) { + if (argv[i] == name) { + return argv[i + 1]; + } + } + return fallback; +} + +std::string normalize_text(const std::string & text) { + std::string out; + out.reserve(text.size()); + for (char ch : text) { + if (std::isalnum(static_cast(ch)) || std::isspace(static_cast(ch))) { + out.push_back(static_cast(std::tolower(static_cast(ch)))); + } + } + return engine::io::trim_ascii_whitespace(std::move(out)); +} + +} // namespace + +int main(int argc, char ** argv) { + const std::filesystem::path model_path = arg_value( + argc, argv, "--model", repo_path("models/vibeasr/vibeasr-lm-i2_s-embed-q6_k.gguf").string()); + const std::filesystem::path spec_override = arg_value( + argc, argv, "--model-spec-override", repo_path("model_specs").string()); + const std::filesystem::path audio_path = arg_value( + argc, argv, "--audio", + repo_path("assets/asr_validation/librispeech/librispeech_test_clean_6930-75918-0000.wav").string()); + const int threads = std::atoi(arg_value(argc, argv, "--threads", "4").c_str()); + + if (!engine::io::is_existing_file(model_path) || !engine::io::is_existing_file(audio_path)) { + std::fprintf( + stderr, + "SKIP: test_vibeasr_asr needs the LM GGUF at '%s' and audio at '%s'.\n" + " Fetch huggingface.co/XsquirrelC/VibeVoice-ASR-BitNet and run\n" + " tools/community_models/convert_vibeasr_gguf.py --in-place on both GGUFs.\n", + model_path.string().c_str(), + audio_path.string().c_str()); + return kExitSkip; + } + + try { + auto registry = engine::runtime::make_default_registry(); + engine::runtime::ModelLoadRequest load_request; + load_request.model_path = model_path; + load_request.model_spec_override = spec_override; + load_request.family_hint = "vibeasr"; + auto model = registry.load(load_request); + + const engine::runtime::TaskSpec task{ + engine::runtime::VoiceTaskKind::Asr, + engine::runtime::RunMode::Offline, + }; + engine::runtime::SessionOptions session_options; + session_options.backend.threads = threads > 0 ? threads : 1; + + auto session = model->create_task_session(task, session_options); + auto * offline = dynamic_cast(session.get()); + if (offline == nullptr) { + std::cerr << "FAIL: VibeASR session is not an IOfflineVoiceTaskSession\n"; + return kExitFail; + } + + const auto wav = engine::audio::read_wav_f32(audio_path); + engine::runtime::AudioBuffer audio; + audio.sample_rate = wav.sample_rate; + audio.channels = wav.channels; + audio.samples = wav.samples; + + offline->prepare(engine::runtime::build_preparation_request(audio)); + + engine::runtime::TaskRequest request; + request.audio_input = audio; + const auto result = offline->run(request); + + if (!result.text_output.has_value()) { + std::cerr << "FAIL: VibeASR produced no text output\n"; + return kExitFail; + } + const std::string actual = result.text_output->text; + std::cout << "transcript: " << actual << "\n"; + std::cout << "expected: " << kExpectedText << "\n"; + + // Raw equality pins punctuation and casing against the reference decode; + // the normalized compare is only there to localize a failure. + if (actual != kExpectedText) { + if (normalize_text(actual) == normalize_text(kExpectedText)) { + std::cerr << "FAIL: transcript differs only in punctuation or casing\n"; + } else { + std::cerr << "FAIL: transcript mismatch\n"; + } + return kExitFail; + } + + std::cout << "PASS: VibeASR end-to-end transcription matches VibeASR.cpp\n"; + return kExitPass; + } catch (const std::exception & error) { + std::cerr << "FAIL: " << error.what() << "\n"; + return kExitFail; + } +} diff --git a/tests/vibeasr/test_vibeasr_vae_encoder.cpp b/tests/vibeasr/test_vibeasr_vae_encoder.cpp index f42384d45..9d10404de 100644 --- a/tests/vibeasr/test_vibeasr_vae_encoder.cpp +++ b/tests/vibeasr/test_vibeasr_vae_encoder.cpp @@ -180,8 +180,8 @@ int main(int argc, char ** argv) { 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"); + "SKIP: test_vibeasr_vae_encoder needs --model and --audio .\n" + " Convert a VibeASR.cpp checkpoint with tools/community_models/convert_vibeasr_gguf.py first.\n"); return kExitSkip; } @@ -192,8 +192,10 @@ int main(int argc, char ** argv) { 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", + // The encoder is a raw-waveform stack: it accepts whatever rate the clip + // carries, and only the frame count and the reported RTF depend on it. + if (wav.sample_rate <= 0) { + std::fprintf(stderr, "SKIP: %s reports sample rate %d\n", audio_path.string().c_str(), wav.sample_rate); return kExitSkip; } @@ -215,7 +217,7 @@ int main(int argc, char ** argv) { 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 double audio_seconds = static_cast(num_samples) / static_cast(wav.sample_rate); const auto acoustic_start = std::chrono::steady_clock::now(); const auto acoustic = runtime.encode_acoustic(wav.samples); diff --git a/tools/community_models/convert_vibeasr_vae.py b/tools/community_models/convert_vibeasr_gguf.py similarity index 73% rename from tools/community_models/convert_vibeasr_vae.py rename to tools/community_models/convert_vibeasr_gguf.py index 6a28d9678..b1af8bb49 100644 --- a/tools/community_models/convert_vibeasr_vae.py +++ b/tools/community_models/convert_vibeasr_gguf.py @@ -1,34 +1,44 @@ #!/usr/bin/env python3 -"""Convert a VibeASR.cpp VAE encoder GGUF into an audio.cpp GGUF package. +"""Convert a VibeASR.cpp 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. +Weights: https://huggingface.co/XsquirrelC/VibeVoice-ASR-BitNet + +Handles both halves of the published package -- the I8_S VAE encoder and the +ternary I2_S language model -- because they need exactly the same fix and +nothing else. VibeASR.cpp ships both already quantized by its own ggml fork, so +there is nothing to re-quantize here. The only thing that differs is the numeric +type id: the VibeASR fork picked 36 (I2_S) and 37 (I8_S), which upstream ggml had +already used for the retired IQ4_NL_4_4 / IQ4_NL_4_8 slots. audio.cpp therefore +registers the same two types at 42 (I8_S) and 43 (I2_S). Tensors of any other +type in the file -- the LM's Q6_K token embedding, its F16 output projection, and +every F32 norm and bias -- are already portable and pass through untouched. + +The on-disk layout is identical either way -- an I8_S tensor is `nelements` int8 +bytes followed by a single padded F32 tensor scale, an I2_S tensor is the same +with 128 ternary codes packed per 32 bytes, and ggml's GGUF writer sizes every +tensor with ggml_nbytes() -- so this tool rewrites the 4-byte type field of each +tensor info and copies everything else through byte for byte. Data offsets, the +data section, and the KV block are untouched. Examples: # inspect a VibeASR GGUF without writing anything - python3 tools/community_models/convert_vibeasr_vae.py \ + python3 tools/community_models/convert_vibeasr_gguf.py \ --input vibeasr-vae-encoder-i8_s.gguf --list - # produce the audio.cpp package - python3 tools/community_models/convert_vibeasr_vae.py \ + # convert a downloaded package where it sits (both halves) + python3 tools/community_models/convert_vibeasr_gguf.py \ + --input VibeVoice-ASR-BitNet/vibeasr-vae-encoder-i8_s.gguf --in-place + python3 tools/community_models/convert_vibeasr_gguf.py \ + --input VibeVoice-ASR-BitNet/vibeasr-lm-i2_s-embed-q6_k.gguf --in-place + + # or write the converted copy somewhere else + python3 tools/community_models/convert_vibeasr_gguf.py \ --input vibeasr-vae-encoder-i8_s.gguf \ --output models/vibeasr/vae_encoder-i8_s.gguf # confirm an already converted package needs no further remapping - python3 tools/community_models/convert_vibeasr_vae.py \ + python3 tools/community_models/convert_vibeasr_gguf.py \ --input models/vibeasr/vae_encoder-i8_s.gguf --check """ import argparse @@ -42,7 +52,7 @@ # 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"} +TYPE_NAMES = {0: "f32", 1: "f16", 8: "q8_0", 14: "q6_k", 42: "i8_s", 43: "i2_s"} # GGUF metadata value type ids. ( @@ -174,12 +184,18 @@ def list_tensors(infos) -> None: 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("--input", type=Path, required=True, help="VibeASR.cpp GGUF (VAE encoder or LM)") parser.add_argument("--output", type=Path, help="audio.cpp GGUF to write") + parser.add_argument("--in-place", action="store_true", help="rewrite --input itself instead of writing a copy") parser.add_argument("--list", action="store_true", help="print the tensor table and exit") parser.add_argument("--check", action="store_true", help="exit non-zero if any tensor still needs remapping") args = parser.parse_args() + if args.in_place: + if args.output is not None: + parser.error("--in-place and --output are mutually exclusive") + args.output = args.input + data = bytearray(args.input.read_bytes()) infos, alignment = parse_tensor_infos(bytes(data)) @@ -196,7 +212,7 @@ def main() -> int: return 0 if args.output is None: - parser.error("--output is required unless --list or --check is given") + parser.error("--output or --in-place is required unless --list or --check is given") # Guard against a double conversion: the fork ids and the audio.cpp ids are # both valid ggml types, so a second pass would silently corrupt nothing but