Function attention_from_heads

src/models/ace_step/diffusion.cpp:130–164

    if (cross_attention_kv.has_value()) {
        k = cross_attention_kv->key;
        v = cross_attention_kv->value;
    } else {
        const auto & kv_source = is_cross ? *encoder_hidden_states : hidden_states;
        k = modules::LinearModule({config.hidden_size, config.num_key_value_heads * dim, false, GGML_PREC_F32})
                .build(ctx, kv_source, {weights.k_weight, std::nullopt});
        v = modules::LinearModule({config.hidden_size, config.num_key_value_heads * dim, false, GGML_PREC_F32})
                .build(ctx, kv_source, {weights.v_weight, std::nullopt});
    }
    q = modules::RMSNormModule({dim, config.rms_norm_eps, true, false})
            .build(ctx, reshape_heads(ctx, q, config.num_attention_heads, dim), {weights.q_norm, std::nullopt});
    if (!cross_attention_kv.has_value()) {
        k = modules::RMSNormModule({dim, config.rms_norm_eps, true, false})
                .build(ctx, reshape_heads(ctx, k, config.num_key_value_heads, dim), {weights.k_norm, std::nullopt});
        v = reshape_heads(ctx, v, config.num_key_value_heads, dim);
    }
    if (!is_cross) {
        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);
    }
    auto q_heads = ensure_contiguous(
        ctx,
        modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q));
    auto k_heads = repeat_kv_heads(
        ctx,
        modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k),
        kv_repeats);
    auto v_heads = repeat_kv_heads(
        ctx,
        modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v),
        kv_repeats);
    k_heads = ensure_contiguous(ctx, k_heads);
    v_heads = ensure_contiguous(ctx, v_heads);
    auto context = attention_from_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask, backend_type);
    context = ensure_contiguous(ctx, context);
    context = core::reshape_tensor(
        ctx,
        context,
        core::TensorShape::from_dims({hidden_states.shape.dims[0], hidden_states.shape.dims[1], config.hidden_size}));