Method PrefillGraph

src/models/ace_step/planner.cpp:1756–1818

        size_t graph_arena_bytes)
        : runtime_(std::move(runtime)),
          prompt_steps_(prompt_steps) {
        if (prompt_steps_ <= 0) {
            throw std::runtime_error("ACE-Step planner prefill requires positive prompt length");
        }
        ggml_init_params params{graph_arena_bytes, nullptr, true};
        ctx_.reset(ggml_init(params));
        if (ctx_ == nullptr) {
            throw std::runtime_error("failed to initialize ACE-Step planner prefill graph context");
        }
        const auto & config = runtime_->assets().config.planner;
        const auto & weights = runtime_->weights();
        core::ModuleBuildContext ctx{ctx_.get(), "ace_step.planner.prefill", runtime_->backend_type()};
        token_ids_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_);
        auto token_ids = core::wrap_tensor(token_ids_, core::TensorShape::from_dims({prompt_steps_}), GGML_TYPE_I32);
        auto x = modules::EmbeddingModule({config.vocab_size, config.hidden_size})
                     .build(ctx, token_ids, weights.token_embedding);
        x = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({1, prompt_steps_, config.hidden_size}));
        positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, prompt_steps_);
        auto positions = core::wrap_tensor(positions_, core::TensorShape::from_dims({prompt_steps_}), GGML_TYPE_I32);
        attention_mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, prompt_steps_, prompt_steps_, 1, 1);
        auto attention_mask = core::wrap_tensor(
            attention_mask_,
            core::TensorShape::from_dims({1, 1, prompt_steps_, prompt_steps_}),
            GGML_TYPE_F16);
        query_mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F32, 1, 1, prompt_steps_, 1);
        auto query_mask = core::wrap_tensor(
            query_mask_,
            core::TensorShape::from_dims({1, prompt_steps_, 1, 1}),
            GGML_TYPE_F32);
        for (size_t layer_index = 0; layer_index < weights.layers.layers.size(); ++layer_index) {
            const auto & layer = weights.layers.layers[layer_index];
            auto out = planner_decoder_layer_batched(
                ctx,
                x,
                positions,
                layer,
                config,
                attention_mask,
                query_mask,
                GGML_TYPE_F32);
            x = out.output;
            keys_.push_back(out.key.tensor);
            values_.push_back(out.value.tensor);
        }
        x = modules::SliceModule({1, prompt_steps_ - 1, 1}).build(ctx, x);
        x = modules::RMSNormModule({config.hidden_size, config.rms_norm_eps, true, false})
                .build(ctx, x, {weights.norm, std::nullopt});
        auto logits = modules::LinearModule({config.hidden_size, config.vocab_size, false})
                          .build(ctx, x, {weights.lm_head, std::nullopt});
        logits_ = logits.tensor;
        ggml_set_output(logits_);
        graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false);
        ggml_build_forward_expand(graph_, logits_);
        for (auto * input : {token_ids_, positions_, attention_mask_, query_mask_}) {
            ggml_set_input(input);
        }
        // The CPU reads every layer's KV outputs after graph execution. They
        // must remain live even after the attention nodes have consumed them.