| 167 | } |
| 168 | |
| 169 | core::TensorValue decoder_layer( |
| 170 | core::ModuleBuildContext & ctx, |
| 171 | const core::TensorValue & input, |
| 172 | const core::TensorValue & positions, |
| 173 | const std::optional<core::TensorValue> & attention_mask, |
| 174 | const TextEncoderLayerWeights & weights, |
| 175 | const AceStepTextEncoderConfig & config) { |
| 176 | const int64_t dim = head_dim(config); |
| 177 | const int64_t kv_repeats = config.num_attention_heads / config.num_key_value_heads; |
| 178 | const modules::LinearModule q_proj( |
| 179 | {config.hidden_size, config.num_attention_heads * dim, false, GGML_PREC_F32}); |
| 180 | const modules::LinearModule k_proj( |
| 181 | {config.hidden_size, config.num_key_value_heads * dim, false, GGML_PREC_F32}); |
| 182 | const modules::LinearModule v_proj( |
| 183 | {config.hidden_size, config.num_key_value_heads * dim, false, GGML_PREC_F32}); |
| 184 | const modules::LinearModule o_proj( |
| 185 | {config.num_attention_heads * dim, config.hidden_size, false, GGML_PREC_F32}); |
| 186 | const modules::RMSNormModule hidden_norm({config.hidden_size, config.rms_norm_eps, true, false}); |
| 187 | const modules::RMSNormModule head_norm({dim, config.rms_norm_eps, true, false}); |
| 188 | const modules::AddModule add; |
| 189 | |
| 190 | auto attn_in = hidden_norm.build(ctx, input, binding::norm_data(ctx, weights.input_norm)); |
| 191 | auto q = q_proj.build(ctx, attn_in, binding::linear_data(ctx, weights.q_proj)); |
| 192 | auto k = k_proj.build(ctx, attn_in, binding::linear_data(ctx, weights.k_proj)); |
| 193 | auto v = v_proj.build(ctx, attn_in, binding::linear_data(ctx, weights.v_proj)); |
| 194 | q = head_norm.build(ctx, reshape_heads(ctx, q, config.num_attention_heads, dim), binding::norm_data(ctx, weights.q_norm)); |
| 195 | k = head_norm.build(ctx, reshape_heads(ctx, k, config.num_key_value_heads, dim), binding::norm_data(ctx, weights.k_norm)); |
| 196 | v = reshape_heads(ctx, v, config.num_key_value_heads, dim); |
| 197 | q = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, q, positions); |
| 198 | k = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, k, positions); |
| 199 | |
| 200 | auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); |
| 201 | auto k_heads = |
| 202 | repeat_kv_heads(ctx, modules::TransposeModule({{0, 2, 1, 3}, k.shape.rank}).build(ctx, k), kv_repeats); |
| 203 | auto v_heads = |
| 204 | repeat_kv_heads(ctx, modules::TransposeModule({{0, 2, 1, 3}, v.shape.rank}).build(ctx, v), kv_repeats); |
| 205 | auto context = attention_from_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask); |
| 206 | context = modules::TransposeModule({{0, 2, 1, 3}, context.shape.rank}).build(ctx, context); |
| 207 | context = ensure_contiguous(ctx, context); |
| 208 | context = core::reshape_tensor( |
| 209 | ctx, |
| 210 | context, |
| 211 | core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], config.num_attention_heads * dim})); |
| 212 | auto x = add.build(ctx, input, o_proj.build(ctx, context, binding::linear_data(ctx, weights.o_proj))); |
| 213 | |
| 214 | auto ff_in = hidden_norm.build(ctx, x, binding::norm_data(ctx, weights.post_norm)); |
| 215 | auto gate = |
| 216 | modules::LinearModule( |
| 217 | {config.hidden_size, config.intermediate_size, false, GGML_PREC_F32}) |
| 218 | .build(ctx, ff_in, binding::linear_data(ctx, weights.gate_proj)); |
| 219 | gate = modules::SiluModule{}.build(ctx, gate); |
| 220 | auto up = |
| 221 | modules::LinearModule( |
| 222 | {config.hidden_size, config.intermediate_size, false, GGML_PREC_F32}) |
| 223 | .build(ctx, ff_in, binding::linear_data(ctx, weights.up_proj)); |
| 224 | auto ff = |
| 225 | modules::LinearModule( |
| 226 | {config.intermediate_size, config.hidden_size, false, GGML_PREC_F32}) |
no test coverage detected