Function build_vibevoice_decoder_layer

src/models/vibevoice/decoder.cpp:1769–1861

}

VibeVoiceDecoderLayerOutputs build_vibevoice_decoder_layer(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & input,
    const core::TensorValue & positions,
    const VibeVoiceDecoderLayerWeights & weights,
    const VibeVoiceDecoderConfig & config,
    core::ConstantTensorCache & constants,
    const std::optional<core::TensorValue> & prefix_key,
    const std::optional<core::TensorValue> & prefix_value,
    const std::optional<core::TensorValue> & attention_mask) {
    if (prefix_key.has_value() != prefix_value.has_value()) {
        throw std::runtime_error("VibeVoice decoder layer requires both prefix key and value or neither");
    }
    const int64_t dim = require_head_dim(config);
    const int64_t kv_repeats = config.num_attention_heads / config.num_key_value_heads;
    const modules::LinearModule q_proj(
        binding::linear_config(config.hidden_size, config.num_attention_heads * dim, true));
    const modules::LinearModule k_proj(
        binding::linear_config(config.hidden_size, config.num_key_value_heads * dim, true));
    const modules::LinearModule v_proj(
        binding::linear_config(config.hidden_size, config.num_key_value_heads * dim, true));
    const modules::LinearModule o_proj(
        binding::linear_config(config.num_attention_heads * dim, config.hidden_size, false));
    const modules::RMSNormModule hidden_norm({config.hidden_size, config.rms_norm_eps, true, false});
    const modules::AddModule add;

    auto attn_in = hidden_norm.build(ctx, input, binding::norm_data(constants, weights.input_norm));
    auto q = q_proj.build(
        ctx,
        attn_in,
        binding::linear_data(constants, weights.self_attention.q_weight, weights.self_attention.q_bias));
    auto k = k_proj.build(
        ctx,
        attn_in,
        binding::linear_data(constants, weights.self_attention.k_weight, weights.self_attention.k_bias));
    auto v = v_proj.build(
        ctx,
        attn_in,
        binding::linear_data(constants, weights.self_attention.v_weight, weights.self_attention.v_bias));
    q = reshape_heads(ctx, q, config.num_attention_heads, dim);
    k = reshape_heads(ctx, k, config.num_key_value_heads, dim);
    v = reshape_heads(ctx, v, config.num_key_value_heads, dim);
    q = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, q, positions);
    k = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, k, positions);
    k = core::ensure_backend_addressable_layout(ctx, k);
    v = core::ensure_backend_addressable_layout(ctx, v);

    auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q);
    auto all_k = prefix_key.has_value() ? modules::ConcatModule({1}).build(ctx, *prefix_key, k) : k;
    auto all_v = prefix_value.has_value() ? modules::ConcatModule({1}).build(ctx, *prefix_value, v) : v;
    core::TensorValue context;
    if (!prefix_key.has_value() && attention_mask.has_value()) {
        q_heads = core::wrap_tensor(ggml_cont(ctx.ggml, q_heads.tensor), q_heads.shape, q_heads.type);
        auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, all_k.shape.rank}).build(ctx, all_k);
        auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, all_v.shape.rank}).build(ctx, all_v);
        context = flash_attention_from_grouped_heads_view_kv(ctx, q_heads, k_heads, v_heads, dim, *attention_mask);
    } else {
        auto k_heads = repeat_kv_heads(