| 117 | const modules::LinearModule o_proj( |
| 118 | binding::linear_config(config.num_attention_heads * dim, config.hidden_size, false)); |
| 119 | const modules::RMSNormModule hidden_norm({config.hidden_size, config.rms_norm_eps, true, false}); |
| 120 | const modules::RMSNormModule head_norm({dim, config.rms_norm_eps, true, false}); |
| 121 | const modules::AddModule add; |
| 122 | |
| 123 | auto attn_in = hidden_norm.build(ctx, input, binding::norm_data(ctx, weights.input_norm)); |
| 124 | auto q = q_proj.build(ctx, attn_in, binding::linear_data(ctx, weights.q_proj)); |
| 125 | auto k = k_proj.build(ctx, attn_in, binding::linear_data(ctx, weights.k_proj)); |
| 126 | auto v = v_proj.build(ctx, attn_in, binding::linear_data(ctx, weights.v_proj)); |
| 127 | q = head_norm.build(ctx, reshape_heads(ctx, q, config.num_attention_heads, dim), binding::norm_data(ctx, weights.q_norm)); |
| 128 | k = head_norm.build(ctx, reshape_heads(ctx, k, config.num_key_value_heads, dim), binding::norm_data(ctx, weights.k_norm)); |
| 129 | v = reshape_heads(ctx, v, config.num_key_value_heads, dim); |
| 130 | q = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, q, positions); |
| 131 | k = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, k, positions); |
| 132 | |
| 133 | auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); |
| 134 | auto k_heads = |
| 135 | repeat_kv_heads(ctx, modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k), kv_repeats); |
| 136 | auto v_heads = |
| 137 | repeat_kv_heads(ctx, modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v), kv_repeats); |
| 138 | auto context = attention_from_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask); |
| 139 | context = ensure_contiguous(ctx, context); |
| 140 | context = core::reshape_tensor( |
| 141 | ctx, |
| 142 | context, |
| 143 | core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], config.num_attention_heads * dim})); |
| 144 | auto x = add.build(ctx, input, o_proj.build(ctx, context, binding::linear_data(ctx, weights.o_proj))); |
| 145 | |
| 146 | auto ff_in = hidden_norm.build(ctx, x, binding::norm_data(ctx, weights.post_norm)); |
| 147 | auto gate = |
| 148 | modules::LinearModule(binding::linear_config( |
| 149 | config.hidden_size, |
| 150 | config.intermediate_size, |
| 151 | false)) |
| 152 | .build(ctx, ff_in, binding::linear_data(ctx, weights.gate_proj)); |
| 153 | gate = modules::SiluModule{}.build(ctx, gate); |
| 154 | auto up = |
| 155 | modules::LinearModule(binding::linear_config( |
| 156 | config.hidden_size, |
| 157 | config.intermediate_size, |
| 158 | false)) |
| 159 | .build(ctx, ff_in, binding::linear_data(ctx, weights.up_proj)); |
| 160 | auto ff = |
| 161 | modules::LinearModule(binding::linear_config( |
| 162 | config.intermediate_size, |
| 163 | config.hidden_size, |
| 164 | false)) |
| 165 | .build(ctx, modules::MulModule{}.build(ctx, gate, up), binding::linear_data(ctx, weights.down_proj)); |
| 166 | return add.build(ctx, x, ff); |
| 167 | } |
| 168 | |
| 169 | std::vector<float> build_padding_attention_mask_values(int64_t tokens, int64_t valid_tokens) { |
| 170 | if (tokens <= 0 || valid_tokens <= 0 || valid_tokens > tokens) { |
| 171 | throw std::runtime_error("ACE-Step padding attention mask requires 0 < valid_tokens <= tokens"); |
| 172 | } |
| 173 | std::vector<float> values(static_cast<size_t>(tokens * tokens), 0.0F); |
| 174 | const float neg_inf = -std::numeric_limits<float>::infinity(); |
| 175 | for (int64_t q = 0; q < valid_tokens; ++q) { |
| 176 | for (int64_t k = valid_tokens; k < tokens; ++k) { |
no test coverage detected