Function same_attention

src/models/stable_audio/same_autoencoder.cpp:206–230

}

core::TensorValue same_attention(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & q_heads,
    const core::TensorValue & k_heads,
    const core::TensorValue & v_heads,
    const core::TensorValue * attention_mask,
    int64_t dim) {
    const auto q = ensure_contiguous(ctx, q_heads);
    const auto k = ensure_contiguous(ctx, k_heads);
    const auto v = ensure_contiguous(ctx, v_heads);
    auto * flash = ggml_flash_attn_ext(
        ctx.ggml,
        q.tensor,
        k.tensor,
        v.tensor,
        attention_mask != nullptr ? attention_mask->tensor : nullptr,
        1.0F / std::sqrt(static_cast<float>(dim)),
        0.0F,
        0.0F);
    ggml_flash_attn_ext_set_prec(flash, GGML_PREC_F32);
    return core::wrap_tensor(
        flash,
        core::TensorShape::from_dims({q.shape.dims[0], q.shape.dims[2], q.shape.dims[1], dim}),
        GGML_TYPE_F32);
}

core::TensorValue same_self_attention(
    core::ModuleBuildContext & ctx,