| 2130 | |
| 2131 | class CfgPrefillGraph { |
| 2132 | public: |
| 2133 | CfgPrefillGraph( |
| 2134 | std::shared_ptr<PlannerWeightsRuntime> runtime, |
| 2135 | int64_t prompt_steps, |
| 2136 | size_t graph_arena_bytes) |
| 2137 | : runtime_(std::move(runtime)), |
| 2138 | prompt_steps_(prompt_steps), |
| 2139 | activation_type_(planner_phase2_activation_type(runtime_->backend_type())) { |
| 2140 | if (prompt_steps_ <= 0) { |
| 2141 | throw std::runtime_error("ACE-Step planner CFG prefill requires positive prompt length"); |
| 2142 | } |
| 2143 | ggml_init_params params{graph_arena_bytes, nullptr, true}; |
| 2144 | ctx_.reset(ggml_init(params)); |
| 2145 | if (ctx_ == nullptr) { |
| 2146 | throw std::runtime_error("failed to initialize ACE-Step planner CFG prefill graph context"); |
| 2147 | } |
| 2148 | const auto & config = runtime_->assets().config.planner; |
| 2149 | const auto & weights = runtime_->weights(); |
| 2150 | core::ModuleBuildContext ctx{ctx_.get(), "ace_step.planner.cfg_prefill", runtime_->backend_type()}; |
| 2151 | token_ids_ = ggml_new_tensor_2d(ctx_.get(), GGML_TYPE_I32, prompt_steps_, 2); |
| 2152 | ggml_set_input(token_ids_); |
| 2153 | auto token_ids = core::wrap_tensor(token_ids_, core::TensorShape::from_dims({2, prompt_steps_}), GGML_TYPE_I32); |
| 2154 | auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size}) |
| 2155 | .build(ctx, token_ids, weights.token_embedding); |
| 2156 | x = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({2, prompt_steps_, config.hidden_size})); |
| 2157 | x = cast_planner_activation(ctx, x, activation_type_); |
| 2158 | positions_ = ggml_new_tensor_2d(ctx_.get(), GGML_TYPE_I32, prompt_steps_, 2); |
| 2159 | ggml_set_input(positions_); |
| 2160 | auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({2, prompt_steps_}), GGML_TYPE_I32); |
| 2161 | attention_mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, prompt_steps_, prompt_steps_, 1, 2); |
| 2162 | ggml_set_input(attention_mask_); |
| 2163 | auto attention_mask = core::wrap_tensor( |
| 2164 | attention_mask_, |
| 2165 | core::TensorShape::from_dims({2, 1, prompt_steps_, prompt_steps_}), |
| 2166 | GGML_TYPE_F16); |
| 2167 | query_mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F32, 1, 1, prompt_steps_, 2); |
| 2168 | ggml_set_input(query_mask_); |
| 2169 | auto query_mask = core::wrap_tensor( |
| 2170 | query_mask_, |
| 2171 | core::TensorShape::from_dims({2, prompt_steps_, 1, 1}), |
| 2172 | GGML_TYPE_F32); |
| 2173 | for (size_t layer_index = 0; layer_index < weights.layers.layers.size(); ++layer_index) { |
| 2174 | const auto & layer = weights.layers.layers[layer_index]; |
| 2175 | auto out = planner_decoder_layer_batched( |
| 2176 | ctx, |
| 2177 | x, |
| 2178 | positions, |
| 2179 | layer, |
| 2180 | config, |
| 2181 | attention_mask, |
| 2182 | query_mask, |
| 2183 | activation_type_); |
| 2184 | x = out.output; |
| 2185 | keys_.push_back(out.key.tensor); |
| 2186 | values_.push_back(out.value.tensor); |
| 2187 | } |
| 2188 | x = modules::SliceModule({1, prompt_steps_ - 1, 1}).build(ctx, x); |
| 2189 | x = modules::RMSNormModule({config.hidden_size, config.rms_norm_eps, true, false}) |
nothing calls this directly
no test coverage detected