Fabio
c36510c21e
Some checks failed
Linux Build / Linux cpu (push) Has been cancelled
Linux Build / Linux cpu-clang (push) Has been cancelled
Linux Build / Linux vulkan (push) Has been cancelled
mac-build / macOS CPU (push) Has been cancelled
Nix Build / Nix (cpu) (push) Has been cancelled
Nix Build / Nix (vulkan) (push) Has been cancelled
windows-build / Windows CPU (push) Has been cancelled
92 lines
3.5 KiB
C++
92 lines
3.5 KiB
C++
#pragma once
|
|
|
|
#include "engine/models/ace_step/assets.h"
|
|
#include "engine/models/ace_step/tokenizer_text.h"
|
|
#include "engine/models/ace_step/types.h"
|
|
#include "engine/framework/core/execution_context.h"
|
|
|
|
#include <array>
|
|
#include <cstdint>
|
|
#include <map>
|
|
#include <memory>
|
|
#include <optional>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace engine::models::ace_step {
|
|
|
|
class Qwen3PlannerWeightsRuntime;
|
|
class Qwen3PlannerPrefillGraph;
|
|
class Qwen3PlannerDecodeGraph;
|
|
class Qwen3PlannerCfgPrefillGraph;
|
|
class Qwen3PlannerCfgDecodeGraph;
|
|
|
|
struct AceStepPlannerPreparedInput {
|
|
std::string formatted_prompt;
|
|
AceStepTokenizedText tokenized_prompt;
|
|
};
|
|
|
|
struct Phase1ConstraintTables {
|
|
std::optional<int32_t> newline_token;
|
|
std::optional<int32_t> backtick_token;
|
|
std::vector<int32_t> forbidden_tokens;
|
|
std::map<std::vector<int32_t>, std::vector<int32_t>> bpm_prefix_map;
|
|
std::map<std::vector<int32_t>, std::vector<int32_t>> duration_prefix_map;
|
|
std::map<std::vector<int32_t>, std::vector<int32_t>> keyscale_prefix_map;
|
|
std::map<std::vector<int32_t>, std::vector<int32_t>> language_prefix_map;
|
|
std::map<std::vector<int32_t>, std::vector<int32_t>> timesig_prefix_map;
|
|
};
|
|
|
|
class AceStepQwen3PlannerRuntime {
|
|
public:
|
|
struct GenerationConfig {
|
|
int64_t max_prompt_tokens = 4096;
|
|
int64_t max_cot_tokens = 512;
|
|
int64_t max_code_tokens = 4032;
|
|
bool prefix_bucketing = true;
|
|
bool fused_projections = false;
|
|
bool f16_decode = false;
|
|
bool independent_cfg_prefill = false;
|
|
};
|
|
|
|
AceStepQwen3PlannerRuntime(
|
|
std::shared_ptr<const AceStepAssets> assets,
|
|
core::ExecutionContext & execution);
|
|
|
|
AceStepQwen3PlannerRuntime(
|
|
std::shared_ptr<const AceStepAssets> assets,
|
|
core::ExecutionContext & execution,
|
|
assets::TensorStorageType weight_storage_type,
|
|
size_t decode_graph_arena_bytes,
|
|
size_t weight_context_bytes,
|
|
GenerationConfig generation);
|
|
~AceStepQwen3PlannerRuntime();
|
|
|
|
AceStepPlannerPreparedInput prepare_prompt(const AceStepRequest & request) const;
|
|
AceStepPlan generate(const AceStepRequest & request, bool generate_audio_codes = true) const;
|
|
void release_graph_workspace() const;
|
|
std::string decode_tokens(const std::vector<int32_t> & token_ids) const;
|
|
AceStepPlan parse_output(const std::string & output_text) const;
|
|
|
|
private:
|
|
std::shared_ptr<const AceStepAssets> assets_;
|
|
AceStepTextTokenizer tokenizer_;
|
|
GenerationConfig generation_;
|
|
std::shared_ptr<Qwen3PlannerWeightsRuntime> weights_runtime_;
|
|
std::unique_ptr<core::ExecutionContext> host_planner_prefill_execution_;
|
|
std::shared_ptr<Qwen3PlannerWeightsRuntime> planner_prefill_weights_runtime_;
|
|
size_t decode_graph_arena_bytes_ = 0;
|
|
mutable std::unique_ptr<Qwen3PlannerPrefillGraph> prefill_graph_;
|
|
mutable std::unique_ptr<Qwen3PlannerDecodeGraph> decode_graph_;
|
|
mutable std::array<std::unique_ptr<Qwen3PlannerPrefillGraph>, 2> code_prefill_graphs_;
|
|
mutable std::unique_ptr<Qwen3PlannerPrefillGraph> unconditional_prefill_graph_;
|
|
mutable std::unique_ptr<Qwen3PlannerDecodeGraph> code_decode_graph_;
|
|
mutable std::unique_ptr<Qwen3PlannerCfgPrefillGraph> cfg_prefill_graph_;
|
|
mutable std::array<std::unique_ptr<Qwen3PlannerCfgDecodeGraph>, 2> code_f16_decode_graphs_;
|
|
mutable std::unique_ptr<Qwen3PlannerCfgDecodeGraph> cot_decode_graph_;
|
|
std::vector<uint8_t> is_audio_code_token_;
|
|
std::vector<int32_t> phase2_candidate_token_ids_;
|
|
Phase1ConstraintTables phase1_constraints_;
|
|
};
|
|
|
|
} // namespace engine::models::ace_step
|