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
54 lines
2.5 KiB
C++
54 lines
2.5 KiB
C++
#pragma once
|
|
|
|
// LFM3-Audio backbone: chunked prefill, one decode step per text token or
|
|
// audio frame, the KV and conv state between them, the tied text head and the
|
|
// audio-frame input embedding.
|
|
|
|
#include "engine/community_models/lfm3_audio/assets.h"
|
|
#include "engine/community_models/lfm3_audio/encoder.h"
|
|
#include "engine/community_models/lfm3_audio/prompt.h"
|
|
#include "engine/community_models/lfm3_audio/weights.h"
|
|
#include "engine/framework/assets/tensor_source.h"
|
|
#include "engine/framework/core/execution_context.h"
|
|
|
|
#include <cstdint>
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
namespace engine::community_models::lfm3_audio {
|
|
|
|
enum class Lfm3Output { Logits, Hidden }; // Logits = [vocab_size], Hidden = final-normed [2048]
|
|
|
|
class Lfm3BackboneRuntime {
|
|
public:
|
|
static constexpr int64_t kPrefillBlock = 256;
|
|
|
|
Lfm3BackboneRuntime(std::shared_ptr<const assets::TensorSource> model, const Lfm3BackboneConfig & config,
|
|
int64_t codebooks, int64_t audio_vocab_size, core::ExecutionContext & execution,
|
|
const Lfm3WeightOptions & options);
|
|
~Lfm3BackboneRuntime();
|
|
|
|
// Resets state, ensures cache capacity L + max_new_steps rounded up to 256, fills <|audio|> rows with `audio` and
|
|
// <|audio_out|> rows with prompt.frames' embeddings in order, and returns the output for the last position.
|
|
// `taps` (tests/dumps only): final-normed hidden rows at the given absolute positions, row-major [P][2048].
|
|
struct Taps {
|
|
std::vector<int64_t> positions;
|
|
std::vector<float> rows;
|
|
};
|
|
std::vector<float> prefill(const Lfm3Prompt & prompt, int32_t audio_id, int32_t audio_out_id,
|
|
const Lfm3AudioEmbeddings & audio, int64_t max_new_steps, Lfm3Output output,
|
|
Taps * taps = nullptr);
|
|
std::vector<float> step_text(int32_t token, Lfm3Output output);
|
|
std::vector<float> step_frame(const std::vector<int32_t> & codes, Lfm3Output output); // sum_k codebook_embedding[k*2017+c_k]
|
|
std::vector<float> logits(const std::vector<float> & hidden); // tied head on a given hidden (prefill dumps / tests)
|
|
// Input embeddings of audio frames, row-major [frames][hidden], summed exactly as prefill and step_frame sum them.
|
|
std::vector<float> embed_frames(const std::vector<std::vector<int32_t>> & frames);
|
|
[[nodiscard]] int64_t position() const noexcept; // next position
|
|
[[nodiscard]] int64_t capacity() const noexcept;
|
|
|
|
private:
|
|
struct Impl;
|
|
std::unique_ptr<Impl> impl_;
|
|
};
|
|
|
|
} // namespace engine::community_models::lfm3_audio
|