| 331 | GGML_TYPE_F32); |
| 332 | } |
| 333 | |
| 334 | modules::QwenDecoderLayerOutputs planner_decoder_layer_batched( |
| 335 | core::ModuleBuildContext & ctx, |
| 336 | const core::TensorValue & input, |
| 337 | const core::TensorValue & positions, |
| 338 | const QwenDecoderLayerWeights & weights, |
| 339 | const AceStepPlannerConfig & config, |
| 340 | const core::TensorValue & attention_mask, |
| 341 | const core::TensorValue & query_mask, |
| 342 | ggml_type activation_type) { |
| 343 | const int64_t dim = planner_attention_head_dim(config); |
| 344 | |
| 345 | auto x_norm = modules::RMSNormModule({config.hidden_size, config.rms_norm_eps, true, weights.input_norm.bias.has_value()}) |
| 346 | .build(ctx, input, weights.input_norm); |
| 347 | auto q = modules::LinearModule({config.hidden_size, config.num_attention_heads * dim, false}) |
| 348 | .build(ctx, x_norm, {weights.self_attention.q_weight, std::nullopt}); |
| 349 | auto k = modules::LinearModule({config.hidden_size, config.num_key_value_heads * dim, false}) |
| 350 | .build(ctx, x_norm, {weights.self_attention.k_weight, std::nullopt}); |
| 351 | auto v = modules::LinearModule({config.hidden_size, config.num_key_value_heads * dim, false}) |
| 352 | .build(ctx, x_norm, {weights.self_attention.v_weight, std::nullopt}); |
| 353 | q = cast_planner_activation(ctx, q, activation_type); |
| 354 | k = cast_planner_activation(ctx, k, activation_type); |
| 355 | v = cast_planner_activation(ctx, v, activation_type); |
| 356 | |
| 357 | q = modules::RMSNormModule({dim, config.rms_norm_eps, true, weights.q_norm.bias.has_value()}) |
| 358 | .build(ctx, reshape_planner_heads(ctx, q, config.num_attention_heads, dim), weights.q_norm); |
| 359 | k = modules::RMSNormModule({dim, config.rms_norm_eps, true, weights.k_norm.bias.has_value()}) |
| 360 | .build(ctx, reshape_planner_heads(ctx, k, config.num_key_value_heads, dim), weights.k_norm); |
| 361 | v = reshape_planner_heads(ctx, v, config.num_key_value_heads, dim); |
| 362 | v = cast_planner_activation(ctx, v, activation_type); |
| 363 | |
| 364 | q = apply_planner_rope_batched(ctx, q, positions, dim, config.rope_theta); |
| 365 | k = apply_planner_rope_batched(ctx, k, positions, dim, config.rope_theta); |
| 366 | q = cast_planner_activation(ctx, q, activation_type); |
| 367 | k = cast_planner_activation(ctx, k, activation_type); |
| 368 | |
| 369 | std::array<int, core::kMaxTensorRank> head_axes = {0, 2, 1, 3}; |
| 370 | auto q_heads = modules::TransposeModule({head_axes, q.shape.rank}).build(ctx, q); |
| 371 | auto k_heads = modules::TransposeModule({head_axes, k.shape.rank}).build(ctx, k); |
| 372 | auto v_heads = modules::TransposeModule({head_axes, v.shape.rank}).build(ctx, v); |
| 373 | q_heads = ensure_planner_contiguous(ctx, q_heads); |
| 374 | k_heads = ensure_planner_contiguous(ctx, k_heads); |
| 375 | v_heads = ensure_planner_contiguous(ctx, v_heads); |
| 376 | |
| 377 | auto context = sdpa_from_planner_grouped_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask); |
| 378 | context = zero_masked_query_rows(ctx, context, query_mask); |
| 379 | context = core::ensure_backend_addressable_layout(ctx, context); |
| 380 | context = core::reshape_tensor( |
| 381 | ctx, |
| 382 | context, |
| 383 | core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], config.num_attention_heads * dim})); |
| 384 | context = cast_planner_activation(ctx, context, activation_type); |
| 385 | |
| 386 | auto attn_out = modules::LinearModule({config.hidden_size, config.hidden_size, false}) |
| 387 | .build(ctx, context, {weights.self_attention.out_weight, std::nullopt}); |
| 388 | attn_out = cast_planner_activation(ctx, attn_out, activation_type); |
| 389 | auto x = modules::AddModule{}.build(ctx, input, attn_out); |
| 390 | x = cast_planner_activation(ctx, x, activation_type); |
no test coverage detected