Function build_attention

src/models/roformer/runtime.cpp:263–307

core::TensorValue reshape_heads(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & input,
    int64_t heads,
    int64_t dim) {
    return core::reshape_tensor(
        ctx,
        ensure_contiguous(ctx, input),
        core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], heads, dim}));
}

core::TensorValue attention_from_heads(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & q_heads,
    const core::TensorValue & k_heads,
    const core::TensorValue & v_heads,
    int64_t dim) {
    if (ctx.backend_type == core::BackendType::Cuda) {
        return modules::ScaledDotProductAttentionModule({
            dim,
            modules::ScaledDotProductAttentionLowering::FlashPreserveViews,
            GGML_PREC_F32,
            modules::AttentionCausality::NonCausal,
        }).build(ctx, q_heads, k_heads, v_heads);
    }

    auto scores = matmul_f32(
        ctx,
        q_heads,
        modules::TransposeModule({{0, 1, 3, 2}, k_heads.shape.rank}).build(ctx, k_heads));
    scores = core::wrap_tensor(
        ggml_scale(ctx.ggml, ensure_contiguous(ctx, scores).tensor, 1.0f / std::sqrt(static_cast<float>(dim))),
        scores.shape,
        GGML_TYPE_F32);
    scores = ensure_contiguous(ctx, scores);
    auto attn = modules::SoftmaxModule{}.build(ctx, scores);
    auto context = matmul_f32(ctx, attn, v_heads);
    return modules::TransposeModule({{0, 2, 1, 3}, context.shape.rank}).build(ctx, context);
}

core::TensorValue build_attention(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & input,
    const core::TensorValue & positions,
    const AttentionWeights & weights,
    const RoformerArchitectureConfig & config) {
    const modules::LinearModule q_proj(binding::linear_config(config.dim, config.heads * config.dim_head, false));
    const modules::LinearModule k_proj(binding::linear_config(config.dim, config.heads * config.dim_head, false));
    const modules::LinearModule v_proj(binding::linear_config(config.dim, config.heads * config.dim_head, false));
    const modules::LinearModule gate_proj(binding::linear_config(config.dim, config.heads, true));