Function encoder_layer

src/models/ace_step/condition_encoder.cpp:119–185

    const modules::RMSNormModule hidden_norm({config.hidden_size, config.rms_norm_eps, true, false});
    const modules::RMSNormModule head_norm({dim, config.rms_norm_eps, true, false});
    const modules::AddModule add;

    auto attn_in = hidden_norm.build(ctx, input, binding::norm_data(ctx, weights.input_norm));
    auto q = q_proj.build(ctx, attn_in, binding::linear_data(ctx, weights.q_proj));
    auto k = k_proj.build(ctx, attn_in, binding::linear_data(ctx, weights.k_proj));
    auto v = v_proj.build(ctx, attn_in, binding::linear_data(ctx, weights.v_proj));
    q = head_norm.build(ctx, reshape_heads(ctx, q, config.num_attention_heads, dim), binding::norm_data(ctx, weights.q_norm));
    k = head_norm.build(ctx, reshape_heads(ctx, k, config.num_key_value_heads, dim), binding::norm_data(ctx, weights.k_norm));
    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);

    auto q_heads = 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);
    auto context = attention_from_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask);
    context = ensure_contiguous(ctx, context);
    context = core::reshape_tensor(
        ctx,
        context,
        core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], config.num_attention_heads * dim}));
    auto x = add.build(ctx, input, o_proj.build(ctx, context, binding::linear_data(ctx, weights.o_proj)));

    auto ff_in = hidden_norm.build(ctx, x, binding::norm_data(ctx, weights.post_norm));
    auto gate =
        modules::LinearModule(binding::linear_config(
                                  config.hidden_size,
                                  config.intermediate_size,
                                  false))
            .build(ctx, ff_in, binding::linear_data(ctx, weights.gate_proj));
    gate = modules::SiluModule{}.build(ctx, gate);
    auto up =
        modules::LinearModule(binding::linear_config(
                                  config.hidden_size,
                                  config.intermediate_size,
                                  false))
            .build(ctx, ff_in, binding::linear_data(ctx, weights.up_proj));
    auto ff =
        modules::LinearModule(binding::linear_config(
                                  config.intermediate_size,
                                  config.hidden_size,
                                  false))
            .build(ctx, modules::MulModule{}.build(ctx, gate, up), binding::linear_data(ctx, weights.down_proj));
    return add.build(ctx, x, ff);
}

std::vector<float> build_padding_attention_mask_values(int64_t tokens, int64_t valid_tokens) {
    if (tokens <= 0 || valid_tokens <= 0 || valid_tokens > tokens) {
        throw std::runtime_error("ACE-Step padding attention mask requires 0 < valid_tokens <= tokens");
    }
    std::vector<float> values(static_cast<size_t>(tokens * tokens), 0.0F);
    const float neg_inf = -std::numeric_limits<float>::infinity();
    for (int64_t q = 0; q < valid_tokens; ++q) {
        for (int64_t k = valid_tokens; k < tokens; ++k) {
            values[static_cast<size_t>(q * tokens + k)] = neg_inf;
        }