| 2796 | } |
| 2797 | |
| 2798 | AceStepPlan AceStepPlannerRuntime::generate(const AceStepRequest & request, bool generate_audio_codes) const { |
| 2799 | const auto total_start = Clock::now(); |
| 2800 | double cot_decode_graph_step_ms = 0.0; |
| 2801 | double code_decode_graph_step_ms = 0.0; |
| 2802 | double code_sampling_ms = 0.0; |
| 2803 | const auto prepare_prefill = [&]( |
| 2804 | const AceStepPlannerPreparedInput & prepared, |
| 2805 | int64_t max_new_tokens, |
| 2806 | std::string_view timing_prefix) { |
| 2807 | const auto & config = assets_->config.planner; |
| 2808 | const auto & prompt_ids = prepared.tokenized_prompt.input_ids; |
| 2809 | const auto & attention_mask = prepared.tokenized_prompt.attention_mask; |
| 2810 | if (static_cast<int64_t>(prompt_ids.size()) + max_new_tokens > config.max_position_embeddings) { |
| 2811 | throw std::runtime_error("ACE-Step planner request exceeds max_position_embeddings"); |
| 2812 | } |
| 2813 | auto prefill_runtime = planner_prefill_weights_runtime_ != nullptr |
| 2814 | ? planner_prefill_weights_runtime_ |
| 2815 | : weights_runtime_; |
| 2816 | const auto graph_prepare_start = Clock::now(); |
| 2817 | const int64_t required_cache_steps = static_cast<int64_t>(prompt_ids.size()) + max_new_tokens; |
| 2818 | if (!decode_graph_ || !decode_graph_->can_run(*weights_runtime_, required_cache_steps)) { |
| 2819 | decode_graph_ = std::make_unique<DecodeGraph>( |
| 2820 | weights_runtime_, |
| 2821 | required_cache_steps, |
| 2822 | decode_graph_arena_bytes_); |
| 2823 | } |
| 2824 | const bool use_metal_prompt_step_prefill = prefill_runtime->backend_type() == core::BackendType::Metal; |
| 2825 | if (!use_metal_prompt_step_prefill && |
| 2826 | (!prefill_graph_ || !prefill_graph_->can_run(*prefill_runtime, static_cast<int64_t>(prompt_ids.size())))) { |
| 2827 | prefill_graph_ = std::make_unique<PrefillGraph>( |
| 2828 | prefill_runtime, |
| 2829 | static_cast<int64_t>(prompt_ids.size()), |
| 2830 | decode_graph_arena_bytes_); |
| 2831 | } |
| 2832 | engine::debug::timing_log_scalar( |
| 2833 | std::string(timing_prefix) + ".prefill.graph.prepare_ms", |
| 2834 | engine::debug::elapsed_ms(graph_prepare_start, Clock::now())); |
| 2835 | engine::debug::trace_log_scalar( |
| 2836 | std::string(timing_prefix) + ".prefill_prompt_tokens", |
| 2837 | prompt_ids.size()); |
| 2838 | engine::debug::trace_log_scalar( |
| 2839 | std::string(timing_prefix) + ".prefill_cache_steps", |
| 2840 | required_cache_steps); |
| 2841 | const auto prefill_run_start = Clock::now(); |
| 2842 | PrefillOutput out; |
| 2843 | if (use_metal_prompt_step_prefill) { |
| 2844 | out = run_decode_graph_prompt_prefill(*decode_graph_, prepared.tokenized_prompt); |
| 2845 | } else { |
| 2846 | out = prefill_graph_->run(prompt_ids, attention_mask); |
| 2847 | } |
| 2848 | engine::debug::timing_log_scalar( |
| 2849 | std::string(timing_prefix) + ".prefill_run_ms", |
| 2850 | engine::debug::elapsed_ms(prefill_run_start, Clock::now())); |
| 2851 | return out; |
| 2852 | }; |
| 2853 | |
| 2854 | const auto decode_cot_phase = [&](const AceStepPlannerPreparedInput & prepared) { |
| 2855 | const auto phase_start = Clock::now(); |
no test coverage detected