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
100 lines
3.5 KiB
C++
100 lines
3.5 KiB
C++
#pragma once
|
|
|
|
#include "engine/framework/assets/tensor_source.h"
|
|
#include "engine/framework/core/execution_context.h"
|
|
|
|
#include <cstddef>
|
|
#include <cstdint>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <vector>
|
|
|
|
namespace engine::models::fish_drama3 {
|
|
|
|
struct Drama3BackboneDecoderConfig {
|
|
int64_t vocab_size = 252288;
|
|
int64_t hidden_size = 2048;
|
|
int64_t intermediate_size = 512;
|
|
int64_t experts = 256;
|
|
int64_t experts_per_token = 8;
|
|
int64_t layers = 40;
|
|
int64_t heads = 16;
|
|
int64_t kv_heads = 2;
|
|
int64_t head_dim = 256;
|
|
int64_t full_attention_interval = 4;
|
|
int64_t linear_conv_kernel_dim = 4;
|
|
int64_t linear_key_head_dim = 128;
|
|
int64_t linear_num_key_heads = 16;
|
|
int64_t linear_num_value_heads = 32;
|
|
int64_t linear_value_head_dim = 128;
|
|
float rms_norm_eps = 1.0e-6F;
|
|
float rope_theta = 10000000.0F;
|
|
float partial_rotary_factor = 0.25F;
|
|
std::string weight_prefix = "model.language_model";
|
|
std::string lm_head_prefix = "lm_head";
|
|
bool tie_word_embeddings = false;
|
|
std::vector<int32_t> lm_head_token_ids;
|
|
// drama3-small: rows of the output head. 0 means the full vocabulary; a speech-only
|
|
// head has codebook_size + 1 rows ordered [EOS, semantic 0..codebook_size-1].
|
|
int64_t lm_head_rows = 0;
|
|
// drama3-small: routed experts kept in each layer (adaptive pruning). Empty means
|
|
// `experts` in every layer; otherwise one entry per layer, each <= `experts`.
|
|
std::vector<int64_t> experts_per_layer;
|
|
bool round_bf16_activations = false;
|
|
bool fused_moe = false;
|
|
ggml_type kv_cache_type = GGML_TYPE_F16;
|
|
int64_t kv_cache_alignment = 1;
|
|
int64_t decode_cache_steps = 2048;
|
|
int64_t cpu_expert_layers = 0;
|
|
bool disk_experts = false;
|
|
bool pin_experts = false;
|
|
int64_t expert_cache_size = 32;
|
|
int64_t disk_expert_read_threads = 1;
|
|
size_t prefill_cache_bytes = 256u * 1024u * 1024u;
|
|
bool profile_prefill_cache = false;
|
|
};
|
|
|
|
struct Drama3BackboneDecoderForwardResult {
|
|
std::vector<float> hidden;
|
|
int64_t steps = 0;
|
|
};
|
|
|
|
class Drama3BackboneDecoderRuntime {
|
|
public:
|
|
Drama3BackboneDecoderRuntime(
|
|
std::shared_ptr<const engine::assets::TensorSource> weights,
|
|
Drama3BackboneDecoderConfig config,
|
|
engine::core::ExecutionContext & execution,
|
|
size_t graph_arena_bytes,
|
|
size_t weight_context_bytes,
|
|
engine::assets::TensorStorageType storage_type);
|
|
~Drama3BackboneDecoderRuntime();
|
|
|
|
Drama3BackboneDecoderRuntime(const Drama3BackboneDecoderRuntime &) = delete;
|
|
Drama3BackboneDecoderRuntime & operator=(const Drama3BackboneDecoderRuntime &) = delete;
|
|
|
|
std::vector<float> token_embedding(const std::vector<int32_t> & token_ids);
|
|
Drama3BackboneDecoderForwardResult forward_embeddings(const std::vector<float> & embeddings, int64_t steps);
|
|
std::vector<float> lm_head(const std::vector<float> & hidden);
|
|
void release_prefill_graphs();
|
|
void release_graphs();
|
|
|
|
class DecodeSession {
|
|
public:
|
|
virtual ~DecodeSession() = default;
|
|
virtual void reset(bool invalidate_experts = false) = 0;
|
|
virtual Drama3BackboneDecoderForwardResult prefill_embeddings(
|
|
const std::vector<float> & embeddings,
|
|
int64_t steps) = 0;
|
|
virtual std::vector<float> run_embedding_step(const std::vector<float> & embedding) = 0;
|
|
virtual int64_t valid_steps() const noexcept = 0;
|
|
};
|
|
|
|
std::unique_ptr<DecodeSession> create_decode_session(int64_t cache_steps);
|
|
|
|
private:
|
|
class Impl;
|
|
std::unique_ptr<Impl> impl_;
|
|
};
|
|
|
|
} // namespace engine::models::fish_drama3
|