From 43a4c30596401e7dd715e5cbc1fd79b549e3edcd Mon Sep 17 00:00:00 2001 From: abeiro Date: Sun, 28 Jun 2026 21:20:02 +0200 Subject: [PATCH 01/27] Change size_t to std::size_t in noise.h Change size_t to std::size_t in noise.h, fill Debian compilation issue --- include/engine/framework/sampling/noise.h | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/engine/framework/sampling/noise.h b/include/engine/framework/sampling/noise.h index 93d1d672..2fab54f5 100644 --- a/include/engine/framework/sampling/noise.h +++ b/include/engine/framework/sampling/noise.h @@ -5,7 +5,7 @@ namespace engine::sampling { -std::vector generate_normal_noise(size_t count, uint32_t seed, float scale = 1.0F); +std::vector generate_normal_noise(std::size_t count, uint32_t seed, float scale = 1.0F); void clamp_noise(std::vector & noise, float min_value, float max_value); } // namespace engine::sampling From 2e5c5e568b252241d39d7d521d10c462ac8f442d Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Fri, 17 Jul 2026 20:45:08 -0400 Subject: [PATCH 02/27] Add Higgs Audio TTS model --- CMakeLists.txt | 9 + include/engine/models/higgs_tts/ar.h | 167 ++ include/engine/models/higgs_tts/assets.h | 65 + include/engine/models/higgs_tts/codebooks.h | 24 + include/engine/models/higgs_tts/codec.h | 136 ++ include/engine/models/higgs_tts/generator.h | 76 + include/engine/models/higgs_tts/loader.h | 33 + include/engine/models/higgs_tts/sampler.h | 54 + include/engine/models/higgs_tts/session.h | 64 + .../engine/models/higgs_tts/tokenizer_text.h | 37 + model_specs/higgs_tts.json | 36 + src/framework/runtime/registry.cpp | 4 +- src/models/higgs_tts/ar.cpp | 1310 +++++++++++++ src/models/higgs_tts/assets.cpp | 171 ++ src/models/higgs_tts/codebooks.cpp | 79 + src/models/higgs_tts/codec.cpp | 1614 +++++++++++++++++ src/models/higgs_tts/generator.cpp | 503 +++++ src/models/higgs_tts/loader.cpp | 146 ++ src/models/higgs_tts/sampler.cpp | 458 +++++ src/models/higgs_tts/session.cpp | 297 +++ src/models/higgs_tts/tokenizer_text.cpp | 94 + 21 files changed, 5375 insertions(+), 2 deletions(-) create mode 100644 include/engine/models/higgs_tts/ar.h create mode 100644 include/engine/models/higgs_tts/assets.h create mode 100644 include/engine/models/higgs_tts/codebooks.h create mode 100644 include/engine/models/higgs_tts/codec.h create mode 100644 include/engine/models/higgs_tts/generator.h create mode 100644 include/engine/models/higgs_tts/loader.h create mode 100644 include/engine/models/higgs_tts/sampler.h create mode 100644 include/engine/models/higgs_tts/session.h create mode 100644 include/engine/models/higgs_tts/tokenizer_text.h create mode 100644 model_specs/higgs_tts.json create mode 100644 src/models/higgs_tts/ar.cpp create mode 100644 src/models/higgs_tts/assets.cpp create mode 100644 src/models/higgs_tts/codebooks.cpp create mode 100644 src/models/higgs_tts/codec.cpp create mode 100644 src/models/higgs_tts/generator.cpp create mode 100644 src/models/higgs_tts/loader.cpp create mode 100644 src/models/higgs_tts/sampler.cpp create mode 100644 src/models/higgs_tts/session.cpp create mode 100644 src/models/higgs_tts/tokenizer_text.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index c678c463..5ca2b830 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -356,6 +356,15 @@ add_library(engine_runtime STATIC src/models/higgs_audio_stt/postprocess.cpp src/models/higgs_audio_stt/session.cpp src/models/higgs_audio_stt/loader.cpp + src/models/higgs_tts/ar.cpp + src/models/higgs_tts/assets.cpp + src/models/higgs_tts/codec.cpp + src/models/higgs_tts/codebooks.cpp + src/models/higgs_tts/generator.cpp + src/models/higgs_tts/loader.cpp + src/models/higgs_tts/sampler.cpp + src/models/higgs_tts/session.cpp + src/models/higgs_tts/tokenizer_text.cpp src/models/irodori_tts/assets.cpp src/models/irodori_tts/codec.cpp src/models/irodori_tts/condition_encoder.cpp diff --git a/include/engine/models/higgs_tts/ar.h b/include/engine/models/higgs_tts/ar.h new file mode 100644 index 00000000..73270860 --- /dev/null +++ b/include/engine/models/higgs_tts/ar.h @@ -0,0 +1,167 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/core/module.h" +#include "engine/framework/modules/attention/qwen_decoder.h" +#include "engine/framework/runtime/kv_cache.h" +#include "engine/models/higgs_tts/assets.h" + +#include +#include +#include +#include + +namespace engine::core { +class BackendWeightStore; +} + +namespace engine::models::higgs_tts { + +struct HiggsQwenDecoderStackWeights { + std::vector layers; +}; + +struct HiggsARWeights { + std::shared_ptr store; + core::TensorValue text_embedding; + core::TensorValue modality_embedding; + HiggsQwenDecoderStackWeights decoder; + core::TensorValue norm; +}; + +HiggsARWeights load_higgs_ar_weights( + const HiggsAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + +class HiggsARRuntime { +public: + HiggsARRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + + const HiggsAssets & assets() const noexcept; + const HiggsARWeights & weights() const noexcept; + ggml_backend_t backend() const noexcept; + core::BackendType backend_type() const noexcept; + int device() const noexcept; + int threads() const noexcept; + +private: + std::shared_ptr assets_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int device_ = 0; + int threads_ = 1; + std::shared_ptr weights_; +}; + +struct HiggsARDecodeInput { + std::vector last_codes; + bool use_last_codes = false; +}; + +struct HiggsARDecodeOutput { + std::vector codebook_logits; +}; + +struct HiggsARDecodeTiming { + double input_upload_ms = 0.0; + double mask_upload_ms = 0.0; + double graph_compute_ms = 0.0; + double output_read_ms = 0.0; + int64_t steps = 0; + + void add(const HiggsARDecodeTiming & other) noexcept; +}; + +struct HiggsARPrefillInput { + std::vector text_tokens; + std::vector fused_code_ids; + std::vector text_gate; + std::vector code_gate; + int64_t steps = 0; +}; + +struct HiggsARPrefillOutput { + HiggsARDecodeOutput output; + runtime::TransformerKVState kv_state; + bool wrote_cache = false; +}; + +class HiggsARKVCache { +public: + HiggsARKVCache(std::shared_ptr runtime, int64_t cache_steps); + ~HiggsARKVCache(); + + HiggsARKVCache(const HiggsARKVCache &) = delete; + HiggsARKVCache & operator=(const HiggsARKVCache &) = delete; + + bool can_run(const HiggsARRuntime & runtime, int64_t required_steps) const; + int64_t cache_steps() const; + int64_t valid_steps() const; + int64_t current_end() const; + void reset(); + void retain_prefix(int64_t prefix_steps); + void import_state(const runtime::TransformerKVState & state); + runtime::TransformerKVState export_state() const; + void advance_after_direct_append(int64_t steps); + const core::TensorValue & key_tensor(size_t layer) const; + const core::TensorValue & value_tensor(size_t layer) const; + +private: + struct Impl; + std::unique_ptr impl_; +}; + +class HiggsARPrefillGraph { +public: + HiggsARPrefillGraph( + std::shared_ptr runtime, + int64_t prompt_steps, + int64_t start_step, + HiggsARKVCache * cache, + size_t graph_arena_bytes); + ~HiggsARPrefillGraph(); + + HiggsARPrefillGraph(const HiggsARPrefillGraph &) = delete; + HiggsARPrefillGraph & operator=(const HiggsARPrefillGraph &) = delete; + + bool matches(const HiggsARRuntime & runtime, int64_t prompt_steps, int64_t start_step) const; + HiggsARPrefillOutput run(const HiggsARPrefillInput & input, int64_t start_step = 0); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +class HiggsARDecodeGraph { +public: + HiggsARDecodeGraph( + std::shared_ptr runtime, + int64_t cache_steps, + HiggsARKVCache & cache, + size_t graph_arena_bytes); + ~HiggsARDecodeGraph(); + + HiggsARDecodeGraph(const HiggsARDecodeGraph &) = delete; + HiggsARDecodeGraph & operator=(const HiggsARDecodeGraph &) = delete; + + bool can_run(const HiggsARRuntime & runtime, int64_t required_steps) const; + int64_t cache_steps() const; + void import_prefill_state(const runtime::TransformerKVState & state); + void begin_decode_run(); + HiggsARDecodeTiming timing() const; + void run_step_into(const HiggsARDecodeInput & input, HiggsARDecodeOutput & output, bool log_timing = false); + +private: + struct Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::higgs_tts diff --git a/include/engine/models/higgs_tts/assets.h b/include/engine/models/higgs_tts/assets.h new file mode 100644 index 00000000..a6899778 --- /dev/null +++ b/include/engine/models/higgs_tts/assets.h @@ -0,0 +1,65 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" + +#include +#include +#include +#include + +namespace engine::assets { +class TensorSource; +} + +namespace engine::models::higgs_tts { + +struct HiggsTextConfig { + std::string model_type; + int64_t vocab_size = 0; + int64_t hidden_size = 0; + int64_t intermediate_size = 0; + int64_t num_hidden_layers = 0; + int64_t num_attention_heads = 0; + int64_t num_key_value_heads = 0; + int64_t head_dim = 0; + int64_t max_position_embeddings = 0; + int64_t bos_token_id = 0; + int64_t eos_token_id = 0; + int64_t pad_token_id = -1; + float rms_norm_eps = 1.0e-6F; + float rope_theta = 1000000.0F; + bool tie_word_embeddings = true; +}; + +struct HiggsAudioEncoderConfig { + std::string model_type; + std::string encoder_type; + int64_t num_codebooks = 0; + int64_t vocab_size = 0; + int64_t out_dim = 0; + int64_t mel_per_sample = 0; + int64_t max_chunk_size = 0; + bool tie_word_embeddings = true; + bool use_delay_pattern = true; +}; + +struct HiggsConfig { + std::string model_type; + std::string architecture; + int64_t hidden_size = 0; + int64_t vocab_size = 0; + int64_t audio_token_id = -100; + int64_t ignore_index = -100; + HiggsTextConfig text; + HiggsAudioEncoderConfig audio; +}; + +struct HiggsAssets { + assets::ResourceBundle resources; + HiggsConfig config; + std::shared_ptr weights; +}; + +std::shared_ptr load_higgs_assets(const std::filesystem::path & model_path); + +} // namespace engine::models::higgs_tts diff --git a/include/engine/models/higgs_tts/codebooks.h b/include/engine/models/higgs_tts/codebooks.h new file mode 100644 index 00000000..5209b8a1 --- /dev/null +++ b/include/engine/models/higgs_tts/codebooks.h @@ -0,0 +1,24 @@ +#pragma once + +#include +#include + +namespace engine::models::higgs_tts { + +constexpr int32_t kHiggsBocId = 1024; +constexpr int32_t kHiggsEocId = 1025; +constexpr int32_t kHiggsStopCode = -1; + +int64_t higgs_delayed_frame_count(int64_t raw_frames, int64_t codebooks); + +std::vector apply_higgs_delay_pattern( + const std::vector & raw_codes, + int64_t raw_frames, + int64_t codebooks); + +std::vector reverse_higgs_delay_pattern( + const std::vector & delayed_codes, + int64_t delayed_frames, + int64_t codebooks); + +} // namespace engine::models::higgs_tts diff --git a/include/engine/models/higgs_tts/codec.h b/include/engine/models/higgs_tts/codec.h new file mode 100644 index 00000000..8e160e91 --- /dev/null +++ b/include/engine/models/higgs_tts/codec.h @@ -0,0 +1,136 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/core/module.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/runtime/session.h" +#include "engine/models/higgs_tts/assets.h" + +#include +#include +#include +#include +#include + +namespace engine::core { +class BackendWeightStore; +} + +namespace engine::models::higgs_tts { + +class HiggsCodecDecodeGraph; +class HiggsCodecEncodeGraph; + +struct HiggsCodecVectorQuantizerWeights { + core::TensorValue codebook; + modules::LinearWeights project_in; + modules::LinearWeights project_out; +}; + +struct HiggsCodecResidualUnitWeights { + modules::Snake1dWeights snake1; + modules::Conv1dWeights conv1; + modules::Snake1dWeights snake2; + modules::Conv1dWeights conv2; +}; + +struct HiggsCodecDecoderBlockWeights { + modules::Snake1dWeights snake; + modules::ConvTranspose1dWeights conv_transpose; + std::vector residual_units; +}; + +struct HiggsCodecEncoderBlockWeights { + modules::Snake1dWeights snake; + modules::Conv1dWeights conv; + std::vector residual_units; +}; + +struct HiggsCodecSemanticResidualUnitWeights { + modules::Conv1dWeights conv1; + modules::Conv1dWeights conv2; +}; + +struct HiggsCodecSemanticEncoderBlockWeights { + std::vector residual_units; + modules::Conv1dWeights conv; +}; + +struct HiggsCodecWeights { + std::shared_ptr store; + std::unordered_map semantic_model; + std::vector quantizers; + modules::Conv1dWeights acoustic_encoder_input; + std::vector acoustic_encoder_blocks; + modules::Snake1dWeights acoustic_encoder_output_snake; + modules::Conv1dWeights acoustic_encoder_output; + modules::Conv1dWeights semantic_encoder_input; + std::vector semantic_encoder_blocks; + modules::LinearWeights codec_project; + modules::LinearWeights acoustic_project; + modules::Conv1dWeights acoustic_decoder_input; + std::vector acoustic_decoder_blocks; + modules::Snake1dWeights acoustic_decoder_output_snake; + modules::Conv1dWeights acoustic_decoder_output; +}; + +struct HiggsCodecDecodeOutput { + int sample_rate = 24000; + int channels = 1; + int64_t samples = 0; + std::vector values; +}; + +struct HiggsCodecEncodeOutput { + std::vector codes; + int64_t frames = 0; + int64_t codebooks = 0; +}; + +class HiggsCodecRuntime { +public: + HiggsCodecRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + size_t decode_graph_arena_bytes, + size_t encode_graph_arena_bytes, + assets::TensorStorageType weight_storage_type); + ~HiggsCodecRuntime(); + + const HiggsCodecWeights & weights() const noexcept; + ggml_backend_t backend() const noexcept; + core::BackendType backend_type() const noexcept; + int threads() const noexcept; + size_t decode_graph_arena_bytes() const noexcept; + size_t encode_graph_arena_bytes() const noexcept; + + HiggsCodecEncodeOutput encode_reference(const runtime::AudioBuffer & audio) const; + HiggsCodecDecodeOutput decode_codes( + const std::vector & codes, + int64_t frames, + int64_t codebooks) const; + +private: + std::shared_ptr assets_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + size_t decode_graph_arena_bytes_ = 0; + size_t encode_graph_arena_bytes_ = 0; + std::shared_ptr weights_; + mutable std::unique_ptr encode_graph_; + mutable std::unique_ptr decode_graph_; +}; + +HiggsCodecWeights load_higgs_codec_decode_weights( + const HiggsAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + +} // namespace engine::models::higgs_tts diff --git a/include/engine/models/higgs_tts/generator.h b/include/engine/models/higgs_tts/generator.h new file mode 100644 index 00000000..2d3733cf --- /dev/null +++ b/include/engine/models/higgs_tts/generator.h @@ -0,0 +1,76 @@ +#pragma once + +#include "engine/models/higgs_tts/ar.h" +#include "engine/models/higgs_tts/codec.h" +#include "engine/models/higgs_tts/sampler.h" +#include "engine/models/higgs_tts/tokenizer_text.h" + +#include +#include +#include +#include +#include + +namespace engine::models::higgs_tts { + +struct HiggsGenerationOptions { + int64_t max_tokens = 1024; + float temperature = 1.0F; + std::optional top_p; + std::optional top_k; + float repetition_penalty = 1.0F; + std::optional seed; +}; + +struct HiggsGenerationRequest { + std::string text; + std::string reference_text; + std::vector reference_codes; + int64_t reference_frames = 0; + int64_t reference_codebooks = 0; + HiggsGenerationOptions options; +}; + +struct HiggsGenerationResult { + HiggsCodecDecodeOutput audio; + std::vector delayed_codes; + int64_t delayed_frames = 0; + std::vector raw_codes; + int64_t raw_frames = 0; +}; + +class HiggsGenerator { +public: + HiggsGenerator(std::shared_ptr assets, + std::shared_ptr ar, + std::shared_ptr codec, + size_t ar_decode_graph_arena_bytes); + + void prepare(const HiggsGenerationRequest & request); + HiggsGenerationResult generate(const HiggsGenerationRequest & request); + +private: + struct ReferencePrefixCache { + std::string reference_text; + std::vector reference_codes; + int64_t reference_frames = 0; + int64_t reference_codebooks = 0; + std::vector delayed_reference_codes; + int64_t delayed_reference_frames = 0; + std::vector prefix_tokens; + int64_t prefix_steps = 0; + }; + + std::shared_ptr assets_; + std::shared_ptr ar_; + std::shared_ptr codec_; + HiggsTextTokenizer tokenizer_; + size_t ar_decode_graph_arena_bytes_ = 0; + std::optional reference_prefix_cache_; + std::optional cuda_sampling_policy_; + std::unique_ptr ar_kv_cache_; + std::unique_ptr prefill_graph_; + std::unique_ptr decode_graph_; +}; + +} // namespace engine::models::higgs_tts diff --git a/include/engine/models/higgs_tts/loader.h b/include/engine/models/higgs_tts/loader.h new file mode 100644 index 00000000..6f17c809 --- /dev/null +++ b/include/engine/models/higgs_tts/loader.h @@ -0,0 +1,33 @@ +#pragma once + +#include "engine/framework/runtime/model.h" +#include "engine/models/higgs_tts/assets.h" + +#include +#include + +namespace engine::models::higgs_tts { + +class HiggsTTSLoadedModel final : public runtime::ILoadedVoiceModel { +public: + HiggsTTSLoadedModel( + runtime::ModelMetadata metadata, + runtime::CapabilitySet capabilities, + std::shared_ptr assets); + + const runtime::ModelMetadata & metadata() const noexcept override; + const runtime::CapabilitySet & capabilities() const noexcept override; + std::unique_ptr create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const override; + +private: + runtime::ModelMetadata metadata_; + runtime::CapabilitySet capabilities_; + std::shared_ptr assets_; +}; + +std::unique_ptr load_higgs_tts_model(const std::filesystem::path & model_path); +std::shared_ptr make_higgs_tts_loader(); + +} // namespace engine::models::higgs_tts diff --git a/include/engine/models/higgs_tts/sampler.h b/include/engine/models/higgs_tts/sampler.h new file mode 100644 index 00000000..5d40a2a2 --- /dev/null +++ b/include/engine/models/higgs_tts/sampler.h @@ -0,0 +1,54 @@ +#pragma once + +#include "engine/framework/sampling/torch_random.h" +#include "engine/models/higgs_tts/codebooks.h" + +#include +#include +#include + +namespace engine::models::higgs_tts { + +constexpr int64_t kHiggsMaxTopK = 1026; + +using HiggsCudaSamplingPolicy = engine::sampling::TorchCudaSamplingPolicy; + +struct HiggsSamplingOptions { + float temperature = 1.0F; + std::optional top_p; + std::optional top_k; + bool has_seed = false; + uint64_t seed = 0; + HiggsCudaSamplingPolicy cuda_policy; +}; + +struct HiggsSamplerState { + int64_t num_codebooks = 0; + int64_t delay_count = 0; + std::optional eoc_countdown; + bool generation_done = false; + int64_t step_count = 0; + std::vector last_codes; +}; + +class HiggsCodebookSampler { +public: + explicit HiggsCodebookSampler(int64_t num_codebooks, int64_t codebook_vocab_size); + + HiggsSamplerState make_state() const; + const std::vector & step(const float * logits, + int64_t logits_count, + HiggsSamplerState & state, + HiggsSamplingOptions & options); + +private: + int64_t num_codebooks_ = 0; + int64_t codebook_vocab_size_ = 0; + std::vector scratch_scores_; + std::vector scratch_probs_; + std::vector scratch_order_; + std::vector scratch_kept_; + std::vector scratch_codes_; +}; + +} // namespace engine::models::higgs_tts diff --git a/include/engine/models/higgs_tts/session.h b/include/engine/models/higgs_tts/session.h new file mode 100644 index 00000000..0898f452 --- /dev/null +++ b/include/engine/models/higgs_tts/session.h @@ -0,0 +1,64 @@ +#pragma once + +#include "engine/framework/runtime/session_base.h" +#include "engine/models/higgs_tts/assets.h" +#include "engine/models/higgs_tts/ar.h" +#include "engine/models/higgs_tts/codec.h" +#include "engine/models/higgs_tts/generator.h" + +#include +#include +#include +#include +#include + +namespace engine::models::higgs_tts { + +class HiggsTTSSession final + : public runtime::RuntimeSessionBase + , public runtime::IOfflineVoiceTaskSession { +public: + HiggsTTSSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets); + + std::string family() const override; + runtime::VoiceTaskKind task_kind() const override; + runtime::RunMode run_mode() const override; + void prepare(const runtime::SessionPreparationRequest & request) override; + runtime::TaskResult run(const runtime::TaskRequest & request) override; + +private: + struct ReferenceCacheEntry { + std::string reference_text; + int sample_rate = 0; + int channels = 0; + uint64_t sample_count = 0; + uint64_t sample_hash = 0; + HiggsCodecEncodeOutput codes; + }; + + HiggsGenerationRequest make_generation_request( + const runtime::TaskRequest & request, + const HiggsCodecEncodeOutput * resolved_reference_codes = nullptr); + const HiggsCodecEncodeOutput & resolve_reference_codes( + const runtime::AudioBuffer & audio, + const std::string & reference_text); + + runtime::TaskSpec task_; + std::shared_ptr assets_; + size_t ar_weight_context_bytes_ = 4096ull * 1024ull * 1024ull; + size_t codec_weight_context_bytes_ = 1536ull * 1024ull * 1024ull; + size_t ar_decode_graph_arena_bytes_ = 512ull * 1024ull * 1024ull; + size_t codec_decode_graph_arena_bytes_ = 128ull * 1024ull * 1024ull; + size_t codec_encode_graph_arena_bytes_ = 256ull * 1024ull * 1024ull; + assets::TensorStorageType ar_weight_storage_type_ = assets::TensorStorageType::Native; + assets::TensorStorageType codec_weight_storage_type_ = assets::TensorStorageType::Native; + std::shared_ptr ar_; + std::shared_ptr codec_; + std::unique_ptr generator_; + std::optional reference_cache_; +}; + +} // namespace engine::models::higgs_tts diff --git a/include/engine/models/higgs_tts/tokenizer_text.h b/include/engine/models/higgs_tts/tokenizer_text.h new file mode 100644 index 00000000..b43f99f1 --- /dev/null +++ b/include/engine/models/higgs_tts/tokenizer_text.h @@ -0,0 +1,37 @@ +#pragma once + +#include "engine/models/higgs_tts/assets.h" + +#include +#include +#include +#include + +namespace engine::models::higgs_tts { + +struct HiggsPromptRequest { + std::string text; + std::string reference_text; + int64_t delayed_reference_tokens = 0; +}; + +struct HiggsPromptEncoding { + std::vector token_ids; + std::vector text_ids; + std::vector reference_text_ids; +}; + +class HiggsTextTokenizer { +public: + struct Impl; + + explicit HiggsTextTokenizer(std::shared_ptr assets); + + std::vector encode(const std::string & text) const; + HiggsPromptEncoding encode_prompt(const HiggsPromptRequest & request) const; + +private: + std::shared_ptr impl_; +}; + +} // namespace engine::models::higgs_tts diff --git a/model_specs/higgs_tts.json b/model_specs/higgs_tts.json new file mode 100644 index 00000000..a28cf9cd --- /dev/null +++ b/model_specs/higgs_tts.json @@ -0,0 +1,36 @@ +{ + "family": "higgs_tts", + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "weights": "$gguf" + }, + "files": { + "config": "model:config.json", + "tokenizer_json": "model:tokenizer.json", + "tokenizer_config": "model:tokenizer_config.json", + "chat_template": "model:chat_template.jinja" + }, + "tensors": { + "weights": "weights:" + } + }, + { + "format": "safetensors", + "roots": { + "model": "." + }, + "files": { + "config": "model:config.json", + "tokenizer_json": "model:tokenizer.json", + "tokenizer_config": "model:tokenizer_config.json", + "chat_template": "model:chat_template.jinja" + }, + "tensors": { + "weights": "model:model.safetensors.index.json" + } + } + ] +} diff --git a/src/framework/runtime/registry.cpp b/src/framework/runtime/registry.cpp index 2818503f..34fadb36 100644 --- a/src/framework/runtime/registry.cpp +++ b/src/framework/runtime/registry.cpp @@ -5,7 +5,6 @@ #include "engine/framework/io/config.h" #include "engine/framework/io/filesystem.h" // Development registry entries from Share/AudioCPP that are not present in this release tree yet: -// #include "engine/models/higgs_tts/loader.h" // #include "engine/models/kokoro_tts/loader.h" // #include "engine/models/parakeet_tdt/loader.h" #include "engine/models/ace_step/loader.h" @@ -14,6 +13,7 @@ #include "engine/models/demucs/loader.h" #include "engine/models/heartmula/loader.h" #include "engine/models/higgs_audio_stt/loader.h" +#include "engine/models/higgs_tts/loader.h" #include "engine/models/hviske_asr/loader.h" #include "engine/models/index_tts2/loader.h" #include "engine/models/irodori_tts/loader.h" @@ -226,7 +226,6 @@ ModelRegistry make_default_registry(const std::optional & const std::vector> available_loaders = { // Development registry entries from Share/AudioCPP that are not present in this release tree yet: // engine::models::kokoro_tts::make_kokoro_tts_loader(), - // engine::models::higgs_tts::make_higgs_tts_loader(), // engine::models::parakeet_tdt::make_parakeet_tdt_loader(), engine::models::ace_step::make_ace_step_loader(), engine::models::demucs::make_htdemucs_loader(), @@ -241,6 +240,7 @@ ModelRegistry make_default_registry(const std::optional & engine::models::vibevoice_asr::make_vibevoice_asr_loader(), engine::models::heartmula::make_heartmula_loader(), engine::models::higgs_audio_stt::make_higgs_audio_stt_loader(), + engine::models::higgs_tts::make_higgs_tts_loader(), engine::models::hviske_asr::make_hviske_asr_loader(), engine::models::irodori_tts::make_irodori_tts_loader(), engine::models::nemotron_asr::make_nemotron_asr_loader(), diff --git a/src/models/higgs_tts/ar.cpp b/src/models/higgs_tts/ar.cpp new file mode 100644 index 00000000..6ffe2c4c --- /dev/null +++ b/src/models/higgs_tts/ar.cpp @@ -0,0 +1,1310 @@ +#include "engine/models/higgs_tts/ar.h" + +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/attention/qwen_causal_decoder.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/optimizations/fast_projection_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::higgs_tts { +namespace { + +namespace modules = engine::modules; +using Clock = std::chrono::steady_clock; +constexpr int64_t kLayerwisePrefillMinSteps = 2048; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +modules::QwenDecoderStackConfig make_higgs_qwen_stack_config(const HiggsTextConfig & config) { + modules::QwenDecoderStackConfig out; + out.hidden_size = config.hidden_size; + out.num_attention_heads = config.num_attention_heads; + out.num_key_value_heads = config.num_key_value_heads; + out.head_dim = config.head_dim; + out.intermediate_size = config.intermediate_size; + out.layers = config.num_hidden_layers; + out.rms_norm_eps = config.rms_norm_eps; + out.rope_theta = config.rope_theta; + out.attention_precision = GGML_PREC_F32; + out.projection_precision = GGML_PREC_DEFAULT; + out.qkv_layout = modules::QwenDecoderQKVLayout::Separate; + out.use_qk_norm = true; + out.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.runtime.static_cache.transpose_context = false; + return out; +} + +class HiggsQwenDecoderComponent { +public: + explicit HiggsQwenDecoderComponent(const HiggsTextConfig & config) + : stack_config_(make_higgs_qwen_stack_config(config)), + layer_config_(modules::qwen_decoder_layer_config_from_stack(stack_config_)), + layer_module_(layer_config_) {} + + modules::QwenDecoderLayerOutputs build_prefill_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const modules::QwenDecoderLayerWeights & weights, + const core::TensorValue & attention_mask) const { + return layer_module_.build(ctx, input, positions, weights, std::nullopt, std::nullopt, attention_mask); + } + + modules::QwenDecoderLayerOutputs build_decode_layer( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const core::TensorValue & input, + const core::TensorValue & positions, + const modules::QwenDecoderLayerWeights & weights, + const core::TensorValue & cache_key, + const core::TensorValue & cache_value, + const core::TensorValue & cache_slot, + const core::TensorValue & attention_mask) const { + return layer_module_.build_with_static_cache_tail( + ctx, + graph, + input, + positions, + weights, + cache_key, + cache_value, + cache_slot, + attention_mask); + } + +private: + modules::QwenDecoderStackConfig stack_config_; + modules::QwenDecoderLayerConfig layer_config_; + modules::QwenDecoderLayerModule layer_module_; +}; + +core::TensorValue higgs_cache_view( + core::ModuleBuildContext & ctx, + const core::TensorValue & cache, + int64_t start, + int64_t steps, + int64_t heads, + int64_t head_dim) { + if (start < 0 || steps <= 0 || start + steps > cache.shape.dims[1]) { + throw std::runtime_error("Higgs TTS AR cache view range is invalid"); + } + return core::wrap_tensor( + ggml_view_4d( + ctx.ggml, + cache.tensor, + head_dim, + heads, + steps, + 1, + cache.tensor->nb[1], + cache.tensor->nb[2], + cache.tensor->nb[3], + static_cast(start) * cache.tensor->nb[2]), + core::TensorShape::from_dims({1, steps, heads, head_dim}), + cache.type); +} + +modules::QwenDecoderLayerWeights load_layer_weights( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const HiggsTextConfig & config, + int64_t layer_index, + assets::TensorStorageType storage_type) { + const std::string prefix = "body.layers." + std::to_string(layer_index); + const int64_t q_out = config.num_attention_heads * config.head_dim; + const int64_t kv_out = config.num_key_value_heads * config.head_dim; + modules::QwenDecoderLayerWeights weights; + weights.input_norm = { + store.load_f32_tensor(source, prefix + ".input_layernorm.weight", {config.hidden_size}), + std::nullopt, + }; + // Keep Q/K/V separate here so Higgs exercises the framework Qwen decoder path. + // A packed-QKV fast path can be evaluated later as a framework-level optimization. + weights.self_attention.q_weight = store.load_tensor( + source, + prefix + ".self_attn.q_proj.weight", + storage_type, + {q_out, config.hidden_size}); + weights.self_attention.k_weight = store.load_tensor( + source, + prefix + ".self_attn.k_proj.weight", + storage_type, + {kv_out, config.hidden_size}); + weights.self_attention.v_weight = store.load_tensor( + source, + prefix + ".self_attn.v_proj.weight", + storage_type, + {kv_out, config.hidden_size}); + weights.self_attention.out_weight = store.load_tensor( + source, + prefix + ".self_attn.o_proj.weight", + storage_type, + {config.hidden_size, q_out}); + weights.q_norm = { + store.load_f32_tensor(source, prefix + ".self_attn.q_norm.weight", {config.head_dim}), + std::nullopt, + }; + weights.k_norm = { + store.load_f32_tensor(source, prefix + ".self_attn.k_norm.weight", {config.head_dim}), + std::nullopt, + }; + weights.post_norm = { + store.load_f32_tensor(source, prefix + ".post_attention_layernorm.weight", {config.hidden_size}), + std::nullopt, + }; + weights.mlp.gate_proj = { + store.load_tensor( + source, + prefix + ".mlp.gate_proj.weight", + storage_type, + {config.intermediate_size, config.hidden_size}), + std::nullopt, + }; + weights.mlp.up_proj = { + store.load_tensor( + source, + prefix + ".mlp.up_proj.weight", + storage_type, + {config.intermediate_size, config.hidden_size}), + std::nullopt, + }; + weights.mlp.down_proj = { + store.load_tensor( + source, + prefix + ".mlp.down_proj.weight", + storage_type, + {config.hidden_size, config.intermediate_size}), + std::nullopt, + }; + return weights; +} + +HiggsQwenDecoderStackWeights load_decoder_weights( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const HiggsTextConfig & config, + assets::TensorStorageType storage_type) { + HiggsQwenDecoderStackWeights weights; + weights.layers.reserve(static_cast(config.num_hidden_layers)); + for (int64_t layer = 0; layer < config.num_hidden_layers; ++layer) { + weights.layers.push_back(load_layer_weights(store, source, config, layer, storage_type)); + } + return weights; +} + +core::TensorValue build_higgs_decode_code_embedding( + core::ModuleBuildContext & ctx, + const HiggsARWeights & weights, + const HiggsConfig & config, + ggml_tensor * fused_code_ids) { + auto code_ids = core::wrap_tensor( + fused_code_ids, + core::TensorShape::from_dims({config.audio.num_codebooks}), + GGML_TYPE_I32); + auto code = modules::EmbeddingModule( + {config.audio.num_codebooks * config.audio.vocab_size, config.text.hidden_size}) + .build(ctx, code_ids, weights.modality_embedding); + code = modules::ReduceSumModule({0}).build(ctx, code); + return core::reshape_tensor(ctx, code, core::TensorShape::from_dims({1, 1, config.text.hidden_size})); +} + +core::TensorValue build_higgs_prefill_input_embedding( + core::ModuleBuildContext & ctx, + const HiggsARWeights & weights, + const HiggsConfig & config, + ggml_tensor * text_tokens, + ggml_tensor * fused_code_ids, + ggml_tensor * text_gate, + ggml_tensor * code_gate, + int64_t steps) { + auto text_ids = core::wrap_tensor(text_tokens, core::TensorShape::from_dims({steps}), GGML_TYPE_I32); + auto text = modules::EmbeddingModule({config.text.vocab_size, config.text.hidden_size}) + .build(ctx, text_ids, weights.text_embedding); + text = core::reshape_tensor(ctx, text, core::TensorShape::from_dims({1, steps, config.text.hidden_size})); + + auto code_ids = core::wrap_tensor( + fused_code_ids, + core::TensorShape::from_dims({steps, config.audio.num_codebooks}), + GGML_TYPE_I32); + auto code = modules::EmbeddingModule( + {config.audio.num_codebooks * config.audio.vocab_size, config.text.hidden_size}) + .build(ctx, code_ids, weights.modality_embedding); + code = modules::ReduceSumModule({1}).build(ctx, code); + code = core::reshape_tensor(ctx, code, core::TensorShape::from_dims({1, steps, config.text.hidden_size})); + + auto text_gate_value = core::wrap_tensor( + text_gate, + core::TensorShape::from_dims({1, steps, 1}), + GGML_TYPE_F32); + auto code_gate_value = core::wrap_tensor( + code_gate, + core::TensorShape::from_dims({1, steps, 1}), + GGML_TYPE_F32); + text_gate_value = core::wrap_tensor( + ggml_repeat(ctx.ggml, text_gate_value.tensor, text.tensor), text.shape, GGML_TYPE_F32); + code_gate_value = core::wrap_tensor( + ggml_repeat(ctx.ggml, code_gate_value.tensor, code.tensor), code.shape, GGML_TYPE_F32); + return modules::AddModule{}.build( + ctx, + modules::MulModule{}.build(ctx, text, text_gate_value), + modules::MulModule{}.build(ctx, code, code_gate_value)); +} + +core::TensorValue build_modality_logits( + core::ModuleBuildContext & ctx, + const core::TensorValue & hidden, + const HiggsARWeights & weights, + const HiggsConfig & config) { + const int64_t out_features = config.audio.num_codebooks * config.audio.vocab_size; + const bool use_fast_projection = + ctx.backend_type == core::BackendType::Cuda && hidden.shape.rank == 3 && + hidden.shape.dims[1] == 1 && out_features % 4 == 0; + auto logits = + use_fast_projection + ? modules::FastPackedProjection4Module({config.text.hidden_size, out_features, GGML_PREC_DEFAULT}) + .build(ctx, hidden, {weights.modality_embedding, std::nullopt}) + : modules::LinearModule({config.text.hidden_size, out_features, false}) + .build(ctx, hidden, {weights.modality_embedding, std::nullopt}); + return core::reshape_tensor( + ctx, + logits, + core::TensorShape::from_dims({config.audio.num_codebooks, config.audio.vocab_size})); +} + +} // namespace + +HiggsARWeights load_higgs_ar_weights( + const HiggsAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) { + if (assets.weights == nullptr) { + throw std::runtime_error("Higgs TTS AR weights require tensor source"); + } + if (backend == nullptr) { + throw std::runtime_error("Higgs TTS AR backend is not initialized"); + } + const auto & config = assets.config; + const auto & source = *assets.weights; + HiggsARWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "higgs_tts.ar.weights", + weight_context_bytes); + weights.text_embedding = weights.store->load_tensor( + source, + "tied.embedding.text_embedding.weight", + weight_storage_type, + {config.text.vocab_size, config.text.hidden_size}); + weights.modality_embedding = weights.store->load_tensor( + source, + "tied.embedding.modality_embeddings.0.embedding.weight", + weight_storage_type, + {config.audio.num_codebooks * config.audio.vocab_size, config.text.hidden_size}); + weights.decoder = load_decoder_weights(*weights.store, source, config.text, weight_storage_type); + weights.norm = weights.store->load_f32_tensor(source, "body.norm.weight", {config.text.hidden_size}); + weights.store->upload(); + return weights; +} + +void HiggsARDecodeTiming::add(const HiggsARDecodeTiming & other) noexcept { + input_upload_ms += other.input_upload_ms; + mask_upload_ms += other.mask_upload_ms; + graph_compute_ms += other.graph_compute_ms; + output_read_ms += other.output_read_ms; + steps += other.steps; +} + +HiggsARRuntime::HiggsARRuntime( + std::shared_ptr assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : assets_(std::move(assets)), + backend_(execution.backend()), + backend_type_(execution.backend_type()), + device_(execution.config().device), + threads_(std::max(1, execution.config().threads)) { + if (assets_ == nullptr) { + throw std::runtime_error("Higgs TTS AR runtime requires assets"); + } + if (assets_->weights == nullptr) { + throw std::runtime_error("Higgs TTS AR runtime requires tensor source"); + } + weights_ = std::make_shared( + load_higgs_ar_weights(*assets_, backend_, backend_type_, weight_context_bytes, weight_storage_type)); +} + +const HiggsAssets & HiggsARRuntime::assets() const noexcept { + return *assets_; +} + +const HiggsARWeights & HiggsARRuntime::weights() const noexcept { + return *weights_; +} + +ggml_backend_t HiggsARRuntime::backend() const noexcept { + return backend_; +} + +core::BackendType HiggsARRuntime::backend_type() const noexcept { + return backend_type_; +} + +int HiggsARRuntime::device() const noexcept { + return device_; +} + +int HiggsARRuntime::threads() const noexcept { + return threads_; +} + +struct HiggsARKVCache::Impl { + Impl(std::shared_ptr input_runtime, int64_t input_cache_steps) + : runtime(std::move(input_runtime)), + cache_steps(input_cache_steps) { + if (runtime == nullptr) { + throw std::runtime_error("Higgs TTS AR KV cache requires runtime"); + } + if (cache_steps <= 0) { + throw std::runtime_error("Higgs TTS AR KV cache requires positive capacity"); + } + const auto & config = runtime->assets().config; + const auto & tensor_weights = runtime->weights(); + const int64_t dim = config.text.head_dim; + cache_layer_count = tensor_weights.decoder.layers.size(); + ggml_init_params params{4 * 1024 * 1024, nullptr, true}; + ctx.reset(ggml_init(params)); + if (ctx == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS AR KV cache context"); + } + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_tts.ar.kv_cache", runtime->backend_type()}; + std::vector key_tensors; + std::vector value_tensors; + key_tensors.reserve(tensor_weights.decoder.layers.size()); + value_tensors.reserve(tensor_weights.decoder.layers.size()); + for (size_t layer = 0; layer < tensor_weights.decoder.layers.size(); ++layer) { + key_tensors.push_back(core::make_tensor( + build_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, cache_steps, config.text.num_key_value_heads, dim}))); + value_tensors.push_back(core::make_tensor( + build_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, cache_steps, config.text.num_key_value_heads, dim}))); + } + cache = runtime::TransformerKVCache( + cache_steps, + config.text.num_key_value_heads * dim, + std::move(key_tensors), + std::move(value_tensors)); + buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), runtime->backend()); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS AR KV cache"); + } + } + + ~Impl() { + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + } + + bool can_run(const HiggsARRuntime & candidate_runtime, int64_t required_steps) const { + return runtime.get() == &candidate_runtime && cache_steps >= required_steps; + } + + void reset() { + runtime::TransformerKVState state; + state.current_end = 0; + state.layers.resize(cache_layer_count); + cache.import_state(state); + } + + void retain_prefix(int64_t prefix_steps) { + cache.retain_prefix(prefix_steps); + } + + void import_state(const runtime::TransformerKVState & state) { + cache.import_state(state); + } + + runtime::TransformerKVState export_state() const { + return cache.export_state(); + } + + void advance_after_direct_append(int64_t steps) { + cache.advance_after_direct_append(steps); + } + + std::shared_ptr runtime; + int64_t cache_steps = 0; + size_t cache_layer_count = 0; + std::unique_ptr ctx; + runtime::TransformerKVCache cache; + ggml_backend_buffer_t buffer = nullptr; +}; + +HiggsARKVCache::HiggsARKVCache(std::shared_ptr runtime, int64_t cache_steps) + : impl_(std::make_unique(std::move(runtime), cache_steps)) {} + +HiggsARKVCache::~HiggsARKVCache() = default; + +bool HiggsARKVCache::can_run(const HiggsARRuntime & runtime, int64_t required_steps) const { + return impl_->can_run(runtime, required_steps); +} + +int64_t HiggsARKVCache::cache_steps() const { + return impl_->cache.cache_steps(); +} + +int64_t HiggsARKVCache::valid_steps() const { + return impl_->cache.valid_steps(); +} + +int64_t HiggsARKVCache::current_end() const { + return impl_->cache.current_end(); +} + +void HiggsARKVCache::reset() { + impl_->reset(); +} + +void HiggsARKVCache::retain_prefix(int64_t prefix_steps) { + impl_->retain_prefix(prefix_steps); +} + +void HiggsARKVCache::import_state(const runtime::TransformerKVState & state) { + impl_->import_state(state); +} + +runtime::TransformerKVState HiggsARKVCache::export_state() const { + return impl_->export_state(); +} + +void HiggsARKVCache::advance_after_direct_append(int64_t steps) { + impl_->advance_after_direct_append(steps); +} + +const core::TensorValue & HiggsARKVCache::key_tensor(size_t layer) const { + return impl_->cache.key_tensor(layer); +} + +const core::TensorValue & HiggsARKVCache::value_tensor(size_t layer) const { + return impl_->cache.value_tensor(layer); +} + +struct HiggsARDecodeGraph::Impl { + Impl( + std::shared_ptr input_runtime, + int64_t input_cache_steps, + HiggsARKVCache & input_cache, + size_t graph_arena_bytes) + : runtime(std::move(input_runtime)), + cache(&input_cache), + cache_steps(input_cache_steps) { + if (runtime == nullptr) { + throw std::runtime_error("Higgs TTS AR decode graph requires runtime"); + } + if (cache == nullptr || !cache->can_run(*runtime, cache_steps)) { + throw std::runtime_error("Higgs TTS AR decode graph requires matching KV cache"); + } + if (cache_steps <= 0) { + throw std::runtime_error("Higgs TTS AR decode graph requires positive cache capacity"); + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx.reset(ggml_init(params)); + if (ctx == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS AR decode graph context"); + } + const auto & config = runtime->assets().config; + const auto & tensor_weights = runtime->weights(); + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_tts.ar.decode", runtime->backend_type()}; + + fused_code_ids = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, config.audio.num_codebooks); + auto x = build_higgs_decode_code_embedding( + build_ctx, + tensor_weights, + config, + fused_code_ids); + + positions = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, 1); + auto positions_value = core::wrap_tensor(positions, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + cache_slot = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I64, 1); + auto cache_slot_value = core::wrap_tensor(cache_slot, core::TensorShape::from_dims({1}), GGML_TYPE_I64); + attention_mask = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, cache_steps, 1, 1, 1); + auto attention_mask_value = core::wrap_tensor( + attention_mask, + core::TensorShape::from_dims({1, 1, 1, cache_steps}), + GGML_TYPE_F16); + + graph = ggml_new_graph_custom(ctx.get(), 65536, false); + const HiggsQwenDecoderComponent decoder(config.text); + for (size_t layer_index = 0; layer_index < tensor_weights.decoder.layers.size(); ++layer_index) { + auto out = decoder.build_decode_layer( + build_ctx, + graph, + x, + positions_value, + tensor_weights.decoder.layers[layer_index], + cache->key_tensor(layer_index), + cache->value_tensor(layer_index), + cache_slot_value, + attention_mask_value); + x = out.output; + } + + x = modules::RMSNormModule({config.text.hidden_size, config.text.rms_norm_eps, true, false}) + .build(build_ctx, x, {tensor_weights.norm, std::nullopt}); + auto logits = build_modality_logits(build_ctx, x, tensor_weights, config); + logits_output = logits.tensor; + ggml_set_output(logits_output); + ggml_build_forward_expand(graph, logits_output); + + buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), runtime->backend()); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS AR decode graph"); + } + core::set_backend_threads(runtime->backend(), runtime->threads()); + fused_code_ids_values.assign(static_cast(config.audio.num_codebooks), 0); + attention_mask_values.assign(static_cast(cache_steps), ggml_fp32_to_fp16(-INFINITY)); + engine::debug::timing_log_scalar( + "higgs_tts.ar.decode.graph.build_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + } + + ~Impl() { + engine::core::release_backend_graph_resources(runtime->backend(), graph); + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + } + + bool can_run(const HiggsARRuntime & candidate_runtime, int64_t required_steps) const { + return runtime.get() == &candidate_runtime && cache != nullptr && cache->can_run(candidate_runtime, required_steps); + } + + int64_t cache_steps_value() const { + return cache_steps; + } + + void import_prefill_state(const runtime::TransformerKVState & state) { + cache->import_state(state); + } + + void reset_timing() noexcept { + input_upload_ms = 0.0; + mask_upload_ms = 0.0; + graph_compute_ms = 0.0; + output_read_ms = 0.0; + steps = 0; + } + + void begin_decode_run() { + reset_timing(); + std::fill(attention_mask_values.begin(), attention_mask_values.end(), ggml_fp32_to_fp16(-INFINITY)); + const int64_t visible_steps = std::min(cache->valid_steps(), cache_steps); + std::fill( + attention_mask_values.begin(), + attention_mask_values.begin() + static_cast(visible_steps), + ggml_fp32_to_fp16(0.0F)); + ggml_backend_tensor_set( + attention_mask, + attention_mask_values.data(), + 0, + attention_mask_values.size() * sizeof(ggml_fp16_t)); + } + + HiggsARDecodeTiming timing() const { + return {input_upload_ms, mask_upload_ms, graph_compute_ms, output_read_ms, steps}; + } + + void run_step_into(const HiggsARDecodeInput & input, HiggsARDecodeOutput & output, bool log_timing) { + const auto & config = runtime->assets().config; + if (cache->valid_steps() >= cache_steps) { + throw std::runtime_error("Higgs TTS AR decode cache exhausted"); + } + if (!input.use_last_codes) { + throw std::runtime_error("Higgs TTS AR decode graph expects code-token steps after prefill"); + } + if (static_cast(input.last_codes.size()) != config.audio.num_codebooks) { + throw std::runtime_error("Higgs TTS AR decode last codebook row shape mismatch"); + } + + auto timing_start = Clock::now(); + for (int64_t codebook = 0; codebook < config.audio.num_codebooks; ++codebook) { + const int32_t code = input.last_codes[static_cast(codebook)]; + if (code < 0 || code >= config.audio.vocab_size) { + throw std::runtime_error("Higgs TTS AR decode codebook token is outside vocabulary"); + } + fused_code_ids_values[static_cast(codebook)] = + static_cast(code + codebook * config.audio.vocab_size); + } + ggml_backend_tensor_set( + fused_code_ids, + fused_code_ids_values.data(), + 0, + fused_code_ids_values.size() * sizeof(int32_t)); + + const int32_t position = static_cast(cache->current_end()); + ggml_backend_tensor_set(positions, &position, 0, sizeof(int32_t)); + const int64_t cache_slot_value = cache->valid_steps(); + ggml_backend_tensor_set(cache_slot, &cache_slot_value, 0, sizeof(int64_t)); + const double input_upload_delta_ms = engine::debug::elapsed_ms(timing_start, Clock::now()); + input_upload_ms += input_upload_delta_ms; + if (log_timing) { + engine::debug::timing_log_scalar("higgs_tts.ar.decode.step0.input_upload_ms", input_upload_delta_ms); + } + timing_start = Clock::now(); + attention_mask_values[static_cast(cache_slot_value)] = ggml_fp32_to_fp16(0.0F); + ggml_backend_tensor_set( + attention_mask, + attention_mask_values.data() + cache_slot_value, + static_cast(cache_slot_value) * sizeof(ggml_fp16_t), + sizeof(ggml_fp16_t)); + const double mask_upload_delta_ms = engine::debug::elapsed_ms(timing_start, Clock::now()); + mask_upload_ms += mask_upload_delta_ms; + if (log_timing) { + engine::debug::timing_log_scalar("higgs_tts.ar.decode.step0.mask_upload_ms", mask_upload_delta_ms); + } + + timing_start = Clock::now(); + const ggml_status status = engine::core::compute_backend_graph(runtime->backend(), graph); + if (log_timing || engine::debug::timing_log_enabled()) { + ggml_backend_synchronize(runtime->backend()); + } + const double graph_compute_delta_ms = engine::debug::elapsed_ms(timing_start, Clock::now()); + graph_compute_ms += graph_compute_delta_ms; + if (log_timing) { + engine::debug::timing_log_scalar("higgs_tts.ar.decode.step0.graph.compute_ms", graph_compute_delta_ms); + } + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS AR decode graph compute failed"); + } + + output.codebook_logits.resize(static_cast(config.audio.num_codebooks * config.audio.vocab_size)); + timing_start = Clock::now(); + ggml_backend_tensor_get( + logits_output, + output.codebook_logits.data(), + 0, + output.codebook_logits.size() * sizeof(float)); + const double output_read_delta_ms = engine::debug::elapsed_ms(timing_start, Clock::now()); + output_read_ms += output_read_delta_ms; + if (log_timing) { + engine::debug::timing_log_scalar("higgs_tts.ar.decode.step0.output_read_ms", output_read_delta_ms); + } + + cache->advance_after_direct_append(1); + ++steps; + } + + std::shared_ptr runtime; + HiggsARKVCache * cache = nullptr; + int64_t cache_steps = 0; + std::unique_ptr ctx; + ggml_tensor * fused_code_ids = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * cache_slot = nullptr; + ggml_tensor * attention_mask = nullptr; + ggml_tensor * logits_output = nullptr; + std::vector fused_code_ids_values; + std::vector attention_mask_values; + ggml_cgraph * graph = nullptr; + ggml_backend_buffer_t buffer = nullptr; + double input_upload_ms = 0.0; + double mask_upload_ms = 0.0; + double graph_compute_ms = 0.0; + double output_read_ms = 0.0; + int64_t steps = 0; +}; + +struct HiggsARPrefillGraph::Impl { + Impl( + std::shared_ptr input_runtime, + int64_t input_prompt_steps, + int64_t input_start_step, + HiggsARKVCache * input_cache, + size_t graph_arena_bytes) + : runtime(std::move(input_runtime)), + target_cache(input_cache), + prompt_steps(input_prompt_steps), + start_step(input_start_step), + run_steps(input_prompt_steps - input_start_step), + prefill_cache_steps(input_prompt_steps), + layerwise(input_prompt_steps >= kLayerwisePrefillMinSteps), + graph_arena_bytes(graph_arena_bytes) { + if (runtime == nullptr) { + throw std::runtime_error("Higgs TTS AR prefill graph requires runtime"); + } + if (prompt_steps <= 0) { + throw std::runtime_error("Higgs TTS AR prefill graph requires positive prompt length"); + } + if (start_step < 0 || start_step >= prompt_steps) { + throw std::runtime_error("Higgs TTS AR prefill graph start step is outside the prompt"); + } + if (start_step != 0) { + throw std::runtime_error("Higgs TTS AR prefill graph requires full prompt prefill"); + } + if (layerwise) { + engine::debug::timing_log_scalar("higgs_tts.ar.prefill.graph.build_ms", 0.0); + return; + } + const auto build_start = Clock::now(); + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx.reset(ggml_init(params)); + if (ctx == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS AR prefill graph context"); + } + const auto & config = runtime->assets().config; + const auto & tensor_weights = runtime->weights(); + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_tts.ar.prefill", runtime->backend_type()}; + + text_tokens = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, run_steps); + fused_code_ids = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_I32, config.audio.num_codebooks, run_steps); + text_gate = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, 1, run_steps, 1); + code_gate = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, 1, run_steps, 1); + auto x = build_higgs_prefill_input_embedding( + build_ctx, + tensor_weights, + config, + text_tokens, + fused_code_ids, + text_gate, + code_gate, + run_steps); + positions = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, run_steps); + auto positions_value = + core::wrap_tensor(positions, core::TensorShape::from_dims({run_steps}), GGML_TYPE_I32); + attention_mask = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, prefill_cache_steps, run_steps, 1, 1); + auto attention_mask_value = core::wrap_tensor( + attention_mask, + core::TensorShape::from_dims({1, 1, run_steps, prefill_cache_steps}), + GGML_TYPE_F16); + graph = ggml_new_graph_custom(ctx.get(), 262144, false); + keys.reserve(tensor_weights.decoder.layers.size()); + values.reserve(tensor_weights.decoder.layers.size()); + const HiggsQwenDecoderComponent decoder(config.text); + for (size_t layer_index = 0; layer_index < tensor_weights.decoder.layers.size(); ++layer_index) { + auto out = decoder.build_prefill_layer( + build_ctx, + x, + positions_value, + tensor_weights.decoder.layers[layer_index], + attention_mask_value); + x = out.output; + if (target_cache != nullptr) { + auto key_dest = higgs_cache_view( + build_ctx, + target_cache->key_tensor(layer_index), + 0, + prompt_steps, + config.text.num_key_value_heads, + config.text.head_dim); + auto value_dest = higgs_cache_view( + build_ctx, + target_cache->value_tensor(layer_index), + 0, + prompt_steps, + config.text.num_key_value_heads, + config.text.head_dim); + ggml_build_forward_expand(graph, ggml_cpy(ctx.get(), out.key.tensor, key_dest.tensor)); + ggml_build_forward_expand(graph, ggml_cpy(ctx.get(), out.value.tensor, value_dest.tensor)); + } else { + keys.push_back(out.key.tensor); + values.push_back(out.value.tensor); + } + } + x = modules::SliceModule({1, run_steps - 1, 1}).build(build_ctx, x); + x = modules::RMSNormModule({config.text.hidden_size, config.text.rms_norm_eps, true, false}) + .build(build_ctx, x, {tensor_weights.norm, std::nullopt}); + auto logits = build_modality_logits(build_ctx, x, tensor_weights, config); + logits_output = logits.tensor; + ggml_set_output(logits_output); + ggml_build_forward_expand(graph, logits_output); + + buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), runtime->backend()); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS AR prefill graph"); + } + text_token_values.assign(static_cast(run_steps), 0); + fused_code_id_values.assign(static_cast(run_steps * config.audio.num_codebooks), 0); + text_gate_values.assign(static_cast(run_steps), 0.0F); + code_gate_values.assign(static_cast(run_steps), 0.0F); + positions_values = modules::qwen_position_ids(run_steps, start_step); + attention_mask_values = modules::qwen_causal_prefill_mask_values(1, run_steps); + engine::debug::timing_log_scalar( + "higgs_tts.ar.prefill.graph.build_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + } + + ~Impl() { + engine::core::release_backend_graph_resources(runtime->backend(), graph); + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + } + + bool matches( + const HiggsARRuntime & candidate_runtime, + int64_t candidate_prompt_steps, + int64_t candidate_start_step) const { + return runtime.get() == &candidate_runtime && + prompt_steps == candidate_prompt_steps && + start_step == candidate_start_step; + } + + struct EmbeddingGraph { + EmbeddingGraph(const HiggsARRuntime & runtime, int64_t steps, size_t arena_bytes) + : runtime(&runtime), steps(steps) { + const auto & config = runtime.assets().config; + ggml_init_params params{arena_bytes, nullptr, true}; + ctx.reset(ggml_init(params)); + if (ctx == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS AR embedding graph context"); + } + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_tts.ar.prefill.embedding", runtime.backend_type()}; + text_tokens = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, steps); + fused_code_ids = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_I32, config.audio.num_codebooks, steps); + text_gate = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, 1, steps, 1); + code_gate = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, 1, steps, 1); + output = build_higgs_prefill_input_embedding( + build_ctx, + runtime.weights(), + config, + text_tokens, + fused_code_ids, + text_gate, + code_gate, + steps) + .tensor; + graph = ggml_new_graph_custom(ctx.get(), 32768, false); + ggml_set_output(output); + ggml_build_forward_expand(graph, output); + buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), runtime.backend()); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS AR embedding graph"); + } + } + + ~EmbeddingGraph() { + engine::core::release_backend_graph_resources(runtime->backend(), graph); + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + } + + std::vector run(const HiggsARPrefillInput & input) { + const auto & config = runtime->assets().config; + ggml_backend_tensor_set(text_tokens, input.text_tokens.data(), 0, input.text_tokens.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + fused_code_ids, + input.fused_code_ids.data(), + 0, + input.fused_code_ids.size() * sizeof(int32_t)); + ggml_backend_tensor_set(text_gate, input.text_gate.data(), 0, input.text_gate.size() * sizeof(float)); + ggml_backend_tensor_set(code_gate, input.code_gate.data(), 0, input.code_gate.size() * sizeof(float)); + core::set_backend_threads(runtime->backend(), runtime->threads()); + const ggml_status status = engine::core::compute_backend_graph(runtime->backend(), graph); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS AR embedding graph compute failed"); + } + std::vector hidden(static_cast(steps * config.text.hidden_size)); + ggml_backend_tensor_get(output, hidden.data(), 0, hidden.size() * sizeof(float)); + return hidden; + } + + const HiggsARRuntime * runtime = nullptr; + int64_t steps = 0; + std::unique_ptr ctx; + ggml_tensor * text_tokens = nullptr; + ggml_tensor * fused_code_ids = nullptr; + ggml_tensor * text_gate = nullptr; + ggml_tensor * code_gate = nullptr; + ggml_tensor * output = nullptr; + ggml_cgraph * graph = nullptr; + ggml_backend_buffer_t buffer = nullptr; + }; + + struct LayerGraph { + LayerGraph( + const HiggsARRuntime & runtime, + const modules::QwenDecoderLayerWeights & layer, + int64_t steps, + size_t arena_bytes) + : runtime(&runtime), steps(steps) { + const auto & config = runtime.assets().config; + ggml_init_params params{arena_bytes, nullptr, true}; + ctx.reset(ggml_init(params)); + if (ctx == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS AR layer prefill graph context"); + } + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_tts.ar.prefill.layer", runtime.backend_type()}; + auto x = core::make_tensor( + build_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, steps, config.text.hidden_size})); + input = x.tensor; + positions = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, steps); + auto positions_value = + core::wrap_tensor(positions, core::TensorShape::from_dims({steps}), GGML_TYPE_I32); + attention_mask = ggml_new_tensor_4d(ctx.get(), GGML_TYPE_F16, steps, steps, 1, 1); + auto attention_mask_value = core::wrap_tensor( + attention_mask, + core::TensorShape::from_dims({1, 1, steps, steps}), + GGML_TYPE_F16); + const HiggsQwenDecoderComponent decoder(config.text); + auto out = decoder.build_prefill_layer( + build_ctx, + x, + positions_value, + layer, + attention_mask_value); + output = out.output.tensor; + key = out.key.tensor; + value = out.value.tensor; + graph = ggml_new_graph_custom(ctx.get(), 65536, false); + ggml_set_output(output); + ggml_build_forward_expand(graph, output); + buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), runtime.backend()); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS AR layer prefill graph"); + } + + const auto position_values = modules::qwen_position_ids(steps); + ggml_backend_tensor_set(positions, position_values.data(), 0, position_values.size() * sizeof(int32_t)); + auto mask = modules::qwen_causal_prefill_mask_values(1, steps); + ggml_backend_tensor_set(attention_mask, mask.data(), 0, mask.size() * sizeof(ggml_fp16_t)); + } + + ~LayerGraph() { + engine::core::release_backend_graph_resources(runtime->backend(), graph); + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + } + + struct Output { + std::vector hidden; + std::vector key; + std::vector value; + }; + + Output run(const std::vector & hidden) { + const auto & config = runtime->assets().config; + const size_t hidden_values = static_cast(steps * config.text.hidden_size); + if (hidden.size() != hidden_values) { + throw std::runtime_error("Higgs TTS AR layer prefill input size mismatch"); + } + ggml_backend_tensor_set(input, hidden.data(), 0, hidden.size() * sizeof(float)); + core::set_backend_threads(runtime->backend(), runtime->threads()); + const ggml_status status = engine::core::compute_backend_graph(runtime->backend(), graph); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS AR layer prefill graph compute failed"); + } + Output out; + out.hidden.resize(hidden_values); + ggml_backend_tensor_get(output, out.hidden.data(), 0, out.hidden.size() * sizeof(float)); + const size_t layer_values = static_cast( + steps * config.text.num_key_value_heads * config.text.head_dim); + out.key.resize(layer_values); + out.value.resize(layer_values); + ggml_backend_tensor_get(key, out.key.data(), 0, out.key.size() * sizeof(float)); + ggml_backend_tensor_get(value, out.value.data(), 0, out.value.size() * sizeof(float)); + return out; + } + + const HiggsARRuntime * runtime = nullptr; + int64_t steps = 0; + std::unique_ptr ctx; + ggml_tensor * input = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * attention_mask = nullptr; + ggml_tensor * output = nullptr; + ggml_tensor * key = nullptr; + ggml_tensor * value = nullptr; + ggml_cgraph * graph = nullptr; + ggml_backend_buffer_t buffer = nullptr; + }; + + struct FinalGraph { + FinalGraph(const HiggsARRuntime & runtime, size_t arena_bytes) : runtime(&runtime) { + const auto & config = runtime.assets().config; + ggml_init_params params{arena_bytes, nullptr, true}; + ctx.reset(ggml_init(params)); + if (ctx == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS AR final prefill graph context"); + } + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_tts.ar.prefill.final", runtime.backend_type()}; + auto x = core::make_tensor( + build_ctx, + GGML_TYPE_F32, + core::TensorShape::from_dims({1, 1, config.text.hidden_size})); + input = x.tensor; + x = modules::RMSNormModule({config.text.hidden_size, config.text.rms_norm_eps, true, false}) + .build(build_ctx, x, {runtime.weights().norm, std::nullopt}); + auto logits = build_modality_logits(build_ctx, x, runtime.weights(), config); + output = logits.tensor; + graph = ggml_new_graph_custom(ctx.get(), 8192, false); + ggml_set_output(output); + ggml_build_forward_expand(graph, output); + buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), runtime.backend()); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS AR final prefill graph"); + } + } + + ~FinalGraph() { + engine::core::release_backend_graph_resources(runtime->backend(), graph); + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + } + + std::vector run(const std::vector & hidden) { + const auto & config = runtime->assets().config; + if (static_cast(hidden.size()) != config.text.hidden_size) { + throw std::runtime_error("Higgs TTS AR final prefill input size mismatch"); + } + ggml_backend_tensor_set(input, hidden.data(), 0, hidden.size() * sizeof(float)); + core::set_backend_threads(runtime->backend(), runtime->threads()); + const ggml_status status = engine::core::compute_backend_graph(runtime->backend(), graph); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS AR final prefill graph compute failed"); + } + std::vector logits(static_cast( + config.audio.num_codebooks * config.audio.vocab_size)); + ggml_backend_tensor_get(output, logits.data(), 0, logits.size() * sizeof(float)); + return logits; + } + + const HiggsARRuntime * runtime = nullptr; + std::unique_ptr ctx; + ggml_tensor * input = nullptr; + ggml_tensor * output = nullptr; + ggml_cgraph * graph = nullptr; + ggml_backend_buffer_t buffer = nullptr; + }; + + HiggsARPrefillOutput run_layerwise(const HiggsARPrefillInput & input) { + const auto & config = runtime->assets().config; + auto hidden = EmbeddingGraph(*runtime, prompt_steps, graph_arena_bytes).run(input); + HiggsARPrefillOutput out; + out.kv_state.current_end = prompt_steps; + out.kv_state.layers.resize(runtime->weights().decoder.layers.size()); + for (size_t layer = 0; layer < runtime->weights().decoder.layers.size(); ++layer) { + LayerGraph graph( + *runtime, + runtime->weights().decoder.layers[layer], + prompt_steps, + graph_arena_bytes); + auto layer_out = graph.run(hidden); + hidden = std::move(layer_out.hidden); + auto & state = out.kv_state.layers[layer]; + state.valid_steps = prompt_steps; + state.key = std::move(layer_out.key); + state.value = std::move(layer_out.value); + } + std::vector last_hidden(static_cast(config.text.hidden_size)); + const auto last_begin = hidden.begin() + + static_cast((prompt_steps - 1) * config.text.hidden_size); + std::copy( + last_begin, + last_begin + static_cast(config.text.hidden_size), + last_hidden.begin()); + out.output.codebook_logits = FinalGraph(*runtime, graph_arena_bytes).run(last_hidden); + return out; + } + + HiggsARPrefillOutput run(const HiggsARPrefillInput & input, int64_t candidate_start_step) { + const auto & config = runtime->assets().config; + if (input.steps != prompt_steps || + static_cast(input.text_tokens.size()) != prompt_steps || + static_cast(input.fused_code_ids.size()) != prompt_steps * config.audio.num_codebooks || + static_cast(input.text_gate.size()) != prompt_steps || + static_cast(input.code_gate.size()) != prompt_steps) { + throw std::runtime_error("Higgs TTS AR prefill graph input shape mismatch"); + } + if (candidate_start_step != start_step) { + throw std::runtime_error("Higgs TTS AR prefill graph start step mismatch"); + } + if (layerwise) { + return run_layerwise(input); + } + for (int64_t step = 0; step < run_steps; ++step) { + const int64_t source_step = start_step + step; + text_token_values[static_cast(step)] = + input.text_tokens[static_cast(source_step)]; + text_gate_values[static_cast(step)] = + input.text_gate[static_cast(source_step)]; + code_gate_values[static_cast(step)] = + input.code_gate[static_cast(source_step)]; + for (int64_t codebook = 0; codebook < config.audio.num_codebooks; ++codebook) { + fused_code_id_values[static_cast(step * config.audio.num_codebooks + codebook)] = + input.fused_code_ids[static_cast(source_step * config.audio.num_codebooks + codebook)]; + } + } + ggml_backend_tensor_set(text_tokens, text_token_values.data(), 0, text_token_values.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + fused_code_ids, + fused_code_id_values.data(), + 0, + fused_code_id_values.size() * sizeof(int32_t)); + ggml_backend_tensor_set(text_gate, text_gate_values.data(), 0, text_gate_values.size() * sizeof(float)); + ggml_backend_tensor_set(code_gate, code_gate_values.data(), 0, code_gate_values.size() * sizeof(float)); + ggml_backend_tensor_set(positions, positions_values.data(), 0, positions_values.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + attention_mask, + attention_mask_values.data(), + 0, + attention_mask_values.size() * sizeof(ggml_fp16_t)); + + core::set_backend_threads(runtime->backend(), runtime->threads()); + const ggml_status status = engine::core::compute_backend_graph(runtime->backend(), graph); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS AR prefill graph compute failed"); + } + HiggsARPrefillOutput out; + out.output.codebook_logits.resize(static_cast(config.audio.num_codebooks * config.audio.vocab_size)); + ggml_backend_tensor_get( + logits_output, + out.output.codebook_logits.data(), + 0, + out.output.codebook_logits.size() * sizeof(float)); + if (target_cache != nullptr) { + target_cache->advance_after_direct_append(prompt_steps); + out.wrote_cache = true; + out.kv_state.current_end = prompt_steps; + return out; + } + out.kv_state.current_end = prompt_steps; + out.kv_state.layers.resize(keys.size()); + const size_t layer_values = static_cast( + prompt_steps * config.text.num_key_value_heads * config.text.head_dim); + for (size_t layer = 0; layer < keys.size(); ++layer) { + auto & state = out.kv_state.layers[layer]; + state.valid_steps = prompt_steps; + state.key.resize(layer_values); + state.value.resize(layer_values); + ggml_backend_tensor_get(keys[layer], state.key.data(), 0, state.key.size() * sizeof(float)); + ggml_backend_tensor_get(values[layer], state.value.data(), 0, state.value.size() * sizeof(float)); + } + return out; + } + + std::shared_ptr runtime; + HiggsARKVCache * target_cache = nullptr; + int64_t prompt_steps = 0; + int64_t start_step = 0; + int64_t run_steps = 0; + int64_t prefill_cache_steps = 0; + bool layerwise = false; + size_t graph_arena_bytes = 0; + std::unique_ptr ctx; + ggml_tensor * text_tokens = nullptr; + ggml_tensor * fused_code_ids = nullptr; + ggml_tensor * text_gate = nullptr; + ggml_tensor * code_gate = nullptr; + ggml_tensor * positions = nullptr; + ggml_tensor * attention_mask = nullptr; + ggml_tensor * logits_output = nullptr; + std::vector keys; + std::vector values; + std::vector text_token_values; + std::vector fused_code_id_values; + std::vector text_gate_values; + std::vector code_gate_values; + std::vector positions_values; + std::vector attention_mask_values; + ggml_cgraph * graph = nullptr; + ggml_backend_buffer_t buffer = nullptr; +}; + +HiggsARPrefillGraph::HiggsARPrefillGraph( + std::shared_ptr runtime, + int64_t prompt_steps, + int64_t start_step, + HiggsARKVCache * cache, + size_t graph_arena_bytes) + : impl_(std::make_unique(std::move(runtime), prompt_steps, start_step, cache, graph_arena_bytes)) {} + +HiggsARPrefillGraph::~HiggsARPrefillGraph() = default; + +bool HiggsARPrefillGraph::matches( + const HiggsARRuntime & runtime, + int64_t prompt_steps, + int64_t start_step) const { + return impl_->matches(runtime, prompt_steps, start_step); +} + +HiggsARPrefillOutput HiggsARPrefillGraph::run(const HiggsARPrefillInput & input, int64_t start_step) { + return impl_->run(input, start_step); +} + + +HiggsARDecodeGraph::HiggsARDecodeGraph( + std::shared_ptr runtime, + int64_t cache_steps, + HiggsARKVCache & cache, + size_t graph_arena_bytes) + : impl_(std::make_unique(std::move(runtime), cache_steps, cache, graph_arena_bytes)) {} + +HiggsARDecodeGraph::~HiggsARDecodeGraph() = default; + +bool HiggsARDecodeGraph::can_run(const HiggsARRuntime & runtime, int64_t required_steps) const { + return impl_->can_run(runtime, required_steps); +} + +int64_t HiggsARDecodeGraph::cache_steps() const { + return impl_->cache_steps_value(); +} + +void HiggsARDecodeGraph::import_prefill_state(const runtime::TransformerKVState & state) { + impl_->import_prefill_state(state); +} + +void HiggsARDecodeGraph::begin_decode_run() { + impl_->begin_decode_run(); +} + +HiggsARDecodeTiming HiggsARDecodeGraph::timing() const { + return impl_->timing(); +} + +void HiggsARDecodeGraph::run_step_into( + const HiggsARDecodeInput & input, + HiggsARDecodeOutput & output, + bool log_timing) { + impl_->run_step_into(input, output, log_timing); +} + +} // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/assets.cpp b/src/models/higgs_tts/assets.cpp new file mode 100644 index 00000000..350be213 --- /dev/null +++ b/src/models/higgs_tts/assets.cpp @@ -0,0 +1,171 @@ +#include "engine/models/higgs_tts/assets.h" + +#include "engine/framework/assets/model_package.h" +#include "engine/framework/io/config.h" +#include "engine/framework/io/json.h" + +#include +#include + +namespace engine::models::higgs_tts { +namespace json = engine::io::json; +namespace { + +constexpr const char * kExpectedArchitecture = "HiggsMultimodalQwen3ForConditionalGeneration"; + +float parse_rope_theta(const engine::io::json::Value & text_config) { + const auto * rope_parameters = text_config.find("rope_parameters"); + if (rope_parameters != nullptr && rope_parameters->is_object()) { + return json::optional_f32(*rope_parameters, "rope_theta", 1000000.0F); + } + const auto * rope_theta = text_config.find("rope_theta"); + if (rope_theta != nullptr && rope_theta->is_number()) { + return rope_theta->as_f32(); + } + return 1000000.0F; +} + +std::string parse_architecture(const engine::io::json::Value & root) { + const auto * architectures = root.find("architectures"); + if (architectures == nullptr || !architectures->is_array() || architectures->as_array().empty()) { + throw std::runtime_error("Higgs TTS config must provide architectures[0]"); + } + return architectures->as_array().front().as_string(); +} + +HiggsTextConfig parse_text_config(const engine::io::json::Value & root) { + HiggsTextConfig config; + config.model_type = json::require_string(root, "model_type"); + if (config.model_type != "qwen3") { + throw std::runtime_error("Higgs TTS text_config.model_type mismatch"); + } + config.vocab_size = json::require_i64(root, "vocab_size"); + config.hidden_size = json::require_i64(root, "hidden_size"); + config.intermediate_size = json::require_i64(root, "intermediate_size"); + config.num_hidden_layers = json::require_i64(root, "num_hidden_layers"); + config.num_attention_heads = json::require_i64(root, "num_attention_heads"); + config.num_key_value_heads = json::require_i64(root, "num_key_value_heads"); + config.head_dim = json::optional_i64(root, "head_dim", config.hidden_size / config.num_attention_heads); + config.max_position_embeddings = json::require_i64(root, "max_position_embeddings"); + config.bos_token_id = json::require_i64(root, "bos_token_id"); + config.eos_token_id = json::require_i64(root, "eos_token_id"); + config.pad_token_id = json::optional_nullable_i64(root, "pad_token_id", -1); + config.rms_norm_eps = json::optional_f32(root, "rms_norm_eps", config.rms_norm_eps); + config.rope_theta = parse_rope_theta(root); + config.tie_word_embeddings = json::optional_bool(root, "tie_word_embeddings", config.tie_word_embeddings); + + engine::io::require_positive(config.vocab_size, "text vocab_size"); + engine::io::require_positive(config.hidden_size, "text hidden_size"); + engine::io::require_positive(config.intermediate_size, "text intermediate_size"); + engine::io::require_positive(config.num_hidden_layers, "text num_hidden_layers"); + engine::io::require_positive(config.num_attention_heads, "text num_attention_heads"); + engine::io::require_positive(config.num_key_value_heads, "text num_key_value_heads"); + engine::io::require_positive(config.head_dim, "text head_dim"); + engine::io::require_positive(config.max_position_embeddings, "text max_position_embeddings"); + engine::io::require_divisible(config.num_attention_heads, config.num_key_value_heads, "text grouped-query attention"); + return config; +} + +HiggsAudioEncoderConfig parse_audio_config(const engine::io::json::Value & root) { + HiggsAudioEncoderConfig config; + config.model_type = json::require_string(root, "model_type"); + if (config.model_type != "higgs_audio_encoder") { + throw std::runtime_error("Higgs TTS audio_encoder_config.model_type mismatch"); + } + config.encoder_type = json::require_string(root, "encoder_type"); + if (config.encoder_type != "discrete") { + throw std::runtime_error("Higgs TTS audio_encoder_config.encoder_type mismatch"); + } + config.num_codebooks = json::require_i64(root, "num_codebooks"); + config.vocab_size = json::require_i64(root, "vocab_size"); + config.out_dim = json::require_i64(root, "out_dim"); + config.mel_per_sample = json::require_i64(root, "mel_per_sample"); + config.max_chunk_size = json::require_i64(root, "max_chunk_size"); + config.tie_word_embeddings = json::optional_bool(root, "tie_word_embeddings", config.tie_word_embeddings); + config.use_delay_pattern = json::optional_bool(root, "use_delay_pattern", config.use_delay_pattern); + + engine::io::require_positive(config.num_codebooks, "audio num_codebooks"); + engine::io::require_positive(config.vocab_size, "audio vocab_size"); + engine::io::require_positive(config.out_dim, "audio out_dim"); + engine::io::require_positive(config.mel_per_sample, "audio mel_per_sample"); + engine::io::require_positive(config.max_chunk_size, "audio max_chunk_size"); + if (!config.tie_word_embeddings) { + throw std::runtime_error("Higgs TTS currently expects tied modality embedding/head weights"); + } + if (!config.use_delay_pattern) { + throw std::runtime_error("Higgs TTS v3 requires delay-pattern audio codebooks"); + } + return config; +} + +HiggsConfig parse_config(const assets::ResourceBundle & resources) { + const auto root = resources.parse_json("config"); + HiggsConfig config; + config.model_type = json::require_string(root, "model_type"); + if (config.model_type != "higgs_multimodal_qwen3") { + throw std::runtime_error("Higgs TTS model_type mismatch"); + } + config.architecture = parse_architecture(root); + if (config.architecture != kExpectedArchitecture) { + throw std::runtime_error("Higgs TTS architecture mismatch"); + } + config.hidden_size = json::require_i64(root, "_hidden_size"); + config.vocab_size = json::require_i64(root, "_vocab_size"); + config.audio_token_id = json::optional_i64(root, "audio_token_id", config.audio_token_id); + config.ignore_index = json::optional_i64(root, "ignore_index", config.ignore_index); + config.text = parse_text_config(root.require("text_config")); + config.audio = parse_audio_config(root.require("audio_encoder_config")); + + if (config.hidden_size != config.text.hidden_size) { + throw std::runtime_error("Higgs TTS _hidden_size must match text hidden_size"); + } + if (config.vocab_size != config.text.vocab_size) { + throw std::runtime_error("Higgs TTS _vocab_size must match text vocab_size"); + } + if (config.audio.out_dim != config.text.hidden_size) { + throw std::runtime_error("Higgs TTS audio out_dim must match text hidden_size"); + } + if (config.audio_token_id != -100) { + throw std::runtime_error("Higgs TTS audio_token_id mismatch"); + } + return config; +} + +void validate_weight_anchors(const HiggsAssets & assets) { + const auto & config = assets.config; + const auto & weights = *assets.weights; + const int64_t hidden = config.text.hidden_size; + const int64_t audio_fused_vocab = config.audio.num_codebooks * config.audio.vocab_size; + assets::require_tensor_shape(weights, "tied.embedding.text_embedding.weight", {config.text.vocab_size, hidden}); + assets::require_tensor_shape(weights, "tied.embedding.modality_embeddings.0.embedding.weight", {audio_fused_vocab, hidden}); + assets::require_tensor_shape(weights, "body.norm.weight", {hidden}); + assets::require_tensor_shape(weights, "body.layers.0.input_layernorm.weight", {hidden}); + assets::require_tensor_shape(weights, "body.layers.0.post_attention_layernorm.weight", {hidden}); + assets::require_tensor_shape(weights, "body.layers.0.self_attn.q_proj.weight", {config.text.num_attention_heads * config.text.head_dim, hidden}); + assets::require_tensor_shape(weights, "body.layers.0.self_attn.k_proj.weight", {config.text.num_key_value_heads * config.text.head_dim, hidden}); + assets::require_tensor_shape(weights, "body.layers.0.self_attn.v_proj.weight", {config.text.num_key_value_heads * config.text.head_dim, hidden}); + assets::require_tensor_shape(weights, "body.layers.0.self_attn.o_proj.weight", {hidden, config.text.num_attention_heads * config.text.head_dim}); + assets::require_tensor_shape(weights, "body.layers.0.self_attn.q_norm.weight", {config.text.head_dim}); + assets::require_tensor_shape(weights, "body.layers.0.self_attn.k_norm.weight", {config.text.head_dim}); + assets::require_tensor_shape(weights, "body.layers.0.mlp.gate_proj.weight", {config.text.intermediate_size, hidden}); + assets::require_tensor_shape(weights, "body.layers.0.mlp.up_proj.weight", {config.text.intermediate_size, hidden}); + assets::require_tensor_shape(weights, "body.layers.0.mlp.down_proj.weight", {hidden, config.text.intermediate_size}); + assets::require_tensor_shape(weights, "tied.embedding.modality_embeddings.0.model.acoustic_encoder.conv1.weight", {64, 1, 7}); + assets::require_tensor_shape(weights, "tied.embedding.modality_embeddings.0.model.acoustic_decoder.conv2.weight", {1, 32, 7}); + assets::require_tensor_shape(weights, "tied.embedding.modality_embeddings.0.model.quantizer.quantizers.0.codebook.embed", {1024, 64}); +} + +} // namespace + +std::shared_ptr load_higgs_assets(const std::filesystem::path & model_path) { + HiggsAssets assets; + assets.resources = assets::load_resource_bundle_from_package_spec( + model_path, + assets::default_model_package_spec_path("higgs_tts")); + assets.config = parse_config(assets.resources); + assets.weights = assets.resources.open_tensor_source("weights"); + validate_weight_anchors(assets); + return std::make_shared(std::move(assets)); +} + +} // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/codebooks.cpp b/src/models/higgs_tts/codebooks.cpp new file mode 100644 index 00000000..9cf0b0b6 --- /dev/null +++ b/src/models/higgs_tts/codebooks.cpp @@ -0,0 +1,79 @@ +#include "engine/models/higgs_tts/codebooks.h" + +#include +#include + +namespace engine::models::higgs_tts { +namespace { + +void require_codebook_matrix( + const std::vector & codes, + int64_t frames, + int64_t codebooks, + const char * label) { + if (frames <= 0 || codebooks <= 0) { + throw std::runtime_error(std::string("Higgs TTS ") + label + " requires positive frames and codebooks"); + } + if (static_cast(codes.size()) != frames * codebooks) { + throw std::runtime_error(std::string("Higgs TTS ") + label + " code matrix shape mismatch"); + } +} + +size_t flat_index(int64_t frame, int64_t codebook, int64_t codebooks) { + return static_cast(frame * codebooks + codebook); +} + +} // namespace + +int64_t higgs_delayed_frame_count(int64_t raw_frames, int64_t codebooks) { + if (raw_frames <= 0 || codebooks <= 0) { + throw std::runtime_error("Higgs TTS delayed frame count requires positive dimensions"); + } + return raw_frames + codebooks - 1; +} + +std::vector apply_higgs_delay_pattern( + const std::vector & raw_codes, + int64_t raw_frames, + int64_t codebooks) { + require_codebook_matrix(raw_codes, raw_frames, codebooks, "delay pattern input"); + const int64_t delayed_frames = higgs_delayed_frame_count(raw_frames, codebooks); + std::vector delayed(static_cast(delayed_frames * codebooks), kHiggsEocId); +#ifdef _OPENMP +#pragma omp parallel for if (raw_frames * codebooks > 1024) +#endif + for (int64_t codebook = 0; codebook < codebooks; ++codebook) { + for (int64_t frame = 0; frame < codebook; ++frame) { + delayed[flat_index(frame, codebook, codebooks)] = kHiggsBocId; + } + for (int64_t frame = 0; frame < raw_frames; ++frame) { + delayed[flat_index(codebook + frame, codebook, codebooks)] = + raw_codes[flat_index(frame, codebook, codebooks)]; + } + } + return delayed; +} + +std::vector reverse_higgs_delay_pattern( + const std::vector & delayed_codes, + int64_t delayed_frames, + int64_t codebooks) { + require_codebook_matrix(delayed_codes, delayed_frames, codebooks, "reverse delay pattern input"); + const int64_t raw_frames = delayed_frames - (codebooks - 1); + if (raw_frames <= 0) { + throw std::runtime_error("Higgs TTS delayed codes must include at least one recoverable raw frame"); + } + std::vector raw(static_cast(raw_frames * codebooks), 0); +#ifdef _OPENMP +#pragma omp parallel for if (raw_frames * codebooks > 1024) +#endif + for (int64_t codebook = 0; codebook < codebooks; ++codebook) { + for (int64_t frame = 0; frame < raw_frames; ++frame) { + raw[flat_index(frame, codebook, codebooks)] = + delayed_codes[flat_index(codebook + frame, codebook, codebooks)]; + } + } + return raw; +} + +} // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/codec.cpp b/src/models/higgs_tts/codec.cpp new file mode 100644 index 00000000..82d94cc8 --- /dev/null +++ b/src/models/higgs_tts/codec.cpp @@ -0,0 +1,1614 @@ +#include "engine/models/higgs_tts/codec.h" + +#include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/resampling.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/structural_modules.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::higgs_tts { +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr const char * kCodecPrefix = "tied.embedding.modality_embeddings.0.model."; +constexpr int64_t kCodecCodebooks = 8; +constexpr int64_t kCodecCodebookSize = 1024; +constexpr int64_t kCodecCodebookDim = 64; +constexpr int64_t kCodecHiddenSize = 1024; +constexpr int64_t kAcousticHiddenSize = 256; +constexpr int64_t kSemanticHiddenSize = 768; +constexpr int64_t kCodecProjectInputSize = kAcousticHiddenSize + kSemanticHiddenSize; +constexpr int64_t kResidualUnitsPerBlock = 3; +constexpr int64_t kSemanticResidualUnitsPerBlock = 2; +constexpr int64_t kSemanticIntermediateSize = 3072; +constexpr int64_t kSemanticAttentionHeads = 12; +constexpr int64_t kSemanticLayers = 12; +constexpr int64_t kSemanticConvLayers = 7; +constexpr int64_t kSemanticSampleRate = 16000; +constexpr int64_t kSemanticPadSamples = 160; +constexpr float kSemanticLayerNormEps = 1.0e-5F; + +const int64_t kUpsampleRatios[] = {8, 5, 4, 2, 3}; +const int64_t kDecoderChannels[] = {1024, 512, 256, 128, 64, 32}; +const int64_t kEncoderChannels[] = {64, 128, 256, 512, 1024, 2048}; +const int64_t kSemanticConvDim[] = {512, 512, 512, 512, 512, 512, 512}; +const int64_t kSemanticConvKernel[] = {10, 3, 3, 3, 3, 2, 2}; +const int64_t kSemanticConvStride[] = {5, 2, 2, 2, 2, 2, 2}; +constexpr int64_t kDecoderBlockCount = 5; +constexpr int64_t kEncoderBlockCount = 5; +constexpr int64_t kSemanticBlockCount = 2; +constexpr int kCodecSampleRate = 24000; +constexpr int64_t kCodecDecodeCapacityBucketFrames = 32; +constexpr int64_t kCodecHopLength = 960; +constexpr int64_t kCodecPadSamples = kCodecHopLength / 2; +constexpr int64_t kCodecDecodeWindowFrames = 128; +constexpr int64_t kCodecDecodeOverlapFrames = 8; +constexpr int64_t kCodecFullDecodeMaxFrames = 512; +constexpr int64_t kResidualDilations[] = {1, 3, 9}; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +std::string codec_name(const std::string & name) { return std::string(kCodecPrefix) + name; } + +modules::LinearWeights load_linear(core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + assets::TensorStorageType storage_type, + int64_t out_features, + int64_t in_features, + bool use_bias) { + modules::LinearWeights weights; + weights.weight = store.load_tensor( + source, codec_name(name + ".weight"), storage_type, {out_features, in_features}); + if (use_bias) { + weights.bias = store.load_f32_tensor(source, codec_name(name + ".bias"), {out_features}); + } + return weights; +} + +modules::Conv1dWeights load_conv1d(core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + assets::TensorStorageType storage_type, + int64_t out_channels, + int64_t in_channels, + int64_t kernel_size, + bool use_bias) { + modules::Conv1dWeights weights; + weights.weight = store.load_tensor(source, + codec_name(name + ".weight"), + storage_type, + {out_channels, in_channels, kernel_size}); + if (use_bias) { + weights.bias = store.load_f32_tensor(source, codec_name(name + ".bias"), {out_channels}); + } + return weights; +} + +modules::ConvTranspose1dWeights load_conv_transpose1d(core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + assets::TensorStorageType storage_type, + int64_t in_channels, + int64_t out_channels, + int64_t kernel_size, + bool use_bias) { + modules::ConvTranspose1dWeights weights; + weights.weight = store.load_tensor(source, + codec_name(name + ".weight"), + storage_type, + {in_channels, out_channels, kernel_size}); + if (use_bias) { + weights.bias = store.load_f32_tensor(source, codec_name(name + ".bias"), {out_channels}); + } + return weights; +} + +modules::Snake1dWeights load_snake(core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & name, + int64_t channels) { + const auto raw = source.require_f32(codec_name(name), std::vector{1, channels, 1}); + return {store.make_f32(core::TensorShape::from_dims({channels}), raw)}; +} + +core::TensorValue require_semantic_tensor(const HiggsCodecWeights & weights, + const std::string & name) { + const auto it = weights.semantic_model.find(name); + if (it == weights.semantic_model.end()) { + throw std::runtime_error("Higgs TTS codec missing semantic tensor: " + name); + } + return it->second; +} + +modules::NormWeights semantic_norm_weights(const HiggsCodecWeights & weights, + const std::string & prefix) { + return modules::NormWeights{require_semantic_tensor(weights, prefix + ".weight"), + require_semantic_tensor(weights, prefix + ".bias")}; +} + +modules::LinearWeights semantic_linear_weights(const HiggsCodecWeights & weights, + const std::string & prefix) { + return modules::LinearWeights{require_semantic_tensor(weights, prefix + ".weight"), + require_semantic_tensor(weights, prefix + ".bias")}; +} + +modules::Conv1dWeights semantic_conv_weights(const HiggsCodecWeights & weights, + const std::string & prefix, + bool use_bias) { + modules::Conv1dWeights out; + out.weight = require_semantic_tensor(weights, prefix + ".weight"); + if (use_bias) { + out.bias = require_semantic_tensor(weights, prefix + ".bias"); + } + return out; +} + +int64_t +conv1d_output_frames(int64_t input_frames, int64_t kernel, int64_t stride, int64_t padding) { + return (input_frames + 2 * padding - kernel) / stride + 1; +} + +int64_t ceil_div(int64_t value, int64_t divisor) { return (value + divisor - 1) / divisor; } + +int64_t semantic_feature_frames(int64_t input_samples) { + int64_t frames = input_samples; + for (int64_t index = 0; index < kSemanticConvLayers; ++index) { + frames = conv1d_output_frames(frames, + kSemanticConvKernel[static_cast(index)], + kSemanticConvStride[static_cast(index)], + 0); + if (frames <= 0) { + throw std::runtime_error("Higgs TTS semantic encoder input is too short"); + } + } + return ceil_div(frames, 2); +} + +int64_t acoustic_encoder_frames(int64_t input_samples) { + int64_t frames = conv1d_output_frames(input_samples, 7, 1, 3); + for (int64_t block = 0; block < kEncoderBlockCount; ++block) { + const int64_t ratio = kUpsampleRatios[static_cast(block)]; + frames = conv1d_output_frames(frames, 2 * ratio, ratio, (ratio + 1) / 2); + if (frames <= 0) { + throw std::runtime_error("Higgs TTS acoustic encoder input is too short"); + } + } + frames = conv1d_output_frames(frames, 3, 1, 1); + if (frames <= 0) { + throw std::runtime_error("Higgs TTS acoustic encoder input is too short"); + } + return frames; +} + +std::vector pad_zeros(const std::vector & input, int64_t left, int64_t right) { + std::vector out(static_cast(left + static_cast(input.size()) + right), + 0.0F); + std::copy(input.begin(), input.end(), out.begin() + left); + return out; +} + +std::vector resample_mono_if_needed(const std::vector & mono, + int source_sample_rate, + int target_sample_rate) { + if (source_sample_rate == target_sample_rate) { + return mono; + } + auto out = + audio::resample_mono_torchaudio_sinc_hann(mono, source_sample_rate, target_sample_rate); + if (out.empty()) { + throw std::runtime_error("Higgs TTS codec resampling produced no samples"); + } + return out; +} + +std::vector prepare_mono_audio(const runtime::AudioBuffer & audio) { + if (audio.sample_rate <= 0) { + throw std::runtime_error("Higgs TTS codec reference audio sample rate must be positive"); + } + if (audio.channels <= 0) { + throw std::runtime_error("Higgs TTS codec reference audio channel count must be positive"); + } + if (audio.samples.empty()) { + throw std::runtime_error("Higgs TTS codec reference audio is empty"); + } + if ((audio.samples.size() % static_cast(audio.channels)) != 0) { + throw std::runtime_error("Higgs TTS codec reference audio sample count is " + "not divisible by channels"); + } + return audio::mixdown_interleaved_to_mono_average(audio.samples, audio.channels); +} + +std::vector prepare_codec_audio_24k(const std::vector & mono, + int source_sample_rate) { + auto audio_24k = resample_mono_if_needed(mono, source_sample_rate, kCodecSampleRate); + if (static_cast(audio_24k.size()) < kCodecSampleRate) { + audio_24k.resize(static_cast(kCodecSampleRate), 0.0F); + } + return audio_24k; +} + +std::vector prepare_semantic_audio_16k(const std::vector & mono, + int source_sample_rate) { + auto semantic_16k = resample_mono_if_needed(mono, source_sample_rate, kSemanticSampleRate); + return pad_zeros(semantic_16k, kSemanticPadSamples, kSemanticPadSamples); +} + +std::vector effective_semantic_pos_conv_weight(const assets::TensorSource & source, + int64_t out_channels, + int64_t in_channels, + int64_t kernel_size) { + const auto g = source.require_f32(codec_name("semantic_model.encoder.pos_conv_embed." + "conv.parametrizations.weight.original0"), + {1, 1, kernel_size}); + const auto v = source.require_f32(codec_name("semantic_model.encoder.pos_conv_embed." + "conv.parametrizations.weight.original1"), + {out_channels, in_channels, kernel_size}); + std::vector weight(v.size()); + for (int64_t k = 0; k < kernel_size; ++k) { + double sum = 0.0; + for (int64_t out = 0; out < out_channels; ++out) { + for (int64_t in = 0; in < in_channels; ++in) { + const size_t index = + static_cast((out * in_channels + in) * kernel_size + k); + sum += static_cast(v[index]) * static_cast(v[index]); + } + } + const double norm = std::sqrt(sum); + if (norm == 0.0) { + throw std::runtime_error("Higgs TTS semantic positional-conv weight norm is zero"); + } + const float scale_value = + static_cast(static_cast(g[static_cast(k)]) / norm); + for (int64_t out = 0; out < out_channels; ++out) { + for (int64_t in = 0; in < in_channels; ++in) { + const size_t index = + static_cast((out * in_channels + in) * kernel_size + k); + weight[index] = v[index] * scale_value; + } + } + } + return weight; +} + +void load_semantic_tensor(HiggsCodecWeights & weights, + const assets::TensorSource & source, + const std::string & name, + const std::vector & shape, + assets::TensorStorageType storage_type) { + weights.semantic_model.emplace( + name, + weights.store->load_tensor( + source, codec_name("semantic_model." + name), storage_type, shape)); +} + +void load_semantic_f32_tensor(HiggsCodecWeights & weights, + const assets::TensorSource & source, + const std::string & name, + const std::vector & shape) { + weights.semantic_model.emplace( + name, weights.store->load_f32_tensor(source, codec_name("semantic_model." + name), shape)); +} + +void load_hubert_semantic_model_weights(HiggsCodecWeights & weights, + const assets::TensorSource & source, + assets::TensorStorageType storage_type) { + for (int64_t layer = 0; layer < kSemanticConvLayers; ++layer) { + const std::string prefix = "feature_extractor.conv_layers." + std::to_string(layer); + load_semantic_tensor(weights, + source, + prefix + ".conv.weight", + {kSemanticConvDim[static_cast(layer)], + layer == 0 ? 1 : kSemanticConvDim[static_cast(layer - 1)], + kSemanticConvKernel[static_cast(layer)]}, + storage_type); + if (layer == 0) { + load_semantic_f32_tensor( + weights, source, prefix + ".layer_norm.weight", {kSemanticConvDim[0]}); + load_semantic_f32_tensor( + weights, source, prefix + ".layer_norm.bias", {kSemanticConvDim[0]}); + } + } + load_semantic_f32_tensor( + weights, source, "feature_projection.layer_norm.weight", {kSemanticConvDim[6]}); + load_semantic_f32_tensor( + weights, source, "feature_projection.layer_norm.bias", {kSemanticConvDim[6]}); + load_semantic_tensor(weights, + source, + "feature_projection.projection.weight", + {kSemanticHiddenSize, kSemanticConvDim[6]}, + storage_type); + load_semantic_f32_tensor( + weights, source, "feature_projection.projection.bias", {kSemanticHiddenSize}); + load_semantic_f32_tensor(weights, source, "encoder.layer_norm.weight", {kSemanticHiddenSize}); + load_semantic_f32_tensor(weights, source, "encoder.layer_norm.bias", {kSemanticHiddenSize}); + weights.semantic_model.emplace( + "encoder.pos_conv_embed.conv.weight", + weights.store->make_f32( + core::TensorShape::from_dims({kSemanticHiddenSize, kSemanticHiddenSize / 16, 128}), + effective_semantic_pos_conv_weight( + source, kSemanticHiddenSize, kSemanticHiddenSize / 16, 128))); + load_semantic_f32_tensor( + weights, source, "encoder.pos_conv_embed.conv.bias", {kSemanticHiddenSize}); + for (int64_t layer = 0; layer < kSemanticLayers; ++layer) { + const std::string prefix = "encoder.layers." + std::to_string(layer); + load_semantic_f32_tensor( + weights, source, prefix + ".layer_norm.weight", {kSemanticHiddenSize}); + load_semantic_f32_tensor( + weights, source, prefix + ".layer_norm.bias", {kSemanticHiddenSize}); + load_semantic_f32_tensor( + weights, source, prefix + ".final_layer_norm.weight", {kSemanticHiddenSize}); + load_semantic_f32_tensor( + weights, source, prefix + ".final_layer_norm.bias", {kSemanticHiddenSize}); + load_semantic_tensor(weights, + source, + prefix + ".attention.q_proj.weight", + {kSemanticHiddenSize, kSemanticHiddenSize}, + storage_type); + load_semantic_f32_tensor( + weights, source, prefix + ".attention.q_proj.bias", {kSemanticHiddenSize}); + load_semantic_tensor(weights, + source, + prefix + ".attention.k_proj.weight", + {kSemanticHiddenSize, kSemanticHiddenSize}, + storage_type); + load_semantic_f32_tensor( + weights, source, prefix + ".attention.k_proj.bias", {kSemanticHiddenSize}); + load_semantic_tensor(weights, + source, + prefix + ".attention.v_proj.weight", + {kSemanticHiddenSize, kSemanticHiddenSize}, + storage_type); + load_semantic_f32_tensor( + weights, source, prefix + ".attention.v_proj.bias", {kSemanticHiddenSize}); + load_semantic_tensor(weights, + source, + prefix + ".attention.out_proj.weight", + {kSemanticHiddenSize, kSemanticHiddenSize}, + storage_type); + load_semantic_f32_tensor( + weights, source, prefix + ".attention.out_proj.bias", {kSemanticHiddenSize}); + load_semantic_tensor(weights, + source, + prefix + ".feed_forward.intermediate_dense.weight", + {kSemanticIntermediateSize, kSemanticHiddenSize}, + storage_type); + load_semantic_f32_tensor(weights, + source, + prefix + ".feed_forward.intermediate_dense.bias", + {kSemanticIntermediateSize}); + load_semantic_tensor(weights, + source, + prefix + ".feed_forward.output_dense.weight", + {kSemanticHiddenSize, kSemanticIntermediateSize}, + storage_type); + load_semantic_f32_tensor( + weights, source, prefix + ".feed_forward.output_dense.bias", {kSemanticHiddenSize}); + } +} + +HiggsCodecResidualUnitWeights load_residual_unit(core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t channels) { + HiggsCodecResidualUnitWeights weights; + weights.snake1 = load_snake(store, source, prefix + ".snake1.alpha", channels); + weights.conv1 = + load_conv1d(store, source, prefix + ".conv1", storage_type, channels, channels, 7, true); + weights.snake2 = load_snake(store, source, prefix + ".snake2.alpha", channels); + weights.conv2 = + load_conv1d(store, source, prefix + ".conv2", storage_type, channels, channels, 1, true); + return weights; +} + +HiggsCodecDecoderBlockWeights load_decoder_block(core::BackendWeightStore & store, + const assets::TensorSource & source, + int64_t block_index, + assets::TensorStorageType storage_type) { + if (block_index < 0 || block_index >= kDecoderBlockCount) { + throw std::runtime_error("Higgs TTS codec decoder block index is out of range"); + } + const int64_t in_channels = kDecoderChannels[static_cast(block_index)]; + const int64_t out_channels = kDecoderChannels[static_cast(block_index + 1)]; + const int64_t ratio = kUpsampleRatios[static_cast(block_index)]; + const std::string prefix = "acoustic_decoder.block." + std::to_string(block_index); + + HiggsCodecDecoderBlockWeights weights; + weights.snake = load_snake(store, source, prefix + ".snake1.alpha", in_channels); + weights.conv_transpose = load_conv_transpose1d(store, + source, + prefix + ".conv_t1", + storage_type, + in_channels, + out_channels, + 2 * ratio, + true); + weights.residual_units.reserve(kResidualUnitsPerBlock); + for (int64_t unit = 0; unit < kResidualUnitsPerBlock; ++unit) { + weights.residual_units.push_back( + load_residual_unit(store, + source, + prefix + ".res_unit" + std::to_string(unit + 1), + storage_type, + out_channels)); + } + return weights; +} + +HiggsCodecEncoderBlockWeights load_encoder_block(core::BackendWeightStore & store, + const assets::TensorSource & source, + int64_t block_index, + assets::TensorStorageType storage_type) { + if (block_index < 0 || block_index >= kEncoderBlockCount) { + throw std::runtime_error("Higgs TTS codec encoder block index is out of range"); + } + const int64_t in_channels = kEncoderChannels[static_cast(block_index)]; + const int64_t out_channels = kEncoderChannels[static_cast(block_index + 1)]; + const int64_t ratio = kUpsampleRatios[static_cast(block_index)]; + const std::string prefix = "acoustic_encoder.block." + std::to_string(block_index); + + HiggsCodecEncoderBlockWeights weights; + weights.snake = load_snake(store, source, prefix + ".snake1.alpha", in_channels); + weights.conv = load_conv1d( + store, source, prefix + ".conv1", storage_type, out_channels, in_channels, 2 * ratio, true); + weights.residual_units.reserve(kResidualUnitsPerBlock); + for (int64_t unit = 0; unit < kResidualUnitsPerBlock; ++unit) { + weights.residual_units.push_back( + load_residual_unit(store, + source, + prefix + ".res_unit" + std::to_string(unit + 1), + storage_type, + in_channels)); + } + return weights; +} + +HiggsCodecSemanticResidualUnitWeights +load_semantic_residual_unit(core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type) { + HiggsCodecSemanticResidualUnitWeights weights; + weights.conv1 = load_conv1d(store, + source, + prefix + ".conv1", + storage_type, + kSemanticHiddenSize, + kSemanticHiddenSize, + 3, + false); + weights.conv2 = load_conv1d(store, + source, + prefix + ".conv2", + storage_type, + kSemanticHiddenSize, + kSemanticHiddenSize, + 1, + false); + return weights; +} + +HiggsCodecSemanticEncoderBlockWeights +load_semantic_encoder_block(core::BackendWeightStore & store, + const assets::TensorSource & source, + int64_t block_index, + assets::TensorStorageType storage_type) { + if (block_index < 0 || block_index >= kSemanticBlockCount) { + throw std::runtime_error("Higgs TTS codec semantic encoder block index is out of range"); + } + const std::string prefix = "encoder_semantic.conv_blocks." + std::to_string(block_index); + HiggsCodecSemanticEncoderBlockWeights weights; + weights.residual_units.reserve(kSemanticResidualUnitsPerBlock); + for (int64_t unit = 0; unit < kSemanticResidualUnitsPerBlock; ++unit) { + weights.residual_units.push_back(load_semantic_residual_unit( + store, source, prefix + ".res_units." + std::to_string(unit), storage_type)); + } + weights.conv = load_conv1d(store, + source, + prefix + ".conv", + storage_type, + kSemanticHiddenSize, + kSemanticHiddenSize, + 3, + true); + return weights; +} + +HiggsCodecVectorQuantizerWeights load_quantizer(core::BackendWeightStore & store, + const assets::TensorSource & source, + int64_t index, + assets::TensorStorageType storage_type) { + const std::string prefix = "quantizer.quantizers." + std::to_string(index); + HiggsCodecVectorQuantizerWeights weights; + weights.codebook = store.load_tensor(source, + codec_name(prefix + ".codebook.embed"), + storage_type, + {kCodecCodebookSize, kCodecCodebookDim}); + weights.project_in = load_linear(store, + source, + prefix + ".project_in", + storage_type, + kCodecCodebookDim, + kCodecHiddenSize, + true); + weights.project_out = load_linear(store, + source, + prefix + ".project_out", + storage_type, + kCodecHiddenSize, + kCodecCodebookDim, + true); + return weights; +} + +core::TensorValue +conv_transpose_with_adjusted_output_padding(core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::ConvTranspose1dWeights & weights, + int64_t in_channels, + int64_t out_channels, + int64_t ratio) { + const int64_t kernel = 2 * ratio; + const int64_t padding = (ratio + 1) / 2; + const int64_t output_padding = ratio % 2; + auto full = modules::ConvTranspose1dModule({ + in_channels, + out_channels, + kernel, + static_cast(ratio), + 0, + 1, + weights.bias.has_value(), + }) + .build(ctx, input, weights); + const int64_t cropped_frames = + (input.shape.dims[2] - 1) * ratio - 2 * padding + kernel + output_padding; + if (cropped_frames <= 0 || cropped_frames > full.shape.dims[2]) { + throw std::runtime_error( + "Higgs TTS codec adjusted ConvTranspose1d output length is invalid"); + } + return modules::SliceModule({2, padding, cropped_frames}).build(ctx, full); +} + +core::TensorValue dac_snake(core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::Snake1dWeights & weights, + int64_t channels) { + const auto input_ready = core::ensure_backend_addressable_layout(ctx, input); + const auto input_f32 = + input_ready.type == GGML_TYPE_F32 + ? input_ready + : core::wrap_tensor(ggml_cast(ctx.ggml, input_ready.tensor, GGML_TYPE_F32), + input_ready.shape, + GGML_TYPE_F32); + auto alpha = + core::reshape_tensor(ctx, weights.alpha, core::TensorShape::from_dims({1, channels, 1})); + auto ax = core::wrap_tensor( + ggml_mul(ctx.ggml, input_f32.tensor, alpha.tensor), input_f32.shape, GGML_TYPE_F32); + auto s = core::wrap_tensor(ggml_sin(ctx.ggml, ax.tensor), input_f32.shape, GGML_TYPE_F32); + auto s2 = + core::wrap_tensor(ggml_mul(ctx.ggml, s.tensor, s.tensor), input_f32.shape, GGML_TYPE_F32); + auto denom = core::wrap_tensor( + ggml_scale_bias(ctx.ggml, alpha.tensor, 1.0F, 1.0e-9F), alpha.shape, GGML_TYPE_F32); + auto periodic = core::wrap_tensor( + ggml_div(ctx.ggml, s2.tensor, denom.tensor), input_f32.shape, GGML_TYPE_F32); + return core::wrap_tensor( + ggml_add(ctx.ggml, input_f32.tensor, periodic.tensor), input_f32.shape, GGML_TYPE_F32); +} + +core::TensorValue residual_unit(core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const HiggsCodecResidualUnitWeights & weights, + int64_t channels, + int64_t dilation) { + auto hidden = dac_snake(ctx, input, weights.snake1, channels); + hidden = modules::Conv1dModule({channels, + channels, + 7, + 1, + static_cast(3 * dilation), + static_cast(dilation), + true}) + .build(ctx, hidden, weights.conv1); + hidden = dac_snake(ctx, hidden, weights.snake2, channels); + hidden = modules::Conv1dModule({channels, channels, 1, 1, 0, 1, true}) + .build(ctx, hidden, weights.conv2); + return modules::AddModule{}.build(ctx, input, hidden); +} + +core::TensorValue contiguous(core::ModuleBuildContext & ctx, const core::TensorValue & value) { + return core::ensure_backend_addressable_layout(ctx, value); +} + +core::TensorValue transpose_bct_btc(core::ModuleBuildContext & ctx, + const core::TensorValue & value) { + return modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, value); +} + +core::TensorValue add_same(core::ModuleBuildContext & ctx, + const core::TensorValue & lhs, + const core::TensorValue & rhs) { + return modules::AddModule{}.build(ctx, lhs, rhs); +} + +core::TensorValue +scale(core::ModuleBuildContext & ctx, const core::TensorValue & value, float factor) { + return core::wrap_tensor( + ggml_scale(ctx.ggml, contiguous(ctx, value).tensor, factor), value.shape, GGML_TYPE_F32); +} + +core::TensorValue group_norm_affine(core::ModuleBuildContext & ctx, + const core::TensorValue & input, + int64_t groups, + float eps, + const modules::NormWeights & weights) { + core::TensorValue output; + if (input.shape.rank == 3 && groups == input.shape.dims[1]) { + auto input_f32 = input.type == GGML_TYPE_F32 + ? input + : core::wrap_tensor(ggml_cast(ctx.ggml, input.tensor, GGML_TYPE_F32), + input.shape, + GGML_TYPE_F32); + auto mean = modules::ReduceMeanModule({2}).build(ctx, input_f32); + auto mean_rep = core::wrap_tensor( + ggml_repeat(ctx.ggml, mean.tensor, input_f32.tensor), input_f32.shape, GGML_TYPE_F32); + auto centered = core::wrap_tensor( + ggml_sub(ctx.ggml, input_f32.tensor, mean_rep.tensor), input_f32.shape, GGML_TYPE_F32); + auto variance = modules::ReduceMeanModule({2}).build( + ctx, modules::MulModule().build(ctx, centered, centered)); + auto stddev = core::wrap_tensor( + ggml_sqrt(ctx.ggml, ggml_scale_bias(ctx.ggml, variance.tensor, 1.0F, eps)), + variance.shape, + GGML_TYPE_F32); + auto stddev_rep = core::wrap_tensor( + ggml_repeat(ctx.ggml, stddev.tensor, input_f32.tensor), input_f32.shape, GGML_TYPE_F32); + output = core::wrap_tensor( + ggml_div(ctx.ggml, centered.tensor, stddev_rep.tensor), input_f32.shape, GGML_TYPE_F32); + } else { + output = core::wrap_tensor( + ggml_group_norm(ctx.ggml, input.tensor, groups, eps), input.shape, GGML_TYPE_F32); + } + if (weights.weight.has_value()) { + auto weight = core::reshape_tensor( + ctx, *weights.weight, core::TensorShape::from_dims({1, input.shape.dims[1], 1})); + auto repeated = core::wrap_tensor( + ggml_repeat(ctx.ggml, weight.tensor, output.tensor), output.shape, GGML_TYPE_F32); + output = core::wrap_tensor( + ggml_mul(ctx.ggml, output.tensor, repeated.tensor), output.shape, GGML_TYPE_F32); + } + if (weights.bias.has_value()) { + auto bias = core::reshape_tensor( + ctx, *weights.bias, core::TensorShape::from_dims({1, input.shape.dims[1], 1})); + auto repeated = core::wrap_tensor( + ggml_repeat(ctx.ggml, bias.tensor, output.tensor), output.shape, GGML_TYPE_F32); + output = core::wrap_tensor( + ggml_add(ctx.ggml, output.tensor, repeated.tensor), output.shape, GGML_TYPE_F32); + } + return output; +} + +core::TensorValue grouped_pos_conv(core::ModuleBuildContext & ctx, + const core::TensorValue & input_bct, + const modules::Conv1dWeights & weights) { + constexpr int64_t groups = 16; + constexpr int64_t channels_per_group = kSemanticHiddenSize / groups; + const auto input_contiguous = contiguous(ctx, input_bct); + core::TensorValue out; + for (int64_t group = 0; group < groups; ++group) { + auto input_group = modules::SliceModule({1, group * channels_per_group, channels_per_group}) + .build(ctx, input_contiguous); + auto weight_group = + modules::SliceModule({0, group * channels_per_group, channels_per_group}) + .build(ctx, weights.weight); + modules::Conv1dWeights group_weights{weight_group, std::nullopt}; + if (weights.bias.has_value()) { + group_weights.bias = + modules::SliceModule({0, group * channels_per_group, channels_per_group}) + .build(ctx, *weights.bias); + } + auto group_out = + modules::Conv1dModule( + {channels_per_group, channels_per_group, 128, 1, 64, 1, weights.bias.has_value()}) + .build(ctx, input_group, group_weights); + out = out.valid() ? modules::ConcatModule({1}).build(ctx, out, group_out) : group_out; + } + out = modules::SliceModule({2, 0, input_bct.shape.dims[2]}).build(ctx, out); + return modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, out); +} + +core::TensorValue semantic_self_attention(core::ModuleBuildContext & ctx, + const core::TensorValue & hidden_btc, + const HiggsCodecWeights & weights, + int64_t layer_index) { + constexpr int64_t head_dim = kSemanticHiddenSize / kSemanticAttentionHeads; + const std::string prefix = "encoder.layers." + std::to_string(layer_index) + ".attention"; + auto q = modules::LinearModule({kSemanticHiddenSize, kSemanticHiddenSize, true, GGML_PREC_F32}) + .build(ctx, hidden_btc, semantic_linear_weights(weights, prefix + ".q_proj")); + auto k = modules::LinearModule({kSemanticHiddenSize, kSemanticHiddenSize, true, GGML_PREC_F32}) + .build(ctx, hidden_btc, semantic_linear_weights(weights, prefix + ".k_proj")); + auto v = modules::LinearModule({kSemanticHiddenSize, kSemanticHiddenSize, true, GGML_PREC_F32}) + .build(ctx, hidden_btc, semantic_linear_weights(weights, prefix + ".v_proj")); + q = core::reshape_tensor(ctx, + contiguous(ctx, q), + core::TensorShape::from_dims({hidden_btc.shape.dims[0], + hidden_btc.shape.dims[1], + kSemanticAttentionHeads, + head_dim})); + k = core::reshape_tensor(ctx, + contiguous(ctx, k), + core::TensorShape::from_dims({hidden_btc.shape.dims[0], + hidden_btc.shape.dims[1], + kSemanticAttentionHeads, + head_dim})); + v = core::reshape_tensor(ctx, + contiguous(ctx, v), + core::TensorShape::from_dims({hidden_btc.shape.dims[0], + hidden_btc.shape.dims[1], + kSemanticAttentionHeads, + head_dim})); + q = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, q); + k = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, k); + v = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, v); + const auto k_t = modules::TransposeModule({{0, 1, 3, 2}, 4}).build(ctx, k); + auto scores = modules::MatMulModule().build(ctx, q, k_t); + scores = scale(ctx, scores, static_cast(1.0 / std::sqrt(static_cast(head_dim)))); + auto attn = core::wrap_tensor( + ggml_soft_max(ctx.ggml, contiguous(ctx, scores).tensor), scores.shape, GGML_TYPE_F32); + auto context = modules::MatMulModule().build(ctx, attn, v); + context = modules::TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, context); + context = core::reshape_tensor( + ctx, + contiguous(ctx, context), + core::TensorShape::from_dims( + {hidden_btc.shape.dims[0], hidden_btc.shape.dims[1], kSemanticHiddenSize})); + return modules::LinearModule({kSemanticHiddenSize, kSemanticHiddenSize, true, GGML_PREC_F32}) + .build(ctx, context, semantic_linear_weights(weights, prefix + ".out_proj")); +} + +core::TensorValue semantic_feed_forward(core::ModuleBuildContext & ctx, + const core::TensorValue & hidden_btc, + const HiggsCodecWeights & weights, + int64_t layer_index) { + const std::string prefix = "encoder.layers." + std::to_string(layer_index) + ".feed_forward"; + auto x = + modules::LinearModule({kSemanticHiddenSize, kSemanticIntermediateSize, true, GGML_PREC_F32}) + .build( + ctx, hidden_btc, semantic_linear_weights(weights, prefix + ".intermediate_dense")); + x = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, x); + return modules::LinearModule( + {kSemanticIntermediateSize, kSemanticHiddenSize, true, GGML_PREC_F32}) + .build(ctx, x, semantic_linear_weights(weights, prefix + ".output_dense")); +} + +core::TensorValue downsample_time_by_2(core::ModuleBuildContext & ctx, + const core::TensorValue & hidden_btc, + int64_t target_frames) { + core::TensorValue out; + for (int64_t frame = 0; frame < target_frames; ++frame) { + auto slice = modules::SliceModule({1, frame * 2, 1}).build(ctx, hidden_btc); + out = out.valid() ? modules::ConcatModule({1}).build(ctx, out, slice) : slice; + } + return out; +} + +struct HiggsCodecEncodeGraphValues { + std::array(kCodecCodebooks)> codes = {}; +}; + +core::TensorValue hubert_hidden_state_mean(core::ModuleBuildContext & ctx, + const core::TensorValue & input_values, + const HiggsCodecWeights & weights, + int64_t target_frames) { + auto hidden = core::reshape_tensor( + ctx, + input_values, + core::TensorShape::from_dims({input_values.shape.dims[0], 1, input_values.shape.dims[1]})); + int64_t in_channels = 1; + for (int64_t index = 0; index < kSemanticConvLayers; ++index) { + const std::string prefix = "feature_extractor.conv_layers." + std::to_string(index); + hidden = modules::Conv1dModule( + {in_channels, + kSemanticConvDim[static_cast(index)], + kSemanticConvKernel[static_cast(index)], + static_cast(kSemanticConvStride[static_cast(index)]), + 0, + 1, + false}) + .build(ctx, hidden, semantic_conv_weights(weights, prefix + ".conv", false)); + if (index == 0) { + hidden = group_norm_affine(ctx, + hidden, + kSemanticConvDim[0], + kSemanticLayerNormEps, + semantic_norm_weights(weights, prefix + ".layer_norm")); + } + hidden = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, hidden); + in_channels = kSemanticConvDim[static_cast(index)]; + } + hidden = transpose_bct_btc(ctx, hidden); + hidden = + modules::LayerNormModule({kSemanticConvDim[6], kSemanticLayerNormEps, true, true}) + .build(ctx, hidden, semantic_norm_weights(weights, "feature_projection.layer_norm")); + hidden = + modules::LinearModule({kSemanticConvDim[6], kSemanticHiddenSize, true, GGML_PREC_F32}) + .build(ctx, hidden, semantic_linear_weights(weights, "feature_projection.projection")); + + auto pos = + grouped_pos_conv(ctx, + transpose_bct_btc(ctx, hidden), + semantic_conv_weights(weights, "encoder.pos_conv_embed.conv", true)); + hidden = add_same(ctx, hidden, transpose_bct_btc(ctx, pos)); + hidden = modules::LayerNormModule({kSemanticHiddenSize, kSemanticLayerNormEps, true, true}) + .build(ctx, hidden, semantic_norm_weights(weights, "encoder.layer_norm")); + + auto sum = hidden; + for (int64_t layer = 0; layer < kSemanticLayers; ++layer) { + const std::string prefix = "encoder.layers." + std::to_string(layer); + const auto attn_residual = hidden; + hidden = semantic_self_attention(ctx, hidden, weights, layer); + hidden = add_same(ctx, attn_residual, hidden); + hidden = modules::LayerNormModule({kSemanticHiddenSize, kSemanticLayerNormEps, true, true}) + .build(ctx, hidden, semantic_norm_weights(weights, prefix + ".layer_norm")); + hidden = add_same(ctx, hidden, semantic_feed_forward(ctx, hidden, weights, layer)); + hidden = + modules::LayerNormModule({kSemanticHiddenSize, kSemanticLayerNormEps, true, true}) + .build(ctx, hidden, semantic_norm_weights(weights, prefix + ".final_layer_norm")); + sum = add_same(ctx, sum, hidden); + } + auto hidden_mean = scale(ctx, sum, 1.0F / static_cast(kSemanticLayers + 1)); + auto features = downsample_time_by_2(ctx, hidden_mean, target_frames); + return features; +} + +core::TensorValue semantic_residual_unit(core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const HiggsCodecSemanticResidualUnitWeights & weights) { + auto hidden = modules::EluModule().build(ctx, input); + hidden = modules::Conv1dModule({kSemanticHiddenSize, kSemanticHiddenSize, 3, 1, 1, 1, false}) + .build(ctx, hidden, weights.conv1); + hidden = modules::EluModule().build(ctx, hidden); + hidden = modules::Conv1dModule({kSemanticHiddenSize, kSemanticHiddenSize, 1, 1, 0, 1, false}) + .build(ctx, hidden, weights.conv2); + return modules::AddModule{}.build(ctx, input, hidden); +} + +core::TensorValue semantic_encoder(core::ModuleBuildContext & ctx, + const core::TensorValue & hidden_btc, + const HiggsCodecWeights & weights) { + auto hidden = transpose_bct_btc(ctx, hidden_btc); + hidden = modules::Conv1dModule({kSemanticHiddenSize, kSemanticHiddenSize, 3, 1, 1, 1, false}) + .build(ctx, hidden, weights.semantic_encoder_input); + for (const auto & block : weights.semantic_encoder_blocks) { + for (const auto & unit : block.residual_units) { + hidden = semantic_residual_unit(ctx, hidden, unit); + } + hidden = modules::Conv1dModule({kSemanticHiddenSize, kSemanticHiddenSize, 3, 1, 1, 1, true}) + .build(ctx, hidden, block.conv); + } + return hidden; +} + +core::TensorValue acoustic_encoder(core::ModuleBuildContext & ctx, + const core::TensorValue & waveform, + const HiggsCodecWeights & weights, + int64_t target_frames) { + auto hidden = core::reshape_tensor( + ctx, waveform, core::TensorShape::from_dims({1, 1, waveform.shape.dims[0]})); + hidden = modules::Conv1dModule({1, kEncoderChannels[0], 7, 1, 3, 1, true}) + .build(ctx, hidden, weights.acoustic_encoder_input); + for (int64_t block = 0; block < kEncoderBlockCount; ++block) { + const int64_t in_channels = kEncoderChannels[static_cast(block)]; + const int64_t out_channels = kEncoderChannels[static_cast(block + 1)]; + const int64_t ratio = kUpsampleRatios[static_cast(block)]; + const auto & block_weights = weights.acoustic_encoder_blocks[static_cast(block)]; + for (size_t unit_index = 0; unit_index < block_weights.residual_units.size(); + ++unit_index) { + hidden = residual_unit(ctx, + hidden, + block_weights.residual_units[unit_index], + in_channels, + kResidualDilations[unit_index]); + } + hidden = dac_snake(ctx, hidden, block_weights.snake, in_channels); + hidden = modules::Conv1dModule({in_channels, + out_channels, + 2 * ratio, + static_cast(ratio), + static_cast((ratio + 1) / 2), + 1, + true}) + .build(ctx, hidden, block_weights.conv); + } + hidden = dac_snake(ctx, hidden, weights.acoustic_encoder_output_snake, kEncoderChannels[5]); + hidden = modules::Conv1dModule({kEncoderChannels[5], kAcousticHiddenSize, 3, 1, 1, 1, true}) + .build(ctx, hidden, weights.acoustic_encoder_output); + if (hidden.shape.dims[2] != target_frames) { + hidden = modules::SliceModule({2, 0, target_frames}).build(ctx, hidden); + } + return hidden; +} + +std::array(kCodecCodebooks)> +quantizer_encode(core::ModuleBuildContext & ctx, + const core::TensorValue & embeddings_bct, + const HiggsCodecWeights & weights) { + auto residual = embeddings_bct; + std::array(kCodecCodebooks)> all_codes{}; + for (int64_t codebook = 0; codebook < kCodecCodebooks; ++codebook) { + auto hidden = transpose_bct_btc(ctx, residual); + hidden = + modules::LinearModule({kCodecHiddenSize, kCodecCodebookDim, true, GGML_PREC_F32}) + .build(ctx, hidden, weights.quantizers[static_cast(codebook)].project_in); + auto flat = core::reshape_tensor( + ctx, + contiguous(ctx, hidden), + core::TensorShape::from_dims({hidden.shape.dims[1], kCodecCodebookDim})); + auto codebook_weight = weights.quantizers[static_cast(codebook)].codebook; + auto codebook_t = modules::TransposeModule({{1, 0, 2, 3}, 2}).build(ctx, codebook_weight); + auto dot = modules::MatMulModule().build(ctx, flat, codebook_t); + auto x2 = + modules::ReduceSumModule({1}).build(ctx, modules::MulModule().build(ctx, flat, flat)); + x2 = core::wrap_tensor( + ggml_repeat(ctx.ggml, x2.tensor, dot.tensor), dot.shape, GGML_TYPE_F32); + const auto codebook_weight_f32 = + codebook_weight.type == GGML_TYPE_F32 + ? codebook_weight + : core::wrap_tensor(ggml_cast(ctx.ggml, codebook_weight.tensor, GGML_TYPE_F32), + codebook_weight.shape, + GGML_TYPE_F32); + auto e2 = modules::ReduceSumModule({1}).build( + ctx, modules::MulModule().build(ctx, codebook_weight_f32, codebook_weight_f32)); + e2 = core::reshape_tensor(ctx, e2, core::TensorShape::from_dims({1, kCodecCodebookSize})); + e2 = core::wrap_tensor( + ggml_repeat(ctx.ggml, e2.tensor, dot.tensor), dot.shape, GGML_TYPE_F32); + auto logits = + core::wrap_tensor(ggml_scale(ctx.ggml, dot.tensor, 2.0F), dot.shape, GGML_TYPE_F32); + logits = core::wrap_tensor( + ggml_sub(ctx.ggml, logits.tensor, x2.tensor), logits.shape, GGML_TYPE_F32); + logits = core::wrap_tensor( + ggml_sub(ctx.ggml, logits.tensor, e2.tensor), logits.shape, GGML_TYPE_F32); + auto ids = core::wrap_tensor(ggml_argmax(ctx.ggml, contiguous(ctx, logits).tensor), + core::TensorShape::from_dims({embeddings_bct.shape.dims[2]}), + GGML_TYPE_I32); + all_codes[static_cast(codebook)] = ids; + auto quantized = modules::EmbeddingModule({kCodecCodebookSize, kCodecCodebookDim}) + .build(ctx, ids, codebook_weight); + quantized = + modules::LinearModule({kCodecCodebookDim, kCodecHiddenSize, true, GGML_PREC_F32}) + .build( + ctx, quantized, weights.quantizers[static_cast(codebook)].project_out); + quantized = core::reshape_tensor( + ctx, + contiguous(ctx, quantized), + core::TensorShape::from_dims({1, embeddings_bct.shape.dims[2], kCodecHiddenSize})); + quantized = transpose_bct_btc(ctx, quantized); + residual = core::wrap_tensor( + ggml_sub(ctx.ggml, residual.tensor, quantized.tensor), residual.shape, GGML_TYPE_F32); + } + return all_codes; +} + +HiggsCodecEncodeGraphValues codec_encode(core::ModuleBuildContext & ctx, + const core::TensorValue & waveform_24k, + const core::TensorValue & semantic_waveform_16k, + const HiggsCodecWeights & weights, + int64_t target_frames) { + HiggsCodecEncodeGraphValues out; + auto semantic = + hubert_hidden_state_mean(ctx, semantic_waveform_16k, weights, target_frames); + semantic = semantic_encoder(ctx, semantic, weights); + if (semantic.shape.dims[2] != target_frames) { + semantic = modules::SliceModule({2, 0, target_frames}).build(ctx, semantic); + } + auto acoustic = acoustic_encoder(ctx, waveform_24k, weights, target_frames); + auto concat = modules::ConcatModule({1}).build(ctx, acoustic, semantic); + auto hidden = transpose_bct_btc(ctx, concat); + hidden = modules::LinearModule({kCodecProjectInputSize, kCodecHiddenSize, true, GGML_PREC_F32}) + .build(ctx, hidden, weights.codec_project); + hidden = transpose_bct_btc(ctx, hidden); + out.codes = quantizer_encode(ctx, hidden, weights); + return out; +} + +core::TensorValue quantizer_decode(core::ModuleBuildContext & ctx, + ggml_tensor * codes, + const HiggsCodecWeights & weights, + int64_t frames) { + auto codes_value = core::wrap_tensor( + codes, core::TensorShape::from_dims({frames, kCodecCodebooks}), GGML_TYPE_I32); + std::vector projected; + projected.reserve(weights.quantizers.size()); + for (int64_t codebook = 0; codebook < kCodecCodebooks; ++codebook) { + auto ids = modules::SliceModule({1, codebook, 1}).build(ctx, codes_value); + ids = core::reshape_tensor(ctx, + core::ensure_backend_addressable_layout(ctx, ids), + core::TensorShape::from_dims({frames})); + auto hidden = + modules::EmbeddingModule({kCodecCodebookSize, kCodecCodebookDim}) + .build(ctx, ids, weights.quantizers[static_cast(codebook)].codebook); + hidden = + modules::LinearModule({kCodecCodebookDim, kCodecHiddenSize, true}) + .build(ctx, hidden, weights.quantizers[static_cast(codebook)].project_out); + projected.push_back(hidden); + } + auto sum = projected.front(); + for (size_t index = 1; index < projected.size(); ++index) { + sum = modules::AddModule{}.build(ctx, sum, projected[index]); + } + return sum; +} + +core::TensorValue acoustic_decoder(core::ModuleBuildContext & ctx, + const core::TensorValue & quantized, + const HiggsCodecWeights & weights) { + auto hidden = modules::LinearModule({kCodecHiddenSize, kAcousticHiddenSize, true}) + .build(ctx, quantized, weights.acoustic_project); + hidden = core::reshape_tensor( + ctx, hidden, core::TensorShape::from_dims({1, hidden.shape.dims[0], kAcousticHiddenSize})); + hidden = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, hidden); + hidden = modules::Conv1dModule({kAcousticHiddenSize, kDecoderChannels[0], 7, 1, 3, 1, true}) + .build(ctx, hidden, weights.acoustic_decoder_input); + for (int64_t block = 0; block < kDecoderBlockCount; ++block) { + const int64_t in_channels = kDecoderChannels[static_cast(block)]; + const int64_t out_channels = kDecoderChannels[static_cast(block + 1)]; + const int64_t ratio = kUpsampleRatios[static_cast(block)]; + const auto & block_weights = weights.acoustic_decoder_blocks[static_cast(block)]; + hidden = dac_snake(ctx, hidden, block_weights.snake, in_channels); + hidden = conv_transpose_with_adjusted_output_padding( + ctx, hidden, block_weights.conv_transpose, in_channels, out_channels, ratio); + for (size_t unit_index = 0; unit_index < block_weights.residual_units.size(); + ++unit_index) { + hidden = residual_unit(ctx, + hidden, + block_weights.residual_units[unit_index], + out_channels, + kResidualDilations[unit_index]); + } + } + hidden = dac_snake(ctx, hidden, weights.acoustic_decoder_output_snake, kDecoderChannels[5]); + return modules::Conv1dModule({kDecoderChannels[5], 1, 7, 1, 3, 1, true}) + .build(ctx, hidden, weights.acoustic_decoder_output); +} + +} // namespace + +class HiggsCodecEncodeGraph { +public: + HiggsCodecEncodeGraph(const HiggsCodecRuntime * runtime, + int64_t acoustic_samples, + int64_t semantic_samples, + int64_t frames) + : runtime_(runtime), acoustic_samples_(acoustic_samples), + semantic_samples_(semantic_samples), frames_(frames) { + if (runtime_ == nullptr) { + throw std::runtime_error("Higgs TTS codec encode graph requires runtime"); + } + if (acoustic_samples_ <= 0 || semantic_samples_ <= 0 || frames_ <= 0) { + throw std::runtime_error("Higgs TTS codec encode graph requires positive dimensions"); + } + const auto build_start = Clock::now(); + ggml_init_params params{runtime_->encode_graph_arena_bytes(), nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS codec encode graph context"); + } + core::ModuleBuildContext build_ctx{ + ctx_.get(), "higgs_tts.codec.encode", runtime_->backend_type()}; + acoustic_input_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_F32, acoustic_samples_); + semantic_input_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_F32, semantic_samples_); + ggml_set_input(acoustic_input_); + ggml_set_input(semantic_input_); + auto acoustic = core::wrap_tensor( + acoustic_input_, core::TensorShape::from_dims({acoustic_samples_}), GGML_TYPE_F32); + auto semantic = core::wrap_tensor( + semantic_input_, core::TensorShape::from_dims({1, semantic_samples_}), GGML_TYPE_F32); + auto encoded = codec_encode( + build_ctx, acoustic, semantic, runtime_->weights(), frames_); + graph_ = ggml_new_graph_custom(ctx_.get(), 262144, false); + for (size_t codebook = 0; codebook < outputs_.size(); ++codebook) { + outputs_[codebook] = encoded.codes[codebook].tensor; + ggml_set_output(outputs_[codebook]); + ggml_build_forward_expand(graph_, outputs_[codebook]); + } + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); + if (buffer_ == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS codec encode graph"); + } + engine::debug::timing_log_scalar("higgs_tts.codec.encode.graph.build_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + } + + ~HiggsCodecEncodeGraph() { + engine::core::release_backend_graph_resources(runtime_->backend(), graph_); + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + } + + bool matches(const HiggsCodecRuntime & runtime, + int64_t acoustic_samples, + int64_t semantic_samples, + int64_t frames) const { + return runtime_ == &runtime && acoustic_samples_ == acoustic_samples && + semantic_samples_ == semantic_samples && frames_ == frames; + } + + HiggsCodecEncodeOutput + run(const std::vector & acoustic, const std::vector & semantic, int64_t frames) { + if (static_cast(acoustic.size()) != acoustic_samples_ || + static_cast(semantic.size()) != semantic_samples_ || frames != frames_) { + throw std::runtime_error("Higgs TTS codec encode graph shape mismatch"); + } + auto timing_start = Clock::now(); + ggml_backend_tensor_set( + acoustic_input_, acoustic.data(), 0, acoustic.size() * sizeof(float)); + ggml_backend_tensor_set( + semantic_input_, semantic.data(), 0, semantic.size() * sizeof(float)); + engine::debug::timing_log_scalar("higgs_tts.codec.encode_input_upload_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + timing_start = Clock::now(); + const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); + engine::debug::timing_log_scalar("higgs_tts.codec.encode.graph.compute_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS codec encode graph compute failed"); + } + timing_start = Clock::now(); + HiggsCodecEncodeOutput out; + out.frames = frames; + out.codebooks = kCodecCodebooks; + out.codes.resize(static_cast(frames * kCodecCodebooks)); + std::vector codebook_codes(static_cast(frames_)); + for (int64_t codebook = 0; codebook < kCodecCodebooks; ++codebook) { + ggml_backend_tensor_get(outputs_[static_cast(codebook)], + codebook_codes.data(), + 0, + codebook_codes.size() * sizeof(int32_t)); + for (int64_t frame = 0; frame < frames_; ++frame) { + out.codes[static_cast(frame * kCodecCodebooks + codebook)] = + codebook_codes[static_cast(frame)]; + } + } + engine::debug::timing_log_scalar("higgs_tts.codec.encode_output_read_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + return out; + } + +private: + const HiggsCodecRuntime * runtime_ = nullptr; + int64_t acoustic_samples_ = 0; + int64_t semantic_samples_ = 0; + int64_t frames_ = 0; + std::unique_ptr ctx_; + ggml_tensor * acoustic_input_ = nullptr; + ggml_tensor * semantic_input_ = nullptr; + std::array(kCodecCodebooks)> outputs_ = {}; + ggml_cgraph * graph_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; +}; + +class HiggsCodecDecodeGraph { +public: + HiggsCodecDecodeGraph(const HiggsCodecRuntime * runtime, int64_t frames) + : runtime_(runtime), capacity_frames_(frames) { + if (runtime_ == nullptr) { + throw std::runtime_error("Higgs TTS codec decode graph requires runtime"); + } + if (capacity_frames_ <= 0) { + throw std::runtime_error("Higgs TTS codec decode graph requires positive frame count"); + } + const auto build_start = Clock::now(); + ggml_init_params params{runtime_->decode_graph_arena_bytes(), nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Higgs TTS codec decode graph context"); + } + core::ModuleBuildContext build_ctx{ + ctx_.get(), "higgs_tts.codec.decode", runtime_->backend_type()}; + codes_ = ggml_new_tensor_2d(ctx_.get(), GGML_TYPE_I32, kCodecCodebooks, capacity_frames_); + frame_mask_ = ggml_new_tensor_2d(ctx_.get(), GGML_TYPE_F32, 1, capacity_frames_); + ggml_set_input(codes_); + ggml_set_input(frame_mask_); + auto hidden = quantizer_decode(build_ctx, codes_, runtime_->weights(), capacity_frames_); + const auto mask = core::wrap_tensor( + frame_mask_, core::TensorShape::from_dims({capacity_frames_, 1}), GGML_TYPE_F32); + hidden = modules::MulModule{}.build( + build_ctx, + hidden, + core::wrap_tensor(ggml_repeat(build_ctx.ggml, mask.tensor, hidden.tensor), + hidden.shape, + GGML_TYPE_F32)); + auto audio = acoustic_decoder(build_ctx, hidden, runtime_->weights()); + output_ = audio.tensor; + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + ggml_build_forward_expand(graph_, output_); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); + if (buffer_ == nullptr) { + throw std::runtime_error("failed to allocate Higgs TTS codec decode graph"); + } + code_scratch_.assign(static_cast(capacity_frames_ * kCodecCodebooks), 0); + frame_mask_values_.assign(static_cast(capacity_frames_), 0.0F); + engine::debug::timing_log_scalar("higgs_tts.codec.decode.graph.build_ms", + engine::debug::elapsed_ms(build_start, Clock::now())); + } + + ~HiggsCodecDecodeGraph() { + engine::core::release_backend_graph_resources(runtime_->backend(), graph_); + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + } + + bool matches(const HiggsCodecRuntime & runtime, int64_t frames) const { + return runtime_ == &runtime && frames <= capacity_frames_; + } + + int64_t capacity_frames() const { return capacity_frames_; } + + HiggsCodecDecodeOutput run(const std::vector & codes, int64_t frames) { + if (frames <= 0 || frames > capacity_frames_) { + throw std::runtime_error("Higgs TTS codec decode frame count exceeds graph capacity"); + } + if (static_cast(codes.size()) != frames * kCodecCodebooks) { + throw std::runtime_error("Higgs TTS codec decode code matrix shape mismatch"); + } + for (const int32_t code : codes) { + if (code < 0 || code >= kCodecCodebookSize) { + throw std::runtime_error("Higgs TTS codec decode code is outside codebook range"); + } + } + std::fill(code_scratch_.begin(), code_scratch_.end(), 0); + for (int64_t frame = 0; frame < frames; ++frame) { + const auto src = codes.begin() + static_cast(frame * kCodecCodebooks); + const auto dst = + code_scratch_.begin() + static_cast(frame * kCodecCodebooks); + std::copy_n(src, static_cast(kCodecCodebooks), dst); + } + std::fill(frame_mask_values_.begin(), frame_mask_values_.end(), 0.0F); + std::fill(frame_mask_values_.begin(), + frame_mask_values_.begin() + static_cast(frames), + 1.0F); + auto timing_start = Clock::now(); + ggml_backend_tensor_set( + codes_, code_scratch_.data(), 0, code_scratch_.size() * sizeof(int32_t)); + ggml_backend_tensor_set( + frame_mask_, frame_mask_values_.data(), 0, frame_mask_values_.size() * sizeof(float)); + engine::debug::timing_log_scalar("higgs_tts.codec.decode_input_upload_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + timing_start = Clock::now(); + const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); + engine::debug::timing_log_scalar("higgs_tts.codec.decode.graph.compute_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Higgs TTS codec decode graph compute failed"); + } + HiggsCodecDecodeOutput out; + out.sample_rate = kCodecSampleRate; + out.channels = 1; + out.samples = frames; + for (int64_t block = 0; block < kDecoderBlockCount; ++block) { + out.samples *= kUpsampleRatios[static_cast(block)]; + } + out.values.resize(static_cast(out.samples)); + timing_start = Clock::now(); + ggml_backend_tensor_get(output_, out.values.data(), 0, out.values.size() * sizeof(float)); + engine::debug::timing_log_scalar("higgs_tts.codec.decode_output_read_ms", + engine::debug::elapsed_ms(timing_start, Clock::now())); + return out; + } + +private: + const HiggsCodecRuntime * runtime_ = nullptr; + int64_t capacity_frames_ = 0; + std::unique_ptr ctx_; + ggml_tensor * codes_ = nullptr; + ggml_tensor * frame_mask_ = nullptr; + ggml_tensor * output_ = nullptr; + std::vector code_scratch_; + std::vector frame_mask_values_; + ggml_cgraph * graph_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; +}; + +HiggsCodecWeights load_higgs_codec_decode_weights(const HiggsAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) { + if (assets.weights == nullptr) { + throw std::runtime_error("Higgs TTS codec weights require tensor source"); + } + if (backend == nullptr) { + throw std::runtime_error("Higgs TTS codec backend is not initialized"); + } + HiggsCodecWeights weights; + weights.store = std::make_shared( + backend, backend_type, "higgs_tts.codec.weights", weight_context_bytes); + const auto & source = *assets.weights; + load_hubert_semantic_model_weights(weights, source, weight_storage_type); + weights.quantizers.reserve(kCodecCodebooks); + for (int64_t index = 0; index < kCodecCodebooks; ++index) { + weights.quantizers.push_back( + load_quantizer(*weights.store, source, index, weight_storage_type)); + } + weights.acoustic_encoder_input = load_conv1d(*weights.store, + source, + "acoustic_encoder.conv1", + weight_storage_type, + kEncoderChannels[0], + 1, + 7, + true); + weights.acoustic_encoder_blocks.reserve(kEncoderBlockCount); + for (int64_t block = 0; block < kEncoderBlockCount; ++block) { + weights.acoustic_encoder_blocks.push_back( + load_encoder_block(*weights.store, source, block, weight_storage_type)); + } + weights.acoustic_encoder_output_snake = + load_snake(*weights.store, source, "acoustic_encoder.snake1.alpha", kEncoderChannels[5]); + weights.acoustic_encoder_output = load_conv1d(*weights.store, + source, + "acoustic_encoder.conv2", + weight_storage_type, + kAcousticHiddenSize, + kEncoderChannels[5], + 3, + true); + weights.semantic_encoder_input = load_conv1d(*weights.store, + source, + "encoder_semantic.conv", + weight_storage_type, + kSemanticHiddenSize, + kSemanticHiddenSize, + 3, + false); + weights.semantic_encoder_blocks.reserve(kSemanticBlockCount); + for (int64_t block = 0; block < kSemanticBlockCount; ++block) { + weights.semantic_encoder_blocks.push_back( + load_semantic_encoder_block(*weights.store, source, block, weight_storage_type)); + } + weights.codec_project = load_linear(*weights.store, + source, + "fc", + weight_storage_type, + kCodecHiddenSize, + kCodecProjectInputSize, + true); + weights.acoustic_project = load_linear(*weights.store, + source, + "fc2", + weight_storage_type, + kAcousticHiddenSize, + kCodecHiddenSize, + true); + weights.acoustic_decoder_input = load_conv1d(*weights.store, + source, + "acoustic_decoder.conv1", + weight_storage_type, + kDecoderChannels[0], + kAcousticHiddenSize, + 7, + true); + weights.acoustic_decoder_blocks.reserve(kDecoderBlockCount); + for (int64_t block = 0; block < kDecoderBlockCount; ++block) { + weights.acoustic_decoder_blocks.push_back( + load_decoder_block(*weights.store, source, block, weight_storage_type)); + } + weights.acoustic_decoder_output_snake = + load_snake(*weights.store, source, "acoustic_decoder.snake1.alpha", kDecoderChannels[5]); + weights.acoustic_decoder_output = load_conv1d(*weights.store, + source, + "acoustic_decoder.conv2", + weight_storage_type, + 1, + kDecoderChannels[5], + 7, + true); + weights.store->upload(); + return weights; +} + +HiggsCodecRuntime::HiggsCodecRuntime(std::shared_ptr assets, + core::ExecutionContext & execution, + size_t weight_context_bytes, + size_t decode_graph_arena_bytes, + size_t encode_graph_arena_bytes, + assets::TensorStorageType weight_storage_type) + : assets_(std::move(assets)), backend_(execution.backend()), + backend_type_(execution.backend_type()), threads_(std::max(1, execution.config().threads)), + decode_graph_arena_bytes_(decode_graph_arena_bytes), + encode_graph_arena_bytes_(encode_graph_arena_bytes), + weights_(nullptr) { + if (assets_ == nullptr) { + throw std::runtime_error("Higgs TTS codec runtime requires assets"); + } + if (assets_->weights == nullptr) { + throw std::runtime_error("Higgs TTS codec runtime requires tensor source"); + } + weights_ = std::make_shared(load_higgs_codec_decode_weights( + *assets_, backend_, backend_type_, weight_context_bytes, weight_storage_type)); + if (decode_graph_arena_bytes_ == 0) { + throw std::runtime_error("Higgs TTS codec decode graph arena bytes must be non-zero"); + } + if (encode_graph_arena_bytes_ == 0) { + throw std::runtime_error("Higgs TTS codec encode graph arena bytes must be non-zero"); + } +} + +HiggsCodecRuntime::~HiggsCodecRuntime() = default; + +const HiggsCodecWeights & HiggsCodecRuntime::weights() const noexcept { return *weights_; } + +ggml_backend_t HiggsCodecRuntime::backend() const noexcept { return backend_; } + +core::BackendType HiggsCodecRuntime::backend_type() const noexcept { return backend_type_; } + +int HiggsCodecRuntime::threads() const noexcept { return threads_; } + +size_t HiggsCodecRuntime::decode_graph_arena_bytes() const noexcept { + return decode_graph_arena_bytes_; +} + +size_t HiggsCodecRuntime::encode_graph_arena_bytes() const noexcept { + return encode_graph_arena_bytes_; +} + +HiggsCodecEncodeOutput +HiggsCodecRuntime::encode_reference(const runtime::AudioBuffer & audio) const { + const auto mono = prepare_mono_audio(audio); + const auto acoustic_24k_base = prepare_codec_audio_24k(mono, audio.sample_rate); + auto semantic_16k = prepare_semantic_audio_16k(mono, audio.sample_rate); + const int64_t frames = semantic_feature_frames(static_cast(semantic_16k.size())); + + std::vector acoustic_24k = acoustic_24k_base; + const int64_t acoustic_frames = + acoustic_encoder_frames(static_cast(acoustic_24k.size())); + if (acoustic_frames != frames) { + acoustic_24k = pad_zeros(acoustic_24k_base, kCodecPadSamples, kCodecPadSamples); + const int64_t padded_frames = + acoustic_encoder_frames(static_cast(acoustic_24k.size())); + if (padded_frames != frames) { + throw std::runtime_error("Higgs TTS codec acoustic and semantic encoder " + "frame counts do not match"); + } + } + + if (encode_graph_ == nullptr || + !encode_graph_->matches(*this, + static_cast(acoustic_24k.size()), + static_cast(semantic_16k.size()), + frames)) { + encode_graph_.reset(); + encode_graph_ = + std::make_unique(this, + static_cast(acoustic_24k.size()), + static_cast(semantic_16k.size()), + frames); + } + engine::debug::trace_log_scalar("higgs_tts.codec.encode.input_frames", frames); + engine::debug::trace_log_f32("higgs_tts.codec.encode.input_acoustic_24k", + {static_cast(acoustic_24k.size())}, + acoustic_24k); + engine::debug::trace_log_f32("higgs_tts.codec.encode.input_semantic_16k", + {static_cast(semantic_16k.size())}, + semantic_16k); + return encode_graph_->run(acoustic_24k, semantic_16k, frames); +} + +HiggsCodecDecodeOutput HiggsCodecRuntime::decode_codes(const std::vector & codes, + int64_t frames, + int64_t codebooks) const { + if (frames <= 0) { + throw std::runtime_error("Higgs TTS codec decode requires positive frame count"); + } + if (codebooks != kCodecCodebooks) { + throw std::runtime_error("Higgs TTS codec decode requires exactly 8 codebooks"); + } + if (static_cast(codes.size()) != frames * codebooks) { + throw std::runtime_error("Higgs TTS codec decode code count mismatch"); + } + engine::debug::trace_log_scalar("higgs_tts.codec.decode.input_frames", frames); + engine::debug::trace_log_scalar("higgs_tts.codec.decode.input_codebooks", codebooks); + engine::debug::trace_log_i32("higgs_tts.codec.decode.input_codes", + {frames, codebooks}, + codes); + + auto run_window = [&](const std::vector & window_codes, + int64_t window_frames, + int64_t min_capacity_frames) -> HiggsCodecDecodeOutput { + const int64_t bucketed_frames = + ((std::max(window_frames, min_capacity_frames) + + kCodecDecodeCapacityBucketFrames - 1) / + kCodecDecodeCapacityBucketFrames) * + kCodecDecodeCapacityBucketFrames; + encode_graph_.reset(); + if (decode_graph_ == nullptr || !decode_graph_->matches(*this, window_frames)) { + decode_graph_.reset(); + decode_graph_ = std::make_unique(this, bucketed_frames); + } + return decode_graph_->run(window_codes, window_frames); + }; + + if (frames <= kCodecFullDecodeMaxFrames) { + return run_window(codes, frames, frames); + } + + HiggsCodecDecodeOutput out; + out.sample_rate = kCodecSampleRate; + out.channels = 1; + out.samples = frames * kCodecHopLength; + out.values.reserve(static_cast(out.samples)); + + std::vector window_codes; + int64_t emitted_frames = 0; + while (emitted_frames < frames) { + const int64_t window_begin = + std::max(0, emitted_frames - kCodecDecodeOverlapFrames); + const int64_t emit_end = + std::min(frames, emitted_frames + kCodecDecodeWindowFrames); + const int64_t window_frames = emit_end - window_begin; + window_codes.resize(static_cast(window_frames * kCodecCodebooks)); + for (int64_t frame = 0; frame < window_frames; ++frame) { + const auto src = + codes.begin() + + static_cast((window_begin + frame) * kCodecCodebooks); + auto dst = + window_codes.begin() + static_cast(frame * kCodecCodebooks); + std::copy_n(src, static_cast(kCodecCodebooks), dst); + } + + const auto window = run_window( + window_codes, + window_frames, + kCodecDecodeWindowFrames + kCodecDecodeOverlapFrames); + const int64_t trim_frames = emitted_frames - window_begin; + const int64_t emit_frames = emit_end - emitted_frames; + const int64_t sample_begin = trim_frames * kCodecHopLength; + const int64_t sample_count = emit_frames * kCodecHopLength; + if (sample_begin < 0 || sample_count <= 0 || + sample_begin + sample_count > static_cast(window.values.size())) { + throw std::runtime_error("Higgs TTS codec decode window produced invalid length"); + } + out.values.insert(out.values.end(), + window.values.begin() + static_cast(sample_begin), + window.values.begin() + + static_cast(sample_begin + sample_count)); + emitted_frames = emit_end; + } + if (static_cast(out.values.size()) != out.samples) { + throw std::runtime_error("Higgs TTS codec chunked decode output length mismatch"); + } + return out; +} + +} // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/generator.cpp b/src/models/higgs_tts/generator.cpp new file mode 100644 index 00000000..51a4dd3b --- /dev/null +++ b/src/models/higgs_tts/generator.cpp @@ -0,0 +1,503 @@ +#include "engine/models/higgs_tts/generator.h" + +#include "engine/framework/debug/profiler.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/sampling/torch_random.h" +#include "engine/models/higgs_tts/codebooks.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::higgs_tts { +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr int64_t kInitialGeneratedCacheSteps = 512; + +void validate_generation_options(const HiggsGenerationOptions & options) { + if (options.max_tokens <= 0) { + throw std::runtime_error("Higgs TTS max_tokens must be positive"); + } + if (!(options.temperature > 0.0F)) { + throw std::runtime_error("Higgs TTS temperature must be positive"); + } + if (options.top_p.has_value() && !(*options.top_p > 0.0F)) { + throw std::runtime_error("Higgs TTS top_p must be positive"); + } + if (options.top_k.has_value() && *options.top_k < 0) { + throw std::runtime_error("Higgs TTS top_k must be non-negative"); + } + if (!(options.repetition_penalty > 0.0F) || !std::isfinite(options.repetition_penalty)) { + throw std::runtime_error("Higgs TTS repetition_penalty must be finite and positive"); + } +} + +size_t flat_index(int64_t frame, int64_t codebook, int64_t codebooks) { + return static_cast(frame * codebooks + codebook); +} + +struct HiggsPromptInput { + std::vector token_ids; + std::vector reference_positions; +}; + +struct HiggsPreparedPrompt { + HiggsPromptInput prompt; + HiggsARPrefillInput ar_input; + int64_t prefix_steps = 0; +}; + +HiggsPromptInput make_prompt_input(const HiggsPromptEncoding & prompt, + int64_t delayed_reference_frames, + const HiggsConfig & config) { + HiggsPromptInput input; + input.token_ids = prompt.token_ids; + input.reference_positions.reserve(static_cast(delayed_reference_frames)); + for (int64_t position = 0; position < static_cast(input.token_ids.size()); + ++position) { + if (input.token_ids[static_cast(position)] == config.audio_token_id) { + input.reference_positions.push_back(static_cast(position)); + input.token_ids[static_cast(position)] = 0; + } + } + if (static_cast(input.reference_positions.size()) != delayed_reference_frames) { + throw std::runtime_error("Higgs TTS prompt audio placeholder count does " + "not match delayed reference codes"); + } + return input; +} + +HiggsPreparedPrompt make_prepared_prompt(const HiggsPromptEncoding & prompt, + const std::vector & delayed_reference_codes, + int64_t delayed_reference_frames, + const HiggsConfig & config) { + HiggsPreparedPrompt prepared; + prepared.prompt = make_prompt_input(prompt, delayed_reference_frames, config); + const int64_t prompt_steps = static_cast(prepared.prompt.token_ids.size()); + prepared.ar_input.steps = prompt_steps; + prepared.ar_input.text_tokens = prepared.prompt.token_ids; + prepared.ar_input.fused_code_ids.assign( + static_cast(prompt_steps * config.audio.num_codebooks), 0); + prepared.ar_input.text_gate.assign(static_cast(prompt_steps), 0.0F); + prepared.ar_input.code_gate.assign(static_cast(prompt_steps), 0.0F); + size_t reference_row = 0; + for (int64_t position = 0; position < prompt_steps; ++position) { + if (reference_row < prepared.prompt.reference_positions.size() && + prepared.prompt.reference_positions[reference_row] == position) { + for (int64_t codebook = 0; codebook < config.audio.num_codebooks; ++codebook) { + const int32_t code = delayed_reference_codes[flat_index( + static_cast(reference_row), codebook, config.audio.num_codebooks)]; + if (code < 0 || code >= config.audio.vocab_size) { + throw std::runtime_error( + "Higgs TTS AR prefill codebook token is outside vocabulary"); + } + prepared.ar_input + .fused_code_ids[flat_index(position, codebook, config.audio.num_codebooks)] = + static_cast(code + codebook * config.audio.vocab_size); + } + prepared.ar_input.code_gate[static_cast(position)] = 1.0F; + ++reference_row; + } else { + const int32_t text_token = prepared.prompt.token_ids[static_cast(position)]; + if (text_token < 0 || text_token >= config.text.vocab_size) { + throw std::runtime_error("Higgs TTS AR prefill text token is outside vocabulary"); + } + prepared.ar_input.text_gate[static_cast(position)] = 1.0F; + } + } + if (reference_row != prepared.prompt.reference_positions.size()) { + throw std::runtime_error( + "Higgs TTS prompt prefill did not consume all reference code rows"); + } + if (!prepared.prompt.reference_positions.empty()) { + prepared.prefix_steps = + static_cast(prepared.prompt.reference_positions.back()) + 1; + } + return prepared; +} + +} // namespace + +HiggsGenerator::HiggsGenerator(std::shared_ptr assets, + std::shared_ptr ar, + std::shared_ptr codec, + size_t ar_decode_graph_arena_bytes) + : assets_([&]() { + if (assets == nullptr) { + throw std::runtime_error("Higgs TTS generator requires assets"); + } + return std::move(assets); + }()), + ar_([&]() { + if (ar == nullptr) { + throw std::runtime_error("Higgs TTS generator requires AR runtime"); + } + return std::move(ar); + }()), + codec_([&]() { + if (codec == nullptr) { + throw std::runtime_error("Higgs TTS generator requires codec runtime"); + } + return std::move(codec); + }()), + tokenizer_(assets_), + ar_decode_graph_arena_bytes_(ar_decode_graph_arena_bytes) { + if (ar_decode_graph_arena_bytes_ == 0) { + throw std::runtime_error("Higgs TTS generator graph arena bytes must be non-zero"); + } +} + +void HiggsGenerator::prepare(const HiggsGenerationRequest & request) { + const auto & config = assets_->config; + const bool has_reference = request.reference_frames > 0 || !request.reference_codes.empty(); + if (has_reference) { + if (request.reference_frames <= 0 || + request.reference_codebooks != config.audio.num_codebooks) { + throw std::runtime_error("Higgs TTS generation requires reference codes " + "shaped [frames, num_codebooks]"); + } + if (static_cast(request.reference_codes.size()) != + request.reference_frames * request.reference_codebooks) { + throw std::runtime_error("Higgs TTS generation reference code count mismatch"); + } + } else if (request.reference_codebooks != 0) { + throw std::runtime_error( + "Higgs TTS generation got reference codebooks without reference codes"); + } + validate_generation_options(request.options); + const int64_t delayed_reference_frames = + has_reference + ? higgs_delayed_frame_count(request.reference_frames, request.reference_codebooks) + : 0; + std::vector delayed_reference_codes; + if (has_reference) { + delayed_reference_codes = apply_higgs_delay_pattern( + request.reference_codes, request.reference_frames, request.reference_codebooks); + } + const HiggsPromptEncoding prompt = tokenizer_.encode_prompt({ + request.text, + request.reference_text, + delayed_reference_frames, + }); + const auto prepared = + make_prepared_prompt(prompt, delayed_reference_codes, delayed_reference_frames, config); + const int64_t prompt_steps = prepared.ar_input.steps; + if (prompt_steps + request.options.max_tokens > config.text.max_position_embeddings) { + throw std::runtime_error("Higgs TTS generation exceeds text model max_position_embeddings"); + } + if (has_reference && prepared.prefix_steps > 0) { + ReferencePrefixCache cache; + cache.reference_text = request.reference_text; + cache.reference_codes = request.reference_codes; + cache.reference_frames = request.reference_frames; + cache.reference_codebooks = request.reference_codebooks; + cache.delayed_reference_codes = std::move(delayed_reference_codes); + cache.delayed_reference_frames = delayed_reference_frames; + cache.prefix_steps = prepared.prefix_steps; + cache.prefix_tokens.assign(prepared.prompt.token_ids.begin(), + prepared.prompt.token_ids.begin() + + static_cast(prepared.prefix_steps)); + reference_prefix_cache_ = std::move(cache); + } else { + reference_prefix_cache_.reset(); + } +} + +HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & request) { + const auto & config = assets_->config; + const bool has_reference = request.reference_frames > 0 || !request.reference_codes.empty(); + if (has_reference) { + if (request.reference_frames <= 0 || + request.reference_codebooks != config.audio.num_codebooks) { + throw std::runtime_error("Higgs TTS generation requires reference codes " + "shaped [frames, num_codebooks]"); + } + if (static_cast(request.reference_codes.size()) != + request.reference_frames * request.reference_codebooks) { + throw std::runtime_error("Higgs TTS generation reference code count mismatch"); + } + } else if (request.reference_codebooks != 0) { + throw std::runtime_error( + "Higgs TTS generation got reference codebooks without reference codes"); + } + validate_generation_options(request.options); + engine::debug::trace_log_scalar("higgs_tts.request.text", request.text); + engine::debug::trace_log_scalar("higgs_tts.request.reference_text", request.reference_text); + engine::debug::trace_log_scalar("higgs_tts.request.text_chars", request.text.size()); + engine::debug::trace_log_scalar("higgs_tts.request.reference_text_chars", + request.reference_text.size()); + engine::debug::trace_log_scalar("higgs_tts.request.max_tokens", + request.options.max_tokens); + engine::debug::trace_log_scalar("higgs_tts.request.temperature", request.options.temperature); + engine::debug::trace_log_scalar("higgs_tts.request.top_p", + request.options.top_p.has_value() ? + std::to_string(*request.options.top_p) : "none"); + engine::debug::trace_log_scalar("higgs_tts.request.top_k", + request.options.top_k.has_value() ? + std::to_string(*request.options.top_k) : "none"); + engine::debug::trace_log_scalar("higgs_tts.request.repetition_penalty", + request.options.repetition_penalty); + engine::debug::trace_log_scalar("higgs_tts.request.has_seed", request.options.seed.has_value()); + engine::debug::trace_log_scalar("higgs_tts.request.seed", + request.options.seed.has_value() ? + std::to_string(*request.options.seed) : "none"); + if (has_reference) { + engine::debug::trace_log_i32("higgs_tts.request.reference_codes", + {request.reference_frames, request.reference_codebooks}, + request.reference_codes); + } + + std::vector delayed_reference_codes_storage; + const std::vector * delayed_reference_codes = &delayed_reference_codes_storage; + int64_t delayed_reference_frames = 0; + const ReferencePrefixCache * matching_reference_cache = nullptr; + if (has_reference) { + if (reference_prefix_cache_.has_value() && + reference_prefix_cache_->reference_text == request.reference_text && + reference_prefix_cache_->reference_frames == request.reference_frames && + reference_prefix_cache_->reference_codebooks == request.reference_codebooks && + reference_prefix_cache_->reference_codes == request.reference_codes) { + matching_reference_cache = &*reference_prefix_cache_; + delayed_reference_codes = &matching_reference_cache->delayed_reference_codes; + delayed_reference_frames = matching_reference_cache->delayed_reference_frames; + } else { + delayed_reference_codes_storage = apply_higgs_delay_pattern( + request.reference_codes, request.reference_frames, request.reference_codebooks); + delayed_reference_frames = + higgs_delayed_frame_count(request.reference_frames, request.reference_codebooks); + } + } + const HiggsPromptEncoding prompt = tokenizer_.encode_prompt({ + request.text, + request.reference_text, + delayed_reference_frames, + }); + engine::debug::trace_log_scalar("higgs_tts.prompt.text_tokens", prompt.text_ids.size()); + engine::debug::trace_log_i32("higgs_tts.prompt.text_ids", + {static_cast(prompt.text_ids.size())}, + prompt.text_ids); + engine::debug::trace_log_scalar("higgs_tts.prompt.reference_text_tokens", + prompt.reference_text_ids.size()); + engine::debug::trace_log_i32("higgs_tts.prompt.reference_text_ids", + {static_cast(prompt.reference_text_ids.size())}, + prompt.reference_text_ids); + const auto prepared = + make_prepared_prompt(prompt, *delayed_reference_codes, delayed_reference_frames, config); + engine::debug::trace_log_scalar("higgs_tts.prompt.tokens", prepared.prompt.token_ids.size()); + engine::debug::trace_log_i32("higgs_tts.prompt.token_ids", + {static_cast(prepared.prompt.token_ids.size())}, + prepared.prompt.token_ids); + engine::debug::trace_log_scalar("higgs_tts.prompt.delayed_reference_rows", + delayed_reference_frames); + engine::debug::trace_log_i32("higgs_tts.prompt.delayed_reference_codes", + {delayed_reference_frames, config.audio.num_codebooks}, + *delayed_reference_codes); + engine::debug::trace_log_i32("higgs_tts.ar.prefill.text_tokens", + {prepared.ar_input.steps}, + prepared.ar_input.text_tokens); + engine::debug::trace_log_i32("higgs_tts.ar.prefill.fused_code_ids", + {prepared.ar_input.steps, config.audio.num_codebooks}, + prepared.ar_input.fused_code_ids); + engine::debug::trace_log_f32("higgs_tts.ar.prefill.text_gate", + {prepared.ar_input.steps}, + prepared.ar_input.text_gate); + engine::debug::trace_log_f32("higgs_tts.ar.prefill.code_gate", + {prepared.ar_input.steps}, + prepared.ar_input.code_gate); + const int64_t prompt_steps = prepared.ar_input.steps; + if (prompt_steps + request.options.max_tokens > config.text.max_position_embeddings) { + throw std::runtime_error("Higgs TTS generation exceeds text model max_position_embeddings"); + } + + const auto prefill_start = Clock::now(); + const bool reference_cache_hit = + has_reference && prepared.prefix_steps > 0 && matching_reference_cache != nullptr && + matching_reference_cache->prefix_steps == prepared.prefix_steps && + static_cast(matching_reference_cache->prefix_tokens.size()) == prepared.prefix_steps && + std::equal(matching_reference_cache->prefix_tokens.begin(), + matching_reference_cache->prefix_tokens.end(), + prepared.prompt.token_ids.begin()); + engine::debug::trace_log_scalar("higgs_tts.generator.reference_prefix_cache_hit", reference_cache_hit); + engine::debug::trace_log_scalar("higgs_tts.generator.reference_prefix_steps", prepared.prefix_steps); + const int64_t max_cache_steps = prompt_steps + request.options.max_tokens; + const int64_t initial_cache_steps = + prompt_steps + std::min(request.options.max_tokens, kInitialGeneratedCacheSteps); + if (ar_kv_cache_ == nullptr || !ar_kv_cache_->can_run(*ar_, initial_cache_steps)) { + decode_graph_.reset(); + ar_kv_cache_ = std::make_unique(ar_, initial_cache_steps); + } + ar_kv_cache_->reset(); + if (prefill_graph_ == nullptr || !prefill_graph_->matches(*ar_, prompt_steps, 0)) { + prefill_graph_.reset(); + prefill_graph_ = std::make_unique( + ar_, prompt_steps, 0, ar_kv_cache_.get(), ar_decode_graph_arena_bytes_); + } + auto prefill_output = prefill_graph_->run(prepared.ar_input, 0); + prefill_graph_.reset(); + + if (decode_graph_ == nullptr || !decode_graph_->can_run(*ar_, ar_kv_cache_->cache_steps())) { + decode_graph_ = std::make_unique( + ar_, ar_kv_cache_->cache_steps(), *ar_kv_cache_, ar_decode_graph_arena_bytes_); + } + if (!prefill_output.wrote_cache) { + decode_graph_->import_prefill_state(prefill_output.kv_state); + } + HiggsARDecodeOutput prefill = std::move(prefill_output.output); + engine::debug::timing_log_scalar("higgs_tts.generator.prefill_ms", + engine::debug::elapsed_ms(prefill_start, Clock::now())); + engine::debug::trace_log_f32("higgs_tts.sampler.prefill_logits", + {config.audio.num_codebooks, config.audio.vocab_size}, + prefill.codebook_logits); + + HiggsCodebookSampler sampler(config.audio.num_codebooks, config.audio.vocab_size); + HiggsSamplerState state = sampler.make_state(); + HiggsSamplingOptions sampling; + sampling.temperature = request.options.temperature; + sampling.top_p = request.options.top_p; + sampling.top_k = request.options.top_k; + // Python accepts repetition_penalty on the public speech request, but the + // Higgs audio-codebook sampler does not consume it. + sampling.has_seed = request.options.seed.has_value(); + sampling.seed = request.options.seed.value_or(runtime::random_u64_seed()); + if (!cuda_sampling_policy_.has_value()) { + cuda_sampling_policy_ = engine::sampling::resolve_torch_cuda_sampling_policy( + ar_->backend_type(), + ar_->device(), + "higgs_tts.cuda_sampling_policy", + "Higgs TTS", + engine::sampling::TorchCudaSamplingPolicyFailureMode::StrictCuda); + } + sampling.cuda_policy = *cuda_sampling_policy_; + engine::debug::trace_log_scalar("higgs_tts.sampler.temperature", sampling.temperature); + engine::debug::trace_log_scalar("higgs_tts.sampler.has_seed", sampling.has_seed); + engine::debug::trace_log_scalar("higgs_tts.sampler.seed", sampling.seed); + engine::debug::trace_log_scalar("higgs_tts.sampler.top_p", + sampling.top_p.has_value() ? std::to_string(*sampling.top_p) + : "none"); + engine::debug::trace_log_scalar("higgs_tts.sampler.top_k", + sampling.top_k.has_value() ? std::to_string(*sampling.top_k) + : "none"); + HiggsGenerationResult result; + result.delayed_codes.reserve( + static_cast(request.options.max_tokens * config.audio.num_codebooks)); + const auto & first_sampled = sampler.step(prefill.codebook_logits.data(), + static_cast(prefill.codebook_logits.size()), + state, + sampling); + engine::debug::trace_log_i32("higgs_tts.sampler.first_output_codes", + {static_cast(first_sampled.size())}, + first_sampled); + result.delayed_codes.insert( + result.delayed_codes.end(), first_sampled.begin(), first_sampled.end()); + result.delayed_frames += 1; + + const auto decode_start = Clock::now(); + HiggsARDecodeOutput decoded; + decoded.codebook_logits.reserve( + static_cast(config.audio.num_codebooks * config.audio.vocab_size)); + bool logged_decode_step_timing = false; + double sampler_total_ms = 0.0; + HiggsARDecodeTiming decode_timing_total; + decode_graph_->begin_decode_run(); + while (!state.generation_done && result.delayed_frames < request.options.max_tokens) { + if (ar_kv_cache_->valid_steps() >= ar_kv_cache_->cache_steps()) { + decode_timing_total.add(decode_graph_->timing()); + const auto kv_state = ar_kv_cache_->export_state(); + const int64_t grown_cache_steps = + std::min(max_cache_steps, + std::max(ar_kv_cache_->cache_steps() * 2, ar_kv_cache_->valid_steps() + 1)); + if (grown_cache_steps <= ar_kv_cache_->cache_steps()) { + throw std::runtime_error("Higgs TTS AR cache cannot grow"); + } + decode_graph_.reset(); + ar_kv_cache_ = std::make_unique(ar_, grown_cache_steps); + ar_kv_cache_->import_state(kv_state); + decode_graph_ = std::make_unique( + ar_, ar_kv_cache_->cache_steps(), *ar_kv_cache_, ar_decode_graph_arena_bytes_); + decode_graph_->begin_decode_run(); + } + HiggsARDecodeInput input; + input.use_last_codes = state.delay_count > 0; + input.last_codes = state.last_codes; + const auto step_start = Clock::now(); + decode_graph_->run_step_into(input, decoded, !logged_decode_step_timing); + if (!logged_decode_step_timing) { + engine::debug::timing_log_scalar("higgs_tts.generator.decode.step0.ar_ms", + engine::debug::elapsed_ms(step_start, Clock::now())); + } + const auto sample_start = Clock::now(); + const auto & sampled = sampler.step(decoded.codebook_logits.data(), + static_cast(decoded.codebook_logits.size()), + state, + sampling); + sampler_total_ms += engine::debug::elapsed_ms(sample_start, Clock::now()); + if (!logged_decode_step_timing) { + engine::debug::timing_log_scalar("higgs_tts.generator.decode.step0.sampler_ms", + sampler_total_ms); + logged_decode_step_timing = true; + } + if (!sampled.empty() && sampled.front() != kHiggsStopCode) { + result.delayed_codes.insert(result.delayed_codes.end(), sampled.begin(), sampled.end()); + result.delayed_frames += 1; + } + } + decode_timing_total.add(decode_graph_->timing()); + engine::debug::timing_log_scalar("higgs_tts.ar.decode.steps", decode_timing_total.steps); + engine::debug::timing_log_scalar("higgs_tts.ar.decode.input_upload_ms", decode_timing_total.input_upload_ms); + engine::debug::timing_log_scalar("higgs_tts.ar.decode.mask_upload_ms", decode_timing_total.mask_upload_ms); + engine::debug::timing_log_scalar("higgs_tts.ar.decode.graph.compute_ms", decode_timing_total.graph_compute_ms); + engine::debug::timing_log_scalar("higgs_tts.ar.decode.output_read_ms", decode_timing_total.output_read_ms); + engine::debug::timing_log_scalar("higgs_tts.generator.decode.sampler_ms", sampler_total_ms); + engine::debug::timing_log_scalar("higgs_tts.generator.decode_ms", + engine::debug::elapsed_ms(decode_start, Clock::now())); + if (!state.generation_done) { + throw std::runtime_error("Higgs TTS generation reached max_tokens before EOC"); + } + + result.raw_codes = reverse_higgs_delay_pattern( + result.delayed_codes, result.delayed_frames, config.audio.num_codebooks); + result.raw_frames = result.delayed_frames - (config.audio.num_codebooks - 1); + engine::debug::trace_log_i32("higgs_tts.generator.delayed_codes", + {result.delayed_frames, config.audio.num_codebooks}, + result.delayed_codes); + const int64_t delayed_head_rows = std::min(result.delayed_frames, 8); + engine::debug::trace_log_i32("higgs_tts.generator.delayed_codes_head8", + {delayed_head_rows, config.audio.num_codebooks}, + std::vector( + result.delayed_codes.begin(), + result.delayed_codes.begin() + + static_cast( + delayed_head_rows * config.audio.num_codebooks))); + const int32_t codec_vocab = static_cast(config.audio.vocab_size - 2); +#ifdef _OPENMP +#pragma omp parallel for if (static_cast(result.raw_codes.size()) > 1024) +#endif + for (int64_t index = 0; index < static_cast(result.raw_codes.size()); ++index) { + int32_t & code = result.raw_codes[static_cast(index)]; + if (code >= codec_vocab) { + code = 0; + } + } + engine::debug::trace_log_i32("higgs_tts.generator.raw_codes_for_codec", + {result.raw_frames, config.audio.num_codebooks}, + result.raw_codes); + const auto codec_start = Clock::now(); + result.audio = + codec_->decode_codes(result.raw_codes, result.raw_frames, config.audio.num_codebooks); + engine::debug::trace_log_f32("higgs_tts.codec.decode.output_audio", + {result.audio.samples}, + result.audio.values); + engine::debug::timing_log_scalar("higgs_tts.generator.codec_decode_ms", + engine::debug::elapsed_ms(codec_start, Clock::now())); + return result; +} + +} // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/loader.cpp b/src/models/higgs_tts/loader.cpp new file mode 100644 index 00000000..a39dc9a8 --- /dev/null +++ b/src/models/higgs_tts/loader.cpp @@ -0,0 +1,146 @@ +#include "engine/models/higgs_tts/loader.h" + +#include "engine/framework/assets/model_package.h" +#include "engine/models/higgs_tts/session.h" + +#include +#include + +namespace engine::models::higgs_tts { +namespace { + +runtime::ModelMetadata metadata(const HiggsAssets &) { + runtime::ModelMetadata out; + out.family = "higgs_tts"; + out.variant = "v3-4b"; + out.description = "Higgs Audio v3 TTS loaded from local SGLang-Omni compatible assets."; + out.config_candidates = { + "config.json", + "tokenizer.json", + "tokenizer_config.json", + "chat_template.jinja", + }; + out.weight_candidates = {"model.safetensors.index.json", "model.gguf"}; + return out; +} + +runtime::CapabilitySet capabilities(const HiggsAssets &) { + runtime::CapabilitySet out; + out.supported_tasks = { + {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, + }; + out.supports_speaker_reference = true; + out.languages = {"Auto"}; + return out; +} + +runtime::ModelCliInterface cli(const HiggsAssets &) { + runtime::ModelCliInterface out; + out.request_options = { + {"max_tokens", "n", "Maximum generated AR tokens."}, + {"temperature", "float", "AR sampling temperature."}, + {"top_k", "n", "AR top-k sampling limit."}, + {"top_p", "float", "AR nucleus sampling probability."}, + {"repetition_penalty", "float", "Accepted for Python API compatibility; Higgs audio sampling does not consume it."}, + {"seed", "n", "Torch RNG seed."}, + {"text_chunk_size", "n", "Long-form text chunk size."}, + }; + out.session_options = { + {"higgs_tts.weight_type", "native|f32|f16|bf16|q8_0", "AR and codec weight storage type."}, + {"higgs_tts.ar_weight_type", "native|f32|f16|bf16|q8_0", "Autoregressive decoder weight storage type."}, + {"higgs_tts.codec_weight_type", "native|f32|f16|bf16|q8_0", "Codec weight storage type."}, + {"higgs_tts.ar_weight_context_mb", "n", "AR weight context size."}, + {"higgs_tts.codec_weight_context_mb", "n", "Codec weight context size."}, + {"higgs_tts.ar_decode_graph_arena_mb", "n", "AR decode graph arena size."}, + {"higgs_tts.codec_decode_graph_arena_mb", "n", "Codec decode graph arena size."}, + {"higgs_tts.codec_encode_graph_arena_mb", "n", "Codec encode graph arena size."}, + }; + return out; +} + +class HiggsTTSLoader final : public runtime::IVoiceModelLoader { +public: + std::string family() const override { + return "higgs_tts"; + } + + bool can_load(const runtime::ModelLoadRequest & request) const override { + if (request.family_hint.has_value() && *request.family_hint != family()) { + return false; + } + try { + const auto package_spec = engine::assets::default_model_package_spec_path(family()); + (void) engine::assets::load_resource_bundle_from_package_spec(request.model_path, package_spec); + return true; + } catch (...) { + return false; + } + } + + runtime::ModelInspection inspect(const runtime::ModelLoadRequest & request) const override { + const auto assets = load_higgs_assets(request.model_path); + runtime::ModelInspection inspection; + inspection.model_root = assets->resources.model_root(); + inspection.metadata = metadata(*assets); + inspection.capabilities = capabilities(*assets); + inspection.cli = cli(*assets); + const auto package_spec = engine::assets::default_model_package_spec_path(family()); + inspection.discovered_configs = runtime::discover_named_assets_from_package_spec( + request.model_path, + package_spec, + engine::assets::ModelPackageResourceKind::Files); + inspection.discovered_weights = runtime::discover_named_assets_from_package_spec( + request.model_path, + package_spec, + engine::assets::ModelPackageResourceKind::Tensors); + return inspection; + } + + std::unique_ptr load(const runtime::ModelLoadRequest & request) const override { + return load_higgs_tts_model(request.model_path); + } +}; + +} // namespace + +HiggsTTSLoadedModel::HiggsTTSLoadedModel( + runtime::ModelMetadata metadata, + runtime::CapabilitySet capabilities, + std::shared_ptr assets) + : metadata_(std::move(metadata)), + capabilities_(std::move(capabilities)), + assets_(std::move(assets)) {} + +const runtime::ModelMetadata & HiggsTTSLoadedModel::metadata() const noexcept { + return metadata_; +} + +const runtime::CapabilitySet & HiggsTTSLoadedModel::capabilities() const noexcept { + return capabilities_; +} + +std::unique_ptr HiggsTTSLoadedModel::create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const { + if (task.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Higgs TTS only supports offline sessions"); + } + if (task.task != runtime::VoiceTaskKind::Tts) { + throw std::runtime_error("Higgs TTS only supports the Tts task"); + } + return std::make_unique(task, options, assets_); +} + +std::unique_ptr load_higgs_tts_model(const std::filesystem::path & model_path) { + auto assets = load_higgs_assets(model_path); + return std::make_unique( + metadata(*assets), + capabilities(*assets), + std::move(assets)); +} + +std::shared_ptr make_higgs_tts_loader() { + return std::make_shared(); +} + +} // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/sampler.cpp b/src/models/higgs_tts/sampler.cpp new file mode 100644 index 00000000..9c6496c3 --- /dev/null +++ b/src/models/higgs_tts/sampler.cpp @@ -0,0 +1,458 @@ +#include "engine/models/higgs_tts/sampler.h" + +#include "engine/framework/sampling/torch_random.h" + +#include +#include +#include +#include +#include +#include + +namespace engine::models::higgs_tts { +namespace { + +constexpr float kGreedyTemperatureThreshold = 1.0e-5F; + +struct SamplerScratch { + std::vector & scores; + std::vector & probs; + std::vector & order; + std::vector & kept; +}; + +void require_cuda_sampling_policy(const HiggsCudaSamplingPolicy & policy) { + if (policy.multiprocessor_count <= 0 || policy.max_threads_per_multiprocessor <= 0) { + throw std::runtime_error("Higgs TTS stochastic sampler requires CUDA " + "sampling device properties"); + } +} + +uint32_t rotl32(uint32_t value, int shift) { + return static_cast((value << shift) | (value >> (32 - shift))); +} + +uint32_t fmix32(uint32_t value) { + value ^= value >> 16; + value *= 0x85EBCA6Bu; + value ^= value >> 13; + value *= 0xC2B2AE35u; + value ^= value >> 16; + return value; +} + +uint32_t murmur3_mix(uint32_t hash, uint32_t key) { + key *= 0xCC9E2D51u; + key = rotl32(key, 15); + key *= 0x1B873593u; + hash ^= key; + hash = rotl32(hash, 13); + hash = hash * 5u + 0xE6546B64u; + return hash; +} + +uint32_t sglang_murmur_hash32(uint64_t seed, uint32_t position, uint32_t column) { + uint32_t hash = 0; + hash = murmur3_mix(hash, static_cast(seed & 0xFFFFFFFFull)); + hash = murmur3_mix(hash, static_cast((seed >> 32) & 0xFFFFFFFFull)); + hash = murmur3_mix(hash, position); + hash = murmur3_mix(hash, column); + hash ^= 16u; + return fmix32(hash); +} + +int32_t argmax_row(const float * logits, int64_t vocab_size) { + int32_t best = 0; + float best_value = logits[0]; + for (int64_t i = 1; i < vocab_size; ++i) { + const float value = logits[i]; + if (value > best_value) { + best_value = value; + best = static_cast(i); + } + } + return best; +} + +float finite_max(const std::vector & scores, const std::vector & candidates) { + float max_score = -std::numeric_limits::infinity(); + if (candidates.empty()) { + for (const float score : scores) { + if (std::isfinite(score)) { + max_score = std::max(max_score, score); + } + } + } else { + for (const int64_t index : candidates) { + const float score = scores[static_cast(index)]; + if (std::isfinite(score)) { + max_score = std::max(max_score, score); + } + } + } + if (!std::isfinite(max_score)) { + throw std::runtime_error("Higgs TTS sampler kept no finite logits"); + } + return max_score; +} + +void scores_to_probs(SamplerScratch & scratch, int64_t vocab_size) { + const float max_score = finite_max(scratch.scores, scratch.kept); + double total = 0.0; + scratch.probs.assign(static_cast(vocab_size), 0.0F); + if (scratch.kept.empty()) { + for (int64_t i = 0; i < vocab_size; ++i) { + const float score = scratch.scores[static_cast(i)]; + if (std::isfinite(score)) { + const float value = + static_cast(std::exp(static_cast(score - max_score))); + scratch.probs[static_cast(i)] = value; + total += static_cast(value); + } + } + } else { + for (const int64_t index : scratch.kept) { + const float value = static_cast(std::exp( + static_cast(scratch.scores[static_cast(index)] - max_score))); + scratch.probs[static_cast(index)] = value; + total += static_cast(value); + } + } + if (!(total > 0.0) || !std::isfinite(total)) { + throw std::runtime_error("Higgs TTS sampler probability mass is invalid"); + } + const float inv_total = static_cast(1.0 / total); + for (float & prob : scratch.probs) { + prob *= inv_total; + } +} + +void renormalize_probs(SamplerScratch & scratch) { + double total = 0.0; + if (scratch.kept.empty()) { + for (const float prob : scratch.probs) { + total += static_cast(prob); + } + } else { + for (const int64_t index : scratch.kept) { + total += static_cast(scratch.probs[static_cast(index)]); + } + } + if (!(total > 0.0) || !std::isfinite(total)) { + throw std::runtime_error("Higgs TTS sampler probability mass is invalid"); + } + const float inv_total = static_cast(1.0 / total); + if (scratch.kept.empty()) { + for (float & prob : scratch.probs) { + prob *= inv_total; + } + } else { + for (const int64_t index : scratch.kept) { + scratch.probs[static_cast(index)] *= inv_total; + } + } +} + +void apply_top_k_to_probs(SamplerScratch & scratch, int64_t vocab_size, int64_t top_k) { + scratch.kept.clear(); + if (top_k <= 0 || top_k >= vocab_size) { + return; + } + scratch.order.resize(static_cast(vocab_size)); + std::iota(scratch.order.begin(), scratch.order.end(), 0); + auto kth = scratch.order.begin() + static_cast(top_k - 1); + std::nth_element( + scratch.order.begin(), kth, scratch.order.end(), [&](int64_t lhs, int64_t rhs) { + return scratch.probs[static_cast(lhs)] > + scratch.probs[static_cast(rhs)]; + }); + const float threshold = scratch.probs[static_cast(*kth)]; + scratch.kept.reserve(static_cast(top_k)); + for (int64_t i = 0; i < vocab_size; ++i) { + auto & prob = scratch.probs[static_cast(i)]; + if (prob < threshold) { + prob = 0.0F; + } else { + scratch.kept.push_back(i); + } + } + renormalize_probs(scratch); +} + +void apply_top_k_to_scores(SamplerScratch & scratch, int64_t vocab_size, int64_t top_k) { + scratch.kept.clear(); + if (top_k <= 0 || top_k >= vocab_size) { + return; + } + scratch.order.resize(static_cast(vocab_size)); + std::iota(scratch.order.begin(), scratch.order.end(), 0); + auto kth = scratch.order.begin() + static_cast(top_k - 1); + std::nth_element( + scratch.order.begin(), kth, scratch.order.end(), [&](int64_t lhs, int64_t rhs) { + return scratch.scores[static_cast(lhs)] > + scratch.scores[static_cast(rhs)]; + }); + const float threshold = scratch.scores[static_cast(*kth)]; + scratch.kept.reserve(static_cast(top_k)); + for (int64_t i = 0; i < vocab_size; ++i) { + if (scratch.scores[static_cast(i)] >= threshold) { + scratch.kept.push_back(i); + } + } +} + +void apply_top_p_to_probs(SamplerScratch & scratch, int64_t vocab_size, float top_p) { + if (!(top_p < 1.0F)) { + return; + } + if (!(top_p > 0.0F)) { + throw std::runtime_error("Higgs TTS sampler top_p must be positive"); + } + if (scratch.kept.empty()) { + scratch.kept.reserve(static_cast(vocab_size)); + for (int64_t i = 0; i < vocab_size; ++i) { + if (scratch.probs[static_cast(i)] > 0.0F) { + scratch.kept.push_back(i); + } + } + } + if (scratch.kept.empty()) { + throw std::runtime_error("Higgs TTS sampler top-p kept no probabilities"); + } + std::stable_sort(scratch.kept.begin(), scratch.kept.end(), [&](int64_t lhs, int64_t rhs) { + return scratch.probs[static_cast(lhs)] > scratch.probs[static_cast(rhs)]; + }); + + float cumulative = 0.0F; + float threshold = scratch.probs[static_cast(scratch.kept.front())]; + for (const int64_t index : scratch.kept) { + threshold = scratch.probs[static_cast(index)]; + cumulative += threshold; + if (cumulative >= top_p) { + break; + } + } + + size_t kept_count = 0; + for (const int64_t index : scratch.kept) { + auto & prob = scratch.probs[static_cast(index)]; + if (prob < threshold) { + prob = 0.0F; + } else { + scratch.kept[kept_count++] = index; + } + } + scratch.kept.resize(kept_count); + renormalize_probs(scratch); +} + +double sglang_gumbel_from_hash(uint32_t hashed) { + constexpr double kUint32Max = static_cast(std::numeric_limits::max()); + const double x = static_cast(hashed) / kUint32Max; + const double log_x = std::max(std::log(x), std::numeric_limits::lowest()); + return -std::log(-log_x); +} + +int32_t sample_seeded_sglang_gumbel(const std::vector & probs, + const std::vector & candidates, + uint64_t seed, + uint64_t position) { + double best_rank = -std::numeric_limits::infinity(); + int32_t best = -1; + const auto sample_one = [&](int64_t i) { + const float prob = probs[static_cast(i)]; + if (!(prob > 0.0F)) { + return; + } + const double logprob = std::log(static_cast(prob)); + const uint32_t hashed = sglang_murmur_hash32( + seed, static_cast(position & 0xFFFFFFFFull), static_cast(i)); + const double rank = logprob + sglang_gumbel_from_hash(hashed); + if (rank > best_rank) { + best_rank = rank; + best = static_cast(i); + } + }; + if (candidates.empty()) { + for (int64_t i = 0; i < static_cast(probs.size()); ++i) { + sample_one(i); + } + } else { + for (const int64_t index : candidates) { + sample_one(index); + } + } + if (best < 0) { + throw std::runtime_error("Higgs TTS sampler failed to select a codebook token"); + } + return best; +} + +int32_t sample_unseeded_torch_multinomial(const std::vector & probs, + const std::vector & candidates, + uint64_t seed, + uint64_t call_index, + const HiggsCudaSamplingPolicy & policy) { + require_cuda_sampling_policy(policy); + const int64_t vocab_size = static_cast(probs.size()); + + double best_rank = -std::numeric_limits::infinity(); + int32_t best = -1; + const auto sample_one = [&](int64_t i) { + const float prob = probs[static_cast(i)]; + if (!(prob > 0.0F)) { + return; + } + const float exponential = engine::sampling::torch_cuda_tensor_iterator_exponential_element( + seed, + static_cast(vocab_size), + static_cast(i), + call_index, + policy.multiprocessor_count, + policy.max_threads_per_multiprocessor); + const double rank = static_cast(prob) / static_cast(exponential); + if (rank > best_rank) { + best_rank = rank; + best = static_cast(i); + } + }; + if (candidates.empty()) { + for (int64_t i = 0; i < vocab_size; ++i) { + sample_one(i); + } + } else { + for (const int64_t index : candidates) { + sample_one(index); + } + } + if (best < 0) { + throw std::runtime_error("Higgs TTS sampler failed to select a codebook token"); + } + return best; +} + +int32_t sample_codebook_row(const float * logits, + int64_t vocab_size, + const HiggsSamplingOptions & options, + uint64_t call_index, + SamplerScratch & scratch) { + if (logits == nullptr || vocab_size <= 0) { + throw std::runtime_error("Higgs TTS sampler requires logits"); + } + if (options.temperature <= kGreedyTemperatureThreshold || + (options.top_k.has_value() && *options.top_k == 1)) { + return argmax_row(logits, vocab_size); + } + if (!(options.temperature > 0.0F) || !std::isfinite(options.temperature)) { + throw std::runtime_error("Higgs TTS sampler temperature must be finite and positive"); + } + + scratch.kept.clear(); + scratch.scores.resize(static_cast(vocab_size)); + for (int64_t i = 0; i < vocab_size; ++i) { + scratch.scores[static_cast(i)] = logits[i] / options.temperature; + } + bool top_k_applied_to_scores = false; + if (options.top_k.has_value()) { + if (*options.top_k < 0) { + throw std::runtime_error("Higgs TTS sampler top_k must be non-negative"); + } + const int64_t top_k = std::min(*options.top_k, vocab_size); + if (top_k > 0 && top_k < vocab_size) { + apply_top_k_to_scores(scratch, vocab_size, top_k); + top_k_applied_to_scores = true; + } + } + scores_to_probs(scratch, vocab_size); + if (options.top_k.has_value() && !top_k_applied_to_scores) { + apply_top_k_to_probs(scratch, vocab_size, std::min(*options.top_k, vocab_size)); + } + if (options.top_p.has_value()) { + apply_top_p_to_probs(scratch, vocab_size, *options.top_p); + } + if (options.has_seed) { + return sample_seeded_sglang_gumbel( + scratch.probs, scratch.kept, options.seed & 0x7FFFFFFFull, call_index); + } + return sample_unseeded_torch_multinomial( + scratch.probs, scratch.kept, options.seed, call_index, options.cuda_policy); +} + +} // namespace + +HiggsCodebookSampler::HiggsCodebookSampler(int64_t num_codebooks, int64_t codebook_vocab_size) + : num_codebooks_(num_codebooks), codebook_vocab_size_(codebook_vocab_size) { + if (num_codebooks_ <= 0 || codebook_vocab_size_ <= 0) { + throw std::runtime_error("Higgs TTS sampler requires positive codebook dimensions"); + } + scratch_scores_.reserve(static_cast(codebook_vocab_size_)); + scratch_probs_.reserve(static_cast(codebook_vocab_size_)); + scratch_order_.reserve(static_cast(codebook_vocab_size_)); + scratch_kept_.reserve(static_cast(codebook_vocab_size_)); + scratch_codes_.reserve(static_cast(num_codebooks_)); +} + +HiggsSamplerState HiggsCodebookSampler::make_state() const { + HiggsSamplerState state; + state.num_codebooks = num_codebooks_; + state.last_codes.assign(static_cast(num_codebooks_), 0); + return state; +} + +const std::vector & HiggsCodebookSampler::step(const float * logits, + int64_t logits_count, + HiggsSamplerState & state, + HiggsSamplingOptions & options) { + if (state.num_codebooks != num_codebooks_) { + throw std::runtime_error("Higgs TTS sampler state codebook count mismatch"); + } + if (logits_count != num_codebooks_ * codebook_vocab_size_) { + throw std::runtime_error("Higgs TTS sampler logits shape mismatch"); + } + + if (state.generation_done) { + scratch_codes_.assign(static_cast(num_codebooks_), kHiggsStopCode); + return scratch_codes_; + } + + scratch_codes_.assign(static_cast(num_codebooks_), 0); + SamplerScratch scratch{scratch_scores_, scratch_probs_, scratch_order_, scratch_kept_}; + for (int64_t codebook = 0; codebook < num_codebooks_; ++codebook) { + const float * row = logits + static_cast(codebook * codebook_vocab_size_); + scratch_codes_[static_cast(codebook)] = + sample_codebook_row(row, + codebook_vocab_size_, + options, + static_cast(state.step_count * num_codebooks_ + codebook), + scratch); + } + + if (state.delay_count < num_codebooks_) { + const int64_t next_codebook = state.delay_count + 1; + if (next_codebook < num_codebooks_) { + for (int64_t codebook = next_codebook; codebook < num_codebooks_; ++codebook) { + scratch_codes_[static_cast(codebook)] = kHiggsBocId; + } + } + state.delay_count += 1; + } else if (state.eoc_countdown.has_value()) { + *state.eoc_countdown -= 1; + if (*state.eoc_countdown <= 0) { + state.generation_done = true; + } + } else if (scratch_codes_.front() == kHiggsEocId) { + if (num_codebooks_ <= 2) { + state.generation_done = true; + } else { + state.eoc_countdown = num_codebooks_ - 2; + } + } + + state.step_count += 1; + if (!state.generation_done) { + state.last_codes = scratch_codes_; + } + return scratch_codes_; +} + +} // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/session.cpp b/src/models/higgs_tts/session.cpp new file mode 100644 index 00000000..7f67cf7c --- /dev/null +++ b/src/models/higgs_tts/session.cpp @@ -0,0 +1,297 @@ +#include "engine/models/higgs_tts/session.h" + +#include "engine/framework/debug/profiler.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/text/chunking.h" + +#include +#include +#include +#include +#include + +namespace engine::models::higgs_tts { +namespace { + +using Clock = std::chrono::steady_clock; + +constexpr int64_t kDefaultTextChunkSize = 4096; + +void validate_matmul_weight_storage(assets::TensorStorageType storage_type, const char * option_name) { + if (storage_type == assets::TensorStorageType::Native || + storage_type == assets::TensorStorageType::F32 || + storage_type == assets::TensorStorageType::F16 || + storage_type == assets::TensorStorageType::BF16 || + storage_type == assets::TensorStorageType::Q8_0) { + return; + } + throw std::runtime_error(std::string(option_name) + " currently supports only native, f32, f16, bf16, and q8_0"); +} + +uint64_t fnv1a_mix(uint64_t hash, const void * data, size_t size) { + const auto * bytes = static_cast(data); + for (size_t i = 0; i < size; ++i) { + hash ^= bytes[i]; + hash *= 1099511628211ull; + } + return hash; +} + +uint64_t hash_audio_samples(const runtime::AudioBuffer & audio) { + uint64_t hash = 1469598103934665603ull; + for (const float sample : audio.samples) { + uint32_t bits = 0; + std::memcpy(&bits, &sample, sizeof(bits)); + hash = fnv1a_mix(hash, &bits, sizeof(bits)); + } + return hash; +} + +const runtime::AudioBuffer * find_reference_audio(const runtime::TaskRequest & request) { + if (request.voice.has_value() + && request.voice->speaker.has_value() + && request.voice->speaker->audio.has_value()) { + return &*request.voice->speaker->audio; + } + if (request.audio_input.has_value()) { + return &*request.audio_input; + } + return nullptr; +} + +HiggsGenerationOptions generation_options_from_request( + const runtime::TaskRequest & request, + const HiggsConfig & config) { + HiggsGenerationOptions options; + options.max_tokens = 1024; + options.temperature = 0.8F; + options.top_p = 0.8F; + options.top_k = 30; + options.repetition_penalty = 1.1F; + if (const auto value = runtime::parse_int_option(request.options, {"max_tokens"})) { + if (*value <= 0) { + throw std::runtime_error("Higgs TTS max_tokens must be positive"); + } + options.max_tokens = *value; + } + if (const auto value = runtime::parse_float_option(request.options, {"temperature"})) { + options.temperature = *value; + } + if (const auto value = runtime::parse_float_option(request.options, {"top_p"})) { + options.top_p = *value; + } + if (const auto value = runtime::parse_int_option(request.options, {"top_k"})) { + options.top_k = *value; + } + if (const auto value = runtime::parse_finite_float_option(request.options, {"repetition_penalty"})) { + options.repetition_penalty = *value; + } + if (const auto value = runtime::parse_u64_option(request.options, {"seed"})) { + options.seed = *value; + } + if (options.max_tokens > config.text.max_position_embeddings) { + throw std::runtime_error("Higgs TTS max_tokens exceeds model max_position_embeddings"); + } + if (!(options.repetition_penalty > 0.0F) || !std::isfinite(options.repetition_penalty)) { + throw std::runtime_error("Higgs TTS repetition_penalty must be finite and positive"); + } + return options; +} + +} // namespace + +HiggsTTSSession::HiggsTTSSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets) + : RuntimeSessionBase(options), + task_(task), + assets_(std::move(assets)) { + if (assets_ == nullptr) { + throw std::runtime_error("Higgs TTS session requires assets"); + } + if (task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Higgs TTS currently supports offline sessions"); + } + if (task_.task != runtime::VoiceTaskKind::Tts) { + throw std::runtime_error("Higgs TTS only supports the Tts task"); + } + if (options.backend.type != core::BackendType::Cuda) { + throw std::runtime_error("Higgs TTS generation requires CUDA backend"); + } + + ar_weight_context_bytes_ = runtime::parse_size_mb_option( + options.options, {"higgs_tts.ar_weight_context_mb"}, ar_weight_context_bytes_); + codec_weight_context_bytes_ = runtime::parse_size_mb_option( + options.options, {"higgs_tts.codec_weight_context_mb"}, codec_weight_context_bytes_); + ar_decode_graph_arena_bytes_ = runtime::parse_size_mb_option( + options.options, {"higgs_tts.ar_decode_graph_arena_mb"}, ar_decode_graph_arena_bytes_); + codec_decode_graph_arena_bytes_ = runtime::parse_size_mb_option( + options.options, {"higgs_tts.codec_decode_graph_arena_mb"}, codec_decode_graph_arena_bytes_); + codec_encode_graph_arena_bytes_ = runtime::parse_size_mb_option( + options.options, {"higgs_tts.codec_encode_graph_arena_mb"}, codec_encode_graph_arena_bytes_); + + if (const auto it = options.options.find("higgs_tts.weight_type"); it != options.options.end()) { + const auto storage_type = assets::parse_tensor_storage_type(it->second); + validate_matmul_weight_storage(storage_type, "higgs_tts.weight_type"); + ar_weight_storage_type_ = storage_type; + codec_weight_storage_type_ = storage_type; + } + if (const auto it = options.options.find("higgs_tts.ar_weight_type"); it != options.options.end()) { + ar_weight_storage_type_ = assets::parse_tensor_storage_type(it->second); + validate_matmul_weight_storage(ar_weight_storage_type_, "higgs_tts.ar_weight_type"); + } + if (const auto it = options.options.find("higgs_tts.codec_weight_type"); it != options.options.end()) { + codec_weight_storage_type_ = assets::parse_tensor_storage_type(it->second); + validate_matmul_weight_storage(codec_weight_storage_type_, "higgs_tts.codec_weight_type"); + } + for (const auto & [key, _] : options.options) { + if (key.rfind("higgs_tts.", 0) == 0 && + key != "higgs_tts.ar_weight_context_mb" && + key != "higgs_tts.codec_weight_context_mb" && + key != "higgs_tts.ar_decode_graph_arena_mb" && + key != "higgs_tts.codec_decode_graph_arena_mb" && + key != "higgs_tts.codec_encode_graph_arena_mb" && + key != "higgs_tts.weight_type" && + key != "higgs_tts.ar_weight_type" && + key != "higgs_tts.codec_weight_type") { + throw std::runtime_error("unknown Higgs TTS session option: " + key); + } + } + + ar_ = std::make_shared( + assets_, + execution_context(), + ar_weight_context_bytes_, + ar_weight_storage_type_); + codec_ = std::make_shared( + assets_, + execution_context(), + codec_weight_context_bytes_, + codec_decode_graph_arena_bytes_, + codec_encode_graph_arena_bytes_, + codec_weight_storage_type_); + generator_ = std::make_unique( + assets_, + ar_, + codec_, + ar_decode_graph_arena_bytes_); +} + +std::string HiggsTTSSession::family() const { + return "higgs_tts"; +} + +runtime::VoiceTaskKind HiggsTTSSession::task_kind() const { + return task_.task; +} + +runtime::RunMode HiggsTTSSession::run_mode() const { + return task_.mode; +} + +void HiggsTTSSession::prepare(const runtime::SessionPreparationRequest & request) { + if (request.text.has_value()) { + runtime::TaskRequest task_request; + task_request.text_input = request.text; + task_request.voice = request.voice; + task_request.options = request.options; + const auto generation_request = make_generation_request(task_request); + generator_->prepare(generation_request); + } + mark_prepared(); +} + +runtime::TaskResult HiggsTTSSession::run(const runtime::TaskRequest & request) { + require_prepared("Higgs TTS run"); + const auto wall_start = Clock::now(); + const int64_t text_chunk_size = + engine::text::parse_text_chunk_size_override(request.options).value_or(kDefaultTextChunkSize); + const auto chunk_requests = runtime::chunk_text_request(request, text_chunk_size); + const std::string reference_text = runtime::find_option(request.options, {"reference_text"}).value_or(""); + const auto * reference_audio = find_reference_audio(request); + const HiggsCodecEncodeOutput * reference_codes = + reference_audio != nullptr ? &resolve_reference_codes(*reference_audio, reference_text) : nullptr; + debug::trace_log_scalar("higgs_tts.text_chunk_size", text_chunk_size); + debug::trace_log_scalar("higgs_tts.text_chunk_count", static_cast(chunk_requests.size())); + + runtime::AudioBuffer merged_audio; + for (const auto & chunk_request : chunk_requests) { + const auto generation_request = make_generation_request(chunk_request, reference_codes); + auto result = generator_->generate(generation_request); + runtime::append_audio_buffer(merged_audio, runtime::AudioBuffer{ + result.audio.sample_rate, + result.audio.channels, + std::move(result.audio.values), + }); + } + + runtime::TaskResult out; + out.audio_output = std::move(merged_audio); + debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start, Clock::now())); + return out; +} + +const HiggsCodecEncodeOutput & HiggsTTSSession::resolve_reference_codes( + const runtime::AudioBuffer & audio, + const std::string & reference_text) { + const uint64_t sample_count = static_cast(audio.samples.size()); + const uint64_t sample_hash = hash_audio_samples(audio); + debug::trace_log_scalar("higgs_tts.reference_audio.sample_rate", audio.sample_rate); + debug::trace_log_scalar("higgs_tts.reference_audio.channels", audio.channels); + debug::trace_log_f32("higgs_tts.reference_audio.samples", + {static_cast(audio.samples.size())}, + audio.samples); + const bool cache_hit = reference_cache_.has_value() + && reference_cache_->reference_text == reference_text + && reference_cache_->sample_rate == audio.sample_rate + && reference_cache_->channels == audio.channels + && reference_cache_->sample_count == sample_count + && reference_cache_->sample_hash == sample_hash; + if (!cache_hit) { + const auto encode_start = Clock::now(); + ReferenceCacheEntry entry; + entry.reference_text = reference_text; + entry.sample_rate = audio.sample_rate; + entry.channels = audio.channels; + entry.sample_count = sample_count; + entry.sample_hash = sample_hash; + entry.codes = codec_->encode_reference(audio); + debug::trace_log_scalar("higgs_tts.reference_codes.frames", entry.codes.frames); + debug::trace_log_scalar("higgs_tts.reference_codes.codebooks", entry.codes.codebooks); + debug::trace_log_i32("higgs_tts.reference_codes.values", + {entry.codes.frames, entry.codes.codebooks}, + entry.codes.codes); + reference_cache_ = std::move(entry); + debug::timing_log_scalar("higgs_tts.codec.encode_reference_ms", engine::debug::elapsed_ms(encode_start)); + } + return reference_cache_->codes; +} + +HiggsGenerationRequest HiggsTTSSession::make_generation_request( + const runtime::TaskRequest & request, + const HiggsCodecEncodeOutput * resolved_reference_codes) { + if (!request.text_input.has_value()) { + throw std::runtime_error("Higgs TTS requires text input"); + } + const std::string reference_text = runtime::find_option(request.options, {"reference_text"}).value_or(""); + + HiggsGenerationRequest out; + out.text = request.text_input->text; + out.reference_text = reference_text; + out.options = generation_options_from_request(request, assets_->config); + if (resolved_reference_codes != nullptr) { + out.reference_codes = resolved_reference_codes->codes; + out.reference_frames = resolved_reference_codes->frames; + out.reference_codebooks = resolved_reference_codes->codebooks; + } else if (const auto * reference_audio = find_reference_audio(request)) { + const auto & reference_codes = resolve_reference_codes(*reference_audio, reference_text); + out.reference_codes = reference_codes.codes; + out.reference_frames = reference_codes.frames; + out.reference_codebooks = reference_codes.codebooks; + } + return out; +} + +} // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/tokenizer_text.cpp b/src/models/higgs_tts/tokenizer_text.cpp new file mode 100644 index 00000000..6c94a403 --- /dev/null +++ b/src/models/higgs_tts/tokenizer_text.cpp @@ -0,0 +1,94 @@ +#include "engine/models/higgs_tts/tokenizer_text.h" + +#include "engine/framework/tokenizers/llama_bpe.h" + +#include +#include +#include +#include + +namespace engine::models::higgs_tts { +namespace { + +int32_t require_token_id(const engine::tokenizers::LlamaBpeTokenizer & tokenizer, const std::string & token) { + const auto id = tokenizer.find_token_id(token); + if (!id.has_value()) { + throw std::runtime_error("Higgs TTS tokenizer missing required token: " + token); + } + return *id; +} + +} // namespace + +struct HiggsTextTokenizer::Impl { + explicit Impl(std::shared_ptr input_assets) + : assets(std::move(input_assets)), + tokenizer(engine::tokenizers::LlamaBpeTokenizerSpec{ + {}, + {}, + assets->resources.require_file("tokenizer_config"), + assets->resources.require_file("tokenizer_json"), + engine::tokenizers::LlamaBpePreTokenizer::Qwen2, + }), + tts_id(require_token_id(tokenizer, "<|tts|>")), + ref_audio_id(require_token_id(tokenizer, "<|ref_audio|>")), + ref_text_id(require_token_id(tokenizer, "<|ref_text|>")), + text_id(require_token_id(tokenizer, "<|text|>")), + audio_id(require_token_id(tokenizer, "<|audio|>")), + audio_placeholder_id(static_cast(assets->config.audio_token_id)) {} + + std::shared_ptr assets; + engine::tokenizers::LlamaBpeTokenizer tokenizer; + int32_t tts_id = 0; + int32_t ref_audio_id = 0; + int32_t ref_text_id = 0; + int32_t text_id = 0; + int32_t audio_id = 0; + int32_t audio_placeholder_id = -100; +}; + +HiggsTextTokenizer::HiggsTextTokenizer(std::shared_ptr assets) + : impl_([&]() { + if (assets == nullptr) { + throw std::runtime_error("Higgs TTS text tokenizer requires assets"); + } + return std::make_shared(std::move(assets)); + }()) {} + +std::vector HiggsTextTokenizer::encode(const std::string & text) const { + return impl_->tokenizer.encode(text, true); +} + +HiggsPromptEncoding HiggsTextTokenizer::encode_prompt(const HiggsPromptRequest & request) const { + if (request.delayed_reference_tokens < 0) { + throw std::runtime_error("Higgs TTS delayed_reference_tokens must be non-negative"); + } + + HiggsPromptEncoding encoding; + encoding.text_ids = encode(request.text); + if (!request.reference_text.empty() && request.delayed_reference_tokens > 0) { + encoding.reference_text_ids = encode(request.reference_text); + } + + encoding.token_ids.push_back(impl_->tts_id); + if (!encoding.reference_text_ids.empty()) { + encoding.token_ids.push_back(impl_->ref_text_id); + encoding.token_ids.insert( + encoding.token_ids.end(), + encoding.reference_text_ids.begin(), + encoding.reference_text_ids.end()); + } + if (request.delayed_reference_tokens > 0) { + encoding.token_ids.push_back(impl_->ref_audio_id); + encoding.token_ids.insert( + encoding.token_ids.end(), + static_cast(request.delayed_reference_tokens), + impl_->audio_placeholder_id); + } + encoding.token_ids.push_back(impl_->text_id); + encoding.token_ids.insert(encoding.token_ids.end(), encoding.text_ids.begin(), encoding.text_ids.end()); + encoding.token_ids.push_back(impl_->audio_id); + return encoding; +} + +} // namespace engine::models::higgs_tts From c0f7caedc7e55ff7549cffb052b7b68735452883 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Fri, 17 Jul 2026 21:32:42 -0400 Subject: [PATCH 03/27] Add conv lowering matrix test --- CMakeLists.txt | 7 + tests/unittests/test_conv_lowering_matrix.cpp | 834 ++++++++++++++++++ 2 files changed, 841 insertions(+) create mode 100644 tests/unittests/test_conv_lowering_matrix.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index 5ca2b830..01d93fea 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -877,6 +877,13 @@ if (ENGINE_BUILD_TESTS) COMMAND depthwise_conv1d_lowering_test ) + add_engine_unittest(conv_lowering_matrix_test tests/unittests/test_conv_lowering_matrix.cpp) + + add_test( + NAME conv_lowering_matrix_test + COMMAND conv_lowering_matrix_test + ) + add_engine_unittest(gguf_tensor_source_test tests/unittests/test_gguf_tensor_source.cpp) target_include_directories(gguf_tensor_source_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tests/unittests) diff --git a/tests/unittests/test_conv_lowering_matrix.cpp b/tests/unittests/test_conv_lowering_matrix.cpp new file mode 100644 index 00000000..0344132e --- /dev/null +++ b/tests/unittests/test_conv_lowering_matrix.cpp @@ -0,0 +1,834 @@ +#include "engine/framework/core/backend.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/streaming_conv_modules.h" +#include "engine/framework/modules/structural_modules.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr size_t kGraphBytes = 512 * 1024 * 1024; +constexpr size_t kGraphNodes = 16384; +constexpr int kWarmupRounds = 1; +constexpr int kMeasureRounds = 5; + +struct DiffStats { + float max_abs = 0.0f; + double mean_abs = 0.0; + double cosine = 1.0; +}; + +struct RunResult { + bool supported = false; + std::string error; + engine::core::TensorShape shape = {}; + std::vector values; + double avg_ms = 0.0; +}; + +std::vector make_patterned_f32(size_t count, float phase, float scale) { + std::vector values(count, 0.0f); + for (size_t i = 0; i < count; ++i) { + const float x = static_cast(i); + values[i] = scale * ( + std::sin(phase + 0.113f * x) + + 0.5f * std::cos(phase * 0.7f + 0.071f * x)); + } + return values; +} + +int64_t conv_out(int64_t input, int64_t kernel, int stride, int padding, int dilation) { + return (input + 2 * padding - dilation * (kernel - 1) - 1) / stride + 1; +} + +int64_t conv_transpose_out(int64_t input, int64_t kernel, int stride, int padding, int dilation) { + return (input - 1) * stride - 2 * padding + dilation * (kernel - 1) + 1; +} + +const char * backend_name(engine::core::BackendType backend_type) { + switch (backend_type) { + case engine::core::BackendType::Cpu: return "cpu"; + case engine::core::BackendType::Cuda: return "cuda"; + case engine::core::BackendType::Vulkan: return "vulkan"; + case engine::core::BackendType::Metal: return "metal"; + default: return "unknown"; + } +} + +bool same_shape(const engine::core::TensorShape & lhs, const engine::core::TensorShape & rhs) { + if (lhs.rank != rhs.rank) { + return false; + } + for (size_t i = 0; i < lhs.rank; ++i) { + if (lhs.dims[i] != rhs.dims[i]) { + return false; + } + } + return true; +} + +DiffStats diff_values(const std::vector & reference, const std::vector & actual) { + if (reference.size() != actual.size()) { + throw std::runtime_error("value count mismatch"); + } + DiffStats stats; + double dot = 0.0; + double ref_norm = 0.0; + double actual_norm = 0.0; + for (size_t i = 0; i < reference.size(); ++i) { + const float diff = std::fabs(reference[i] - actual[i]); + stats.max_abs = std::max(stats.max_abs, diff); + stats.mean_abs += diff; + dot += static_cast(reference[i]) * static_cast(actual[i]); + ref_norm += static_cast(reference[i]) * static_cast(reference[i]); + actual_norm += static_cast(actual[i]) * static_cast(actual[i]); + } + stats.mean_abs /= static_cast(reference.size()); + if (ref_norm > 0.0 && actual_norm > 0.0) { + stats.cosine = dot / (std::sqrt(ref_norm) * std::sqrt(actual_norm)); + } + return stats; +} + +engine::core::TensorValue add_bias_3d( + engine::core::ModuleBuildContext & ctx, + const engine::core::TensorValue & output, + int64_t channels, + const std::optional & bias) { + if (!bias.has_value()) { + return output; + } + auto output_contiguous = engine::core::ensure_backend_addressable_layout(ctx, output); + auto bias_view = engine::core::reshape_tensor(ctx, *bias, engine::core::TensorShape::from_dims({1, channels, 1})); + auto repeated = engine::core::wrap_tensor( + ggml_repeat(ctx.ggml, bias_view.tensor, output_contiguous.tensor), + output.shape, + GGML_TYPE_F32); + return engine::core::wrap_tensor(ggml_add(ctx.ggml, output_contiguous.tensor, repeated.tensor), output.shape, GGML_TYPE_F32); +} + +engine::core::TensorValue add_bias_4d( + engine::core::ModuleBuildContext & ctx, + const engine::core::TensorValue & output, + int64_t channels, + const std::optional & bias) { + if (!bias.has_value()) { + return output; + } + auto output_contiguous = engine::core::ensure_backend_addressable_layout(ctx, output); + auto bias_view = engine::core::reshape_tensor(ctx, *bias, engine::core::TensorShape::from_dims({1, channels, 1, 1})); + auto repeated = engine::core::wrap_tensor( + ggml_repeat(ctx.ggml, bias_view.tensor, output_contiguous.tensor), + output.shape, + GGML_TYPE_F32); + return engine::core::wrap_tensor(ggml_add(ctx.ggml, output_contiguous.tensor, repeated.tensor), output.shape, GGML_TYPE_F32); +} + +engine::core::TensorValue view_batch_matrix( + engine::core::ModuleBuildContext & ctx, + const engine::core::TensorValue & input, + int64_t batch_index, + int64_t channels, + int64_t frames) { + auto * view = ggml_view_2d( + ctx.ggml, + input.tensor, + frames, + channels, + input.tensor->nb[1], + static_cast(batch_index) * input.tensor->nb[2]); + return engine::core::wrap_tensor(view, engine::core::TensorShape::from_dims({channels, frames}), input.type); +} + +class GraphRunner { +public: + GraphRunner(const char * name, engine::core::BackendType backend_type) : backend_type_(backend_type) { + backend_ = engine::core::init_backend({backend_type, 0, 8}); + engine::core::set_backend_threads(backend_, 8); + ggml_init_params params{}; + params.mem_size = kGraphBytes; + params.mem_buffer = nullptr; + params.no_alloc = true; + ggml_ = ggml_init(params); + if (ggml_ == nullptr) { + throw std::runtime_error("failed to initialize ggml test context"); + } + ctx_.ggml = ggml_; + ctx_.module_instance_name = name; + ctx_.backend_type = backend_type; + } + + ~GraphRunner() { + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + if (ggml_ != nullptr) { + ggml_free(ggml_); + } + if (backend_ != nullptr) { + ggml_backend_free(backend_); + } + } + + engine::core::TensorValue make_f32(const engine::core::TensorShape & shape) { + return engine::core::make_tensor(ctx_, GGML_TYPE_F32, shape); + } + + engine::core::ModuleBuildContext & ctx() noexcept { return ctx_; } + + RunResult run( + const engine::core::TensorValue & output, + const std::vector>> & writes) { + ggml_cgraph * graph = ggml_new_graph_custom(ggml_, kGraphNodes, false); + ggml_build_forward_expand(graph, output.tensor); + engine::core::validate_backend_graph_supported(backend_, graph, "conv_lowering_matrix"); + buffer_ = ggml_backend_alloc_ctx_tensors(ggml_, backend_); + if (buffer_ == nullptr) { + throw std::runtime_error("failed to allocate backend tensors"); + } + for (const auto & write : writes) { + engine::core::write_tensor_f32(write.first, write.second); + } + for (int i = 0; i < kWarmupRounds; ++i) { + if (ggml_backend_graph_compute(backend_, graph) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("warmup graph compute failed"); + } + } + double total_ms = 0.0; + for (int i = 0; i < kMeasureRounds; ++i) { + const auto start = std::chrono::steady_clock::now(); + if (ggml_backend_graph_compute(backend_, graph) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("graph compute failed"); + } + const auto end = std::chrono::steady_clock::now(); + total_ms += std::chrono::duration(end - start).count(); + } + RunResult result; + result.supported = true; + result.shape = output.shape; + result.avg_ms = total_ms / static_cast(kMeasureRounds); + engine::core::read_tensor_f32_into(output.tensor, result.values); + return result; + } + +private: + engine::core::BackendType backend_type_; + ggml_backend_t backend_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; + ggml_context * ggml_ = nullptr; + engine::core::ModuleBuildContext ctx_{}; +}; + +template +RunResult run_guarded( + const char * label, + engine::core::BackendType backend_type, + Fn && fn) { + try { + GraphRunner runner(label, backend_type); + return fn(runner); + } catch (const std::exception & ex) { + RunResult result; + result.supported = false; + result.error = ex.what(); + return result; + } +} + +struct Conv1dCase { + const char * name; + int64_t batch; + int64_t in_channels; + int64_t out_channels; + int64_t frames; + int64_t kernel; + int stride; + int padding; + int dilation; + bool bias; +}; + +RunResult run_conv1d(const Conv1dCase & c, const char * candidate, engine::core::BackendType backend_type) { + return run_guarded(candidate, backend_type, [&](GraphRunner & runner) { + const auto input_shape = engine::core::TensorShape::from_dims({c.batch, c.in_channels, c.frames}); + const auto weight_shape = engine::core::TensorShape::from_dims({c.out_channels, c.in_channels, c.kernel}); + const auto bias_shape = engine::core::TensorShape::from_dims({c.out_channels}); + auto input = runner.make_f32(input_shape); + auto weight = runner.make_f32(weight_shape); + std::optional bias = c.bias ? std::optional(runner.make_f32(bias_shape)) : std::nullopt; + + engine::core::TensorValue output; + if (std::string(candidate) == "native") { + const auto output_shape = engine::core::TensorShape::from_dims( + {c.batch, c.out_channels, conv_out(c.frames, c.kernel, c.stride, c.padding, c.dilation)}); + if (c.batch == 1) { + output = engine::core::wrap_tensor( + ggml_conv_1d(runner.ctx().ggml, weight.tensor, input.tensor, c.stride, c.padding, c.dilation), + output_shape, + GGML_TYPE_F32); + } else { + for (int64_t batch = 0; batch < c.batch; ++batch) { + auto batch_input = view_batch_matrix(runner.ctx(), input, batch, c.in_channels, c.frames); + auto batch_output = engine::core::wrap_tensor( + ggml_conv_1d(runner.ctx().ggml, weight.tensor, batch_input.tensor, c.stride, c.padding, c.dilation), + engine::core::TensorShape::from_dims({1, c.out_channels, output_shape.dims[2]}), + GGML_TYPE_F32); + output = output.valid() ? engine::modules::ConcatModule({0}).build(runner.ctx(), output, batch_output) : batch_output; + } + } + output = add_bias_3d(runner.ctx(), output, c.out_channels, bias); + } else if (std::string(candidate) == "conv2d_normal") { + auto x4 = engine::core::reshape_tensor(runner.ctx(), input, engine::core::TensorShape::from_dims({c.batch, c.in_channels, 1, c.frames})); + auto w4 = engine::core::reshape_tensor(runner.ctx(), weight, engine::core::TensorShape::from_dims({c.out_channels, c.in_channels, 1, c.kernel})); + auto y4 = engine::core::wrap_tensor( + ggml_conv_2d(runner.ctx().ggml, w4.tensor, x4.tensor, c.stride, 1, c.padding, 0, c.dilation, 1), + engine::core::TensorShape::from_dims({c.batch, c.out_channels, 1, conv_out(c.frames, c.kernel, c.stride, c.padding, c.dilation)}), + GGML_TYPE_F32); + y4 = add_bias_4d(runner.ctx(), y4, c.out_channels, bias); + output = engine::core::reshape_tensor(runner.ctx(), y4, engine::core::TensorShape::from_dims({c.batch, c.out_channels, y4.shape.dims[3]})); + } else if (std::string(candidate) == "conv2d_direct") { + auto x4 = engine::core::reshape_tensor(runner.ctx(), input, engine::core::TensorShape::from_dims({c.batch, c.in_channels, 1, c.frames})); + auto w4 = engine::core::reshape_tensor(runner.ctx(), weight, engine::core::TensorShape::from_dims({c.out_channels, c.in_channels, 1, c.kernel})); + auto y4 = engine::core::wrap_tensor( + ggml_conv_2d_direct(runner.ctx().ggml, w4.tensor, x4.tensor, c.stride, 1, c.padding, 0, c.dilation, 1), + engine::core::TensorShape::from_dims({c.batch, c.out_channels, 1, conv_out(c.frames, c.kernel, c.stride, c.padding, c.dilation)}), + GGML_TYPE_F32); + y4 = add_bias_4d(runner.ctx(), y4, c.out_channels, bias); + output = engine::core::reshape_tensor(runner.ctx(), y4, engine::core::TensorShape::from_dims({c.batch, c.out_channels, y4.shape.dims[3]})); + } else { + throw std::runtime_error("unknown conv1d candidate"); + } + + std::vector>> writes; + writes.push_back({input, make_patterned_f32(static_cast(input_shape.num_elements()), 0.19f, 0.031f)}); + writes.push_back({weight, make_patterned_f32(static_cast(weight_shape.num_elements()), 0.47f, 0.017f)}); + if (bias) { + writes.push_back({*bias, make_patterned_f32(static_cast(bias_shape.num_elements()), 0.83f, 0.011f)}); + } + return runner.run(output, writes); + }); +} + +struct Conv2dCase { + const char * name; + int64_t batch; + int64_t in_channels; + int64_t out_channels; + int64_t height; + int64_t width; + int64_t kernel_h; + int64_t kernel_w; + int stride_h; + int stride_w; + int padding_h; + int padding_w; + int dilation_h; + int dilation_w; + bool bias; +}; + +RunResult run_conv2d(const Conv2dCase & c, const char * candidate, engine::core::BackendType backend_type) { + return run_guarded(candidate, backend_type, [&](GraphRunner & runner) { + const auto input_shape = engine::core::TensorShape::from_dims({c.batch, c.in_channels, c.height, c.width}); + const auto weight_shape = engine::core::TensorShape::from_dims({c.out_channels, c.in_channels, c.kernel_h, c.kernel_w}); + const auto bias_shape = engine::core::TensorShape::from_dims({c.out_channels}); + auto input = runner.make_f32(input_shape); + auto weight = runner.make_f32(weight_shape); + std::optional bias = c.bias ? std::optional(runner.make_f32(bias_shape)) : std::nullopt; + + engine::core::TensorValue output; + const auto output_shape = engine::core::TensorShape::from_dims({ + c.batch, + c.out_channels, + conv_out(c.height, c.kernel_h, c.stride_h, c.padding_h, c.dilation_h), + conv_out(c.width, c.kernel_w, c.stride_w, c.padding_w, c.dilation_w), + }); + if (std::string(candidate) == "im2col_matmul") { + output = engine::core::wrap_tensor( + ggml_conv_2d(runner.ctx().ggml, weight.tensor, input.tensor, c.stride_w, c.stride_h, c.padding_w, c.padding_h, c.dilation_w, c.dilation_h), + output_shape, + GGML_TYPE_F32); + output = add_bias_4d(runner.ctx(), output, c.out_channels, bias); + } else if (std::string(candidate) == "direct") { + output = engine::core::wrap_tensor( + ggml_conv_2d_direct(runner.ctx().ggml, weight.tensor, input.tensor, c.stride_w, c.stride_h, c.padding_w, c.padding_h, c.dilation_w, c.dilation_h), + output_shape, + GGML_TYPE_F32); + output = add_bias_4d(runner.ctx(), output, c.out_channels, bias); + } else { + throw std::runtime_error("unknown conv2d candidate"); + } + std::vector>> writes; + writes.push_back({input, make_patterned_f32(static_cast(input_shape.num_elements()), 0.21f, 0.021f)}); + writes.push_back({weight, make_patterned_f32(static_cast(weight_shape.num_elements()), 0.51f, 0.013f)}); + if (bias) { + writes.push_back({*bias, make_patterned_f32(static_cast(bias_shape.num_elements()), 0.91f, 0.009f)}); + } + return runner.run(output, writes); + }); +} + +struct Depthwise1dCase { + const char * name; + int64_t batch; + int64_t channels; + int64_t frames; + int64_t kernel; + int stride; + int padding; + int dilation; + bool bias; +}; + +RunResult run_depthwise1d(const Depthwise1dCase & c, const char * candidate, engine::core::BackendType backend_type) { + return run_guarded(candidate, backend_type, [&](GraphRunner & runner) { + const auto input_shape = engine::core::TensorShape::from_dims({c.batch, c.channels, c.frames}); + const auto weight_shape = engine::core::TensorShape::from_dims({c.channels, 1, c.kernel}); + const auto bias_shape = engine::core::TensorShape::from_dims({c.channels}); + auto input = runner.make_f32(input_shape); + auto weight = runner.make_f32(weight_shape); + std::optional bias = c.bias ? std::optional(runner.make_f32(bias_shape)) : std::nullopt; + engine::core::TensorValue output; + if (std::string(candidate) == "dw2d_direct") { + auto x4 = engine::core::reshape_tensor(runner.ctx(), input, engine::core::TensorShape::from_dims({c.batch, c.channels, 1, c.frames})); + auto w4 = engine::core::reshape_tensor(runner.ctx(), weight, engine::core::TensorShape::from_dims({c.channels, 1, 1, c.kernel})); + auto y4 = engine::core::wrap_tensor( + ggml_conv_2d_dw_direct(runner.ctx().ggml, w4.tensor, x4.tensor, c.stride, 1, c.padding, 0, c.dilation, 1), + engine::core::TensorShape::from_dims({c.batch, c.channels, 1, conv_out(c.frames, c.kernel, c.stride, c.padding, c.dilation)}), + GGML_TYPE_F32); + y4 = add_bias_4d(runner.ctx(), y4, c.channels, bias); + output = engine::core::reshape_tensor(runner.ctx(), y4, engine::core::TensorShape::from_dims({c.batch, c.channels, y4.shape.dims[3]})); + } else if (std::string(candidate) == "native_1d_dw") { + if (c.batch != 1) { + throw std::runtime_error("native ggml_conv_1d_dw asserts for batched rank-3 input; slice batch first"); + } + output = engine::core::wrap_tensor( + ggml_conv_1d_dw(runner.ctx().ggml, weight.tensor, input.tensor, c.stride, c.padding, c.dilation), + engine::core::TensorShape::from_dims({c.batch, c.channels, conv_out(c.frames, c.kernel, c.stride, c.padding, c.dilation)}), + GGML_TYPE_F32); + output = add_bias_3d(runner.ctx(), output, c.channels, bias); + } else { + throw std::runtime_error("unknown depthwise1d candidate"); + } + std::vector>> writes; + writes.push_back({input, make_patterned_f32(static_cast(input_shape.num_elements()), 0.23f, 0.025f)}); + writes.push_back({weight, make_patterned_f32(static_cast(weight_shape.num_elements()), 0.53f, 0.015f)}); + if (bias) { + writes.push_back({*bias, make_patterned_f32(static_cast(bias_shape.num_elements()), 0.93f, 0.007f)}); + } + return runner.run(output, writes); + }); +} + +struct Pointwise1dCase { + const char * name; + int64_t batch; + int64_t in_channels; + int64_t out_channels; + int64_t frames; + bool bias; +}; + +RunResult run_pointwise1d(const Pointwise1dCase & c, const char * candidate, engine::core::BackendType backend_type) { + return run_guarded(candidate, backend_type, [&](GraphRunner & runner) { + const auto input_shape = engine::core::TensorShape::from_dims({c.batch, c.in_channels, c.frames}); + const auto weight_shape = engine::core::TensorShape::from_dims({c.out_channels, c.in_channels, 1}); + const auto bias_shape = engine::core::TensorShape::from_dims({c.out_channels}); + auto input = runner.make_f32(input_shape); + auto weight = runner.make_f32(weight_shape); + std::optional bias = c.bias ? std::optional(runner.make_f32(bias_shape)) : std::nullopt; + + engine::core::TensorValue output; + if (std::string(candidate) == "conv1d_kernel1") { + const auto output_shape = engine::core::TensorShape::from_dims({c.batch, c.out_channels, c.frames}); + if (c.batch == 1) { + output = engine::core::wrap_tensor( + ggml_conv_1d(runner.ctx().ggml, weight.tensor, input.tensor, 1, 0, 1), + output_shape, + GGML_TYPE_F32); + } else { + for (int64_t batch = 0; batch < c.batch; ++batch) { + auto batch_input = view_batch_matrix(runner.ctx(), input, batch, c.in_channels, c.frames); + auto batch_output = engine::core::wrap_tensor( + ggml_conv_1d(runner.ctx().ggml, weight.tensor, batch_input.tensor, 1, 0, 1), + engine::core::TensorShape::from_dims({1, c.out_channels, c.frames}), + GGML_TYPE_F32); + output = output.valid() ? engine::modules::ConcatModule({0}).build(runner.ctx(), output, batch_output) : batch_output; + } + } + output = add_bias_3d(runner.ctx(), output, c.out_channels, bias); + } else if (std::string(candidate) == "linear_matmul") { + auto x = engine::modules::TransposeModule({{0, 2, 1, 3}, 3}).build(runner.ctx(), input); + x = engine::core::ensure_backend_addressable_layout(runner.ctx(), x); + auto matrix = engine::core::reshape_tensor(runner.ctx(), x, engine::core::TensorShape::from_dims({c.batch * c.frames, c.in_channels})); + auto w2 = engine::core::reshape_tensor(runner.ctx(), weight, engine::core::TensorShape::from_dims({c.out_channels, c.in_channels})); + auto projected = engine::core::wrap_tensor( + ggml_mul_mat(runner.ctx().ggml, w2.tensor, matrix.tensor), + engine::core::TensorShape::from_dims({c.batch * c.frames, c.out_channels}), + GGML_TYPE_F32); + if (bias) { + projected = engine::core::wrap_tensor(ggml_add(runner.ctx().ggml, projected.tensor, bias->tensor), projected.shape, GGML_TYPE_F32); + } + auto y = engine::core::reshape_tensor(runner.ctx(), projected, engine::core::TensorShape::from_dims({c.batch, c.frames, c.out_channels})); + output = engine::modules::TransposeModule({{0, 2, 1, 3}, 3}).build(runner.ctx(), y); + } else { + throw std::runtime_error("unknown pointwise1d candidate"); + } + std::vector>> writes; + writes.push_back({input, make_patterned_f32(static_cast(input_shape.num_elements()), 0.24f, 0.027f)}); + writes.push_back({weight, make_patterned_f32(static_cast(weight_shape.num_elements()), 0.54f, 0.014f)}); + if (bias) { + writes.push_back({*bias, make_patterned_f32(static_cast(bias_shape.num_elements()), 0.94f, 0.007f)}); + } + return runner.run(output, writes); + }); +} + +RunResult run_depthwise2d(const Conv2dCase & c, const char * candidate, engine::core::BackendType backend_type) { + return run_guarded(candidate, backend_type, [&](GraphRunner & runner) { + const auto input_shape = engine::core::TensorShape::from_dims({c.batch, c.in_channels, c.height, c.width}); + const auto weight_shape = engine::core::TensorShape::from_dims({c.in_channels, 1, c.kernel_h, c.kernel_w}); + const auto bias_shape = engine::core::TensorShape::from_dims({c.in_channels}); + auto input = runner.make_f32(input_shape); + auto weight = runner.make_f32(weight_shape); + std::optional bias = c.bias ? std::optional(runner.make_f32(bias_shape)) : std::nullopt; + const auto output_shape = engine::core::TensorShape::from_dims({ + c.batch, + c.in_channels, + conv_out(c.height, c.kernel_h, c.stride_h, c.padding_h, c.dilation_h), + conv_out(c.width, c.kernel_w, c.stride_w, c.padding_w, c.dilation_w), + }); + engine::core::TensorValue output; + if (std::string(candidate) == "direct") { + output = engine::core::wrap_tensor( + ggml_conv_2d_dw_direct(runner.ctx().ggml, weight.tensor, input.tensor, c.stride_w, c.stride_h, c.padding_w, c.padding_h, c.dilation_w, c.dilation_h), + output_shape, + GGML_TYPE_F32); + output = add_bias_4d(runner.ctx(), output, c.in_channels, bias); + } else if (std::string(candidate) == "im2col_matmul") { + output = engine::core::wrap_tensor( + ggml_conv_2d_dw(runner.ctx().ggml, weight.tensor, input.tensor, c.stride_w, c.stride_h, c.padding_w, c.padding_h, c.dilation_w, c.dilation_h), + output_shape, + GGML_TYPE_F32); + output = add_bias_4d(runner.ctx(), output, c.in_channels, bias); + } else { + throw std::runtime_error("unknown depthwise2d candidate"); + } + std::vector>> writes; + writes.push_back({input, make_patterned_f32(static_cast(input_shape.num_elements()), 0.25f, 0.023f)}); + writes.push_back({weight, make_patterned_f32(static_cast(weight_shape.num_elements()), 0.55f, 0.012f)}); + if (bias) { + writes.push_back({*bias, make_patterned_f32(static_cast(bias_shape.num_elements()), 0.95f, 0.008f)}); + } + return runner.run(output, writes); + }); +} + +struct ConvTranspose1dCase { + const char * name; + int64_t batch; + int64_t in_channels; + int64_t out_channels; + int64_t frames; + int64_t kernel; + int stride; + int padding; + int dilation; + bool bias; +}; + +engine::core::TensorValue build_conv_transpose_native( + engine::core::ModuleBuildContext & ctx, + const ConvTranspose1dCase & c, + const engine::core::TensorValue & input, + const engine::core::TensorValue & weight, + const std::optional & bias) { + if (c.padding != 0 || c.dilation != 1) { + throw std::runtime_error("native ggml_conv_transpose_1d supports only padding=0 and dilation=1"); + } + engine::core::TensorValue output; + for (int64_t batch = 0; batch < c.batch; ++batch) { + auto matrix = view_batch_matrix(ctx, input, batch, c.in_channels, c.frames); + auto batch_out = engine::core::wrap_tensor( + ggml_conv_transpose_1d(ctx.ggml, weight.tensor, matrix.tensor, c.stride, c.padding, c.dilation), + engine::core::TensorShape::from_dims({1, c.out_channels, conv_transpose_out(c.frames, c.kernel, c.stride, c.padding, c.dilation)}), + GGML_TYPE_F32); + output = output.valid() ? engine::modules::ConcatModule({0}).build(ctx, output, batch_out) : batch_out; + } + return add_bias_3d(ctx, output, c.out_channels, bias); +} + +engine::core::TensorValue build_conv_transpose_col2im( + engine::core::ModuleBuildContext & ctx, + const ConvTranspose1dCase & c, + const engine::core::TensorValue & input, + const engine::core::TensorValue & weight, + const std::optional & bias) { + if (c.dilation != 1) { + throw std::runtime_error("col2im lowering currently supports only dilation=1"); + } + auto * weight_perm = ggml_reshape_2d( + ctx.ggml, + ggml_cont(ctx.ggml, ggml_permute(ctx.ggml, weight.tensor, 1, 2, 0, 3)), + c.in_channels, + c.kernel * c.out_channels); + ggml_tensor * bias_matrix = nullptr; + if (c.bias) { + if (!bias.has_value()) { + throw std::runtime_error("missing bias"); + } + bias_matrix = ggml_reshape_2d(ctx.ggml, bias->tensor, 1, c.out_channels); + } + engine::core::TensorValue output; + for (int64_t batch = 0; batch < c.batch; ++batch) { + auto * batch_input = ggml_view_2d( + ctx.ggml, + input.tensor, + input.tensor->ne[0], + input.tensor->ne[1], + input.tensor->nb[1], + static_cast(batch) * input.tensor->nb[2]); + auto * transposed_input = ggml_cont(ctx.ggml, ggml_transpose(ctx.ggml, batch_input)); + auto * columns = ggml_mul_mat(ctx.ggml, weight_perm, transposed_input); + auto * batch_output = ggml_col2im_1d(ctx.ggml, columns, c.stride, static_cast(c.out_channels), c.padding); + if (bias_matrix != nullptr) { + batch_output = ggml_add(ctx.ggml, batch_output, bias_matrix); + } + auto batch_value = engine::core::wrap_tensor( + ggml_reshape_3d(ctx.ggml, batch_output, batch_output->ne[0], batch_output->ne[1], 1), + engine::core::TensorShape::from_dims({1, c.out_channels, batch_output->ne[0]}), + GGML_TYPE_F32); + output = output.valid() ? engine::modules::ConcatModule({0}).build(ctx, output, batch_value) : batch_value; + } + return output; +} + +RunResult run_conv_transpose1d(const ConvTranspose1dCase & c, const char * candidate, engine::core::BackendType backend_type) { + if (backend_type == engine::core::BackendType::Cpu && std::string(candidate) == "matmul_col2im") { + RunResult result; + result.supported = false; + result.error = "current ggml CPU backend aborts for COL2IM_1D"; + return result; + } + return run_guarded(candidate, backend_type, [&](GraphRunner & runner) { + const auto input_shape = engine::core::TensorShape::from_dims({c.batch, c.in_channels, c.frames}); + const auto weight_shape = engine::core::TensorShape::from_dims({c.in_channels, c.out_channels, c.kernel}); + const auto bias_shape = engine::core::TensorShape::from_dims({c.out_channels}); + auto input = runner.make_f32(input_shape); + auto weight = runner.make_f32(weight_shape); + std::optional bias = c.bias ? std::optional(runner.make_f32(bias_shape)) : std::nullopt; + engine::core::TensorValue output; + if (std::string(candidate) == "native_direct") { + output = build_conv_transpose_native(runner.ctx(), c, input, weight, bias); + } else if (std::string(candidate) == "matmul_col2im") { + output = build_conv_transpose_col2im(runner.ctx(), c, input, weight, bias); + } else { + throw std::runtime_error("unknown conv_transpose1d candidate"); + } + std::vector>> writes; + writes.push_back({input, make_patterned_f32(static_cast(input_shape.num_elements()), 0.27f, 0.019f)}); + writes.push_back({weight, make_patterned_f32(static_cast(weight_shape.num_elements()), 0.57f, 0.011f)}); + if (bias) { + writes.push_back({*bias, make_patterned_f32(static_cast(bias_shape.num_elements()), 0.97f, 0.006f)}); + } + return runner.run(output, writes); + }); +} + +struct MatrixRow { + std::string module; + std::string case_name; + std::string candidate; + engine::core::BackendType backend; + RunResult result; + std::optional diff; +}; + +void print_row(const MatrixRow & row) { + std::cout << "| " << row.module + << " | " << row.case_name + << " | " << row.candidate + << " | " << backend_name(row.backend) + << " | "; + if (!row.result.supported) { + std::string error = row.result.error; + std::replace(error.begin(), error.end(), '|', '/'); + std::cout << "unsupported | - | - | - | - | " << error << " |\n"; + return; + } + std::cout << "ok | " << row.result.shape.to_string() + << " | " << std::fixed << std::setprecision(4) << row.result.avg_ms + << " | "; + if (row.diff.has_value()) { + std::cout << std::scientific << std::setprecision(3) << row.diff->max_abs + << " | " << row.diff->mean_abs + << " | " << std::fixed << std::setprecision(9) << row.diff->cosine << " |\n"; + } else { + std::cout << "- | - | - |\n"; + } +} + +void add_result( + std::vector & rows, + const std::string & module, + const std::string & case_name, + const std::string & candidate, + engine::core::BackendType backend, + const RunResult & result, + const RunResult & reference) { + MatrixRow row{module, case_name, candidate, backend, result, std::nullopt}; + if (result.supported && reference.supported) { + if (!same_shape(reference.shape, result.shape)) { + row.result.supported = false; + row.result.error = "shape mismatch vs reference " + reference.shape.to_string(); + } else { + row.diff = diff_values(reference.values, result.values); + } + } + rows.push_back(std::move(row)); +} + +bool backend_available(engine::core::BackendType backend_type) { + try { + GraphRunner runner("conv_lowering_matrix.probe", backend_type); + return true; + } catch (...) { + return false; + } +} + +} // namespace + +int main() { + try { + std::vector backends = {engine::core::BackendType::Cpu}; + if (backend_available(engine::core::BackendType::Cuda)) { + backends.push_back(engine::core::BackendType::Cuda); + } else { + std::cout << "[SKIP] cuda backend unavailable\n"; + } + if (backend_available(engine::core::BackendType::Vulkan)) { + backends.push_back(engine::core::BackendType::Vulkan); + } else { + std::cout << "[SKIP] vulkan backend unavailable\n"; + } + + std::vector rows; + + const std::vector conv1d_cases = { + {"citrinet_like_large_regular", 1, 80, 256, 256, 11, 1, 5, 1, true}, + {"bigvgan_like_resblock", 1, 192, 192, 384, 7, 1, 3, 1, true}, + {"batched_stride_regular", 2, 64, 128, 160, 5, 2, 2, 1, true}, + {"dilated_regular", 1, 128, 128, 192, 3, 1, 2, 2, false}, + }; + const std::vector conv1d_candidates = {"native", "conv2d_normal", "conv2d_direct"}; + for (const auto & c : conv1d_cases) { + const auto reference = run_conv1d(c, "native", engine::core::BackendType::Cpu); + for (const auto backend : backends) { + for (const auto & candidate : conv1d_candidates) { + add_result(rows, "Conv1dModule", c.name, candidate, backend, run_conv1d(c, candidate.c_str(), backend), reference); + } + } + } + + const std::vector conv2d_cases = { + {"spectrogram_small_kernel", 1, 64, 128, 20, 160, 3, 3, 1, 1, 1, 1, 1, 1, true}, + {"conv1d_lowered_shape", 1, 256, 256, 1, 384, 1, 7, 1, 1, 0, 3, 1, 1, true}, + {"batched_feature_map", 2, 32, 64, 12, 96, 3, 5, 1, 2, 1, 2, 1, 1, false}, + }; + const std::vector conv2d_candidates = {"im2col_matmul", "direct"}; + for (const auto & c : conv2d_cases) { + const auto reference = run_conv2d(c, "im2col_matmul", engine::core::BackendType::Cpu); + for (const auto backend : backends) { + for (const auto & candidate : conv2d_candidates) { + add_result(rows, "Conv2dModule", c.name, candidate, backend, run_conv2d(c, candidate.c_str(), backend), reference); + } + } + } + + const std::vector depthwise1d_cases = { + {"conformer_like_depthwise", 1, 256, 192, 31, 1, 15, 1, true}, + {"tokenizer_stride_depthwise", 2, 96, 256, 7, 2, 3, 1, true}, + {"dilated_depthwise", 1, 128, 160, 5, 1, 4, 2, false}, + }; + const std::vector depthwise1d_candidates = {"dw2d_direct", "native_1d_dw"}; + for (const auto & c : depthwise1d_cases) { + const auto reference = run_depthwise1d(c, "dw2d_direct", engine::core::BackendType::Cpu); + for (const auto backend : backends) { + for (const auto & candidate : depthwise1d_candidates) { + add_result(rows, "DepthwiseConv1dModule", c.name, candidate, backend, run_depthwise1d(c, candidate.c_str(), backend), reference); + } + } + } + + const std::vector pointwise1d_cases = { + {"conformer_projection", 1, 256, 512, 192, true}, + {"batched_token_projection", 2, 192, 384, 160, true}, + {"vocoder_channel_mix", 1, 192, 192, 384, false}, + }; + const std::vector pointwise1d_candidates = {"conv1d_kernel1", "linear_matmul"}; + for (const auto & c : pointwise1d_cases) { + const auto reference = run_pointwise1d(c, "conv1d_kernel1", engine::core::BackendType::Cpu); + for (const auto backend : backends) { + for (const auto & candidate : pointwise1d_candidates) { + add_result(rows, "PointwiseConv1dModule", c.name, candidate, backend, run_pointwise1d(c, candidate.c_str(), backend), reference); + } + } + } + + const std::vector depthwise2d_cases = { + {"depthwise_1d_lowered_shape", 1, 192, 192, 1, 384, 1, 7, 1, 1, 0, 3, 1, 1, true}, + {"image_depthwise_small", 1, 64, 64, 24, 80, 3, 3, 1, 1, 1, 1, 1, 1, true}, + }; + const std::vector depthwise2d_candidates = {"direct", "im2col_matmul"}; + for (const auto & c : depthwise2d_cases) { + const auto reference = run_depthwise2d(c, "direct", engine::core::BackendType::Cpu); + for (const auto backend : backends) { + for (const auto & candidate : depthwise2d_candidates) { + add_result(rows, "DepthwiseConv2dModule", c.name, candidate, backend, run_depthwise2d(c, candidate.c_str(), backend), reference); + } + } + } + + const std::vector conv_transpose_cases = { + {"qwen3_like_stride5_padding0", 1, 256, 128, 96, 10, 5, 0, 1, true}, + {"vocoder_stride2_padding1", 1, 192, 96, 192, 4, 2, 1, 1, true}, + {"batched_stride2_no_bias", 2, 128, 128, 96, 2, 2, 0, 1, false}, + {"dilated_unsupported_probe", 1, 64, 64, 96, 3, 2, 0, 2, true}, + }; + const std::vector conv_transpose_candidates = {"native_direct", "matmul_col2im"}; + for (const auto & c : conv_transpose_cases) { + const auto reference = run_conv_transpose1d(c, "native_direct", engine::core::BackendType::Cpu); + for (const auto backend : backends) { + for (const auto & candidate : conv_transpose_candidates) { + add_result(rows, "ConvTranspose1dModule", c.name, candidate, backend, run_conv_transpose1d(c, candidate.c_str(), backend), reference); + } + } + } + + std::cout << "| module | case | candidate | backend | status | shape | avg_ms | max_abs_vs_cpu_ref | mean_abs_vs_cpu_ref | cosine_vs_cpu_ref |\n"; + std::cout << "| --- | --- | --- | --- | --- | --- | ---: | ---: | ---: | ---: |\n"; + for (const auto & row : rows) { + print_row(row); + } + } catch (const std::exception & ex) { + std::cerr << "[FAIL] " << ex.what() << '\n'; + return 1; + } + return 0; +} From b266f351aad4f9e20694c06f557fbc1ce6d0217e Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Sat, 18 Jul 2026 08:47:36 -0400 Subject: [PATCH 04/27] Cache Qwen3 TTS talker prefill state --- src/models/qwen3_tts/talker.cpp | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/src/models/qwen3_tts/talker.cpp b/src/models/qwen3_tts/talker.cpp index 57a45b59..0570b944 100644 --- a/src/models/qwen3_tts/talker.cpp +++ b/src/models/qwen3_tts/talker.cpp @@ -1628,6 +1628,7 @@ class Qwen3TalkerStepRuntime::Impl { !talker_prefill_equal(*cached_prompt_prefill_, request)) { cached_prompt_state_ = build_prompt_state(request, weights_->assets().config, weights_->weights()); cached_prompt_prefill_ = request; + cached_prefill_output_.reset(); } const auto & state = *cached_prompt_state_; const auto prompt_state_end = Clock::now(); @@ -1637,9 +1638,13 @@ class Qwen3TalkerStepRuntime::Impl { throw std::runtime_error("Qwen3 talker prompt exceeds step runtime capacity"); } const auto prefill_start = Clock::now(); - auto prefill_output = run_prefill_embeddings_with_state(state.prompt, prompt_steps); + const bool prefill_cache_hit = cached_prefill_output_.has_value(); + if (!prefill_cache_hit) { + cached_prefill_output_ = run_prefill_embeddings_with_state(state.prompt, prompt_steps); + } + const auto & prefill_output = *cached_prefill_output_; const auto prefill_end = Clock::now(); - auto current = std::move(prefill_output.result); + auto current = prefill_output.result; double code_predictor_build_ms = 0.0; if (code_predictor_graph_ == nullptr) { const auto build_start = Clock::now(); @@ -1649,12 +1654,14 @@ class Qwen3TalkerStepRuntime::Impl { double cached_step_build_ms = 0.0; double import_prefill_state_ms = 0.0; int64_t cached_step_capacity = 0; - auto cached_state = std::move(prefill_output.state); + runtime::TransformerKVState exported_cached_state; + const runtime::TransformerKVState * cached_state = &prefill_output.state; bool cached_graph_has_state = false; auto ensure_cached_step_capacity = [&](int64_t required_capacity) { if (cached_step_graph_ != nullptr && cached_graph_has_state && !cached_step_graph_->can_run(*weights_, required_capacity)) { - cached_state = cached_step_graph_->export_state(); + exported_cached_state = cached_step_graph_->export_state(); + cached_state = &exported_cached_state; cached_graph_has_state = false; } if (cached_step_graph_ == nullptr || !cached_step_graph_->can_run(*weights_, required_capacity)) { @@ -1676,7 +1683,7 @@ class Qwen3TalkerStepRuntime::Impl { } if (!cached_graph_has_state) { const auto import_start = Clock::now(); - cached_step_graph_->import_prefill_state(cached_state); + cached_step_graph_->import_prefill_state(*cached_state); import_prefill_state_ms += engine::debug::elapsed_ms(import_start, Clock::now()); cached_graph_has_state = true; } @@ -1761,6 +1768,7 @@ class Qwen3TalkerStepRuntime::Impl { out.decoder_input_codes.frames += out.generated_codes.frames; debug::timing_log_scalar("qwen3_tts.talker.prompt_state_ms", engine::debug::elapsed_ms(prompt_state_start, prompt_state_end)); debug::timing_log_scalar("qwen3_tts.talker.prefill_ms", engine::debug::elapsed_ms(prefill_start, prefill_end)); + debug::timing_log_scalar("qwen3_tts.talker.prefill_cache.hit", prefill_cache_hit); debug::timing_log_scalar("qwen3_tts.talker.code_predictor_build_ms", code_predictor_build_ms); debug::timing_log_scalar("qwen3_tts.talker.cached_step_build_ms", cached_step_build_ms); debug::timing_log_scalar("qwen3_tts.talker.import_prefill_state_ms", import_prefill_state_ms); @@ -1813,6 +1821,7 @@ class Qwen3TalkerStepRuntime::Impl { std::unique_ptr code_predictor_graph_; std::optional cached_prompt_prefill_; std::optional cached_prompt_state_; + std::optional cached_prefill_output_; }; Qwen3TalkerStepRuntime::Qwen3TalkerStepRuntime(std::unique_ptr impl) : impl_(std::move(impl)) { From 9ba3c4b82e3cc0f9cc1f45baaf672d34e40386bb Mon Sep 17 00:00:00 2001 From: mirek190 Date: Sun, 19 Jul 2026 10:44:14 +0100 Subject: [PATCH 05/27] Optimize Higgs Audio v3 inference - pack QKV and gate/up projections and use fused SwiGLU in the shared Qwen decoder path - use grouped FlashAttention, F16 bucketed KV caches, direct set-rows updates, and CUDA-friendly RoPE/view/set-rows graphs - retain and reuse cloned-reference KV prefixes so repeated server requests prefill only their text suffix - add graph and F16 KV correctness coverage plus fixed-seed request-level CUDA benchmark/parity tooling - document the optimized Higgs runtime controls and reproducible validation flow --- CMakeLists.txt | 11 + docs/tts.md | 6 +- include/engine/framework/core/backend.h | 2 + .../modules/attention/qwen_causal_decoder.h | 5 + .../modules/attention/qwen_decoder.h | 1 + include/engine/models/higgs_tts/ar.h | 1 + include/engine/models/higgs_tts/generator.h | 1 + src/framework/core/backend.cpp | 20 +- .../modules/attention/qwen_causal_decoder.cpp | 30 ++ .../modules/attention/qwen_decoder.cpp | 122 +++-- .../modules/optimizations/fast_kv_modules.cpp | 31 +- src/framework/runtime/kv_cache.cpp | 38 +- src/models/higgs_tts/ar.cpp | 148 +++-- src/models/higgs_tts/generator.cpp | 61 ++- tests/higgs_tts/.gitignore | 2 + tests/higgs_tts/README.md | 38 ++ tests/higgs_tts/compare_warmbench_results.py | 160 ++++++ .../higgs_tts/higgs_tts_cuda_bench_cases.json | 35 ++ .../higgs_tts/higgs_tts_cuda_mixed_cases.json | 44 ++ .../higgs_tts/higgs_tts_cuda_perf_cases.json | 57 ++ tests/higgs_tts/higgs_tts_warm_bench.cpp | 51 +- tests/higgs_tts/run_cuda_performance.ps1 | 69 +++ .../test_qwen_decoder_packed_projections.cpp | 504 ++++++++++++++++++ 23 files changed, 1321 insertions(+), 116 deletions(-) create mode 100644 tests/higgs_tts/.gitignore create mode 100644 tests/higgs_tts/README.md create mode 100644 tests/higgs_tts/compare_warmbench_results.py create mode 100644 tests/higgs_tts/higgs_tts_cuda_bench_cases.json create mode 100644 tests/higgs_tts/higgs_tts_cuda_mixed_cases.json create mode 100644 tests/higgs_tts/higgs_tts_cuda_perf_cases.json create mode 100644 tests/higgs_tts/run_cuda_performance.ps1 create mode 100644 tests/unittests/test_qwen_decoder_packed_projections.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index cd0b9a6b..18dbec27 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -709,6 +709,7 @@ if (ENGINE_BUILD_WARMBENCH) add_engine_warmbench(chatterbox_warm_bench tests/chatterbox/chatterbox_warm_bench.cpp) add_engine_warmbench(citrinet_asr_warm_bench tests/citrinet_asr/citrinet_asr_warm_bench.cpp) add_engine_warmbench(higgs_audio_stt_warm_bench tests/higgs_audio_stt/higgs_audio_stt_warm_bench.cpp) + add_engine_warmbench(higgs_tts_warm_bench tests/higgs_tts/higgs_tts_warm_bench.cpp) add_engine_warmbench(hviske_asr_warm_bench tests/hviske_asr/hviske_asr_warm_bench.cpp) add_engine_warmbench(index_tts2_warm_bench tests/index_tts2/index_tts2_warm_bench.cpp) add_engine_warmbench(irodori_tts_warm_bench tests/irodori_tts/irodori_tts_warm_bench.cpp) @@ -857,6 +858,16 @@ if (ENGINE_BUILD_TESTS) COMMAND encoder_module_test ) + add_engine_unittest( + qwen_decoder_packed_projection_test + tests/unittests/test_qwen_decoder_packed_projections.cpp + ) + + add_test( + NAME qwen_decoder_packed_projection_test + COMMAND qwen_decoder_packed_projection_test + ) + add_engine_unittest(conv_transpose_fast_path_test tests/unittests/test_conv_transpose_fast_path.cpp) add_test( diff --git a/docs/tts.md b/docs/tts.md index d0ff8eee..6a3a6768 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -368,9 +368,9 @@ audiocpp_cli --task tts --family higgs_tts --model models/higgs-audio-v3-tts-4b | `--text-chunk-size` | integer chars | `512` | Long-form chunk size. | | `--max-tokens` | integer | `1024` | Maximum generated AR tokens per chunk. | | `--temperature` | float | `0.8` | AR sampling temperature. | -| `--top-k` | integer | `30` | AR top-k sampling limit. | -| `--top-p` | float | `0.8` | AR nucleus sampling limit. | -| `--repetition-penalty` | float | `1.1` | AR repetition penalty. | +| `--top-k` | integer | `30` | AR top-k sampling limit. The narrower default is less prone to premature EOC than the Python client's `50`. | +| `--top-p` | float | `0.8` | AR nucleus sampling limit. The Python client's unfiltered equivalent is `1.0`. | +| `--repetition-penalty` | float | `1.1` | Accepted for Python API compatibility; Higgs audio-code sampling does not consume it. | ## IndexTTS2 diff --git a/include/engine/framework/core/backend.h b/include/engine/framework/core/backend.h index ac8c0277..4597e357 100644 --- a/include/engine/framework/core/backend.h +++ b/include/engine/framework/core/backend.h @@ -74,6 +74,8 @@ void write_tensor_i32(const TensorValue & tensor, const int32_t * values, size_t void write_tensor_i32(const TensorValue & tensor, const std::vector & values); void read_tensor_f32_into(const ggml_tensor * tensor, std::vector & values); std::vector read_tensor_f32(const ggml_tensor * tensor); +void read_tensor_f16_into(const ggml_tensor * tensor, std::vector & values); +std::vector read_tensor_f16(const ggml_tensor * tensor); void read_tensor_i32_into(const ggml_tensor * tensor, std::vector & values); std::vector read_tensor_i32(const ggml_tensor * tensor); diff --git a/include/engine/framework/modules/attention/qwen_causal_decoder.h b/include/engine/framework/modules/attention/qwen_causal_decoder.h index ffc51353..0ee5ebe7 100644 --- a/include/engine/framework/modules/attention/qwen_causal_decoder.h +++ b/include/engine/framework/modules/attention/qwen_causal_decoder.h @@ -79,6 +79,11 @@ std::vector qwen_position_ids(int64_t steps, int64_t offset = 0); std::vector qwen_causal_prefill_mask_values(int64_t batch_size, int64_t steps); +std::vector qwen_causal_suffix_mask_values( + int64_t batch_size, + int64_t query_steps, + int64_t prefix_steps); + void write_qwen_causal_prefill_mask( ggml_tensor * tensor, int64_t batch_size, diff --git a/include/engine/framework/modules/attention/qwen_decoder.h b/include/engine/framework/modules/attention/qwen_decoder.h index 968d6d60..0cfaf6bc 100644 --- a/include/engine/framework/modules/attention/qwen_decoder.h +++ b/include/engine/framework/modules/attention/qwen_decoder.h @@ -84,6 +84,7 @@ struct QwenDecoderLayerConfig { struct QwenMLPWeights { LinearWeights gate_proj; LinearWeights up_proj; + std::optional gate_up_proj; LinearWeights down_proj; }; diff --git a/include/engine/models/higgs_tts/ar.h b/include/engine/models/higgs_tts/ar.h index 73270860..fd47ba3f 100644 --- a/include/engine/models/higgs_tts/ar.h +++ b/include/engine/models/higgs_tts/ar.h @@ -28,6 +28,7 @@ struct HiggsARWeights { core::TensorValue modality_embedding; HiggsQwenDecoderStackWeights decoder; core::TensorValue norm; + bool packed_qkv = false; }; HiggsARWeights load_higgs_ar_weights( diff --git a/include/engine/models/higgs_tts/generator.h b/include/engine/models/higgs_tts/generator.h index 2d3733cf..4eb3e761 100644 --- a/include/engine/models/higgs_tts/generator.h +++ b/include/engine/models/higgs_tts/generator.h @@ -67,6 +67,7 @@ class HiggsGenerator { HiggsTextTokenizer tokenizer_; size_t ar_decode_graph_arena_bytes_ = 0; std::optional reference_prefix_cache_; + bool reference_kv_ready_ = false; std::optional cuda_sampling_policy_; std::unique_ptr ar_kv_cache_; std::unique_ptr prefill_graph_; diff --git a/src/framework/core/backend.cpp b/src/framework/core/backend.cpp index b34cdea3..255d59b4 100644 --- a/src/framework/core/backend.cpp +++ b/src/framework/core/backend.cpp @@ -527,10 +527,22 @@ void read_tensor_f32_into(const ggml_tensor * tensor, std::vector & value read_tensor_typed_into(tensor, GGML_TYPE_F32, values); } -std::vector read_tensor_f32(const ggml_tensor * tensor) { - return read_tensor_typed(tensor, GGML_TYPE_F32); -} - +std::vector read_tensor_f32(const ggml_tensor * tensor) { + return read_tensor_typed(tensor, GGML_TYPE_F32); +} + +void read_tensor_f16_into(const ggml_tensor * tensor, std::vector & values) { + const auto fp16_values = read_tensor_typed(tensor, GGML_TYPE_F16); + values.resize(fp16_values.size()); + ggml_fp16_to_fp32_row(fp16_values.data(), values.data(), static_cast(values.size())); +} + +std::vector read_tensor_f16(const ggml_tensor * tensor) { + std::vector values; + read_tensor_f16_into(tensor, values); + return values; +} + void read_tensor_i32_into(const ggml_tensor * tensor, std::vector & values) { read_tensor_typed_into(tensor, GGML_TYPE_I32, values); } diff --git a/src/framework/modules/attention/qwen_causal_decoder.cpp b/src/framework/modules/attention/qwen_causal_decoder.cpp index 229243c1..63c5fcd8 100644 --- a/src/framework/modules/attention/qwen_causal_decoder.cpp +++ b/src/framework/modules/attention/qwen_causal_decoder.cpp @@ -175,6 +175,36 @@ std::vector qwen_causal_prefill_mask_values(int64_t batch_size, int return out; } +std::vector qwen_causal_suffix_mask_values( + int64_t batch_size, + int64_t query_steps, + int64_t prefix_steps) { + if (batch_size <= 0) { + throw std::runtime_error("qwen_causal_suffix_mask_values requires positive batch size"); + } + validate_steps(query_steps, "qwen_causal_suffix_mask_values"); + if (prefix_steps < 0) { + throw std::runtime_error("qwen_causal_suffix_mask_values requires non-negative prefix steps"); + } + const int64_t key_steps = prefix_steps + query_steps; + const auto masked = ggml_fp32_to_fp16(-INFINITY); + const auto visible = ggml_fp32_to_fp16(0.0F); + std::vector one(static_cast(query_steps * key_steps), masked); + for (int64_t row = 0; row < query_steps; ++row) { + const size_t row_offset = static_cast(row * key_steps); + std::fill_n( + one.begin() + static_cast(row_offset), + prefix_steps + row + 1, + visible); + } + std::vector out; + out.reserve(static_cast(batch_size) * one.size()); + for (int64_t batch = 0; batch < batch_size; ++batch) { + out.insert(out.end(), one.begin(), one.end()); + } + return out; +} + void write_qwen_causal_prefill_mask( ggml_tensor * tensor, int64_t batch_size, diff --git a/src/framework/modules/attention/qwen_decoder.cpp b/src/framework/modules/attention/qwen_decoder.cpp index 886f1c03..56b7c601 100644 --- a/src/framework/modules/attention/qwen_decoder.cpp +++ b/src/framework/modules/attention/qwen_decoder.cpp @@ -332,35 +332,75 @@ core::TensorValue build_mlp( const core::TensorValue & input, const QwenDecoderLayerConfig & config, const QwenMLPWeights & weights) { - auto gate = LinearModule( - { - config.hidden_size, - config.intermediate_size, - weights.gate_proj.bias.has_value(), - config.projection_precision, - }) - .build(ctx, input, require_linear(weights.gate_proj, false, "QwenMLPWeights.gate_proj")); - if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { - gate = activation_cast(ctx, gate, config.activation_cast); - } - gate = SiluModule{}.build(ctx, gate); - if (config.activation_cast.enabled && config.activation_cast.after_mlp_silu) { - gate = activation_cast(ctx, gate, config.activation_cast); - } - auto up = LinearModule( - { - config.hidden_size, - config.intermediate_size, - weights.up_proj.bias.has_value(), - config.projection_precision, - }) - .build(ctx, input, require_linear(weights.up_proj, false, "QwenMLPWeights.up_proj")); - if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { - up = activation_cast(ctx, up, config.activation_cast); - } - auto gated = MulModule{}.build(ctx, gate, up); - if (config.activation_cast.enabled && config.activation_cast.after_mlp_mul) { - gated = activation_cast(ctx, gated, config.activation_cast); + core::TensorValue gate; + core::TensorValue up; + std::optional packed_gate_up; + if (weights.gate_up_proj.has_value()) { + auto gate_up = LinearModule( + { + config.hidden_size, + config.intermediate_size * 2, + weights.gate_up_proj->bias.has_value(), + config.projection_precision, + }) + .build( + ctx, + input, + require_linear(*weights.gate_up_proj, false, "QwenMLPWeights.gate_up_proj")); + packed_gate_up = gate_up; + gate = SliceModule({2, 0, config.intermediate_size}).build(ctx, gate_up); + up = SliceModule({2, config.intermediate_size, config.intermediate_size}).build(ctx, gate_up); + } else { + gate = LinearModule( + { + config.hidden_size, + config.intermediate_size, + weights.gate_proj.bias.has_value(), + config.projection_precision, + }) + .build(ctx, input, require_linear(weights.gate_proj, false, "QwenMLPWeights.gate_proj")); + up = LinearModule( + { + config.hidden_size, + config.intermediate_size, + weights.up_proj.bias.has_value(), + config.projection_precision, + }) + .build(ctx, input, require_linear(weights.up_proj, false, "QwenMLPWeights.up_proj")); + } + const bool can_use_fused_swiglu = + !config.activation_cast.enabled || + (!config.activation_cast.after_mlp_projection && + !config.activation_cast.after_mlp_silu && + !config.activation_cast.after_mlp_mul); + core::TensorValue gated; + if (can_use_fused_swiglu && packed_gate_up.has_value()) { + gated = core::wrap_tensor( + ggml_swiglu(ctx.ggml, packed_gate_up->tensor), + core::TensorShape::from_dims({ + input.shape.dims[0], + input.shape.dims[1], + config.intermediate_size, + }), + packed_gate_up->type); + } else if (can_use_fused_swiglu) { + gated = core::wrap_tensor( + ggml_swiglu_split(ctx.ggml, gate.tensor, up.tensor), + gate.shape, + gate.type); + } else { + if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { + gate = activation_cast(ctx, gate, config.activation_cast); + up = activation_cast(ctx, up, config.activation_cast); + } + gate = SiluModule{}.build(ctx, gate); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_silu) { + gate = activation_cast(ctx, gate, config.activation_cast); + } + gated = MulModule{}.build(ctx, gate, up); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_mul) { + gated = activation_cast(ctx, gated, config.activation_cast); + } } auto down = LinearModule( { @@ -441,8 +481,22 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build( v = core::ensure_backend_addressable_layout(ctx, v); auto q_heads = TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); - auto all_k = prefix_key.has_value() ? ConcatModule({1}).build(ctx, *prefix_key, k) : k; - auto all_v = prefix_value.has_value() ? ConcatModule({1}).build(ctx, *prefix_value, v) : v; + auto attention_prefix_key = prefix_key; + auto attention_prefix_value = prefix_value; + if (attention_prefix_key.has_value() && attention_prefix_key->type != k.type) { + attention_prefix_key = core::wrap_tensor( + ggml_cast(ctx.ggml, attention_prefix_key->tensor, k.type), + attention_prefix_key->shape, + k.type); + } + if (attention_prefix_value.has_value() && attention_prefix_value->type != v.type) { + attention_prefix_value = core::wrap_tensor( + ggml_cast(ctx.ggml, attention_prefix_value->tensor, v.type), + attention_prefix_value->shape, + v.type); + } + auto all_k = attention_prefix_key.has_value() ? ConcatModule({1}).build(ctx, *attention_prefix_key, k) : k; + auto all_v = attention_prefix_value.has_value() ? ConcatModule({1}).build(ctx, *attention_prefix_value, v) : v; core::TensorValue context; if (!prefix_key.has_value() && attention_mask.has_value() && config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGroupedViewKV) { @@ -457,8 +511,10 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build( dim, *attention_mask, config_.attention_precision); - } else if (!prefix_key.has_value() && attention_mask.has_value() && - config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGrouped) { + } else if (attention_mask.has_value() && + (config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGrouped || + (prefix_key.has_value() && + config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGroupedViewKV))) { q_heads = core::wrap_tensor(ggml_cont(ctx.ggml, q_heads.tensor), q_heads.shape, q_heads.type); auto k_heads = TransposeModule({{0, 2, 1, 3}, all_k.shape.rank}).build(ctx, all_k); auto v_heads = TransposeModule({{0, 2, 1, 3}, all_v.shape.rank}).build(ctx, all_v); diff --git a/src/framework/modules/optimizations/fast_kv_modules.cpp b/src/framework/modules/optimizations/fast_kv_modules.cpp index ce9e08eb..64389763 100644 --- a/src/framework/modules/optimizations/fast_kv_modules.cpp +++ b/src/framework/modules/optimizations/fast_kv_modules.cpp @@ -51,8 +51,8 @@ core::TensorValue FastKVSetRowsModule::build( if (row_index.shape.rank != 1 || (row_index.shape.dims[0] != 1 && row_index.shape.dims[0] != batch)) { throw std::runtime_error("FastKVSetRowsModule row_index must have shape {1} or {batch}"); } - if (cache.type != GGML_TYPE_F32 || row.type != GGML_TYPE_F32) { - throw std::runtime_error("FastKVSetRowsModule requires f32 cache and row tensors"); + if ((cache.type != GGML_TYPE_F32 && cache.type != GGML_TYPE_F16) || row.type != GGML_TYPE_F32) { + throw std::runtime_error("FastKVSetRowsModule requires an f32/f16 cache and an f32 row tensor"); } if (row_index.type != GGML_TYPE_I32 && row_index.type != GGML_TYPE_I64) { throw std::runtime_error("FastKVSetRowsModule requires i32 or i64 row_index tensor"); @@ -69,16 +69,39 @@ core::TensorValue FastKVSetRowsModule::build( } auto flat_cache = core::reshape_tensor(ctx, cache, core::TensorShape::from_dims({steps, row_elems})); auto contiguous_row = tensor_layout::ensure_contiguous_layout_if_needed(ctx, row); - auto flat_row = core::reshape_tensor(ctx, contiguous_row, core::TensorShape::from_dims({1, row_elems})); + auto flat_row = core::wrap_tensor( + ggml_view_2d( + ctx.ggml, + contiguous_row.tensor, + row_elems, + 1, + contiguous_row.tensor->nb[2], + 0), + core::TensorShape::from_dims({1, row_elems}), + row.type); ggml_tensor * updated = ggml_set_rows(ctx.ggml, flat_cache.tensor, flat_row.tensor, row_index.tensor); + // ggml_set_rows src[2] is only a legacy dependency anchor for the + // destination. Point it at the underlying cache so the metadata-only + // flatten does not interrupt CUDA's ROPE -> VIEW -> SET_ROWS fusion. + updated->src[2] = cache.tensor; auto flat_updated = core::wrap_tensor(updated, flat_cache.shape, cache.type); return core::reshape_tensor(ctx, flat_updated, cache.shape); } auto flat_cache = core::reshape_tensor(ctx, cache, core::TensorShape::from_dims({batch * steps, row_elems})); auto contiguous_row = tensor_layout::ensure_contiguous_layout_if_needed(ctx, row); - auto flat_row = core::reshape_tensor(ctx, contiguous_row, core::TensorShape::from_dims({batch, row_elems})); + auto flat_row = core::wrap_tensor( + ggml_view_2d( + ctx.ggml, + contiguous_row.tensor, + row_elems, + batch, + contiguous_row.tensor->nb[3], + 0), + core::TensorShape::from_dims({batch, row_elems}), + row.type); ggml_tensor * updated = ggml_set_rows(ctx.ggml, flat_cache.tensor, flat_row.tensor, row_index.tensor); + updated->src[2] = cache.tensor; auto flat_updated = core::wrap_tensor(updated, flat_cache.shape, cache.type); return core::reshape_tensor(ctx, flat_updated, cache.shape); } diff --git a/src/framework/runtime/kv_cache.cpp b/src/framework/runtime/kv_cache.cpp index 24f9925c..17e236fa 100644 --- a/src/framework/runtime/kv_cache.cpp +++ b/src/framework/runtime/kv_cache.cpp @@ -9,6 +9,30 @@ namespace engine::runtime { +namespace { + +void write_cache_tensor(const core::TensorValue & tensor, const std::vector & values) { + if (tensor.type == GGML_TYPE_F32) { + core::write_tensor_f32(tensor, values); + } else if (tensor.type == GGML_TYPE_F16) { + core::write_tensor_f16(tensor, values); + } else { + throw std::runtime_error("TransformerKVCache supports only f32 and f16 cache tensors"); + } +} + +std::vector read_cache_tensor(const core::TensorValue & tensor) { + if (tensor.type == GGML_TYPE_F32) { + return core::read_tensor_f32(tensor.tensor); + } + if (tensor.type == GGML_TYPE_F16) { + return core::read_tensor_f16(tensor.tensor); + } + throw std::runtime_error("TransformerKVCache supports only f32 and f16 cache tensors"); +} + +} // namespace + TransformerKVCache::TransformerKVCache( int64_t cache_steps, int64_t step_elems, @@ -68,8 +92,8 @@ void TransformerKVCache::import_state(const TransformerKVState & state) { std::copy(source.key.begin(), source.key.end(), cache.import_key_scratch.begin()); std::copy(source.value.begin(), source.value.end(), cache.import_value_scratch.begin()); } - core::write_tensor_f32(cache.key_tensor, cache.import_key_scratch); - core::write_tensor_f32(cache.value_tensor, cache.import_value_scratch); + write_cache_tensor(cache.key_tensor, cache.import_key_scratch); + write_cache_tensor(cache.value_tensor, cache.import_value_scratch); } } } @@ -85,8 +109,8 @@ TransformerKVState TransformerKVCache::export_state() const { if (keep_elems == 0) { continue; } - const auto key_values = core::read_tensor_f32(layers_[layer].key_tensor.tensor); - const auto value_values = core::read_tensor_f32(layers_[layer].value_tensor.tensor); + const auto key_values = read_cache_tensor(layers_[layer].key_tensor); + const auto value_values = read_cache_tensor(layers_[layer].value_tensor); out.key.assign(key_values.begin(), key_values.begin() + static_cast(keep_elems)); out.value.assign(value_values.begin(), value_values.begin() + static_cast(keep_elems)); } @@ -141,11 +165,11 @@ void TransformerKVCache::trace_log_state(const std::string & name, int64_t num_h return; } const size_t keep_elems = static_cast(valid_steps_ * step_elems_); - const auto first_key = core::read_tensor_f32(layers_.front().key_tensor.tensor); + const auto first_key = read_cache_tensor(layers_.front().key_tensor); std::vector first_key_keep(first_key.begin(), first_key.begin() + static_cast(keep_elems)); debug::trace_log_f32(name + ".layer0.key", {1, valid_steps_, num_heads, head_dim}, first_key_keep); if (layers_.size() > 1) { - const auto last_key = core::read_tensor_f32(layers_.back().key_tensor.tensor); + const auto last_key = read_cache_tensor(layers_.back().key_tensor); std::vector last_key_keep(last_key.begin(), last_key.begin() + static_cast(keep_elems)); debug::trace_log_f32(name + ".layer_last.key", {1, valid_steps_, num_heads, head_dim}, last_key_keep); } @@ -175,7 +199,7 @@ core::TensorValue view_transformer_kv_cache_steps( cache.tensor->nb[3], static_cast(start) * cache.tensor->nb[2]), core::TensorShape::from_dims({1, steps, heads, head_dim}), - GGML_TYPE_F32); + cache.type); } } // namespace engine::runtime diff --git a/src/models/higgs_tts/ar.cpp b/src/models/higgs_tts/ar.cpp index 6ffe2c4c..922d4830 100644 --- a/src/models/higgs_tts/ar.cpp +++ b/src/models/higgs_tts/ar.cpp @@ -62,18 +62,25 @@ modules::QwenDecoderStackConfig make_higgs_qwen_stack_config(const HiggsTextConf class HiggsQwenDecoderComponent { public: - explicit HiggsQwenDecoderComponent(const HiggsTextConfig & config) + HiggsQwenDecoderComponent(const HiggsTextConfig & config, bool packed_qkv) : stack_config_(make_higgs_qwen_stack_config(config)), layer_config_(modules::qwen_decoder_layer_config_from_stack(stack_config_)), - layer_module_(layer_config_) {} + layer_module_([&] { + layer_config_.qkv_layout = packed_qkv + ? modules::QwenDecoderQKVLayout::PackedQKV + : modules::QwenDecoderQKVLayout::Separate; + return layer_config_; + }()) {} modules::QwenDecoderLayerOutputs build_prefill_layer( core::ModuleBuildContext & ctx, const core::TensorValue & input, const core::TensorValue & positions, const modules::QwenDecoderLayerWeights & weights, - const core::TensorValue & attention_mask) const { - return layer_module_.build(ctx, input, positions, weights, std::nullopt, std::nullopt, attention_mask); + const core::TensorValue & attention_mask, + const std::optional & prefix_key = std::nullopt, + const std::optional & prefix_value = std::nullopt) const { + return layer_module_.build(ctx, input, positions, weights, prefix_key, prefix_value, attention_mask); } modules::QwenDecoderLayerOutputs build_decode_layer( @@ -144,23 +151,33 @@ modules::QwenDecoderLayerWeights load_layer_weights( store.load_f32_tensor(source, prefix + ".input_layernorm.weight", {config.hidden_size}), std::nullopt, }; - // Keep Q/K/V separate here so Higgs exercises the framework Qwen decoder path. - // A packed-QKV fast path can be evaluated later as a framework-level optimization. - weights.self_attention.q_weight = store.load_tensor( - source, - prefix + ".self_attn.q_proj.weight", - storage_type, - {q_out, config.hidden_size}); - weights.self_attention.k_weight = store.load_tensor( - source, - prefix + ".self_attn.k_proj.weight", - storage_type, - {kv_out, config.hidden_size}); - weights.self_attention.v_weight = store.load_tensor( - source, - prefix + ".self_attn.v_proj.weight", - storage_type, - {kv_out, config.hidden_size}); + { + const auto q = source.require_tensor( + prefix + ".self_attn.q_proj.weight", + storage_type, + {q_out, config.hidden_size}); + const auto k = source.require_tensor( + prefix + ".self_attn.k_proj.weight", + storage_type, + {kv_out, config.hidden_size}); + const auto v = source.require_tensor( + prefix + ".self_attn.v_proj.weight", + storage_type, + {kv_out, config.hidden_size}); + if (q.type != k.type || q.type != v.type) { + throw std::runtime_error("Higgs TTS packed QKV weights require matching storage types"); + } + std::vector packed; + packed.reserve(q.bytes.size() + k.bytes.size() + v.bytes.size()); + packed.insert(packed.end(), q.bytes.begin(), q.bytes.end()); + packed.insert(packed.end(), k.bytes.begin(), k.bytes.end()); + packed.insert(packed.end(), v.bytes.begin(), v.bytes.end()); + weights.self_attention.qkv_weight = store.make_tensor( + core::TensorShape::from_dims({q_out + 2 * kv_out, config.hidden_size}), + q.type, + packed.data(), + packed.size()); + } weights.self_attention.out_weight = store.load_tensor( source, prefix + ".self_attn.o_proj.weight", @@ -178,22 +195,31 @@ modules::QwenDecoderLayerWeights load_layer_weights( store.load_f32_tensor(source, prefix + ".post_attention_layernorm.weight", {config.hidden_size}), std::nullopt, }; - weights.mlp.gate_proj = { - store.load_tensor( - source, + { + const auto gate = source.require_tensor( prefix + ".mlp.gate_proj.weight", storage_type, - {config.intermediate_size, config.hidden_size}), - std::nullopt, - }; - weights.mlp.up_proj = { - store.load_tensor( - source, + {config.intermediate_size, config.hidden_size}); + const auto up = source.require_tensor( prefix + ".mlp.up_proj.weight", storage_type, - {config.intermediate_size, config.hidden_size}), - std::nullopt, - }; + {config.intermediate_size, config.hidden_size}); + if (gate.type != up.type) { + throw std::runtime_error("Higgs TTS packed gate/up weights require matching storage types"); + } + std::vector packed; + packed.reserve(gate.bytes.size() + up.bytes.size()); + packed.insert(packed.end(), gate.bytes.begin(), gate.bytes.end()); + packed.insert(packed.end(), up.bytes.begin(), up.bytes.end()); + weights.mlp.gate_up_proj = modules::LinearWeights{ + store.make_tensor( + core::TensorShape::from_dims({config.intermediate_size * 2, config.hidden_size}), + gate.type, + packed.data(), + packed.size()), + std::nullopt, + }; + } weights.mlp.down_proj = { store.load_tensor( source, @@ -331,6 +357,7 @@ HiggsARWeights load_higgs_ar_weights( {config.audio.num_codebooks * config.audio.vocab_size, config.text.hidden_size}); weights.decoder = load_decoder_weights(*weights.store, source, config.text, weight_storage_type); weights.norm = weights.store->load_f32_tensor(source, "body.norm.weight", {config.text.hidden_size}); + weights.packed_qkv = true; weights.store->upload(); return weights; } @@ -414,11 +441,11 @@ struct HiggsARKVCache::Impl { for (size_t layer = 0; layer < tensor_weights.decoder.layers.size(); ++layer) { key_tensors.push_back(core::make_tensor( build_ctx, - GGML_TYPE_F32, + GGML_TYPE_F16, core::TensorShape::from_dims({1, cache_steps, config.text.num_key_value_heads, dim}))); value_tensors.push_back(core::make_tensor( build_ctx, - GGML_TYPE_F32, + GGML_TYPE_F16, core::TensorShape::from_dims({1, cache_steps, config.text.num_key_value_heads, dim}))); } cache = runtime::TransformerKVCache( @@ -568,7 +595,7 @@ struct HiggsARDecodeGraph::Impl { GGML_TYPE_F16); graph = ggml_new_graph_custom(ctx.get(), 65536, false); - const HiggsQwenDecoderComponent decoder(config.text); + const HiggsQwenDecoderComponent decoder(config.text, tensor_weights.packed_qkv); for (size_t layer_index = 0; layer_index < tensor_weights.decoder.layers.size(); ++layer_index) { auto out = decoder.build_decode_layer( build_ctx, @@ -761,7 +788,7 @@ struct HiggsARPrefillGraph::Impl { start_step(input_start_step), run_steps(input_prompt_steps - input_start_step), prefill_cache_steps(input_prompt_steps), - layerwise(input_prompt_steps >= kLayerwisePrefillMinSteps), + layerwise(input_prompt_steps >= kLayerwisePrefillMinSteps && input_start_step == 0), graph_arena_bytes(graph_arena_bytes) { if (runtime == nullptr) { throw std::runtime_error("Higgs TTS AR prefill graph requires runtime"); @@ -772,8 +799,8 @@ struct HiggsARPrefillGraph::Impl { if (start_step < 0 || start_step >= prompt_steps) { throw std::runtime_error("Higgs TTS AR prefill graph start step is outside the prompt"); } - if (start_step != 0) { - throw std::runtime_error("Higgs TTS AR prefill graph requires full prompt prefill"); + if (start_step > 0 && target_cache == nullptr) { + throw std::runtime_error("Higgs TTS AR suffix prefill requires a target KV cache"); } if (layerwise) { engine::debug::timing_log_scalar("higgs_tts.ar.prefill.graph.build_ms", 0.0); @@ -813,28 +840,48 @@ struct HiggsARPrefillGraph::Impl { graph = ggml_new_graph_custom(ctx.get(), 262144, false); keys.reserve(tensor_weights.decoder.layers.size()); values.reserve(tensor_weights.decoder.layers.size()); - const HiggsQwenDecoderComponent decoder(config.text); + const HiggsQwenDecoderComponent decoder(config.text, tensor_weights.packed_qkv); for (size_t layer_index = 0; layer_index < tensor_weights.decoder.layers.size(); ++layer_index) { + std::optional prefix_key; + std::optional prefix_value; + if (start_step > 0) { + prefix_key = higgs_cache_view( + build_ctx, + target_cache->key_tensor(layer_index), + 0, + start_step, + config.text.num_key_value_heads, + config.text.head_dim); + prefix_value = higgs_cache_view( + build_ctx, + target_cache->value_tensor(layer_index), + 0, + start_step, + config.text.num_key_value_heads, + config.text.head_dim); + } auto out = decoder.build_prefill_layer( build_ctx, x, positions_value, tensor_weights.decoder.layers[layer_index], - attention_mask_value); + attention_mask_value, + prefix_key, + prefix_value); x = out.output; if (target_cache != nullptr) { auto key_dest = higgs_cache_view( build_ctx, target_cache->key_tensor(layer_index), - 0, - prompt_steps, + start_step, + run_steps, config.text.num_key_value_heads, config.text.head_dim); auto value_dest = higgs_cache_view( build_ctx, target_cache->value_tensor(layer_index), - 0, - prompt_steps, + start_step, + run_steps, config.text.num_key_value_heads, config.text.head_dim); ggml_build_forward_expand(graph, ggml_cpy(ctx.get(), out.key.tensor, key_dest.tensor)); @@ -861,7 +908,7 @@ struct HiggsARPrefillGraph::Impl { text_gate_values.assign(static_cast(run_steps), 0.0F); code_gate_values.assign(static_cast(run_steps), 0.0F); positions_values = modules::qwen_position_ids(run_steps, start_step); - attention_mask_values = modules::qwen_causal_prefill_mask_values(1, run_steps); + attention_mask_values = modules::qwen_causal_suffix_mask_values(1, run_steps, start_step); engine::debug::timing_log_scalar( "higgs_tts.ar.prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); @@ -982,7 +1029,7 @@ struct HiggsARPrefillGraph::Impl { attention_mask, core::TensorShape::from_dims({1, 1, steps, steps}), GGML_TYPE_F16); - const HiggsQwenDecoderComponent decoder(config.text); + const HiggsQwenDecoderComponent decoder(config.text, runtime.weights().packed_qkv); auto out = decoder.build_prefill_layer( build_ctx, x, @@ -1157,6 +1204,11 @@ struct HiggsARPrefillGraph::Impl { if (candidate_start_step != start_step) { throw std::runtime_error("Higgs TTS AR prefill graph start step mismatch"); } + if (start_step > 0 && + (target_cache == nullptr || target_cache->valid_steps() < start_step || + target_cache->current_end() != start_step)) { + throw std::runtime_error("Higgs TTS AR suffix prefill requires the retained prefix in KV cache"); + } if (layerwise) { return run_layerwise(input); } @@ -1201,7 +1253,7 @@ struct HiggsARPrefillGraph::Impl { 0, out.output.codebook_logits.size() * sizeof(float)); if (target_cache != nullptr) { - target_cache->advance_after_direct_append(prompt_steps); + target_cache->advance_after_direct_append(run_steps); out.wrote_cache = true; out.kv_state.current_end = prompt_steps; return out; diff --git a/src/models/higgs_tts/generator.cpp b/src/models/higgs_tts/generator.cpp index 51a4dd3b..058fa34d 100644 --- a/src/models/higgs_tts/generator.cpp +++ b/src/models/higgs_tts/generator.cpp @@ -20,7 +20,21 @@ namespace { using Clock = std::chrono::steady_clock; -constexpr int64_t kInitialGeneratedCacheSteps = 512; +constexpr int64_t kInitialGeneratedCacheSteps = 128; +constexpr int64_t kMinimumCacheBucketSteps = 128; + +int64_t bucketed_initial_cache_steps(int64_t prompt_steps, int64_t max_tokens) { + const int64_t maximum = prompt_steps + max_tokens; + const int64_t required = prompt_steps + std::min(max_tokens, kInitialGeneratedCacheSteps); + int64_t bucket = kMinimumCacheBucketSteps; + while (bucket < required && bucket <= maximum / 2) { + bucket *= 2; + } + if (bucket < required) { + bucket = required; + } + return std::min(bucket, maximum); +} void validate_generation_options(const HiggsGenerationOptions & options) { if (options.max_tokens <= 0) { @@ -205,9 +219,21 @@ void HiggsGenerator::prepare(const HiggsGenerationRequest & request) { cache.prefix_tokens.assign(prepared.prompt.token_ids.begin(), prepared.prompt.token_ids.begin() + static_cast(prepared.prefix_steps)); + const bool same_reference = + reference_prefix_cache_.has_value() && + reference_prefix_cache_->reference_text == cache.reference_text && + reference_prefix_cache_->reference_codes == cache.reference_codes && + reference_prefix_cache_->reference_frames == cache.reference_frames && + reference_prefix_cache_->reference_codebooks == cache.reference_codebooks && + reference_prefix_cache_->prefix_steps == cache.prefix_steps && + reference_prefix_cache_->prefix_tokens == cache.prefix_tokens; reference_prefix_cache_ = std::move(cache); + if (!same_reference) { + reference_kv_ready_ = false; + } } else { reference_prefix_cache_.reset(); + reference_kv_ready_ = false; } } @@ -328,20 +354,38 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re engine::debug::trace_log_scalar("higgs_tts.generator.reference_prefix_cache_hit", reference_cache_hit); engine::debug::trace_log_scalar("higgs_tts.generator.reference_prefix_steps", prepared.prefix_steps); const int64_t max_cache_steps = prompt_steps + request.options.max_tokens; - const int64_t initial_cache_steps = - prompt_steps + std::min(request.options.max_tokens, kInitialGeneratedCacheSteps); - if (ar_kv_cache_ == nullptr || !ar_kv_cache_->can_run(*ar_, initial_cache_steps)) { + const int64_t initial_cache_steps = bucketed_initial_cache_steps(prompt_steps, request.options.max_tokens); + const bool cache_rebuild = + ar_kv_cache_ == nullptr || !ar_kv_cache_->can_run(*ar_, initial_cache_steps) || + ar_kv_cache_->cache_steps() != initial_cache_steps; + if (cache_rebuild) { decode_graph_.reset(); ar_kv_cache_ = std::make_unique(ar_, initial_cache_steps); + reference_kv_ready_ = false; + } + const bool reference_kv_cache_hit = + reference_cache_hit && reference_kv_ready_ && + ar_kv_cache_->valid_steps() >= prepared.prefix_steps; + const int64_t prefill_start_step = reference_kv_cache_hit ? prepared.prefix_steps : 0; + if (reference_kv_cache_hit) { + ar_kv_cache_->retain_prefix(prefill_start_step); + } else { + ar_kv_cache_->reset(); } - ar_kv_cache_->reset(); - if (prefill_graph_ == nullptr || !prefill_graph_->matches(*ar_, prompt_steps, 0)) { + engine::debug::trace_log_scalar("higgs_tts.generator.reference_kv_cache_hit", reference_kv_cache_hit); + engine::debug::trace_log_scalar("higgs_tts.generator.prefill_start_step", prefill_start_step); + engine::debug::trace_log_scalar("higgs_tts.generator.prefill_run_steps", prompt_steps - prefill_start_step); + engine::debug::trace_log_scalar("higgs_tts.generator.kv_cache_steps", initial_cache_steps); + engine::debug::trace_log_scalar("higgs_tts.generator.kv_cache_rebuild", cache_rebuild); + if (prefill_graph_ == nullptr || + !prefill_graph_->matches(*ar_, prompt_steps, prefill_start_step)) { prefill_graph_.reset(); prefill_graph_ = std::make_unique( - ar_, prompt_steps, 0, ar_kv_cache_.get(), ar_decode_graph_arena_bytes_); + ar_, prompt_steps, prefill_start_step, ar_kv_cache_.get(), ar_decode_graph_arena_bytes_); } - auto prefill_output = prefill_graph_->run(prepared.ar_input, 0); + auto prefill_output = prefill_graph_->run(prepared.ar_input, prefill_start_step); prefill_graph_.reset(); + reference_kv_ready_ = reference_cache_hit; if (decode_graph_ == nullptr || !decode_graph_->can_run(*ar_, ar_kv_cache_->cache_steps())) { decode_graph_ = std::make_unique( @@ -420,6 +464,7 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re decode_graph_.reset(); ar_kv_cache_ = std::make_unique(ar_, grown_cache_steps); ar_kv_cache_->import_state(kv_state); + engine::debug::trace_log_scalar("higgs_tts.generator.kv_cache_grown_steps", grown_cache_steps); decode_graph_ = std::make_unique( ar_, ar_kv_cache_->cache_steps(), *ar_kv_cache_, ar_decode_graph_arena_bytes_); decode_graph_->begin_decode_run(); diff --git a/tests/higgs_tts/.gitignore b/tests/higgs_tts/.gitignore new file mode 100644 index 00000000..b45f5596 --- /dev/null +++ b/tests/higgs_tts/.gitignore @@ -0,0 +1,2 @@ +results/ +__pycache__/ diff --git a/tests/higgs_tts/README.md b/tests/higgs_tts/README.md new file mode 100644 index 00000000..393a1e63 --- /dev/null +++ b/tests/higgs_tts/README.md @@ -0,0 +1,38 @@ +# Higgs Audio v3 TTS tests + +The focused framework unit test validates that packed QKV/gate-up projections +match the separate projections, suffix causal masks are correct, F16 KV writes +preserve their values, and the decode graph exposes the intended CUDA paths: +grouped FlashAttention, packed SwiGLU, direct KV updates, and the +`ROPE -> VIEW -> SET_ROWS` fusion pattern. + +```powershell +cmake --build build/windows-cuda-release --config Release --target qwen_decoder_packed_projection_test higgs_tts_warm_bench -j 8 +ctest --test-dir build/windows-cuda-release -C Release -R qwen_decoder_packed_projection_test --output-on-failure +``` + +Run the fixed-seed, five-request CUDA benchmark and save every generated WAV: + +```powershell +tests/higgs_tts/run_cuda_performance.ps1 ` + -Model ../models/higgs-audio-v3-tts-4b_Q8/higgs-audio-v3-tts-4b_Q8.gguf ` + -Label candidate +``` + +Compare a candidate run with a prior result directory request by request: + +```powershell +tests/higgs_tts/run_cuda_performance.ps1 ` + -Model ../models/higgs-audio-v3-tts-4b_Q8/higgs-audio-v3-tts-4b_Q8.gguf ` + -Label candidate ` + -Baseline tests/higgs_tts/results/baseline +``` + +The comparison reports frame counts, wall time, RTF, speedup, waveform cosine, +and 80-band log-mel cosine. Result WAVs, logs, and JSON reports are written below +`tests/higgs_tts/results/`, which is intentionally ignored by Git. +The comparison helper requires Python 3 with NumPy. + +Add `-RequireSameFrames` when comparing paths that are expected to be +deterministic and numerically identical. Sampled or mixed-precision paths still +report their frame drift and similarity metrics without hiding the results. diff --git a/tests/higgs_tts/compare_warmbench_results.py b/tests/higgs_tts/compare_warmbench_results.py new file mode 100644 index 00000000..3facf3a5 --- /dev/null +++ b/tests/higgs_tts/compare_warmbench_results.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +"""Compare Higgs warmbench runs request by request. + +Each result directory is expected to contain timing.log and audio/audio_N.wav, +as emitted by higgs_tts_warm_bench. The report includes exact frame counts, +wall time, RTF, speedup, waveform cosine, and log-mel cosine per request. +""" + +from __future__ import annotations + +import argparse +import json +import math +import wave +from pathlib import Path + +import numpy as np + + +def read_wav(path: Path) -> tuple[int, np.ndarray]: + with wave.open(str(path), "rb") as wav: + if wav.getsampwidth() != 2: + raise ValueError(f"{path}: expected PCM16 WAV") + channels = wav.getnchannels() + sample_rate = wav.getframerate() + samples = np.frombuffer(wav.readframes(wav.getnframes()), dtype=" 1: + samples = samples.reshape(-1, channels).mean(axis=1) + return sample_rate, samples / 32768.0 + + +def cosine(a: np.ndarray, b: np.ndarray) -> float: + count = min(a.size, b.size) + if count == 0: + return math.nan + a = a[:count].astype(np.float64, copy=False) + b = b[:count].astype(np.float64, copy=False) + denom = np.linalg.norm(a) * np.linalg.norm(b) + return float(np.dot(a, b) / denom) if denom > 0.0 else math.nan + + +def hz_to_mel(hz: np.ndarray | float) -> np.ndarray | float: + return 2595.0 * np.log10(1.0 + np.asarray(hz) / 700.0) + + +def mel_to_hz(mel: np.ndarray | float) -> np.ndarray | float: + return 700.0 * (np.power(10.0, np.asarray(mel) / 2595.0) - 1.0) + + +def log_mel(samples: np.ndarray, sample_rate: int, n_fft: int = 1024, hop: int = 256, bands: int = 80) -> np.ndarray: + if samples.size < n_fft: + samples = np.pad(samples, (0, n_fft - samples.size)) + frame_count = 1 + (samples.size - n_fft) // hop + shape = (frame_count, n_fft) + strides = (samples.strides[0] * hop, samples.strides[0]) + frames = np.lib.stride_tricks.as_strided(samples, shape=shape, strides=strides) + spectrum = np.abs(np.fft.rfft(frames * np.hanning(n_fft), axis=1)) ** 2 + + mel_points = np.linspace(hz_to_mel(0.0), hz_to_mel(sample_rate / 2.0), bands + 2) + bins = np.floor((n_fft + 1) * mel_to_hz(mel_points) / sample_rate).astype(np.int64) + bins = np.clip(bins, 0, spectrum.shape[1] - 1) + filters = np.zeros((bands, spectrum.shape[1]), dtype=np.float64) + for band in range(bands): + left, center, right = bins[band : band + 3] + if center > left: + filters[band, left:center] = np.arange(center - left) / (center - left) + if right > center: + filters[band, center:right] = np.arange(right - center, 0, -1) / (right - center) + return np.log(np.maximum(spectrum @ filters.T, 1.0e-10)).astype(np.float32) + + +def timings(path: Path) -> dict[int, float]: + result: dict[int, float] = {} + for line in (path / "timing.log").read_text(encoding="utf-8").splitlines(): + prefix = "higgs_tts.cpp.request_" + if not line.startswith(prefix) or ".wall_ms=" not in line: + continue + index_text, value = line[len(prefix) :].split(".wall_ms=", 1) + result[int(index_text)] = float(value) + return result + + +def audio_files(path: Path) -> dict[int, Path]: + result: dict[int, Path] = {} + for wav_path in (path / "audio").glob("audio_*.wav"): + result[int(wav_path.stem.removeprefix("audio_"))] = wav_path + return result + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--baseline", type=Path, required=True) + parser.add_argument("--candidate", type=Path, required=True) + parser.add_argument("--output", type=Path) + parser.add_argument("--require-same-frames", action="store_true") + parser.add_argument("--min-wav-cosine", type=float) + parser.add_argument("--min-logmel-cosine", type=float) + args = parser.parse_args() + + baseline_times = timings(args.baseline) + candidate_times = timings(args.candidate) + baseline_audio = audio_files(args.baseline) + candidate_audio = audio_files(args.candidate) + indices = sorted(set(baseline_times) & set(candidate_times) & set(baseline_audio) & set(candidate_audio)) + if not indices: + raise ValueError("no matching warmbench requests were found") + + failed = False + requests = [] + for index in indices: + baseline_rate, baseline_samples = read_wav(baseline_audio[index]) + candidate_rate, candidate_samples = read_wav(candidate_audio[index]) + if baseline_rate != candidate_rate: + raise ValueError(f"request {index}: sample rates differ") + same_frames = baseline_samples.size == candidate_samples.size + wav_cosine = cosine(baseline_samples, candidate_samples) + baseline_mel = log_mel(baseline_samples, baseline_rate) + candidate_mel = log_mel(candidate_samples, candidate_rate) + mel_frames = min(baseline_mel.shape[0], candidate_mel.shape[0]) + logmel_cosine = cosine(baseline_mel[:mel_frames].reshape(-1), candidate_mel[:mel_frames].reshape(-1)) + baseline_duration_sec = baseline_samples.size / baseline_rate + candidate_duration_sec = candidate_samples.size / candidate_rate + baseline_ms = baseline_times[index] + candidate_ms = candidate_times[index] + baseline_rtf = baseline_ms / 1000.0 / baseline_duration_sec + candidate_rtf = candidate_ms / 1000.0 / candidate_duration_sec + request = { + "request_index": index, + "baseline_frames": int(baseline_samples.size), + "candidate_frames": int(candidate_samples.size), + "same_frames": same_frames, + "baseline_wall_ms": baseline_ms, + "candidate_wall_ms": candidate_ms, + "wall_speedup": baseline_ms / candidate_ms, + "baseline_rtf": baseline_rtf, + "candidate_rtf": candidate_rtf, + "rtf_speedup": baseline_rtf / candidate_rtf, + "wav_cosine": wav_cosine, + "logmel_cosine": logmel_cosine, + } + requests.append(request) + failed = failed or (args.require_same_frames and not same_frames) + failed = failed or (args.min_wav_cosine is not None and wav_cosine < args.min_wav_cosine) + failed = failed or (args.min_logmel_cosine is not None and logmel_cosine < args.min_logmel_cosine) + print( + f"request={index} frames={baseline_samples.size}/{candidate_samples.size} " + f"wall_ms={baseline_ms:.3f}/{candidate_ms:.3f} " + f"rtf={baseline_rtf:.4f}/{candidate_rtf:.4f} speedup={baseline_rtf / candidate_rtf:.3f}x " + f"wav_cos={wav_cosine:.8f} mel_cos={logmel_cosine:.8f}" + ) + + payload = {"baseline": str(args.baseline), "candidate": str(args.candidate), "requests": requests} + if args.output: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8") + return 1 if failed else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/higgs_tts/higgs_tts_cuda_bench_cases.json b/tests/higgs_tts/higgs_tts_cuda_bench_cases.json new file mode 100644 index 00000000..cc4ff314 --- /dev/null +++ b/tests/higgs_tts/higgs_tts_cuda_bench_cases.json @@ -0,0 +1,35 @@ +[ + { + "id": "clone_prefix_first", + "text": "Hello. This is the first retained prefix test.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 1234 + }, + { + "id": "clone_prefix_second", + "text": "The second request should reuse the cloned voice prefix.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 2234 + }, + { + "id": "clone_prefix_third", + "text": "A third short sentence checks stable repeated generation.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 3234 + } +] diff --git a/tests/higgs_tts/higgs_tts_cuda_mixed_cases.json b/tests/higgs_tts/higgs_tts_cuda_mixed_cases.json new file mode 100644 index 00000000..345b8f3c --- /dev/null +++ b/tests/higgs_tts/higgs_tts_cuda_mixed_cases.json @@ -0,0 +1,44 @@ +[ + { + "id": "clone_first", + "text": "This cloned request prepares the reusable reference prefix.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 6134 + }, + { + "id": "clone_reuse", + "text": "This cloned request reuses the same reference prefix.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 6234 + }, + { + "id": "unconditioned", + "text": "This request intentionally generates an unconditioned random voice.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 6334 + }, + { + "id": "clone_after_unconditioned", + "text": "The cloned voice is rebuilt safely after the unconditioned request.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 256, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 6434 + } +] diff --git a/tests/higgs_tts/higgs_tts_cuda_perf_cases.json b/tests/higgs_tts/higgs_tts_cuda_perf_cases.json new file mode 100644 index 00000000..83403b80 --- /dev/null +++ b/tests/higgs_tts/higgs_tts_cuda_perf_cases.json @@ -0,0 +1,57 @@ +[ + { + "id": "control_room_update", + "text": "The control room confirmed the overnight checks and the field team can restart the survey.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 1234 + }, + { + "id": "lab_briefing", + "text": "The lab briefing is ready. Please confirm the calibration notes before the afternoon check-in.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 2234 + }, + { + "id": "dispatch_note", + "text": "Dispatch logged the revised route. The north access road is clear and the receiver test can begin after lunch.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 3234 + }, + { + "id": "short_status", + "text": "All systems are ready for the next test.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 4234 + }, + { + "id": "weather_report", + "text": "Light rain is expected this evening, but tomorrow morning should remain calm and clear.", + "reference_audio": "../SAMPLES/EN_2.wav", + "reference_text": "If you actually care about security.", + "max_tokens": 512, + "temperature": 1.0, + "top_p": 0.95, + "top_k": 50, + "seed": 5234 + } +] diff --git a/tests/higgs_tts/higgs_tts_warm_bench.cpp b/tests/higgs_tts/higgs_tts_warm_bench.cpp index 665e3881..73c400b9 100644 --- a/tests/higgs_tts/higgs_tts_warm_bench.cpp +++ b/tests/higgs_tts/higgs_tts_warm_bench.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -106,13 +107,14 @@ void set_optional_option( } } -engine::runtime::AudioBuffer read_reference_audio(const engine::io::json::Value & object) { +std::optional read_reference_audio( + const engine::io::json::Value & object) { auto reference_path = optional_string(object, "reference_audio"); if (reference_path.empty()) { reference_path = optional_string(object, "voice_ref"); } if (reference_path.empty()) { - throw std::runtime_error("Higgs TTS warmbench request missing field: reference_audio"); + return std::nullopt; } const auto wav = engine::audio::read_wav_f32(resolve_path(reference_path)); return engine::runtime::AudioBuffer{wav.sample_rate, wav.channels, wav.samples}; @@ -121,10 +123,12 @@ engine::runtime::AudioBuffer read_reference_audio(const engine::io::json::Value engine::runtime::TaskRequest make_request(const engine::io::json::Value & object) { engine::runtime::TaskRequest request; request.text_input = engine::runtime::Transcript{required_string(object, "text"), ""}; - request.voice = engine::runtime::VoiceCondition{}; - request.voice->speaker = engine::runtime::VoiceReference{}; - request.voice->speaker->audio = read_reference_audio(object); - request.options["reference_text"] = required_string(object, "reference_text"); + if (auto reference_audio = read_reference_audio(object); reference_audio.has_value()) { + request.voice = engine::runtime::VoiceCondition{}; + request.voice->speaker = engine::runtime::VoiceReference{}; + request.voice->speaker->audio = std::move(*reference_audio); + request.options["reference_text"] = required_string(object, "reference_text"); + } set_optional_option(request, object, "max_tokens", "max_tokens"); set_optional_option(request, object, "temperature", "temperature"); set_optional_option(request, object, "top_p", "top_p"); @@ -205,10 +209,18 @@ engine::io::json::Value step_json( if (!audio_path.empty()) { stem.emplace("audio", string(audio_path.string())); } + const auto & audio = *result.audio_output; + const double frames = static_cast( + audio.samples.size() / static_cast(std::max(1, audio.channels))); + const double duration_sec = audio.sample_rate > 0 ? frames / audio.sample_rate : 0.0; + const double rtf = duration_sec > 0.0 ? wall_ms / 1000.0 / duration_sec : 0.0; return engine::io::json::Value::make_object({ {"request_index", number(static_cast(request_index))}, {"stems", engine::io::json::Value::make_array({engine::io::json::Value::make_object(std::move(stem))})}, - {"metrics", engine::io::json::Value::make_object({{"wall_ms", number(wall_ms)}})}, + {"metrics", engine::io::json::Value::make_object({ + {"wall_ms", number(wall_ms)}, + {"rtf", number(rtf)}, + })}, }); } @@ -238,7 +250,19 @@ int main(int argc, char ** argv) { const int threads = int_arg(argc, argv, "--threads", 8); const int warmup = int_arg(argc, argv, "--warmup", 0); const int iterations = int_arg(argc, argv, "--iterations", 1); - const std::string request_sequence_json = arg_value(argc, argv, "--request-sequence-json", ""); + std::string request_sequence_json = arg_value(argc, argv, "--request-sequence-json", ""); + const std::filesystem::path request_sequence_file = + arg_value(argc, argv, "--request-sequence-file", ""); + if (request_sequence_json.empty() && !request_sequence_file.empty()) { + std::ifstream input(request_sequence_file, std::ios::binary); + if (!input) { + throw std::runtime_error( + "failed to open Higgs TTS request sequence: " + request_sequence_file.string()); + } + request_sequence_json.assign( + std::istreambuf_iterator(input), + std::istreambuf_iterator()); + } const std::filesystem::path output_dir = arg_value(argc, argv, "--output-dir", ""); const std::filesystem::path timing_path = arg_value(argc, argv, "--timing-file", "/tmp/higgs_tts_warm_bench_timing.log"); @@ -308,7 +332,16 @@ int main(int argc, char ** argv) { } timing_lines.push_back( "higgs_tts.cpp.request_" + std::to_string(request_index) + ".wall_ms=" + std::to_string(wall_ms)); - std::cout << "higgs_tts.cpp.wall_ms=" << wall_ms << "\n"; + const auto & audio = *last_result.audio_output; + const double frames = static_cast( + audio.samples.size() / static_cast(std::max(1, audio.channels))); + const double duration_sec = audio.sample_rate > 0 ? frames / audio.sample_rate : 0.0; + const double rtf = duration_sec > 0.0 ? wall_ms / 1000.0 / duration_sec : 0.0; + timing_lines.push_back( + "higgs_tts.cpp.request_" + std::to_string(request_index) + ".rtf=" + std::to_string(rtf)); + std::cout << "higgs_tts.cpp.request=" << request_index + << " wall_ms=" << wall_ms + << " rtf=" << rtf << "\n"; steps.push_back(step_json(last_result, static_cast(request_index), wall_ms, audio_path)); } diff --git a/tests/higgs_tts/run_cuda_performance.ps1 b/tests/higgs_tts/run_cuda_performance.ps1 new file mode 100644 index 00000000..456a3549 --- /dev/null +++ b/tests/higgs_tts/run_cuda_performance.ps1 @@ -0,0 +1,69 @@ +param( + [Parameter(Mandatory = $true)] + [string]$Model, + [string]$BuildDir = "build/windows-cuda-release", + [string]$Label = (Get-Date -Format "yyyyMMdd-HHmmss"), + [string]$Baseline = "", + [switch]$RequireSameFrames, + [int]$Device = 0, + [int]$Threads = 8, + [int]$Warmup = 1, + [int]$Iterations = 1 +) + +$ErrorActionPreference = "Stop" +$RepoRoot = (Resolve-Path (Join-Path $PSScriptRoot "../..")).Path +$ModelPath = (Resolve-Path $Model).Path +$Bench = Join-Path $RepoRoot "$BuildDir/bin/higgs_tts_warm_bench.exe" +$Cases = Join-Path $PSScriptRoot "higgs_tts_cuda_perf_cases.json" +$ResultDir = Join-Path $PSScriptRoot "results/$Label" + +if (-not (Test-Path -LiteralPath $Bench)) { + throw "Warmbench binary does not exist: $Bench" +} + +New-Item -ItemType Directory -Force (Join-Path $ResultDir "audio") | Out-Null +$PreviousErrorActionPreference = $ErrorActionPreference +$ErrorActionPreference = "Continue" +& $Bench ` + --model $ModelPath ` + --backend cuda ` + --device $Device ` + --threads $Threads ` + --warmup $Warmup ` + --iterations $Iterations ` + --request-sequence-file $Cases ` + --output-dir (Join-Path $ResultDir "audio") ` + --timing-file (Join-Path $ResultDir "timing.log") 2>&1 | + ForEach-Object { + if ($_ -is [System.Management.Automation.ErrorRecord]) { + $_.Exception.Message + } else { + $_ + } + } | + Tee-Object -FilePath (Join-Path $ResultDir "console.log") +$BenchExitCode = $LASTEXITCODE +$ErrorActionPreference = $PreviousErrorActionPreference +if ($BenchExitCode -ne 0) { + exit $BenchExitCode +} + +if ($Baseline) { + $BaselinePath = (Resolve-Path $Baseline).Path + $CompareArgs = @( + (Join-Path $PSScriptRoot "compare_warmbench_results.py"), + "--baseline", $BaselinePath, + "--candidate", $ResultDir, + "--output", (Join-Path $ResultDir "comparison.json") + ) + if ($RequireSameFrames) { + $CompareArgs += "--require-same-frames" + } + python @CompareArgs + if ($LASTEXITCODE -ne 0) { + exit $LASTEXITCODE + } +} + +Write-Host "Higgs CUDA performance artifacts: $ResultDir" diff --git a/tests/unittests/test_qwen_decoder_packed_projections.cpp b/tests/unittests/test_qwen_decoder_packed_projections.cpp new file mode 100644 index 00000000..c34c4535 --- /dev/null +++ b/tests/unittests/test_qwen_decoder_packed_projections.cpp @@ -0,0 +1,504 @@ +#include "engine/framework/core/backend.h" +#include "engine/framework/modules/attention/qwen_causal_decoder.h" +#include "engine/framework/modules/attention/qwen_decoder.h" +#include "engine/framework/modules/optimizations/fast_kv_modules.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr size_t kGraphBytes = 16 * 1024 * 1024; +constexpr size_t kGraphNodes = 4096; + +std::vector patterned(size_t count, float phase, float scale) { + std::vector values(count); + for (size_t i = 0; i < count; ++i) { + const float x = static_cast(i); + values[i] = scale * (std::sin(phase + 0.19f * x) + 0.35f * std::cos(phase + 0.07f * x)); + } + return values; +} + +void require_allclose( + const std::vector & actual, + const std::vector & expected, + float tolerance, + const std::string & label) { + if (actual.size() != expected.size()) { + throw std::runtime_error(label + " size mismatch"); + } + for (size_t i = 0; i < actual.size(); ++i) { + const float diff = std::fabs(actual[i] - expected[i]); + if (diff > tolerance) { + std::ostringstream message; + message << label << " mismatch at " << i << ": expected " << expected[i] + << ", got " << actual[i] << ", diff=" << diff; + throw std::runtime_error(message.str()); + } + } +} + +struct LayerResult { + std::vector output; + std::vector key; + std::vector value; +}; + +LayerResult run_layer(bool packed) { + constexpr int64_t batch = 1; + constexpr int64_t steps = 3; + constexpr int64_t hidden = 8; + constexpr int64_t heads = 2; + constexpr int64_t kv_heads = 1; + constexpr int64_t head_dim = 4; + constexpr int64_t intermediate = 12; + constexpr int64_t q_out = heads * head_dim; + constexpr int64_t kv_out = kv_heads * head_dim; + + engine::core::BackendConfig backend_config{engine::core::BackendType::Cpu, 0, 4}; + ggml_backend_t backend = engine::core::init_backend(backend_config); + if (backend == nullptr) { + throw std::runtime_error("failed to initialize CPU backend"); + } + + ggml_init_params params{kGraphBytes, nullptr, true}; + ggml_context * ggml = ggml_init(params); + if (ggml == nullptr) { + ggml_backend_free(backend); + throw std::runtime_error("failed to initialize GGML context"); + } + + ggml_backend_buffer_t buffer = nullptr; + try { + engine::core::ModuleBuildContext ctx{ggml, "qwen_packed_projection_test", engine::core::BackendType::Cpu}; + auto make_f32 = [&](std::initializer_list dims) { + return engine::core::make_tensor(ctx, GGML_TYPE_F32, engine::core::TensorShape::from_dims(dims)); + }; + + auto input = make_f32({batch, steps, hidden}); + auto positions = engine::core::make_tensor( + ctx, + GGML_TYPE_I32, + engine::core::TensorShape::from_dims({steps})); + + engine::modules::QwenDecoderLayerWeights weights; + weights.input_norm = {make_f32({hidden}), std::nullopt}; + weights.post_norm = {make_f32({hidden}), std::nullopt}; + weights.self_attention.out_weight = make_f32({hidden, hidden}); + weights.mlp.down_proj = {make_f32({hidden, intermediate}), std::nullopt}; + + const auto q_values = patterned(static_cast(q_out * hidden), 0.1f, 0.12f); + const auto k_values = patterned(static_cast(kv_out * hidden), 0.5f, 0.10f); + const auto v_values = patterned(static_cast(kv_out * hidden), 0.9f, 0.08f); + const auto gate_values = patterned(static_cast(intermediate * hidden), 1.3f, 0.11f); + const auto up_values = patterned(static_cast(intermediate * hidden), 1.7f, 0.09f); + + if (packed) { + weights.self_attention.qkv_weight = make_f32({q_out + 2 * kv_out, hidden}); + weights.mlp.gate_up_proj = engine::modules::LinearWeights{ + make_f32({intermediate * 2, hidden}), + std::nullopt, + }; + } else { + weights.self_attention.q_weight = make_f32({q_out, hidden}); + weights.self_attention.k_weight = make_f32({kv_out, hidden}); + weights.self_attention.v_weight = make_f32({kv_out, hidden}); + weights.mlp.gate_proj = {make_f32({intermediate, hidden}), std::nullopt}; + weights.mlp.up_proj = {make_f32({intermediate, hidden}), std::nullopt}; + } + + engine::modules::QwenDecoderLayerConfig config; + config.hidden_size = hidden; + config.num_attention_heads = heads; + config.num_key_value_heads = kv_heads; + config.head_dim = head_dim; + config.intermediate_size = intermediate; + config.rms_norm_eps = 1e-5f; + config.qkv_layout = packed + ? engine::modules::QwenDecoderQKVLayout::PackedQKV + : engine::modules::QwenDecoderQKVLayout::Separate; + config.use_qk_norm = false; + config.runtime.attention.prefill_mode = engine::modules::QwenDecoderAttentionMode::ManualRepeat; + + const auto outputs = engine::modules::QwenDecoderLayerModule(config).build( + ctx, + input, + positions, + weights); + + ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); + ggml_build_forward_expand(graph, outputs.output.tensor); + buffer = ggml_backend_alloc_ctx_tensors(ggml, backend); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate test tensors"); + } + + engine::core::write_tensor_f32(input, patterned(static_cast(batch * steps * hidden), 2.1f, 0.20f)); + engine::core::write_tensor_i32(positions, {0, 1, 2}); + engine::core::write_tensor_f32(*weights.input_norm.weight, patterned(hidden, 0.3f, 0.7f)); + engine::core::write_tensor_f32(*weights.post_norm.weight, patterned(hidden, 0.7f, 0.8f)); + engine::core::write_tensor_f32( + weights.self_attention.out_weight, + patterned(static_cast(hidden * hidden), 1.1f, 0.10f)); + engine::core::write_tensor_f32( + weights.mlp.down_proj.weight, + patterned(static_cast(hidden * intermediate), 1.9f, 0.10f)); + + if (packed) { + std::vector qkv_values; + qkv_values.reserve(q_values.size() + k_values.size() + v_values.size()); + qkv_values.insert(qkv_values.end(), q_values.begin(), q_values.end()); + qkv_values.insert(qkv_values.end(), k_values.begin(), k_values.end()); + qkv_values.insert(qkv_values.end(), v_values.begin(), v_values.end()); + engine::core::write_tensor_f32(*weights.self_attention.qkv_weight, qkv_values); + + std::vector gate_up_values; + gate_up_values.reserve(gate_values.size() + up_values.size()); + gate_up_values.insert(gate_up_values.end(), gate_values.begin(), gate_values.end()); + gate_up_values.insert(gate_up_values.end(), up_values.begin(), up_values.end()); + engine::core::write_tensor_f32(weights.mlp.gate_up_proj->weight, gate_up_values); + } else { + engine::core::write_tensor_f32(weights.self_attention.q_weight, q_values); + engine::core::write_tensor_f32(weights.self_attention.k_weight, k_values); + engine::core::write_tensor_f32(weights.self_attention.v_weight, v_values); + engine::core::write_tensor_f32(weights.mlp.gate_proj.weight, gate_values); + engine::core::write_tensor_f32(weights.mlp.up_proj.weight, up_values); + } + + ggml_backend_graph_compute(backend, graph); + LayerResult result; + engine::core::read_tensor_f32_into(outputs.output.tensor, result.output); + engine::core::read_tensor_f32_into(outputs.key.tensor, result.key); + engine::core::read_tensor_f32_into(outputs.value.tensor, result.value); + + ggml_backend_buffer_free(buffer); + buffer = nullptr; + ggml_free(ggml); + ggml_backend_free(backend); + return result; + } catch (...) { + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + ggml_free(ggml); + ggml_backend_free(backend); + throw; + } +} + +void test_packed_qkv_and_gate_up_match_separate_projections() { + const auto separate = run_layer(false); + const auto packed = run_layer(true); + require_allclose(packed.output, separate.output, 2.0e-5f, "decoder output"); + require_allclose(packed.key, separate.key, 2.0e-5f, "decoder key"); + require_allclose(packed.value, separate.value, 2.0e-5f, "decoder value"); +} + +void test_suffix_causal_mask() { + const auto values = engine::modules::qwen_causal_suffix_mask_values(2, 3, 2); + if (values.size() != 30) { + throw std::runtime_error("suffix causal mask size mismatch"); + } + const std::vector expected{ + true, true, true, false, false, + true, true, true, true, false, + true, true, true, true, true, + }; + for (int batch = 0; batch < 2; ++batch) { + for (size_t i = 0; i < expected.size(); ++i) { + const float actual = ggml_fp16_to_fp32(values[static_cast(batch) * expected.size() + i]); + if ((expected[i] && actual != 0.0F) || (!expected[i] && !std::isinf(actual))) { + throw std::runtime_error("suffix causal mask visibility mismatch"); + } + } + } +} + +void test_f16_kv_set_rows() { + engine::core::BackendConfig backend_config{engine::core::BackendType::Cpu, 0, 4}; + ggml_backend_t backend = engine::core::init_backend(backend_config); + if (backend == nullptr) { + throw std::runtime_error("failed to initialize CPU backend"); + } + + ggml_init_params params{kGraphBytes, nullptr, true}; + ggml_context * ggml = ggml_init(params); + if (ggml == nullptr) { + ggml_backend_free(backend); + throw std::runtime_error("failed to initialize GGML context"); + } + + ggml_backend_buffer_t buffer = nullptr; + try { + engine::core::ModuleBuildContext ctx{ggml, "f16_kv_set_rows_test", engine::core::BackendType::Cpu}; + const auto cache = engine::core::make_tensor( + ctx, + GGML_TYPE_F16, + engine::core::TensorShape::from_dims({1, 3, 1, 2})); + const auto row = engine::core::make_tensor( + ctx, + GGML_TYPE_F32, + engine::core::TensorShape::from_dims({1, 1, 1, 2})); + const auto row_index = engine::core::make_tensor( + ctx, + GGML_TYPE_I64, + engine::core::TensorShape::from_dims({1})); + const auto output = engine::modules::FastKVSetRowsModule{}.build(ctx, cache, row, row_index); + + ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); + ggml_build_forward_expand(graph, output.tensor); + buffer = ggml_backend_alloc_ctx_tensors(ggml, backend); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate f16 KV test tensors"); + } + + engine::core::write_tensor_f16(cache, std::vector(6, 0.0F)); + engine::core::write_tensor_f32(row, {1.25F, -2.5F}); + const int64_t index = 1; + ggml_backend_tensor_set(row_index.tensor, &index, 0, sizeof(index)); + if (ggml_backend_graph_compute(backend, graph) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("f16 KV set-rows graph compute failed"); + } + + const auto values = engine::core::read_tensor_f16(output.tensor); + require_allclose(values, {0.0F, 0.0F, 1.25F, -2.5F, 0.0F, 0.0F}, 1.0e-3F, "f16 KV cache"); + + ggml_backend_buffer_free(buffer); + buffer = nullptr; + ggml_free(ggml); + ggml_backend_free(backend); + } catch (...) { + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + ggml_free(ggml); + ggml_backend_free(backend); + throw; + } +} + +void test_f16_kv_set_rows_batched() { + engine::core::BackendConfig backend_config{engine::core::BackendType::Cpu, 0, 4}; + ggml_backend_t backend = engine::core::init_backend(backend_config); + if (backend == nullptr) { + throw std::runtime_error("failed to initialize CPU backend"); + } + + ggml_init_params params{kGraphBytes, nullptr, true}; + ggml_context * ggml = ggml_init(params); + if (ggml == nullptr) { + ggml_backend_free(backend); + throw std::runtime_error("failed to initialize GGML context"); + } + + ggml_backend_buffer_t buffer = nullptr; + try { + engine::core::ModuleBuildContext ctx{ggml, "f16_kv_set_rows_batched_test", engine::core::BackendType::Cpu}; + const auto cache = engine::core::make_tensor( + ctx, + GGML_TYPE_F16, + engine::core::TensorShape::from_dims({2, 3, 1, 2})); + const auto rows = engine::core::make_tensor( + ctx, + GGML_TYPE_F32, + engine::core::TensorShape::from_dims({2, 1, 1, 2})); + const auto row_indices = engine::core::make_tensor( + ctx, + GGML_TYPE_I64, + engine::core::TensorShape::from_dims({2})); + const auto output = engine::modules::FastKVSetRowsModule{}.build(ctx, cache, rows, row_indices); + + ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); + ggml_build_forward_expand(graph, output.tensor); + buffer = ggml_backend_alloc_ctx_tensors(ggml, backend); + if (buffer == nullptr) { + throw std::runtime_error("failed to allocate batched F16 KV test tensors"); + } + + engine::core::write_tensor_f16(cache, std::vector(12, 0.0F)); + engine::core::write_tensor_f32(rows, {1.25F, -2.5F, 3.5F, -4.5F}); + const std::vector indices{1, 4}; + ggml_backend_tensor_set(row_indices.tensor, indices.data(), 0, indices.size() * sizeof(int64_t)); + if (ggml_backend_graph_compute(backend, graph) != GGML_STATUS_SUCCESS) { + throw std::runtime_error("batched F16 KV set-rows graph compute failed"); + } + + const auto values = engine::core::read_tensor_f16(output.tensor); + require_allclose( + values, + {0.0F, 0.0F, 1.25F, -2.5F, 0.0F, 0.0F, + 0.0F, 0.0F, 3.5F, -4.5F, 0.0F, 0.0F}, + 1.0e-3F, + "batched f16 KV cache"); + + ggml_backend_buffer_free(buffer); + buffer = nullptr; + ggml_free(ggml); + ggml_backend_free(backend); + } catch (...) { + if (buffer != nullptr) { + ggml_backend_buffer_free(buffer); + } + ggml_free(ggml); + ggml_backend_free(backend); + throw; + } +} + +int count_graph_op(ggml_cgraph * graph, ggml_op op) { + int count = 0; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + const ggml_tensor * node = ggml_graph_node(graph, i); + count += node != nullptr && node->op == op ? 1 : 0; + } + return count; +} + +bool graph_contains_sequence(ggml_cgraph * graph, std::initializer_list ops) { + if (ops.size() == 0 || static_cast(ops.size()) > ggml_graph_n_nodes(graph)) { + return false; + } + for (int start = 0; start + static_cast(ops.size()) <= ggml_graph_n_nodes(graph); ++start) { + bool matches = true; + int offset = 0; + for (const ggml_op op : ops) { + const ggml_tensor * node = ggml_graph_node(graph, start + offset++); + matches = matches && node != nullptr && node->op == op; + } + if (matches) { + return true; + } + } + return false; +} + +std::string graph_ops(ggml_cgraph * graph) { + std::ostringstream out; + for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) { + const ggml_tensor * node = ggml_graph_node(graph, i); + if (i != 0) { + out << ','; + } + out << (node != nullptr ? ggml_op_name(node->op) : "null"); + } + return out.str(); +} + +void test_higgs_decode_graph_exposes_cuda_fast_paths() { + constexpr int64_t hidden = 8; + constexpr int64_t heads = 2; + constexpr int64_t kv_heads = 1; + constexpr int64_t head_dim = 4; + constexpr int64_t intermediate = 12; + constexpr int64_t cache_steps = 8; + constexpr int64_t qkv_out = heads * head_dim + 2 * kv_heads * head_dim; + + ggml_init_params params{kGraphBytes, nullptr, true}; + ggml_context * ggml = ggml_init(params); + if (ggml == nullptr) { + throw std::runtime_error("failed to initialize Higgs decode graph test context"); + } + + try { + engine::core::ModuleBuildContext ctx{ggml, "higgs_decode_fast_path_test", engine::core::BackendType::Cuda}; + auto make_tensor = [&](ggml_type type, std::initializer_list dims) { + return engine::core::make_tensor(ctx, type, engine::core::TensorShape::from_dims(dims)); + }; + + const auto input = make_tensor(GGML_TYPE_F32, {1, 1, hidden}); + const auto positions = make_tensor(GGML_TYPE_I32, {1}); + const auto cache_key = make_tensor(GGML_TYPE_F16, {1, cache_steps, kv_heads, head_dim}); + const auto cache_value = make_tensor(GGML_TYPE_F16, {1, cache_steps, kv_heads, head_dim}); + const auto cache_slot = make_tensor(GGML_TYPE_I64, {1}); + const auto attention_mask = make_tensor(GGML_TYPE_F16, {1, 1, 1, cache_steps}); + + engine::modules::QwenDecoderLayerWeights weights; + weights.input_norm = {make_tensor(GGML_TYPE_F32, {hidden}), std::nullopt}; + weights.q_norm = {make_tensor(GGML_TYPE_F32, {head_dim}), std::nullopt}; + weights.k_norm = {make_tensor(GGML_TYPE_F32, {head_dim}), std::nullopt}; + weights.post_norm = {make_tensor(GGML_TYPE_F32, {hidden}), std::nullopt}; + weights.self_attention.qkv_weight = make_tensor(GGML_TYPE_F32, {qkv_out, hidden}); + weights.self_attention.out_weight = make_tensor(GGML_TYPE_F32, {hidden, hidden}); + weights.mlp.gate_up_proj = engine::modules::LinearWeights{ + make_tensor(GGML_TYPE_F32, {intermediate * 2, hidden}), + std::nullopt, + }; + weights.mlp.down_proj = { + make_tensor(GGML_TYPE_F32, {hidden, intermediate}), + std::nullopt, + }; + + engine::modules::QwenDecoderLayerConfig config; + config.hidden_size = hidden; + config.num_attention_heads = heads; + config.num_key_value_heads = kv_heads; + config.head_dim = head_dim; + config.intermediate_size = intermediate; + config.qkv_layout = engine::modules::QwenDecoderQKVLayout::PackedQKV; + config.use_qk_norm = true; + config.runtime.attention.static_mode = + engine::modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + config.runtime.static_cache.update_mode = + engine::modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + + ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); + const auto outputs = engine::modules::QwenDecoderLayerModule(config).build_with_static_cache_tail( + ctx, + graph, + input, + positions, + weights, + cache_key, + cache_value, + cache_slot, + attention_mask); + ggml_build_forward_expand(graph, outputs.output.tensor); + + if (count_graph_op(graph, GGML_OP_FLASH_ATTN_EXT) != 1) { + throw std::runtime_error("Higgs decode graph must contain one grouped FlashAttention op"); + } + if (count_graph_op(graph, GGML_OP_SET_ROWS) != 2) { + throw std::runtime_error("Higgs decode graph must update both F16 KV caches with set-rows"); + } + if (count_graph_op(graph, GGML_OP_GLU) != 1) { + throw std::runtime_error("Higgs decode graph must contain one packed SwiGLU op"); + } + if (count_graph_op(graph, GGML_OP_REPEAT) != 0) { + throw std::runtime_error("grouped FlashAttention must not materialize repeated KV heads"); + } + if (!graph_contains_sequence(graph, {GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS})) { + throw std::runtime_error( + "Higgs key-cache update must expose CUDA RoPE/view/set-rows fusion; graph=" + + graph_ops(graph)); + } + + ggml_free(ggml); + } catch (...) { + ggml_free(ggml); + throw; + } +} + +} // namespace + +int main() { + try { + test_packed_qkv_and_gate_up_match_separate_projections(); + test_suffix_causal_mask(); + test_f16_kv_set_rows(); + test_f16_kv_set_rows_batched(); + test_higgs_decode_graph_exposes_cuda_fast_paths(); + std::cout << "qwen_decoder_packed_projection_test: ok\n"; + return 0; + } catch (const std::exception & ex) { + std::cerr << "qwen_decoder_packed_projection_test: failed: " << ex.what() << "\n"; + return 1; + } +} From 2a7fbd5a8b1bc46657d19fd2a2efb6b706550ddd Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Mon, 20 Jul 2026 14:21:56 -0400 Subject: [PATCH 06/27] Add Fish Audio model integration --- CMakeLists.txt | 8 + include/engine/models/fish_audio/ar.h | 31 + include/engine/models/fish_audio/assets.h | 24 + include/engine/models/fish_audio/codec.h | 33 + include/engine/models/fish_audio/generator.h | 40 + include/engine/models/fish_audio/loader.h | 33 + .../engine/models/fish_audio/prompt_builder.h | 19 + include/engine/models/fish_audio/session.h | 58 + .../engine/models/fish_audio/tokenizer_text.h | 28 + include/engine/models/fish_audio/types.h | 100 ++ model_specs/fish_audio.json | 42 + src/framework/runtime/registry.cpp | 2 + src/models/fish_audio/ar.cpp | 1371 +++++++++++++++++ src/models/fish_audio/assets.cpp | 127 ++ src/models/fish_audio/codec.cpp | 1130 ++++++++++++++ src/models/fish_audio/generator.cpp | 69 + src/models/fish_audio/loader.cpp | 135 ++ src/models/fish_audio/prompt_builder.cpp | 101 ++ src/models/fish_audio/session.cpp | 432 ++++++ src/models/fish_audio/tokenizer_text.cpp | 69 + src/models/higgs_tts/ar.cpp | 1 - .../audiocpp_cli/audiocpp_cli_path_cases.json | 47 + 22 files changed, 3899 insertions(+), 1 deletion(-) create mode 100644 include/engine/models/fish_audio/ar.h create mode 100644 include/engine/models/fish_audio/assets.h create mode 100644 include/engine/models/fish_audio/codec.h create mode 100644 include/engine/models/fish_audio/generator.h create mode 100644 include/engine/models/fish_audio/loader.h create mode 100644 include/engine/models/fish_audio/prompt_builder.h create mode 100644 include/engine/models/fish_audio/session.h create mode 100644 include/engine/models/fish_audio/tokenizer_text.h create mode 100644 include/engine/models/fish_audio/types.h create mode 100644 model_specs/fish_audio.json create mode 100644 src/models/fish_audio/ar.cpp create mode 100644 src/models/fish_audio/assets.cpp create mode 100644 src/models/fish_audio/codec.cpp create mode 100644 src/models/fish_audio/generator.cpp create mode 100644 src/models/fish_audio/loader.cpp create mode 100644 src/models/fish_audio/prompt_builder.cpp create mode 100644 src/models/fish_audio/session.cpp create mode 100644 src/models/fish_audio/tokenizer_text.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index ca237bf3..fd235a43 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -348,6 +348,14 @@ add_library(engine_runtime STATIC src/models/voxtral_realtime/text_decoder.cpp src/models/voxtral_realtime/session.cpp src/models/voxtral_realtime/loader.cpp + src/models/fish_audio/ar.cpp + src/models/fish_audio/assets.cpp + src/models/fish_audio/codec.cpp + src/models/fish_audio/generator.cpp + src/models/fish_audio/loader.cpp + src/models/fish_audio/prompt_builder.cpp + src/models/fish_audio/session.cpp + src/models/fish_audio/tokenizer_text.cpp src/models/heartmula/assets.cpp src/models/heartmula/codec.cpp src/models/heartmula/generator.cpp diff --git a/include/engine/models/fish_audio/ar.h b/include/engine/models/fish_audio/ar.h new file mode 100644 index 00000000..db0e0a0c --- /dev/null +++ b/include/engine/models/fish_audio/ar.h @@ -0,0 +1,31 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/models/fish_audio/assets.h" +#include "engine/models/fish_audio/types.h" + +#include + +namespace engine::models::fish_audio { + +class FishAudioARRuntime { +public: + FishAudioARRuntime( + std::shared_ptr assets, + core::BackendConfig backend, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type); + ~FishAudioARRuntime(); + + FishAudioCodes generate(const FishAudioPrompt & prompt, const FishAudioGenerationOptions & options); + void release_runtime_graphs(); + +private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/assets.h b/include/engine/models/fish_audio/assets.h new file mode 100644 index 00000000..338ce866 --- /dev/null +++ b/include/engine/models/fish_audio/assets.h @@ -0,0 +1,24 @@ +#pragma once + +#include "engine/framework/assets/resource_bundle.h" +#include "engine/models/fish_audio/types.h" + +#include +#include + +namespace engine::assets { +class TensorSource; +} + +namespace engine::models::fish_audio { + +struct FishAudioAssets { + assets::ResourceBundle resources; + FishAudioConfig config; + std::shared_ptr model_weights; + std::shared_ptr codec_weights; +}; + +std::shared_ptr load_fish_audio_assets(const std::filesystem::path & model_path); + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/codec.h b/include/engine/models/fish_audio/codec.h new file mode 100644 index 00000000..5427fef4 --- /dev/null +++ b/include/engine/models/fish_audio/codec.h @@ -0,0 +1,33 @@ +#pragma once + +#include "engine/framework/assets/tensor_source.h" +#include "engine/framework/core/backend.h" +#include "engine/models/fish_audio/assets.h" +#include "engine/models/fish_audio/types.h" + +#include + +namespace engine::models::fish_audio { + +class FishAudioCodecRuntime { +public: + FishAudioCodecRuntime( + std::shared_ptr assets, + core::BackendConfig backend, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType matmul_weight_storage_type, + assets::TensorStorageType conv_weight_storage_type); + ~FishAudioCodecRuntime(); + + FishAudioCodes encode_reference(const runtime::AudioBuffer & audio); + runtime::AudioBuffer decode(const FishAudioCodes & codes); + void release_runtime_graphs(); + +private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/generator.h b/include/engine/models/fish_audio/generator.h new file mode 100644 index 00000000..374c9aac --- /dev/null +++ b/include/engine/models/fish_audio/generator.h @@ -0,0 +1,40 @@ +#pragma once + +#include "engine/models/fish_audio/ar.h" +#include "engine/models/fish_audio/codec.h" +#include "engine/models/fish_audio/prompt_builder.h" +#include "engine/models/fish_audio/tokenizer_text.h" + +#include +#include + +namespace engine::models::fish_audio { + +struct FishAudioGenerationResult { + runtime::AudioBuffer audio; + FishAudioCodes codes; +}; + +class FishAudioGenerator { +public: + FishAudioGenerator( + std::shared_ptr assets, + std::unique_ptr ar, + std::unique_ptr codec); + ~FishAudioGenerator(); + + FishAudioCodes encode_reference(const runtime::AudioBuffer & audio); + FishAudioGenerationResult generate( + const FishAudioRequest & request, + const std::optional & reference_codes, + bool mem_saver); + +private: + std::shared_ptr assets_; + FishAudioTextTokenizer tokenizer_; + FishAudioPromptBuilder prompt_builder_; + std::unique_ptr ar_; + std::unique_ptr codec_; +}; + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/loader.h b/include/engine/models/fish_audio/loader.h new file mode 100644 index 00000000..65e49f4e --- /dev/null +++ b/include/engine/models/fish_audio/loader.h @@ -0,0 +1,33 @@ +#pragma once + +#include "engine/framework/runtime/model.h" +#include "engine/models/fish_audio/assets.h" + +#include +#include + +namespace engine::models::fish_audio { + +class FishAudioLoadedModel final : public runtime::ILoadedVoiceModel { +public: + FishAudioLoadedModel( + runtime::ModelMetadata metadata, + runtime::CapabilitySet capabilities, + std::shared_ptr assets); + + const runtime::ModelMetadata & metadata() const noexcept override; + const runtime::CapabilitySet & capabilities() const noexcept override; + std::unique_ptr create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const override; + +private: + runtime::ModelMetadata metadata_; + runtime::CapabilitySet capabilities_; + std::shared_ptr assets_; +}; + +std::unique_ptr load_fish_audio_model(const std::filesystem::path & model_path); +std::shared_ptr make_fish_audio_loader(); + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/prompt_builder.h b/include/engine/models/fish_audio/prompt_builder.h new file mode 100644 index 00000000..825e42bf --- /dev/null +++ b/include/engine/models/fish_audio/prompt_builder.h @@ -0,0 +1,19 @@ +#pragma once + +#include "engine/models/fish_audio/tokenizer_text.h" +#include "engine/models/fish_audio/types.h" + +namespace engine::models::fish_audio { + +class FishAudioPromptBuilder { +public: + FishAudioPromptBuilder(std::shared_ptr assets, FishAudioTextTokenizer tokenizer); + + FishAudioPrompt build(const FishAudioRequest & request, const std::optional & reference_codes) const; + +private: + std::shared_ptr assets_; + FishAudioTextTokenizer tokenizer_; +}; + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/session.h b/include/engine/models/fish_audio/session.h new file mode 100644 index 00000000..40e6bb5f --- /dev/null +++ b/include/engine/models/fish_audio/session.h @@ -0,0 +1,58 @@ +#pragma once + +#include "engine/framework/runtime/cache_slots.h" +#include "engine/framework/runtime/session_base.h" +#include "engine/models/fish_audio/assets.h" +#include "engine/models/fish_audio/generator.h" + +#include +#include +#include +#include +#include + +namespace engine::models::fish_audio { + +class FishAudioSession final : public runtime::RuntimeSessionBase, public runtime::IOfflineVoiceTaskSession { +public: + FishAudioSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets); + ~FishAudioSession() override; + + std::string family() const override; + runtime::VoiceTaskKind task_kind() const override; + runtime::RunMode run_mode() const override; + void prepare(const runtime::SessionPreparationRequest & request) override; + runtime::TaskResult run(const runtime::TaskRequest & request) override; + +private: + struct ReferenceCacheKey { + std::string source_id; + int sample_rate = 0; + int channels = 0; + uint64_t sample_count = 0; + uint64_t sample_hash = 0; + }; + + struct ReferenceCacheKeyEqual { + bool operator()(const ReferenceCacheKey & lhs, const ReferenceCacheKey & rhs) const; + }; + + struct ReferenceCacheEntry { + FishAudioCodes codes; + }; + + FishAudioRequest make_request(const runtime::TaskRequest & request) const; + const FishAudioCodes & resolve_reference_codes(const FishAudioReference & reference); + + runtime::TaskSpec task_; + std::shared_ptr assets_; + std::unique_ptr generator_; + std::optional defaults_; + runtime::CacheSlots reference_cache_; + std::optional uncached_reference_; +}; + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/tokenizer_text.h b/include/engine/models/fish_audio/tokenizer_text.h new file mode 100644 index 00000000..b457fb2f --- /dev/null +++ b/include/engine/models/fish_audio/tokenizer_text.h @@ -0,0 +1,28 @@ +#pragma once + +#include "engine/models/fish_audio/assets.h" + +#include +#include +#include +#include + +namespace engine::models::fish_audio { + +class FishAudioTextTokenizer { +public: + struct Impl; + + explicit FishAudioTextTokenizer(std::shared_ptr assets); + + std::vector encode(const std::string & text) const; + int32_t token_id(const std::string & token) const; + int32_t im_end_id() const noexcept; + int32_t semantic_begin_id() const noexcept; + int32_t semantic_end_id() const noexcept; + +private: + std::shared_ptr impl_; +}; + +} // namespace engine::models::fish_audio diff --git a/include/engine/models/fish_audio/types.h b/include/engine/models/fish_audio/types.h new file mode 100644 index 00000000..3c05e94a --- /dev/null +++ b/include/engine/models/fish_audio/types.h @@ -0,0 +1,100 @@ +#pragma once + +#include "engine/framework/runtime/session.h" + +#include +#include +#include +#include + +namespace engine::models::fish_audio { + +struct FishAudioGenerationOptions { + int64_t max_new_tokens = 1024; + int64_t chunk_length = 200; + float top_p = 0.8F; + int top_k = 30; + float temperature = 0.8F; + uint32_t seed = 1234; +}; + +struct FishAudioReference { + std::optional audio = std::nullopt; + std::string text; + std::string cache_id; +}; + +struct FishAudioRequest { + std::string text; + std::optional reference = std::nullopt; + FishAudioGenerationOptions generation; +}; + +struct FishAudioCodes { + std::vector codes; + int64_t codebooks = 0; + int64_t frames = 0; +}; + +struct FishAudioPrompt { + std::vector matrix; + int64_t codebook_rows = 0; + int64_t steps = 0; + std::string text; +}; + +struct FishAudioTextConfig { + int64_t vocab_size = 0; + int64_t n_layer = 0; + int64_t dim = 0; + int64_t intermediate_size = 0; + int64_t n_head = 0; + int64_t n_local_heads = 0; + int64_t head_dim = 0; + int64_t max_seq_len = 0; + float rope_base = 1000000.0F; + float norm_eps = 1.0e-6F; + bool tie_word_embeddings = true; + bool attention_qk_norm = true; +}; + +struct FishAudioFastConfig { + int64_t vocab_size = 0; + int64_t num_codebooks = 0; + int64_t n_layer = 0; + int64_t dim = 0; + int64_t intermediate_size = 0; + int64_t n_head = 0; + int64_t n_local_heads = 0; + int64_t head_dim = 0; + int64_t max_seq_len = 0; + float rope_base = 1000000.0F; + float norm_eps = 1.0e-6F; + bool tie_word_embeddings = false; + bool attention_qk_norm = false; +}; + +struct FishAudioCodecConfig { + int sample_rate = 44100; + int64_t semantic_codebook_size = 4096; + int64_t residual_codebook_size = 1024; + int64_t quantizer_codebooks = 9; + int64_t total_codebooks = 10; + int64_t codebook_dim = 8; + int64_t latent_dim = 1024; + int64_t frame_length = 2048; +}; + +struct FishAudioConfig { + std::string model_type; + std::string torch_dtype; + int64_t semantic_start_token_id = 0; + int64_t semantic_end_token_id = 0; + int64_t im_end_token_id = 0; + bool norm_fastlayer_input = false; + FishAudioTextConfig text; + FishAudioFastConfig fast; + FishAudioCodecConfig codec; +}; + +} // namespace engine::models::fish_audio diff --git a/model_specs/fish_audio.json b/model_specs/fish_audio.json new file mode 100644 index 00000000..73e7a869 --- /dev/null +++ b/model_specs/fish_audio.json @@ -0,0 +1,42 @@ +{ + "family": "fish_audio", + "sources": [ + { + "format": "gguf", + "roots": { + "model": ".", + "weights": "$gguf" + }, + "files": { + "config": "model:config.json", + "tokenizer_config": "model:tokenizer_config.json", + "tokenizer_json": "model:tokenizer.json" + }, + "tensors": { + "model_weights": { + "source": "weights:", + "prefix": "model_weights" + }, + "codec_weights": { + "source": "weights:", + "prefix": "codec_weights" + } + } + }, + { + "format": "safetensors", + "roots": { + "model": "." + }, + "files": { + "config": "model:config.json", + "tokenizer_config": "model:tokenizer_config.json", + "tokenizer_json": "model:tokenizer.json" + }, + "tensors": { + "model_weights": "model:model_audio_cpp.safetensors.index.json", + "codec_weights": "model:codec.safetensors" + } + } + ] +} diff --git a/src/framework/runtime/registry.cpp b/src/framework/runtime/registry.cpp index 3a93eec4..78bdcce2 100644 --- a/src/framework/runtime/registry.cpp +++ b/src/framework/runtime/registry.cpp @@ -11,6 +11,7 @@ #include "engine/models/chatterbox/loader.h" #include "engine/models/citrinet_asr/session.h" #include "engine/models/demucs/loader.h" +#include "engine/models/fish_audio/loader.h" #include "engine/models/heartmula/loader.h" #include "engine/models/higgs_audio_stt/loader.h" #include "engine/models/higgs_tts/loader.h" @@ -257,6 +258,7 @@ ModelRegistry make_default_registry(const std::optional & engine::models::vibevoice::make_vibevoice_loader(), engine::models::vibevoice_asr::make_vibevoice_asr_loader(), engine::models::voxtral_realtime::make_voxtral_realtime_loader(), + engine::models::fish_audio::make_fish_audio_loader(), engine::models::heartmula::make_heartmula_loader(), engine::models::higgs_audio_stt::make_higgs_audio_stt_loader(), engine::models::higgs_tts::make_higgs_tts_loader(), diff --git a/src/models/fish_audio/ar.cpp b/src/models/fish_audio/ar.cpp new file mode 100644 index 00000000..91704cd9 --- /dev/null +++ b/src/models/fish_audio/ar.cpp @@ -0,0 +1,1371 @@ +#include "engine/models/fish_audio/ar.h" + +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/modules/attention/qwen_causal_decoder.h" +#include "engine/framework/modules/attention/qwen_decoder.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" +#include "engine/framework/sampling/torch_random.h" + +#include "../common/constant_tensor_cache.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::fish_audio { +namespace { + +namespace binding = engine::modules::binding; +using Clock = std::chrono::steady_clock; + +constexpr int64_t kRasWindow = 10; +constexpr float kRasHighTemperature = 1.0F; +constexpr float kRasHighTopP = 0.9F; + +struct FishARProfile { + double graph_build_prefill_ms = 0.0; + double graph_build_step_ms = 0.0; + double graph_build_fast_ms = 0.0; + double slow_embedding_ms = 0.0; + double fast_embedding_ms = 0.0; + double prefill_input_upload_ms = 0.0; + double prefill_graph_ms = 0.0; + double prefill_output_read_ms = 0.0; + double prefill_state_read_ms = 0.0; + double step_input_upload_ms = 0.0; + double step_mask_upload_ms = 0.0; + double step_graph_ms = 0.0; + double step_output_read_ms = 0.0; + double fast_input_upload_ms = 0.0; + double fast_mask_upload_ms = 0.0; + double fast_graph_ms = 0.0; + double fast_output_read_ms = 0.0; + double import_prefill_state_ms = 0.0; + double sample_bias_ms = 0.0; + double sample_main_ms = 0.0; + double sample_high_ms = 0.0; + double sample_fast_ms = 0.0; + int64_t prefill_runs = 0; + int64_t step_runs = 0; + int64_t fast_runs = 0; + int64_t generated_frames = 0; +}; + +struct SampleCandidate { + int32_t index = 0; + float probability = 0.0F; +}; + +struct SampleDistribution { + size_t source_size = 0; + std::vector candidates; +}; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct FishLayerWeights { + assets::TensorDataF32 input_norm; + core::TensorValue q_proj; + core::TensorValue k_proj; + core::TensorValue v_proj; + core::TensorValue o_proj; + std::optional q_norm; + std::optional k_norm; + assets::TensorDataF32 post_norm; + core::TensorValue gate_proj; + core::TensorValue up_proj; + core::TensorValue down_proj; +}; + +struct FishARWeights { + std::shared_ptr store; + assets::TensorData text_embedding_host; + assets::TensorData codebook_embedding_host; + assets::TensorData fast_embedding_host; + core::TensorValue text_embedding; + std::vector slow_layers; + assets::TensorDataF32 slow_norm; + std::vector fast_layers; + assets::TensorDataF32 fast_norm; + core::TensorValue fast_output; +}; + +struct SlowForwardOutput { + std::vector logits; + std::vector hidden; +}; + +struct SlowPrefillOutput { + SlowForwardOutput forward; + runtime::TransformerKVState state; +}; + +modules::QwenDecoderActivationCastPolicy fish_activation_cast_policy() { + modules::QwenDecoderActivationCastPolicy policy; + policy.enabled = true; + policy.type = GGML_TYPE_BF16; + policy.after_input_norm = true; + policy.after_qkv_projection = true; + policy.after_qk_norm = true; + policy.after_rope = true; + policy.after_static_cache_update = true; + policy.after_attention = true; + policy.after_attention_output = true; + policy.after_residual = true; + policy.after_ffn_norm = true; + policy.after_mlp_projection = true; + policy.after_mlp_silu = true; + policy.after_mlp_mul = true; + policy.after_output = true; + return policy; +} + +modules::QwenCausalDecoderConfig make_slow_decoder_config(const FishAudioTextConfig & config) { + modules::QwenCausalDecoderConfig out; + out.stack.hidden_size = config.dim; + out.stack.num_attention_heads = config.n_head; + out.stack.num_key_value_heads = config.n_local_heads; + out.stack.head_dim = config.head_dim; + out.stack.intermediate_size = config.intermediate_size; + out.stack.layers = config.n_layer; + out.stack.rms_norm_eps = config.norm_eps; + out.stack.rope_theta = config.rope_base; + out.stack.rope_type = GGML_ROPE_TYPE_NORMAL; + out.stack.attention_precision = GGML_PREC_F32; + out.stack.qkv_layout = modules::QwenDecoderQKVLayout::Separate; + out.stack.use_qk_norm = config.attention_qk_norm; + out.stack.activation_cast = fish_activation_cast_policy(); + out.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + out.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + out.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.logits_size = config.vocab_size; + out.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; + out.lm_head_precision = GGML_PREC_F32; + return out; +} + +modules::QwenCausalDecoderConfig make_fast_decoder_config(const FishAudioFastConfig & config) { + modules::QwenCausalDecoderConfig out; + out.stack.hidden_size = config.dim; + out.stack.num_attention_heads = config.n_head; + out.stack.num_key_value_heads = config.n_local_heads; + out.stack.head_dim = config.head_dim; + out.stack.intermediate_size = config.intermediate_size; + out.stack.layers = config.n_layer; + out.stack.rms_norm_eps = config.norm_eps; + out.stack.rope_theta = config.rope_base; + out.stack.rope_type = GGML_ROPE_TYPE_NORMAL; + out.stack.attention_precision = GGML_PREC_F32; + out.stack.qkv_layout = modules::QwenDecoderQKVLayout::Separate; + out.stack.use_qk_norm = config.attention_qk_norm; + out.stack.activation_cast = fish_activation_cast_policy(); + out.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + out.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + out.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.logits_size = config.vocab_size; + out.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; + out.lm_head_precision = GGML_PREC_F32; + return out; +} + +modules::QwenDecoderLayerWeights bind_layer( + common::ConstantTensorCache & constants, + const FishLayerWeights & weights, + bool use_qk_norm) { + modules::QwenDecoderLayerWeights out; + out.input_norm = binding::norm_data(constants, weights.input_norm); + out.self_attention.q_weight = weights.q_proj; + out.self_attention.k_weight = weights.k_proj; + out.self_attention.v_weight = weights.v_proj; + out.self_attention.out_weight = weights.o_proj; + if (use_qk_norm) { + if (!weights.q_norm.has_value() || !weights.k_norm.has_value()) { + throw std::runtime_error("Fish Audio q/k norm weights are missing"); + } + out.q_norm = binding::norm_data(constants, *weights.q_norm); + out.k_norm = binding::norm_data(constants, *weights.k_norm); + } + out.post_norm = binding::norm_data(constants, weights.post_norm); + out.mlp.gate_proj = binding::linear_data(constants, weights.gate_proj); + out.mlp.up_proj = binding::linear_data(constants, weights.up_proj); + out.mlp.down_proj = binding::linear_data(constants, weights.down_proj); + return out; +} + +modules::QwenCausalDecoderWeights bind_slow_weights( + common::ConstantTensorCache & constants, + const FishARWeights & weights, + const FishAudioTextConfig & config) { + modules::QwenCausalDecoderWeights out; + out.stack.layers.reserve(weights.slow_layers.size()); + for (const auto & layer : weights.slow_layers) { + out.stack.layers.push_back(bind_layer(constants, layer, config.attention_qk_norm)); + } + out.final_norm = binding::norm_data(constants, weights.slow_norm); + out.lm_head = binding::linear_data(constants, weights.text_embedding); + return out; +} + +modules::QwenDecoderLayerWeights bind_fast_layer( + common::ConstantTensorCache & constants, + const FishLayerWeights & weights, + const FishAudioFastConfig & config) { + return bind_layer(constants, weights, config.attention_qk_norm); +} + +void copy_tensor_row_to_f32(const assets::TensorData & table, int64_t row, int64_t width, float * out) { + if (row < 0 || width <= 0 || table.shape.rank != 2 || table.shape.dims[1] != width || + row >= table.shape.dims[0]) { + throw std::runtime_error("Fish Audio embedding row lookup shape mismatch"); + } + const size_t row_bytes = ggml_row_size(table.type, width); + const size_t offset = static_cast(row) * row_bytes; + if (offset + row_bytes > table.bytes.size()) { + throw std::runtime_error("Fish Audio embedding row lookup exceeded tensor storage"); + } + const auto * bytes = reinterpret_cast(table.bytes.data()) + offset; + if (table.type == GGML_TYPE_F32) { + std::memcpy(out, bytes, static_cast(width) * sizeof(float)); + } else if (table.type == GGML_TYPE_F16) { + ggml_fp16_to_fp32_row(reinterpret_cast(bytes), out, width); + } else if (table.type == GGML_TYPE_BF16) { + ggml_bf16_to_fp32_row(reinterpret_cast(bytes), out, width); + } else { + throw std::runtime_error("Fish Audio host embedding lookup requires f32/f16/bf16 native embeddings"); + } +} + +std::vector lookup_row(const assets::TensorData & table, int64_t row, int64_t width) { + std::vector out(static_cast(width), 0.0F); + copy_tensor_row_to_f32(table, row, width, out.data()); + return out; +} + +void add_row(const assets::TensorData & table, int64_t row, int64_t width, std::vector & out) { + std::vector tmp(static_cast(width), 0.0F); + copy_tensor_row_to_f32(table, row, width, tmp.data()); + for (int64_t i = 0; i < width; ++i) { + out[static_cast(i)] += tmp[static_cast(i)]; + } +} + +bool is_semantic_token(const FishAudioConfig & config, int32_t token) { + return token >= config.semantic_start_token_id && token <= config.semantic_end_token_id; +} + +std::vector build_slow_embeddings( + const FishAudioConfig & config, + const FishARWeights & weights, + const int32_t * matrix, + int64_t steps) { + const int64_t rows = config.fast.num_codebooks + 1; + const int64_t hidden = config.text.dim; + std::vector out(static_cast(steps * hidden), 0.0F); + const float semantic_scale = 1.0F / std::sqrt(static_cast(rows)); + for (int64_t step = 0; step < steps; ++step) { + const int32_t token = matrix[step]; + auto row = lookup_row(weights.text_embedding_host, token, hidden); + if (is_semantic_token(config, token)) { + for (int64_t codebook = 0; codebook < config.fast.num_codebooks; ++codebook) { + const int32_t code = matrix[(codebook + 1) * steps + step]; + add_row( + weights.codebook_embedding_host, + codebook * config.fast.vocab_size + code, + hidden, + row); + } + for (float & value : row) { + value *= semantic_scale; + } + } + std::copy(row.begin(), row.end(), out.begin() + static_cast(step * hidden)); + } + return out; +} + +std::vector build_slow_embedding_for_frame( + const FishAudioConfig & config, + const FishARWeights & weights, + const std::vector & frame) { + if (static_cast(frame.size()) != config.fast.num_codebooks + 1) { + throw std::runtime_error("Fish Audio frame size mismatch"); + } + return build_slow_embeddings(config, weights, frame.data(), 1); +} + +std::vector build_fast_embedding( + const FishAudioConfig & config, + const FishARWeights & weights, + int32_t code) { + return lookup_row(weights.fast_embedding_host, code, config.fast.dim); +} + +FishLayerWeights load_layer( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + int64_t hidden, + int64_t heads, + int64_t kv_heads, + int64_t head_dim, + int64_t intermediate, + bool qk_norm, + assets::TensorStorageType storage_type) { + FishLayerWeights w; + w.input_norm = source.require_f32_tensor(prefix + ".attention_norm.weight", {hidden}); + w.q_proj = store.load_tensor(source, prefix + ".attention.q_proj.weight", storage_type, {heads * head_dim, hidden}); + w.k_proj = store.load_tensor(source, prefix + ".attention.k_proj.weight", storage_type, {kv_heads * head_dim, hidden}); + w.v_proj = store.load_tensor(source, prefix + ".attention.v_proj.weight", storage_type, {kv_heads * head_dim, hidden}); + w.o_proj = store.load_tensor(source, prefix + ".attention.wo.weight", storage_type, {hidden, heads * head_dim}); + if (qk_norm) { + w.q_norm = source.require_f32_tensor(prefix + ".attention.q_norm.weight", {head_dim}); + w.k_norm = source.require_f32_tensor(prefix + ".attention.k_norm.weight", {head_dim}); + } + w.post_norm = source.require_f32_tensor(prefix + ".ffn_norm.weight", {hidden}); + w.gate_proj = store.load_tensor(source, prefix + ".feed_forward.w1.weight", storage_type, {intermediate, hidden}); + w.down_proj = store.load_tensor(source, prefix + ".feed_forward.w2.weight", storage_type, {hidden, intermediate}); + w.up_proj = store.load_tensor(source, prefix + ".feed_forward.w3.weight", storage_type, {intermediate, hidden}); + return w; +} + +FishARWeights load_ar_weights( + const FishAudioAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType storage_type) { + const auto & source = *assets.model_weights; + const auto & config = assets.config; + FishARWeights weights; + weights.store = std::make_shared( + backend, + backend_type, + "fish_audio.ar.weights", + weight_context_bytes); + weights.text_embedding_host = source.require_tensor( + "embeddings.weight", + assets::TensorStorageType::Native, + {config.text.vocab_size, config.text.dim}); + weights.codebook_embedding_host = source.require_tensor( + "codebook_embeddings.weight", + assets::TensorStorageType::Native, + {config.fast.vocab_size * config.fast.num_codebooks, config.text.dim}); + weights.fast_embedding_host = source.require_tensor( + "fast_embeddings.weight", + assets::TensorStorageType::Native, + {config.fast.vocab_size, config.fast.dim}); + weights.text_embedding = weights.store->load_tensor( + source, + "embeddings.weight", + storage_type, + {config.text.vocab_size, config.text.dim}); + weights.slow_layers.reserve(static_cast(config.text.n_layer)); + for (int64_t i = 0; i < config.text.n_layer; ++i) { + weights.slow_layers.push_back(load_layer( + *weights.store, + source, + "layers." + std::to_string(i), + config.text.dim, + config.text.n_head, + config.text.n_local_heads, + config.text.head_dim, + config.text.intermediate_size, + config.text.attention_qk_norm, + storage_type)); + } + weights.slow_norm = source.require_f32_tensor("norm.weight", {config.text.dim}); + weights.fast_layers.reserve(static_cast(config.fast.n_layer)); + for (int64_t i = 0; i < config.fast.n_layer; ++i) { + weights.fast_layers.push_back(load_layer( + *weights.store, + source, + "fast_layers." + std::to_string(i), + config.fast.dim, + config.fast.n_head, + config.fast.n_local_heads, + config.fast.head_dim, + config.fast.intermediate_size, + config.fast.attention_qk_norm, + storage_type)); + } + weights.fast_norm = source.require_f32_tensor("fast_norm.weight", {config.fast.dim}); + weights.fast_output = weights.store->load_tensor( + source, + "fast_output.weight", + storage_type, + {config.fast.vocab_size, config.fast.dim}); + weights.store->upload(); + return weights; +} + +struct SampleState { + uint64_t seed = 0; + uint64_t call_index = 0; + std::vector previous_main; +}; + +SampleDistribution logits_to_distribution( + const std::vector & logits, + float temperature, + float top_p, + int top_k) { + if (logits.empty()) { + throw std::runtime_error("Fish Audio sampling requires non-empty logits"); + } + std::vector order; + order.reserve(logits.size()); + float max_logit = -std::numeric_limits::infinity(); + for (size_t i = 0; i < logits.size(); ++i) { + const float logit = logits[i]; + if (!std::isfinite(logit)) { + continue; + } + order.push_back(static_cast(i)); + max_logit = std::max(max_logit, logit); + } + double denom = 0.0; + for (const int32_t index : order) { + denom += std::exp(logits[static_cast(index)] - max_logit); + } + if (denom <= 0.0) { + throw std::runtime_error("Fish Audio sampling logits produced zero probability mass"); + } + const size_t candidate_count = std::min(order.size(), static_cast(std::max(top_k, 1))); + const auto by_logit_desc = [&](int32_t lhs, int32_t rhs) { + return logits[static_cast(lhs)] > logits[static_cast(rhs)]; + }; + if (candidate_count < order.size()) { + std::partial_sort(order.begin(), order.begin() + static_cast(candidate_count), order.end(), by_logit_desc); + order.resize(candidate_count); + } else { + std::sort(order.begin(), order.end(), by_logit_desc); + } + double cumulative = 0.0; + std::vector kept; + kept.reserve(candidate_count); + for (size_t i = 0; i < order.size(); ++i) { + const int32_t index = order[i]; + const float logit = logits[static_cast(index)]; + const float prob = static_cast(std::exp(logit - max_logit) / denom); + cumulative += prob; + const bool remove = cumulative > static_cast(top_p) && i != 0; + if (!remove) { + kept.push_back({index, 0.0F}); + } + } + float filtered_max = -std::numeric_limits::infinity(); + const float temperature_scale = std::max(temperature, 1.0e-5F); + for (const auto & candidate : kept) { + filtered_max = std::max(filtered_max, logits[static_cast(candidate.index)] / temperature_scale); + } + double filtered_denom = 0.0; + for (auto & candidate : kept) { + candidate.probability = + std::exp(logits[static_cast(candidate.index)] / temperature_scale - filtered_max); + filtered_denom += candidate.probability; + } + if (filtered_denom <= 0.0) { + throw std::runtime_error("Fish Audio sampling filter produced zero probability mass"); + } + for (auto & candidate : kept) { + candidate.probability = static_cast(static_cast(candidate.probability) / filtered_denom); + } + return {logits.size(), std::move(kept)}; +} + +int32_t sample_from_logits( + const std::vector & logits, + float temperature, + float top_p, + int top_k, + SampleState & state, + const sampling::TorchCudaSamplingPolicy & policy) { + const auto distribution = logits_to_distribution(logits, temperature, top_p, top_k); + const uint64_t call_index = state.call_index++; + int32_t best = 0; + double best_score = -std::numeric_limits::infinity(); + for (const auto & candidate : distribution.candidates) { + if (!(candidate.probability > 0.0F)) { + continue; + } + const float exponential = sampling::torch_cuda_tensor_iterator_exponential_element( + state.seed, + static_cast(distribution.source_size), + static_cast(candidate.index), + call_index, + policy.multiprocessor_count, + policy.max_threads_per_multiprocessor); + const float uniform = std::exp(-exponential); + const float uniform_bf16 = ggml_bf16_to_fp32(ggml_fp32_to_bf16(uniform)); + const float exponential_bf16 = ggml_bf16_to_fp32(ggml_fp32_to_bf16(-std::log(uniform_bf16))); + const double score = static_cast(candidate.probability) / static_cast(exponential_bf16); + if (score > best_score) { + best_score = score; + best = candidate.index; + } + } + return best; +} + +std::vector apply_semantic_bias( + const FishAudioConfig & config, + int32_t im_end_id, + const std::vector & logits) { + std::vector out(logits.size(), -std::numeric_limits::infinity()); + const int64_t begin = std::max(0, config.semantic_start_token_id); + const int64_t end = std::min(static_cast(logits.size()) - 1, config.semantic_end_token_id); + for (int64_t i = begin; i <= end; ++i) { + out[static_cast(i)] = logits[static_cast(i)]; + } + if (im_end_id >= 0 && static_cast(im_end_id) < logits.size()) { + out[static_cast(im_end_id)] = logits[static_cast(im_end_id)]; + } + return out; +} + +core::TensorValue make_fish_causal_mask( + core::ModuleBuildContext &, + common::ConstantTensorCache & constants, + int64_t steps) { + auto values = modules::qwen_causal_prefill_mask_values(1, steps); + return constants.make_tensor( + core::TensorShape::from_dims({1, 1, steps, steps}), + GGML_TYPE_F16, + values.data(), + values.size() * sizeof(ggml_fp16_t)); +} + +struct FishCausalDecoderOutputs { + core::TensorValue hidden; + core::TensorValue logits; + modules::QwenDecoderStackState state; +}; + +FishCausalDecoderOutputs build_fish_causal_decoder( + core::ModuleBuildContext & ctx, + common::ConstantTensorCache & constants, + const core::TensorValue & input, + const core::TensorValue & positions, + const modules::QwenCausalDecoderWeights & weights, + const modules::QwenCausalDecoderConfig & config, + bool norm_fastlayer_input) { + auto mask = make_fish_causal_mask(ctx, constants, input.shape.dims[1]); + auto x = input; + modules::QwenDecoderStackState state; + state.layers.reserve(weights.stack.layers.size()); + const auto layer_config = modules::qwen_decoder_layer_config_from_stack(config.stack); + const modules::QwenDecoderLayerModule layer_module(layer_config); + for (const auto & layer : weights.stack.layers) { + auto out = layer_module.build(ctx, x, positions, layer, std::nullopt, std::nullopt, mask); + x = out.output; + auto state_key = core::wrap_tensor(ggml_dup(ctx.ggml, out.key.tensor), out.key.shape, out.key.type); + auto state_value = core::wrap_tensor(ggml_dup(ctx.ggml, out.value.tensor), out.value.shape, out.value.type); + state.layers.push_back({state_key, state_value}); + } + auto hidden_sequence = modules::RMSNormModule({config.stack.hidden_size, config.stack.rms_norm_eps, true, false}) + .build(ctx, x, weights.final_norm); + const int64_t steps = hidden_sequence.shape.dims[1]; + auto fast_hidden_source = norm_fastlayer_input ? hidden_sequence : x; + auto hidden = modules::SliceModule({1, steps - 1, 1}).build(ctx, fast_hidden_source); + auto logits = modules::LinearModule({config.stack.hidden_size, config.logits_size, false, config.lm_head_precision}) + .build(ctx, modules::SliceModule({1, steps - 1, 1}).build(ctx, hidden_sequence), weights.lm_head); + auto hidden_out = core::wrap_tensor(ggml_dup(ctx.ggml, hidden.tensor), hidden.shape, hidden.type); + auto logits_out = core::wrap_tensor(ggml_dup(ctx.ggml, logits.tensor), logits.shape, logits.type); + return {hidden_out, logits_out, std::move(state)}; +} + +struct FishStaticDecoderOutputs { + core::TensorValue hidden; + core::TensorValue logits; + runtime::TransformerKVCache cache; +}; + +FishStaticDecoderOutputs build_fish_static_decoder( + core::ModuleBuildContext & ctx, + ggml_cgraph * graph, + const core::TensorValue & input, + const core::TensorValue & positions, + const modules::QwenCausalDecoderWeights & weights, + const modules::QwenCausalDecoderConfig & config, + int64_t cache_steps, + const core::TensorValue & attention_mask, + const core::TensorValue & cache_slot, + bool norm_fastlayer_input) { + auto decoder = modules::QwenCausalDecoderModule(config).build_static_cache_tail( + ctx, + graph, + input, + positions, + weights, + cache_steps, + attention_mask, + cache_slot); + auto fast_hidden = norm_fastlayer_input ? decoder.hidden : decoder.sequence; + return { + fast_hidden, + decoder.logits, + std::move(decoder.cache), + }; +} + +} // namespace + +class FishARWeightsRuntime { +public: + FishARWeightsRuntime( + std::shared_ptr assets, + core::BackendConfig backend_config, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : assets_(std::move(assets)), + threads_(threads), + graph_arena_bytes_(graph_arena_bytes) { + if (assets_ == nullptr) { + throw std::runtime_error("Fish Audio AR weights runtime requires assets"); + } + backend_config.threads = threads_; + backend_ = core::init_backend(backend_config); + backend_type_ = core::backend_type(backend_); + weights_ = std::make_shared( + load_ar_weights(*assets_, backend_, backend_type_, weight_context_bytes, weight_storage_type)); + slow_step_constants_ = std::make_unique( + backend_, + threads_, + "fish_audio.ar.step.constants", + 256ull * 1024ull * 1024ull); + fast_constants_ = std::make_unique( + backend_, + threads_, + "fish_audio.ar.fast.constants", + 256ull * 1024ull * 1024ull); + } + + ~FishARWeightsRuntime() { + fast_constants_.reset(); + slow_step_constants_.reset(); + weights_.reset(); + if (backend_ != nullptr) { + ggml_backend_free(backend_); + } + } + + FishARWeightsRuntime(const FishARWeightsRuntime &) = delete; + FishARWeightsRuntime & operator=(const FishARWeightsRuntime &) = delete; + + const FishAudioAssets & assets() const noexcept { + return *assets_; + } + + const FishARWeights & weights() const noexcept { + return *weights_; + } + + int threads() const noexcept { + return threads_; + } + + size_t graph_arena_bytes() const noexcept { + return graph_arena_bytes_; + } + + ggml_backend_t backend() const noexcept { + return backend_; + } + + core::BackendType backend_type() const noexcept { + return backend_type_; + } + + common::ConstantTensorCache & slow_step_constants() const noexcept { + return *slow_step_constants_; + } + + common::ConstantTensorCache & fast_constants() const noexcept { + return *fast_constants_; + } + +private: + std::shared_ptr assets_; + std::shared_ptr weights_; + int threads_ = 1; + size_t graph_arena_bytes_ = 0; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + std::unique_ptr slow_step_constants_; + std::unique_ptr fast_constants_; +}; + +class FishAudioARRuntime::Impl { +public: + Impl( + std::shared_ptr assets, + core::BackendConfig backend_config, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : runtime_(std::make_shared( + std::move(assets), + backend_config, + threads, + graph_arena_bytes, + weight_context_bytes, + weight_storage_type)), + sampling_policy_(sampling::resolve_torch_cuda_sampling_policy( + runtime_->backend_type(), + backend_config.device, + "fish_audio.ar.cuda_sampling_policy", + "Fish Audio", + sampling::TorchCudaSamplingPolicyFailureMode::StrictCuda)) {} + + ~Impl() { + step_graph_.reset(); + prefill_graph_.reset(); + fast_graph_.reset(); + runtime_.reset(); + } + + FishAudioCodes generate(const FishAudioPrompt & prompt, const FishAudioGenerationOptions & options) { + FishARProfile profile; + const auto & assets = runtime_->assets(); + const auto & weights = runtime_->weights(); + if (prompt.codebook_rows != assets.config.fast.num_codebooks + 1 || + static_cast(prompt.matrix.size()) != prompt.codebook_rows * prompt.steps) { + throw std::runtime_error("Fish Audio AR prompt shape mismatch"); + } + const int64_t max_new_tokens = std::min(options.max_new_tokens, assets.config.text.max_seq_len - prompt.steps); + if (max_new_tokens <= 0) { + throw std::runtime_error("Fish Audio prompt leaves no room for generated tokens"); + } + ensure_prefill_graph(prompt.steps, profile); + ensure_step_graph(prompt.steps + max_new_tokens, profile); + ensure_fast_graph(profile); + SampleState sample; + sample.seed = options.seed; + sample.previous_main.assign(static_cast(kRasWindow), 0); + auto timing_start = Clock::now(); + auto embeddings = build_slow_embeddings(assets.config, weights, prompt.matrix.data(), prompt.steps); + profile.slow_embedding_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + auto prefill = prefill_graph_->run(embeddings, profile); + std::vector generated_frame_major; + generated_frame_major.reserve(static_cast(max_new_tokens * assets.config.fast.num_codebooks)); + auto frame = sample_frame(prefill.forward.logits, prefill.forward.hidden, options, sample, false, profile); + if (frame.front() == im_end_id()) { + log_profile(profile); + return FishAudioCodes{{}, assets.config.fast.num_codebooks, 0}; + } + append_frame(generated_frame_major, frame); + ++profile.generated_frames; + timing_start = Clock::now(); + step_graph_->import_state(prefill.state); + profile.import_prefill_state_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + bool ended_by_im_end = false; + for (int64_t step = 1; step < max_new_tokens; ++step) { + timing_start = Clock::now(); + const auto input = build_slow_embedding_for_frame(assets.config, weights, frame); + profile.slow_embedding_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + auto step_out = step_graph_->run(input, profile); + frame = sample_frame(step_out.logits, step_out.hidden, options, sample, true, profile); + if (frame.front() == im_end_id()) { + ended_by_im_end = true; + break; + } + append_frame(generated_frame_major, frame); + ++profile.generated_frames; + } + if (!ended_by_im_end && !generated_frame_major.empty()) { + generated_frame_major.resize(generated_frame_major.size() - static_cast(assets.config.fast.num_codebooks)); + --profile.generated_frames; + } + FishAudioCodes out; + out.codebooks = assets.config.fast.num_codebooks; + out.frames = static_cast(generated_frame_major.size()) / out.codebooks; + out.codes.assign(static_cast(out.codebooks * out.frames), 0); + for (int64_t frame_index = 0; frame_index < out.frames; ++frame_index) { + for (int64_t codebook = 0; codebook < out.codebooks; ++codebook) { + out.codes[static_cast(codebook * out.frames + frame_index)] = + generated_frame_major[static_cast(frame_index * out.codebooks + codebook)]; + } + } + log_profile(profile); + return out; + } + + void release_runtime_graphs() { + step_graph_.reset(); + prefill_graph_.reset(); + fast_graph_.reset(); + } + +private: + class PrefillGraph { + public: + PrefillGraph(std::shared_ptr runtime, int64_t steps) + : runtime_(std::move(runtime)), + steps_(steps) { + ggml_init_params params{runtime_->graph_arena_bytes(), nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio AR prefill context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "fish_audio.ar.prefill", runtime_->backend_type()}; + const auto & assets = runtime_->assets(); + const auto & config = assets.config.text; + auto input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, steps_, config.dim})); + input_ = input.tensor; + positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, steps_); + auto positions_value = core::wrap_tensor(positions_, core::TensorShape::from_dims({steps_}), GGML_TYPE_I32); + constants_ = std::make_unique( + runtime_->backend(), + runtime_->threads(), + "fish_audio.ar.prefill.constants", + 256ull * 1024ull * 1024ull); + constants_->begin_graph(); + auto decoder = build_fish_causal_decoder( + ctx, + *constants_, + input, + positions_value, + bind_slow_weights(*constants_, runtime_->weights(), config), + make_slow_decoder_config(config), + assets.config.norm_fastlayer_input); + for (const auto & layer : decoder.state.layers) { + if (!layer.key.has_value() || !layer.value.has_value()) { + throw std::runtime_error("Fish Audio prefill decoder did not produce K/V state"); + } + keys_.push_back(layer.key->tensor); + values_.push_back(layer.value->tensor); + } + hidden_ = decoder.hidden.tensor; + logits_ = decoder.logits.tensor; + ggml_set_output(hidden_); + for (ggml_tensor * key : keys_) { + ggml_set_output(key); + } + for (ggml_tensor * value : values_) { + ggml_set_output(value); + } + ggml_set_output(logits_); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + ggml_build_forward_expand(graph_, logits_); + ggml_build_forward_expand(graph_, hidden_); + for (ggml_tensor * key : keys_) { + ggml_build_forward_expand(graph_, key); + } + for (ggml_tensor * value : values_) { + ggml_build_forward_expand(graph_, value); + } + constants_->finish_graph(); + constants_->ensure_uploaded(); + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + throw std::runtime_error("failed to allocate Fish Audio AR prefill graph"); + } + auto positions = modules::qwen_position_ids(steps_); + ggml_backend_tensor_set(positions_, positions.data(), 0, positions.size() * sizeof(int32_t)); + } + + ~PrefillGraph() { + core::release_backend_graph_resources(runtime_->backend(), graph_); + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + } + + SlowPrefillOutput run(const std::vector & embeddings, FishARProfile & profile) { + const auto & config = runtime_->assets().config.text; + if (static_cast(embeddings.size()) != steps_ * config.dim) { + throw std::runtime_error("Fish Audio prefill embedding size mismatch"); + } + ++profile.prefill_runs; + auto timing_start = Clock::now(); + ggml_backend_tensor_set(input_, embeddings.data(), 0, embeddings.size() * sizeof(float)); + profile.prefill_input_upload_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + timing_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(runtime_->backend(), graph_, nullptr, "fish_audio.ar.prefill"); + ggml_backend_synchronize(runtime_->backend()); + profile.prefill_graph_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Fish Audio AR prefill graph compute failed"); + } + SlowPrefillOutput out; + out.forward.logits.resize(static_cast(config.vocab_size)); + out.forward.hidden.resize(static_cast(config.dim)); + timing_start = Clock::now(); + ggml_backend_tensor_get(logits_, out.forward.logits.data(), 0, out.forward.logits.size() * sizeof(float)); + ggml_backend_tensor_get(hidden_, out.forward.hidden.data(), 0, out.forward.hidden.size() * sizeof(float)); + profile.prefill_output_read_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + out.state.current_end = steps_; + out.state.layers.resize(keys_.size()); + const size_t values_per_layer = static_cast(steps_ * config.n_local_heads * config.head_dim); + timing_start = Clock::now(); + for (size_t i = 0; i < keys_.size(); ++i) { + out.state.layers[i].valid_steps = steps_; + out.state.layers[i].key.resize(values_per_layer); + out.state.layers[i].value.resize(values_per_layer); + ggml_backend_tensor_get(keys_[i], out.state.layers[i].key.data(), 0, values_per_layer * sizeof(float)); + ggml_backend_tensor_get(values_[i], out.state.layers[i].value.data(), 0, values_per_layer * sizeof(float)); + } + profile.prefill_state_read_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + return out; + } + + int64_t steps() const noexcept { return steps_; } + + private: + std::shared_ptr runtime_; + int64_t steps_ = 0; + std::unique_ptr ctx_; + ggml_tensor * input_ = nullptr; + ggml_tensor * positions_ = nullptr; + ggml_tensor * hidden_ = nullptr; + ggml_tensor * logits_ = nullptr; + std::vector keys_; + std::vector values_; + ggml_cgraph * graph_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + std::unique_ptr constants_; + }; + + class StepGraph { + public: + StepGraph(std::shared_ptr runtime, int64_t cache_steps) + : runtime_(std::move(runtime)), + cache_steps_(cache_steps) { + ggml_init_params params{runtime_->graph_arena_bytes(), nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio AR step context"); + } + const auto & assets = runtime_->assets(); + const auto & config = assets.config.text; + core::ModuleBuildContext ctx{ctx_.get(), "fish_audio.ar.step", runtime_->backend_type()}; + auto input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, config.dim})); + input_ = input.tensor; + position_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + cache_slot_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, cache_steps_, 1, 1, 1); + auto position_value = core::wrap_tensor(position_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + auto cache_slot_value = core::wrap_tensor(cache_slot_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + auto mask_value = core::wrap_tensor(mask_, core::TensorShape::from_dims({1, 1, 1, cache_steps_}), GGML_TYPE_F16); + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + auto & constants = runtime_->slow_step_constants(); + constants.begin_graph(); + auto decoder = build_fish_static_decoder( + ctx, + graph_, + input, + position_value, + bind_slow_weights(constants, runtime_->weights(), config), + make_slow_decoder_config(config), + cache_steps_, + mask_value, + cache_slot_value, + assets.config.norm_fastlayer_input); + cache_ = std::move(decoder.cache); + hidden_ = decoder.hidden.tensor; + logits_ = decoder.logits.tensor; + ggml_set_output(hidden_); + ggml_set_output(logits_); + ggml_build_forward_expand(graph_, logits_); + ggml_build_forward_expand(graph_, hidden_); + constants.finish_graph(); + constants.ensure_uploaded(); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); + if (buffer_ == nullptr) { + throw std::runtime_error("failed to allocate Fish Audio AR step tensors"); + } + mask_scratch_.assign(static_cast(cache_steps_), ggml_fp32_to_fp16(-INFINITY)); + } + + ~StepGraph() { + core::release_backend_graph_resources(runtime_->backend(), graph_); + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + } + + int64_t cache_steps() const noexcept { return cache_steps_; } + + void import_state(const runtime::TransformerKVState & state) { + cache_.import_state(state); + const auto masked = ggml_fp32_to_fp16(-INFINITY); + const auto visible = ggml_fp32_to_fp16(0.0F); + std::fill(mask_scratch_.begin(), mask_scratch_.end(), masked); + for (int64_t i = 0; i < cache_.valid_steps(); ++i) { + mask_scratch_[static_cast(i)] = visible; + } + ggml_backend_tensor_set(mask_, mask_scratch_.data(), 0, mask_scratch_.size() * sizeof(ggml_fp16_t)); + } + + SlowForwardOutput run(const std::vector & embedding, FishARProfile & profile) { + const auto & config = runtime_->assets().config.text; + if (static_cast(embedding.size()) != config.dim) { + throw std::runtime_error("Fish Audio step embedding size mismatch"); + } + if (cache_.valid_steps() >= cache_steps_) { + throw std::runtime_error("Fish Audio step cache exceeds capacity"); + } + ++profile.step_runs; + auto timing_start = Clock::now(); + const int32_t pos = static_cast(cache_.current_end()); + ggml_backend_tensor_set(position_, &pos, 0, sizeof(pos)); + const int32_t cache_slot = static_cast(cache_.valid_steps()); + ggml_backend_tensor_set(cache_slot_, &cache_slot, 0, sizeof(cache_slot)); + const auto visible = ggml_fp32_to_fp16(0.0F); + mask_scratch_[static_cast(cache_.valid_steps())] = visible; + ggml_backend_tensor_set( + mask_, + &visible, + static_cast(cache_.valid_steps()) * sizeof(ggml_fp16_t), + sizeof(ggml_fp16_t)); + profile.step_mask_upload_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + timing_start = Clock::now(); + ggml_backend_tensor_set(input_, embedding.data(), 0, embedding.size() * sizeof(float)); + profile.step_input_upload_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + timing_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(runtime_->backend(), graph_, nullptr, "fish_audio.ar.step"); + ggml_backend_synchronize(runtime_->backend()); + profile.step_graph_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Fish Audio AR step graph compute failed"); + } + cache_.advance_after_direct_append(1); + SlowForwardOutput out; + out.logits.resize(static_cast(config.vocab_size)); + out.hidden.resize(static_cast(config.dim)); + timing_start = Clock::now(); + ggml_backend_tensor_get(logits_, out.logits.data(), 0, out.logits.size() * sizeof(float)); + ggml_backend_tensor_get(hidden_, out.hidden.data(), 0, out.hidden.size() * sizeof(float)); + profile.step_output_read_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + return out; + } + + private: + std::shared_ptr runtime_; + int64_t cache_steps_ = 0; + std::unique_ptr ctx_; + ggml_tensor * input_ = nullptr; + ggml_tensor * position_ = nullptr; + ggml_tensor * cache_slot_ = nullptr; + ggml_tensor * mask_ = nullptr; + ggml_tensor * hidden_ = nullptr; + ggml_tensor * logits_ = nullptr; + runtime::TransformerKVCache cache_; + std::vector mask_scratch_; + ggml_cgraph * graph_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; + }; + + class FastGraph { + public: + explicit FastGraph(std::shared_ptr runtime) + : runtime_(std::move(runtime)) { + ggml_init_params params{runtime_->graph_arena_bytes(), nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio fast AR context"); + } + const auto & config = runtime_->assets().config.fast; + const auto & weights = runtime_->weights(); + core::ModuleBuildContext ctx{ctx_.get(), "fish_audio.ar.fast", runtime_->backend_type()}; + auto input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, config.dim})); + input_ = input.tensor; + position_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); + mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, config.num_codebooks, 1, 1, 1); + auto position_value = core::wrap_tensor(position_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); + auto mask_value = core::wrap_tensor(mask_, core::TensorShape::from_dims({1, 1, 1, config.num_codebooks}), GGML_TYPE_F16); + graph_ = ggml_new_graph_custom(ctx_.get(), 32768, false); + auto & constants = runtime_->fast_constants(); + constants.begin_graph(); + modules::QwenCausalDecoderWeights decoder_weights; + decoder_weights.stack.layers.reserve(weights.fast_layers.size()); + for (const auto & layer : weights.fast_layers) { + decoder_weights.stack.layers.push_back(bind_fast_layer(constants, layer, config)); + } + decoder_weights.final_norm = binding::norm_data(constants, weights.fast_norm); + decoder_weights.lm_head = binding::linear_data(constants, weights.fast_output); + auto decoder = modules::QwenCausalDecoderModule(make_fast_decoder_config(config)) + .build_static_cache_tail( + ctx, + graph_, + input, + position_value, + decoder_weights, + config.num_codebooks, + mask_value, + position_value); + for (size_t layer = 0; layer < weights.fast_layers.size(); ++layer) { + cache_keys_.push_back(decoder.cache.key_tensor(layer).tensor); + cache_values_.push_back(decoder.cache.value_tensor(layer).tensor); + } + logits_ = decoder.logits.tensor; + ggml_set_output(logits_); + ggml_build_forward_expand(graph_, logits_); + constants.finish_graph(); + constants.ensure_uploaded(); + buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); + if (buffer_ == nullptr) { + throw std::runtime_error("failed to allocate Fish Audio fast AR graph"); + } + mask_scratch_.assign(static_cast(config.num_codebooks), ggml_fp32_to_fp16(-INFINITY)); + } + + ~FastGraph() { + core::release_backend_graph_resources(runtime_->backend(), graph_); + if (buffer_ != nullptr) { + ggml_backend_buffer_free(buffer_); + } + } + + std::vector run(const std::vector & input, int64_t position, FishARProfile & profile) { + const auto & config = runtime_->assets().config.fast; + if (static_cast(input.size()) != config.dim) { + throw std::runtime_error("Fish Audio fast AR input size mismatch"); + } + ++profile.fast_runs; + auto timing_start = Clock::now(); + const int32_t pos = static_cast(position); + ggml_backend_tensor_set(position_, &pos, 0, sizeof(pos)); + const auto visible = ggml_fp32_to_fp16(0.0F); + if (position == 0) { + std::fill(mask_scratch_.begin(), mask_scratch_.end(), ggml_fp32_to_fp16(-INFINITY)); + mask_scratch_[0] = visible; + ggml_backend_tensor_set(mask_, mask_scratch_.data(), 0, mask_scratch_.size() * sizeof(ggml_fp16_t)); + } else { + mask_scratch_[static_cast(position)] = visible; + ggml_backend_tensor_set( + mask_, + &visible, + static_cast(position) * sizeof(ggml_fp16_t), + sizeof(ggml_fp16_t)); + } + profile.fast_mask_upload_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + timing_start = Clock::now(); + ggml_backend_tensor_set(input_, input.data(), 0, input.size() * sizeof(float)); + profile.fast_input_upload_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + core::set_backend_threads(runtime_->backend(), runtime_->threads()); + timing_start = Clock::now(); + const ggml_status status = core::compute_backend_graph(runtime_->backend(), graph_, nullptr, "fish_audio.ar.fast"); + ggml_backend_synchronize(runtime_->backend()); + profile.fast_graph_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Fish Audio fast AR graph compute failed"); + } + std::vector logits(static_cast(config.vocab_size), 0.0F); + timing_start = Clock::now(); + ggml_backend_tensor_get(logits_, logits.data(), 0, logits.size() * sizeof(float)); + profile.fast_output_read_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + return logits; + } + + private: + std::shared_ptr runtime_; + std::unique_ptr ctx_; + ggml_tensor * input_ = nullptr; + ggml_tensor * position_ = nullptr; + ggml_tensor * mask_ = nullptr; + ggml_tensor * logits_ = nullptr; + std::vector cache_keys_; + std::vector cache_values_; + std::vector mask_scratch_; + ggml_cgraph * graph_ = nullptr; + ggml_backend_buffer_t buffer_ = nullptr; + }; + + void ensure_prefill_graph(int64_t steps, FishARProfile & profile) { + if (!prefill_graph_ || prefill_graph_->steps() != steps) { + const auto build_start = Clock::now(); + prefill_graph_ = std::make_unique(runtime_, steps); + profile.graph_build_prefill_ms += engine::debug::elapsed_ms(build_start, Clock::now()); + } + } + + void ensure_step_graph(int64_t cache_steps, FishARProfile & profile) { + if (!step_graph_ || step_graph_->cache_steps() < cache_steps) { + const auto build_start = Clock::now(); + step_graph_ = std::make_unique(runtime_, cache_steps); + profile.graph_build_step_ms += engine::debug::elapsed_ms(build_start, Clock::now()); + } + } + + void ensure_fast_graph(FishARProfile & profile) { + if (!fast_graph_) { + const auto build_start = Clock::now(); + fast_graph_ = std::make_unique(runtime_); + profile.graph_build_fast_ms += engine::debug::elapsed_ms(build_start, Clock::now()); + } + } + + int32_t im_end_id() const { + return static_cast(runtime_->assets().config.im_end_token_id); + } + + void append_frame(std::vector & out, const std::vector & frame) const { + if (static_cast(frame.size()) != runtime_->assets().config.fast.num_codebooks + 1) { + throw std::runtime_error("Fish Audio generated frame shape mismatch"); + } + out.insert(out.end(), frame.begin() + 1, frame.end()); + } + + std::vector sample_frame( + const std::vector & slow_logits, + const std::vector & slow_hidden, + const FishAudioGenerationOptions & options, + SampleState & sample, + bool apply_ras, + FishARProfile & profile) { + const auto & config = runtime_->assets().config; + const auto & weights = runtime_->weights(); + auto timing_start = Clock::now(); + const auto biased = apply_semantic_bias(config, im_end_id(), slow_logits); + profile.sample_bias_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + timing_start = Clock::now(); + // Upstream Python accepts repetition_penalty on the request but does not apply it in this generation path. + int32_t main_token = sample_from_logits( + biased, + options.temperature, + options.top_p, + options.top_k, + sample, + sampling_policy_); + profile.sample_main_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + timing_start = Clock::now(); + const int32_t high_token = sample_from_logits( + biased, + kRasHighTemperature, + kRasHighTopP, + options.top_k, + sample, + sampling_policy_); + profile.sample_high_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + if (apply_ras && is_semantic_token(config, main_token) && + std::find(sample.previous_main.begin(), sample.previous_main.end(), main_token) != sample.previous_main.end()) { + main_token = high_token; + } + std::rotate(sample.previous_main.begin(), sample.previous_main.begin() + 1, sample.previous_main.end()); + sample.previous_main.back() = main_token; + + std::vector frame(static_cast(config.fast.num_codebooks + 1), 0); + frame[0] = main_token; + if (!is_semantic_token(config, main_token)) { + return frame; + } + const auto fast0_logits = fast_graph_->run(slow_hidden, 0, profile); + int32_t code = std::clamp( + main_token - static_cast(config.semantic_start_token_id), + 0, + static_cast(config.fast.vocab_size - 1)); + frame[1] = code; + for (int64_t codebook = 1; codebook < config.fast.num_codebooks; ++codebook) { + timing_start = Clock::now(); + const auto embedding = build_fast_embedding(config, weights, code); + profile.fast_embedding_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + const auto logits = fast_graph_->run(embedding, codebook, profile); + timing_start = Clock::now(); + code = sample_from_logits( + logits, + options.temperature, + options.top_p, + options.top_k, + sample, + sampling_policy_); + profile.sample_fast_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + frame[static_cast(codebook + 1)] = code; + } + return frame; + } + + void log_profile(const FishARProfile & profile) const { + engine::debug::timing_log_scalar("fish_audio.ar.profile.graph_build_prefill_ms", profile.graph_build_prefill_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.graph_build_step_ms", profile.graph_build_step_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.graph_build_fast_ms", profile.graph_build_fast_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.slow_embedding_ms", profile.slow_embedding_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.fast_embedding_ms", profile.fast_embedding_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.prefill_input_upload_ms", profile.prefill_input_upload_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.prefill_graph_ms", profile.prefill_graph_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.prefill_output_read_ms", profile.prefill_output_read_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.prefill_state_read_ms", profile.prefill_state_read_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.step_input_upload_ms", profile.step_input_upload_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.step_mask_upload_ms", profile.step_mask_upload_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.step_graph_ms", profile.step_graph_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.step_output_read_ms", profile.step_output_read_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.fast_input_upload_ms", profile.fast_input_upload_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.fast_mask_upload_ms", profile.fast_mask_upload_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.fast_graph_ms", profile.fast_graph_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.fast_output_read_ms", profile.fast_output_read_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.import_prefill_state_ms", profile.import_prefill_state_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.sample_bias_ms", profile.sample_bias_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.sample_main_ms", profile.sample_main_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.sample_high_ms", profile.sample_high_ms); + engine::debug::timing_log_scalar("fish_audio.ar.profile.sample_fast_ms", profile.sample_fast_ms); + engine::debug::trace_log_scalar("fish_audio.ar.profile.prefill_runs", profile.prefill_runs); + engine::debug::trace_log_scalar("fish_audio.ar.profile.step_runs", profile.step_runs); + engine::debug::trace_log_scalar("fish_audio.ar.profile.fast_runs", profile.fast_runs); + engine::debug::trace_log_scalar("fish_audio.ar.profile.generated_frames", profile.generated_frames); + } + + std::shared_ptr runtime_; + sampling::TorchCudaSamplingPolicy sampling_policy_; + std::unique_ptr prefill_graph_; + std::unique_ptr step_graph_; + std::unique_ptr fast_graph_; +}; + +FishAudioARRuntime::FishAudioARRuntime( + std::shared_ptr assets, + core::BackendConfig backend, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType weight_storage_type) + : impl_(std::make_unique( + std::move(assets), + backend, + threads, + graph_arena_bytes, + weight_context_bytes, + weight_storage_type)) {} + +FishAudioARRuntime::~FishAudioARRuntime() = default; + +FishAudioCodes FishAudioARRuntime::generate( + const FishAudioPrompt & prompt, + const FishAudioGenerationOptions & options) { + return impl_->generate(prompt, options); +} + +void FishAudioARRuntime::release_runtime_graphs() { + impl_->release_runtime_graphs(); +} + +} // namespace engine::models::fish_audio diff --git a/src/models/fish_audio/assets.cpp b/src/models/fish_audio/assets.cpp new file mode 100644 index 00000000..7870a300 --- /dev/null +++ b/src/models/fish_audio/assets.cpp @@ -0,0 +1,127 @@ +#include "engine/models/fish_audio/assets.h" + +#include "engine/framework/assets/model_package.h" +#include "engine/framework/io/config.h" +#include "engine/framework/io/json.h" + +#include +#include + +namespace engine::models::fish_audio { +namespace json = engine::io::json; +namespace { + +FishAudioTextConfig parse_text_config(const json::Value & value) { + if (json::optional_string(value, "model_type", "") != "fish_qwen3") { + throw std::runtime_error("Fish Audio text_config.model_type mismatch"); + } + FishAudioTextConfig config; + config.vocab_size = json::require_i64(value, "vocab_size"); + config.n_layer = json::require_i64(value, "n_layer"); + config.dim = json::require_i64(value, "dim"); + config.intermediate_size = json::require_i64(value, "intermediate_size"); + config.n_head = json::require_i64(value, "n_head"); + config.n_local_heads = json::optional_i64(value, "n_local_heads", config.n_head); + config.head_dim = json::require_i64(value, "head_dim"); + config.max_seq_len = json::require_i64(value, "max_seq_len"); + config.rope_base = json::optional_f32(value, "rope_base", config.rope_base); + config.norm_eps = json::optional_f32(value, "norm_eps", config.norm_eps); + config.tie_word_embeddings = json::optional_bool(value, "tie_word_embeddings", config.tie_word_embeddings); + config.attention_qk_norm = json::optional_bool(value, "attention_qk_norm", config.attention_qk_norm); + engine::io::require_positive(config.vocab_size, "text vocab_size"); + engine::io::require_positive(config.n_layer, "text n_layer"); + engine::io::require_positive(config.dim, "text dim"); + engine::io::require_positive(config.intermediate_size, "text intermediate_size"); + engine::io::require_positive(config.n_head, "text n_head"); + engine::io::require_positive(config.n_local_heads, "text n_local_heads"); + engine::io::require_positive(config.head_dim, "text head_dim"); + engine::io::require_positive(config.max_seq_len, "text max_seq_len"); + engine::io::require_divisible(config.n_head, config.n_local_heads, "text n_head / n_local_heads"); + return config; +} + +FishAudioFastConfig parse_fast_config(const json::Value & value) { + if (json::optional_string(value, "model_type", "") != "fish_qwen3_audio_decoder") { + throw std::runtime_error("Fish Audio audio_decoder_config.model_type mismatch"); + } + FishAudioFastConfig config; + config.vocab_size = json::require_i64(value, "vocab_size"); + config.num_codebooks = json::require_i64(value, "num_codebooks"); + config.n_layer = json::require_i64(value, "n_layer"); + config.dim = json::require_i64(value, "dim"); + config.intermediate_size = json::require_i64(value, "intermediate_size"); + config.n_head = json::require_i64(value, "n_head"); + config.n_local_heads = json::optional_i64(value, "n_local_heads", config.n_head); + config.head_dim = json::require_i64(value, "head_dim"); + config.max_seq_len = json::optional_i64(value, "max_seq_len", config.num_codebooks + 1); + config.rope_base = json::optional_f32(value, "rope_base", config.rope_base); + config.norm_eps = json::optional_f32(value, "norm_eps", config.norm_eps); + config.tie_word_embeddings = json::optional_bool(value, "tie_word_embeddings", config.tie_word_embeddings); + config.attention_qk_norm = json::optional_bool(value, "attention_qk_norm", config.attention_qk_norm); + engine::io::require_positive(config.vocab_size, "fast vocab_size"); + engine::io::require_positive(config.num_codebooks, "fast num_codebooks"); + engine::io::require_positive(config.n_layer, "fast n_layer"); + engine::io::require_positive(config.dim, "fast dim"); + engine::io::require_positive(config.intermediate_size, "fast intermediate_size"); + engine::io::require_positive(config.n_head, "fast n_head"); + engine::io::require_positive(config.n_local_heads, "fast n_local_heads"); + engine::io::require_positive(config.head_dim, "fast head_dim"); + engine::io::require_divisible(config.n_head, config.n_local_heads, "fast n_head / n_local_heads"); + return config; +} + +FishAudioConfig parse_config(const assets::ResourceBundle & resources) { + const auto root = resources.parse_json("config"); + FishAudioConfig config; + config.model_type = json::optional_string(root, "model_type", ""); + if (config.model_type != "fish_qwen3_omni") { + throw std::runtime_error("Fish Audio model_type mismatch"); + } + config.torch_dtype = json::optional_string(root, "torch_dtype", config.torch_dtype); + config.semantic_start_token_id = json::require_i64(root, "semantic_start_token_id"); + config.semantic_end_token_id = json::require_i64(root, "semantic_end_token_id"); + config.im_end_token_id = json::require_i64(root, "eos_token_id"); + config.norm_fastlayer_input = json::optional_bool(root, "norm_fastlayer_input", config.model_type == "fish_qwen3_omni"); + config.text = parse_text_config(root.require("text_config")); + config.fast = parse_fast_config(root.require("audio_decoder_config")); + config.codec.total_codebooks = config.fast.num_codebooks; + if (config.fast.dim != config.text.dim) { + throw std::runtime_error("Fish Audio fast dim must match text dim for S2-Pro"); + } + if (!config.text.tie_word_embeddings) { + throw std::runtime_error("Fish Audio S2-Pro expects tied text embeddings"); + } + if (config.semantic_start_token_id <= 0 || config.semantic_end_token_id < config.semantic_start_token_id) { + throw std::runtime_error("Fish Audio semantic token range is invalid"); + } + return config; +} + +void validate_weight_anchors(const FishAudioAssets & assets) { + assets.model_weights->require_metadata("embeddings.weight"); + assets.model_weights->require_metadata("codebook_embeddings.weight"); + assets.model_weights->require_metadata("layers.0.attention.q_proj.weight"); + assets.model_weights->require_metadata("layers.0.attention.k_proj.weight"); + assets.model_weights->require_metadata("layers.0.attention.v_proj.weight"); + assets.model_weights->require_metadata("fast_layers.0.attention.q_proj.weight"); + assets.model_weights->require_metadata("fast_embeddings.weight"); + assets.model_weights->require_metadata("fast_output.weight"); + assets.codec_weights->require_metadata("quantizer.semantic_quantizer.quantizers.0.codebook.weight"); + assets.codec_weights->require_metadata("decoder.model.0.conv.weight"); +} + +} // namespace + +std::shared_ptr load_fish_audio_assets(const std::filesystem::path & model_path) { + FishAudioAssets assets; + assets.resources = assets::load_resource_bundle_from_package_spec( + model_path, + assets::default_model_package_spec_path("fish_audio")); + assets.config = parse_config(assets.resources); + assets.model_weights = assets.resources.open_tensor_source("model_weights"); + assets.codec_weights = assets.resources.open_tensor_source("codec_weights"); + validate_weight_anchors(assets); + return std::make_shared(std::move(assets)); +} + +} // namespace engine::models::fish_audio diff --git a/src/models/fish_audio/codec.cpp b/src/models/fish_audio/codec.cpp new file mode 100644 index 00000000..0c22cf88 --- /dev/null +++ b/src/models/fish_audio/codec.cpp @@ -0,0 +1,1130 @@ +#include "engine/models/fish_audio/codec.h" + +#include "engine/framework/audio/conversion.h" +#include "engine/framework/audio/resampling.h" +#include "engine/framework/core/backend.h" +#include "engine/framework/core/backend_weight_store.h" +#include "engine/framework/debug/trace.h" +#include "engine/framework/core/execution_context.h" +#include "engine/framework/modules/activation_modules.h" +#include "engine/framework/modules/attention_modules.h" +#include "engine/framework/modules/conditioning_modules.h" +#include "engine/framework/modules/conv_modules.h" +#include "engine/framework/modules/linear_module.h" +#include "engine/framework/modules/lookup_modules.h" +#include "engine/framework/modules/norm_modules.h" +#include "engine/framework/modules/positional_modules.h" +#include "engine/framework/modules/primitive_modules.h" +#include "engine/framework/modules/streaming_conv_modules.h" +#include "engine/framework/modules/structural_modules.h" +#include "engine/framework/modules/weight_binding.h" + +#include "../common/constant_tensor_cache.h" + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::fish_audio { +namespace { + +namespace binding = engine::modules::binding; + +constexpr int64_t kCodecDim = 1024; +constexpr int64_t kCodecTransformerHeads = 16; +constexpr int64_t kCodecHeadDim = 64; +constexpr int64_t kCodecIntermediate = 3072; +constexpr int64_t kCodecTransformerLayers = 8; +constexpr float kCodecNormEps = 1.0e-5F; +constexpr float kConvNextNormEps = 1.0e-6F; +constexpr float kCodecRopeTheta = 10000.0F; + +struct GgmlContextDeleter { + void operator()(ggml_context * ctx) const noexcept { + if (ctx != nullptr) { + ggml_free(ctx); + } + } +}; + +struct GgmlGallocrDeleter { + void operator()(ggml_gallocr_t alloc) const noexcept { + if (alloc != nullptr) { + ggml_gallocr_free(alloc); + } + } +}; + +std::vector dims_vector(const core::TensorShape & shape) { + std::vector out; + out.reserve(shape.rank); + for (size_t i = 0; i < shape.rank; ++i) { + out.push_back(shape.dims[i]); + } + return out; +} + +std::vector prepare_codec_mono( + const runtime::AudioBuffer & audio, + int target_sample_rate_hz) { + auto mono = engine::audio::mixdown_interleaved_to_mono_average(audio.samples, audio.channels); + if (audio.sample_rate != target_sample_rate_hz) { + mono = engine::audio::resample_mono_torchaudio_sinc_hann( + mono, + audio.sample_rate, + target_sample_rate_hz); + } + return mono; +} + +struct CodecTransformerLayerWeights { + modules::NormWeights attention_norm; + modules::AttentionWeights attention; + modules::LayerScaleWeights attention_scale; + modules::NormWeights ffn_norm; + modules::QwenMLPWeights feed_forward; + modules::LayerScaleWeights ffn_scale; +}; + +struct CodecTransformerWeights { + std::vector layers; + modules::NormWeights norm; +}; + +struct ResidualUnitWeights { + modules::Snake1dWeights snake1; + modules::Conv1dWeights conv1; + modules::Snake1dWeights snake2; + modules::Conv1dWeights conv2; +}; + +struct EncoderBlockWeights { + ResidualUnitWeights residual1; + ResidualUnitWeights residual3; + ResidualUnitWeights residual9; + modules::Snake1dWeights snake; + modules::Conv1dWeights conv; + std::optional transformer; +}; + +struct DecoderBlockWeights { + modules::Snake1dWeights snake; + modules::ConvTranspose1dWeights conv; + ResidualUnitWeights residual1; + ResidualUnitWeights residual3; + ResidualUnitWeights residual9; +}; + +struct ConvNeXtBlockWeights { + modules::DepthwiseConv1dWeights dwconv; + modules::NormWeights norm; + modules::LinearWeights pwconv1; + modules::LinearWeights pwconv2; + modules::LayerScaleWeights gamma; +}; + +struct QuantizerUnitWeights { + modules::Conv1dWeights in_proj; + modules::Conv1dWeights out_proj; + core::TensorValue codebook; + core::TensorValue normalized_codebook; +}; + +struct FishCodecWeights { + std::shared_ptr store; + modules::Conv1dWeights encoder_first; + std::vector encoder_blocks; + modules::Snake1dWeights encoder_final_snake; + modules::Conv1dWeights encoder_final; + + std::vector> downsample; + CodecTransformerWeights pre_module; + QuantizerUnitWeights semantic_quantizer; + std::vector residual_quantizers; + CodecTransformerWeights post_module; + std::vector> upsample; + + modules::Conv1dWeights decoder_first; + std::vector decoder_blocks; + modules::Snake1dWeights decoder_final_snake; + modules::Conv1dWeights decoder_final; +}; + +int64_t ceil_div(int64_t a, int64_t b) { + return (a + b - 1) / b; +} + +std::vector normalized_rows(const std::vector & values, int64_t rows, int64_t cols) { + if (static_cast(values.size()) != rows * cols) { + throw std::runtime_error("Fish Audio normalized_rows shape mismatch"); + } + std::vector out(values.size(), 0.0F); + for (int64_t row = 0; row < rows; ++row) { + double sum = 0.0; + for (int64_t col = 0; col < cols; ++col) { + const float value = values[static_cast(row * cols + col)]; + sum += static_cast(value) * static_cast(value); + } + const float inv = sum > 0.0 ? static_cast(1.0 / std::sqrt(sum)) : 0.0F; + for (int64_t col = 0; col < cols; ++col) { + const size_t index = static_cast(row * cols + col); + out[index] = values[index] * inv; + } + } + return out; +} + +core::TensorValue slice_frames(core::ModuleBuildContext & ctx, const core::TensorValue & input, int64_t start, int64_t frames) { + if (frames <= 0) { + throw std::runtime_error("Fish Audio codec slice_frames requires positive frames"); + } + return modules::SliceModule({2, start, frames}).build(ctx, input); +} + +core::TensorValue zero_prefix_like(core::ModuleBuildContext & ctx, const core::TensorValue & input, int64_t frames) { + if (frames <= 0) { + throw std::runtime_error("Fish Audio zero_prefix_like requires positive frames"); + } + auto first = modules::SliceModule({2, 0, 1}).build(ctx, input); + first = core::wrap_tensor(ggml_scale(ctx.ggml, first.tensor, 0.0F), first.shape, GGML_TYPE_F32); + return modules::RepeatModule({core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], frames})}) + .build(ctx, first); +} + +core::TensorValue zero_suffix_like(core::ModuleBuildContext & ctx, const core::TensorValue & input, int64_t frames) { + if (frames <= 0) { + throw std::runtime_error("Fish Audio zero_suffix_like requires positive frames"); + } + auto last = modules::SliceModule({2, input.shape.dims[2] - 1, 1}).build(ctx, input); + last = core::wrap_tensor(ggml_scale(ctx.ggml, last.tensor, 0.0F), last.shape, GGML_TYPE_F32); + return modules::RepeatModule({core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], frames})}) + .build(ctx, last); +} + +int64_t extra_padding_for_conv1d(int64_t frames, int64_t effective_kernel, int64_t stride, int64_t left_pad) { + const double n_frames = (static_cast(frames - effective_kernel + left_pad) / static_cast(stride)) + 1.0; + const int64_t ideal_length = + (static_cast(std::ceil(n_frames)) - 1) * stride + (effective_kernel - left_pad); + return ideal_length - frames; +} + +core::TensorValue causal_pad(core::ModuleBuildContext & ctx, const core::TensorValue & input, int64_t left_pad, int64_t right_pad) { + if (left_pad < 0) { + throw std::runtime_error("Fish Audio causal conv requires non-negative left padding"); + } + if (right_pad < 0) { + throw std::runtime_error("Fish Audio causal conv requires non-negative right padding"); + } + if (left_pad == 0 && right_pad == 0) { + return input; + } + auto out = input; + if (left_pad > 0) { + out = modules::ConcatModule({2}).build(ctx, zero_prefix_like(ctx, input, left_pad), out); + } + if (right_pad > 0) { + out = modules::ConcatModule({2}).build(ctx, out, zero_suffix_like(ctx, input, right_pad)); + } + return out; +} + +core::TensorValue causal_conv1d( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::Conv1dWeights & weights, + int64_t in_channels, + int64_t out_channels, + int64_t kernel, + int stride, + int dilation, + bool use_bias) { + const int64_t effective_kernel = (kernel - 1) * dilation + 1; + const int64_t left_pad = effective_kernel - stride; + const int64_t right_pad = extra_padding_for_conv1d(input.shape.dims[2], effective_kernel, stride, left_pad); + auto padded = causal_pad(ctx, input, left_pad, right_pad); + return modules::Conv1dModule({ + in_channels, + out_channels, + kernel, + stride, + 0, + dilation, + use_bias, + }).build(ctx, padded, weights); +} + +core::TensorValue causal_depthwise_conv1d( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::DepthwiseConv1dWeights & weights, + int64_t channels, + int64_t kernel, + int stride, + int dilation, + bool use_bias) { + const int64_t effective_kernel = (kernel - 1) * dilation + 1; + const int64_t left_pad = effective_kernel - stride; + const int64_t right_pad = extra_padding_for_conv1d(input.shape.dims[2], effective_kernel, stride, left_pad); + auto padded = causal_pad(ctx, input, left_pad, right_pad); + return modules::DepthwiseConv1dModule({ + channels, + kernel, + stride, + 0, + dilation, + use_bias, + }).build(ctx, padded, weights); +} + +core::TensorValue causal_conv_transpose1d( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::ConvTranspose1dWeights & weights, + int64_t in_channels, + int64_t out_channels, + int64_t kernel, + int stride, + bool use_bias) { + auto out = modules::ConvTranspose1dModule({ + in_channels, + out_channels, + kernel, + stride, + 0, + 1, + use_bias, + }).build(ctx, input, weights); + const int64_t pad = kernel - stride; + const int64_t padding_right = static_cast(std::ceil(static_cast(pad))); + const int64_t padding_left = pad - padding_right; + return slice_frames(ctx, out, padding_left, out.shape.dims[2] - padding_left - padding_right); +} + +core::TensorValue l2_normalize_last(core::ModuleBuildContext & ctx, const core::TensorValue & input) { + auto squared = modules::MulModule{}.build(ctx, input, input); + auto sum = modules::ReduceSumModule({static_cast(input.shape.rank - 1)}).build(ctx, squared); + auto shifted = core::wrap_tensor(ggml_scale_bias(ctx.ggml, sum.tensor, 1.0F, 1.0e-12F), sum.shape, GGML_TYPE_F32); + auto denom = modules::SqrtModule{}.build(ctx, shifted); + auto repeated = modules::RepeatModule({input.shape}).build(ctx, denom); + return core::wrap_tensor(ggml_div(ctx.ggml, input.tensor, repeated.tensor), input.shape, GGML_TYPE_F32); +} + +core::TensorValue build_mlp( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const modules::QwenMLPWeights & weights) { + auto gate = modules::LinearModule({kCodecDim, kCodecIntermediate, false, GGML_PREC_F32}) + .build(ctx, input, weights.gate_proj); + gate = modules::SiluModule{}.build(ctx, gate); + auto up = modules::LinearModule({kCodecDim, kCodecIntermediate, false, GGML_PREC_F32}) + .build(ctx, input, weights.up_proj); + auto hidden = modules::MulModule{}.build(ctx, gate, up); + return modules::LinearModule({kCodecIntermediate, kCodecDim, false, GGML_PREC_F32}) + .build(ctx, hidden, weights.down_proj); +} + +core::TensorValue reshape_heads(core::ModuleBuildContext & ctx, const core::TensorValue & input) { + const auto contiguous = core::ensure_backend_addressable_layout(ctx, input); + return core::reshape_tensor( + ctx, + contiguous, + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], kCodecTransformerHeads, kCodecHeadDim})); +} + +core::TensorValue attention_from_heads( + core::ModuleBuildContext & ctx, + const core::TensorValue & q_heads, + const core::TensorValue & k_heads, + const core::TensorValue & v_heads, + const core::TensorValue & attention_mask) { + auto q = modules::TransposeModule({{0, 2, 1, 3}, q_heads.shape.rank}).build(ctx, q_heads); + auto k = modules::TransposeModule({{0, 2, 1, 3}, k_heads.shape.rank}).build(ctx, k_heads); + auto v = modules::TransposeModule({{0, 2, 1, 3}, v_heads.shape.rank}).build(ctx, v_heads); + q = core::wrap_tensor(ggml_cont(ctx.ggml, q.tensor), q.shape, q.type); + k = core::wrap_tensor(ggml_cont(ctx.ggml, k.tensor), k.shape, k.type); + v = core::wrap_tensor(ggml_cont(ctx.ggml, v.tensor), v.shape, v.type); + auto * flash = ggml_flash_attn_ext( + ctx.ggml, + q.tensor, + k.tensor, + v.tensor, + attention_mask.tensor, + 1.0F / std::sqrt(static_cast(kCodecHeadDim)), + 0.0F, + 0.0F); + ggml_flash_attn_ext_set_prec(flash, GGML_PREC_F32); + return core::wrap_tensor( + flash, + core::TensorShape::from_dims({q.shape.dims[0], q.shape.dims[2], q.shape.dims[1], kCodecHeadDim}), + GGML_TYPE_F32); +} + +core::TensorValue build_transformer_layer( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const core::TensorValue & positions, + const core::TensorValue & attention_mask, + const CodecTransformerLayerWeights & weights) { + auto normed = modules::RMSNormModule({kCodecDim, kCodecNormEps, true, false}).build(ctx, input, weights.attention_norm); + auto q = modules::LinearModule({kCodecDim, kCodecDim, false, GGML_PREC_F32}) + .build(ctx, normed, {weights.attention.q_weight, std::nullopt}); + auto k = modules::LinearModule({kCodecDim, kCodecDim, false, GGML_PREC_F32}) + .build(ctx, normed, {weights.attention.k_weight, std::nullopt}); + auto v = modules::LinearModule({kCodecDim, kCodecDim, false, GGML_PREC_F32}) + .build(ctx, normed, {weights.attention.v_weight, std::nullopt}); + q = modules::RoPEModule({kCodecHeadDim, GGML_ROPE_TYPE_NORMAL, kCodecRopeTheta}).build(ctx, reshape_heads(ctx, q), positions); + k = modules::RoPEModule({kCodecHeadDim, GGML_ROPE_TYPE_NORMAL, kCodecRopeTheta}).build(ctx, reshape_heads(ctx, k), positions); + v = reshape_heads(ctx, v); + auto context = attention_from_heads(ctx, q, k, v, attention_mask); + context = core::ensure_backend_addressable_layout(ctx, context); + context = core::reshape_tensor( + ctx, + context, + core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], kCodecDim})); + auto attn = modules::LinearModule({kCodecDim, kCodecDim, false, GGML_PREC_F32}) + .build(ctx, context, {weights.attention.out_weight, std::nullopt}); + attn = modules::LayerScaleModule{}.build(ctx, attn, weights.attention_scale); + auto hidden = modules::AddModule{}.build(ctx, input, attn); + auto ffn_in = modules::RMSNormModule({kCodecDim, kCodecNormEps, true, false}).build(ctx, hidden, weights.ffn_norm); + auto ff = build_mlp(ctx, ffn_in, weights.feed_forward); + ff = modules::LayerScaleModule{}.build(ctx, ff, weights.ffn_scale); + return modules::AddModule{}.build(ctx, hidden, ff); +} + +core::TensorValue make_positions( + core::ModuleBuildContext &, + common::ConstantTensorCache & constants, + int64_t frames) { + std::vector values(static_cast(frames)); + for (int64_t i = 0; i < frames; ++i) { + values[static_cast(i)] = static_cast(i); + } + return constants.make_tensor(core::TensorShape::from_dims({frames}), GGML_TYPE_I32, values.data(), values.size() * sizeof(int32_t)); +} + +core::TensorValue make_causal_mask( + core::ModuleBuildContext &, + common::ConstantTensorCache & constants, + int64_t frames, + int64_t window_size) { + std::vector values(static_cast(frames * frames), ggml_fp32_to_fp16(-std::numeric_limits::infinity())); + for (int64_t row = 0; row < frames; ++row) { + const int64_t begin = window_size > 0 ? std::max(0, row - window_size + 1) : 0; + for (int64_t col = begin; col <= row; ++col) { + values[static_cast(row * frames + col)] = ggml_fp32_to_fp16(0.0F); + } + } + return constants.make_tensor(core::TensorShape::from_dims({frames, frames}), GGML_TYPE_F16, values.data(), values.size() * sizeof(ggml_fp16_t)); +} + +core::TensorValue build_window_transformer( + core::ModuleBuildContext & ctx, + common::ConstantTensorCache & constants, + const core::TensorValue & input_bct, + const CodecTransformerWeights & weights, + int64_t window_size) { + auto x = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, input_bct); + auto positions = make_positions(ctx, constants, x.shape.dims[1]); + auto mask = make_causal_mask(ctx, constants, x.shape.dims[1], window_size); + for (const auto & layer : weights.layers) { + x = build_transformer_layer(ctx, x, positions, mask, layer); + } + x = modules::RMSNormModule({kCodecDim, kCodecNormEps, true, false}).build(ctx, x, weights.norm); + return modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, x); +} + +core::TensorValue build_residual_unit( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const ResidualUnitWeights & weights, + int64_t channels, + int dilation) { + auto y = modules::Snake1dModule({channels}).build(ctx, input, weights.snake1); + y = causal_conv1d(ctx, y, weights.conv1, channels, channels, 7, 1, dilation, true); + y = modules::Snake1dModule({channels}).build(ctx, y, weights.snake2); + y = causal_conv1d(ctx, y, weights.conv2, channels, channels, 1, 1, 1, true); + core::TensorValue x = input; + if (x.shape.dims[2] != y.shape.dims[2]) { + x = slice_frames(ctx, x, 0, y.shape.dims[2]); + } + return modules::AddModule{}.build(ctx, x, y); +} + +core::TensorValue build_convnext( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const ConvNeXtBlockWeights & weights, + int64_t channels) { + auto y = causal_depthwise_conv1d(ctx, input, weights.dwconv, channels, 7, 1, 1, true); + y = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, y); + y = modules::LayerNormModule({channels, kConvNextNormEps, true, true}).build(ctx, y, weights.norm); + y = modules::LinearModule({channels, channels * 4, true, GGML_PREC_F32}).build(ctx, y, weights.pwconv1); + y = modules::GeluModule({modules::GeluApproximation::ExactErf}).build(ctx, y); + y = modules::LinearModule({channels * 4, channels, true, GGML_PREC_F32}).build(ctx, y, weights.pwconv2); + y = modules::LayerScaleModule{}.build(ctx, y, weights.gamma); + y = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, y); + core::TensorValue x = input; + if (x.shape.dims[2] != y.shape.dims[2]) { + x = slice_frames(ctx, x, 0, y.shape.dims[2]); + } + return modules::AddModule{}.build(ctx, x, y); +} + +core::TensorValue build_encoder( + core::ModuleBuildContext & ctx, + common::ConstantTensorCache & constants, + const core::TensorValue & input, + const FishCodecWeights & weights) { + auto x = causal_conv1d(ctx, input, weights.encoder_first, 1, 64, 7, 1, 1, true); + int64_t channels = 64; + const int strides[] = {2, 4, 8, 8}; + for (size_t index = 0; index < weights.encoder_blocks.size(); ++index) { + const auto & block = weights.encoder_blocks[index]; + x = build_residual_unit(ctx, x, block.residual1, channels, 1); + x = build_residual_unit(ctx, x, block.residual3, channels, 3); + x = build_residual_unit(ctx, x, block.residual9, channels, 9); + x = modules::Snake1dModule({channels}).build(ctx, x, block.snake); + x = causal_conv1d(ctx, x, block.conv, channels, channels * 2, 2 * strides[index], strides[index], 1, true); + channels *= 2; + if (block.transformer.has_value()) { + x = build_window_transformer(ctx, constants, x, *block.transformer, 512); + } + } + x = modules::Snake1dModule({channels}).build(ctx, x, weights.encoder_final_snake); + return causal_conv1d(ctx, x, weights.encoder_final, channels, kCodecDim, 3, 1, 1, true); +} + +core::TensorValue build_decoder( + core::ModuleBuildContext & ctx, + const core::TensorValue & input, + const FishCodecWeights & weights) { + auto x = causal_conv1d(ctx, input, weights.decoder_first, kCodecDim, 1536, 7, 1, 1, true); + int64_t channels = 1536; + const int strides[] = {8, 8, 4, 2}; + for (size_t index = 0; index < weights.decoder_blocks.size(); ++index) { + const auto & block = weights.decoder_blocks[index]; + x = modules::Snake1dModule({channels}).build(ctx, x, block.snake); + x = causal_conv_transpose1d(ctx, x, block.conv, channels, channels / 2, 2 * strides[index], strides[index], true); + channels /= 2; + x = build_residual_unit(ctx, x, block.residual1, channels, 1); + x = build_residual_unit(ctx, x, block.residual3, channels, 3); + x = build_residual_unit(ctx, x, block.residual9, channels, 9); + } + x = modules::Snake1dModule({channels}).build(ctx, x, weights.decoder_final_snake); + x = causal_conv1d(ctx, x, weights.decoder_final, channels, 1, 7, 1, 1, true); + return modules::TanhModule{}.build(ctx, x); +} + +core::TensorValue build_quantizer_out( + core::ModuleBuildContext & ctx, + const core::TensorValue & ids_bt, + const QuantizerUnitWeights & weights, + int64_t codebook_size) { + auto emb_btd = modules::CodebookLookupModule({codebook_size, 8}).build(ctx, ids_bt, weights.codebook); + auto emb_bdt = modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, emb_btd); + return modules::Conv1dModule({8, kCodecDim, 1, 1, 0, 1, true}).build(ctx, emb_bdt, weights.out_proj); +} + +core::TensorValue build_decode_quantizer( + core::ModuleBuildContext & ctx, + common::ConstantTensorCache & constants, + const std::vector & code_inputs, + const FishCodecWeights & weights) { + auto latent = build_quantizer_out(ctx, code_inputs[0], weights.semantic_quantizer, 4096); + for (size_t index = 0; index < weights.residual_quantizers.size(); ++index) { + auto residual = build_quantizer_out(ctx, code_inputs[index + 1], weights.residual_quantizers[index], 1024); + latent = modules::AddModule{}.build(ctx, latent, residual); + } + latent = build_window_transformer(ctx, constants, latent, weights.post_module, 128); + for (const auto & stage : weights.upsample) { + latent = causal_conv_transpose1d(ctx, latent, stage.first, kCodecDim, kCodecDim, 2, 2, true); + latent = build_convnext(ctx, latent, stage.second, kCodecDim); + } + return latent; +} + +core::TensorValue build_encode_quantizer( + core::ModuleBuildContext & ctx, + common::ConstantTensorCache & constants, + const core::TensorValue & encoder_latent, + const FishCodecWeights & weights, + std::vector & code_outputs, + std::vector> & trace_outputs) { + auto x = encoder_latent; + for (const auto & stage : weights.downsample) { + x = causal_conv1d(ctx, x, stage.first, kCodecDim, kCodecDim, 2, 2, 1, true); + x = build_convnext(ctx, x, stage.second, kCodecDim); + } + trace_outputs.push_back({"fish_audio.codec.after_downsample", x}); + x = build_window_transformer(ctx, constants, x, weights.pre_module, 128); + trace_outputs.push_back({"fish_audio.codec.after_pre_module", x}); + + auto residual = x; + auto quantize_one = [&](const QuantizerUnitWeights & quantizer, int64_t codebook_size) { + auto projected = modules::Conv1dModule({kCodecDim, 8, 1, 1, 0, 1, true}).build(ctx, residual, quantizer.in_proj); + auto projected_btd = l2_normalize_last(ctx, modules::TransposeModule({{0, 2, 1, 3}, 3}).build(ctx, projected)); + auto logits = modules::LinearModule({8, codebook_size, false, GGML_PREC_F32}) + .build(ctx, projected_btd, {quantizer.normalized_codebook, std::nullopt}); + auto flat_logits = core::reshape_tensor( + ctx, + core::ensure_backend_addressable_layout(ctx, logits), + core::TensorShape::from_dims({logits.shape.dims[1], codebook_size})); + auto * ids_raw = ggml_argmax(ctx.ggml, flat_logits.tensor); + ggml_set_output(ids_raw); + code_outputs.push_back(ids_raw); + auto ids = core::reshape_tensor( + ctx, + core::wrap_tensor(ids_raw, core::TensorShape::from_dims({logits.shape.dims[1]}), GGML_TYPE_I32), + core::TensorShape::from_dims({1, logits.shape.dims[1]})); + auto quantized = build_quantizer_out(ctx, ids, quantizer, codebook_size); + residual = core::wrap_tensor(ggml_sub(ctx.ggml, residual.tensor, quantized.tensor), residual.shape, GGML_TYPE_F32); + }; + quantize_one(weights.semantic_quantizer, 4096); + for (const auto & quantizer : weights.residual_quantizers) { + quantize_one(quantizer, 1024); + } + return x; +} + +modules::Snake1dWeights load_snake(core::BackendWeightStore & store, const assets::TensorSource & source, const std::string & name, int64_t channels) { + return {store.make_f32( + core::TensorShape::from_dims({channels}), + source.require_f32(name + ".alpha", {1, channels, 1}))}; +} + +CodecTransformerWeights load_transformer( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t layers) { + CodecTransformerWeights out; + out.layers.reserve(static_cast(layers)); + for (int64_t layer = 0; layer < layers; ++layer) { + const std::string layer_prefix = prefix + ".layers." + std::to_string(layer); + CodecTransformerLayerWeights weights; + weights.attention_norm = binding::norm_weight_from_source(store, source, layer_prefix + ".attention_norm", kCodecDim); + weights.attention.q_weight = store.load_tensor(source, layer_prefix + ".attention.q_proj.weight", storage_type, {kCodecDim, kCodecDim}); + weights.attention.k_weight = store.load_tensor(source, layer_prefix + ".attention.k_proj.weight", storage_type, {kCodecDim, kCodecDim}); + weights.attention.v_weight = store.load_tensor(source, layer_prefix + ".attention.v_proj.weight", storage_type, {kCodecDim, kCodecDim}); + weights.attention.out_weight = store.load_tensor(source, layer_prefix + ".attention.wo.weight", storage_type, {kCodecDim, kCodecDim}); + weights.attention_scale = binding::layer_scale_from_named_source(store, source, layer_prefix + ".attention_layer_scale.gamma"); + weights.ffn_norm = binding::norm_weight_from_source(store, source, layer_prefix + ".ffn_norm", kCodecDim); + weights.feed_forward.gate_proj.weight = store.load_tensor(source, layer_prefix + ".feed_forward.w1.weight", storage_type, {kCodecIntermediate, kCodecDim}); + weights.feed_forward.down_proj.weight = store.load_tensor(source, layer_prefix + ".feed_forward.w2.weight", storage_type, {kCodecDim, kCodecIntermediate}); + weights.feed_forward.up_proj.weight = store.load_tensor(source, layer_prefix + ".feed_forward.w3.weight", storage_type, {kCodecIntermediate, kCodecDim}); + weights.ffn_scale = binding::layer_scale_from_named_source(store, source, layer_prefix + ".ffn_layer_scale.gamma"); + out.layers.push_back(std::move(weights)); + } + out.norm = binding::norm_weight_from_source(store, source, prefix + ".norm", kCodecDim); + return out; +} + +ResidualUnitWeights load_residual_unit( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t channels) { + ResidualUnitWeights out; + out.snake1 = load_snake(store, source, prefix + ".block.0", channels); + out.conv1 = binding::conv1d_from_named_source( + store, + source, + prefix + ".block.1.conv.weight", + prefix + ".block.1.conv.bias", + storage_type); + out.snake2 = load_snake(store, source, prefix + ".block.2", channels); + out.conv2 = binding::conv1d_from_named_source( + store, + source, + prefix + ".block.3.conv.weight", + prefix + ".block.3.conv.bias", + storage_type); + return out; +} + +ConvNeXtBlockWeights load_convnext( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t channels) { + ConvNeXtBlockWeights out; + out.dwconv = binding::depthwise_conv1d_from_source( + store, + source, + prefix + ".dwconv.conv", + storage_type, + channels, + 7, + true); + out.norm = binding::norm_from_source(store, source, prefix + ".norm", channels); + out.pwconv1 = binding::linear_from_source(store, source, prefix + ".pwconv1", storage_type, channels * 4, channels, true); + out.pwconv2 = binding::linear_from_source(store, source, prefix + ".pwconv2", storage_type, channels, channels * 4, true); + out.gamma = binding::layer_scale_from_named_source(store, source, prefix + ".gamma"); + return out; +} + +QuantizerUnitWeights load_quantizer_unit( + core::BackendWeightStore & store, + const assets::TensorSource & source, + const std::string & prefix, + assets::TensorStorageType storage_type, + int64_t codebook_size) { + QuantizerUnitWeights out; + out.in_proj = binding::conv1d_from_named_source( + store, + source, + prefix + ".in_proj.weight", + prefix + ".in_proj.bias", + storage_type); + out.out_proj = binding::conv1d_from_named_source( + store, + source, + prefix + ".out_proj.weight", + prefix + ".out_proj.bias", + storage_type); + const auto codebook = source.require_f32(prefix + ".codebook.weight", {codebook_size, 8}); + out.codebook = store.make_from_f32(core::TensorShape::from_dims({codebook_size, 8}), storage_type, codebook); + out.normalized_codebook = store.make_from_f32( + core::TensorShape::from_dims({codebook_size, 8}), + storage_type, + normalized_rows(codebook, codebook_size, 8)); + return out; +} + +std::shared_ptr load_weights( + const FishAudioAssets & assets, + ggml_backend_t backend, + core::BackendType backend_type, + size_t weight_context_bytes, + assets::TensorStorageType matmul_storage_type, + assets::TensorStorageType conv_storage_type) { + auto weights = std::make_shared(); + weights->store = std::make_shared(backend, backend_type, "Fish Audio codec", weight_context_bytes); + auto & store = *weights->store; + const auto & source = *assets.codec_weights; + + weights->encoder_first = binding::conv1d_from_named_source( + store, + source, + "encoder.block.0.conv.weight", + "encoder.block.0.conv.bias", + conv_storage_type); + int64_t encoder_channels = 64; + for (int64_t block_index = 0; block_index < 4; ++block_index) { + const std::string prefix = "encoder.block." + std::to_string(block_index + 1) + ".block"; + EncoderBlockWeights block; + block.residual1 = load_residual_unit(store, source, prefix + ".0", conv_storage_type, encoder_channels); + block.residual3 = load_residual_unit(store, source, prefix + ".1", conv_storage_type, encoder_channels); + block.residual9 = load_residual_unit(store, source, prefix + ".2", conv_storage_type, encoder_channels); + block.snake = load_snake(store, source, prefix + ".3", encoder_channels); + block.conv = binding::conv1d_from_named_source( + store, + source, + prefix + ".4.conv.weight", + prefix + ".4.conv.bias", + conv_storage_type); + encoder_channels *= 2; + if (block_index == 3) { + block.transformer = load_transformer(store, source, prefix + ".5", matmul_storage_type, 4); + } + weights->encoder_blocks.push_back(std::move(block)); + } + weights->encoder_final_snake = load_snake(store, source, "encoder.block.5", kCodecDim); + weights->encoder_final = binding::conv1d_from_named_source( + store, + source, + "encoder.block.6.conv.weight", + "encoder.block.6.conv.bias", + conv_storage_type); + + for (int64_t i = 0; i < 2; ++i) { + const std::string prefix = "quantizer.downsample." + std::to_string(i); + weights->downsample.push_back({ + binding::conv1d_from_named_source( + store, + source, + prefix + ".0.conv.weight", + prefix + ".0.conv.bias", + conv_storage_type), + load_convnext(store, source, prefix + ".1", matmul_storage_type, kCodecDim), + }); + } + weights->pre_module = load_transformer(store, source, "quantizer.pre_module", matmul_storage_type, kCodecTransformerLayers); + weights->semantic_quantizer = load_quantizer_unit(store, source, "quantizer.semantic_quantizer.quantizers.0", matmul_storage_type, 4096); + for (int64_t i = 0; i < assets.config.codec.quantizer_codebooks; ++i) { + weights->residual_quantizers.push_back( + load_quantizer_unit(store, source, "quantizer.quantizer.quantizers." + std::to_string(i), matmul_storage_type, 1024)); + } + weights->post_module = load_transformer(store, source, "quantizer.post_module", matmul_storage_type, kCodecTransformerLayers); + for (int64_t i = 0; i < 2; ++i) { + const std::string prefix = "quantizer.upsample." + std::to_string(i); + weights->upsample.push_back({ + binding::conv_transpose1d_from_named_source( + store, + source, + prefix + ".0.conv.weight", + prefix + ".0.conv.bias", + conv_storage_type), + load_convnext(store, source, prefix + ".1", matmul_storage_type, kCodecDim), + }); + } + + weights->decoder_first = binding::conv1d_from_named_source( + store, + source, + "decoder.model.0.conv.weight", + "decoder.model.0.conv.bias", + conv_storage_type); + int64_t decoder_channels = 1536; + for (int64_t block_index = 0; block_index < 4; ++block_index) { + const std::string prefix = "decoder.model." + std::to_string(block_index + 1) + ".block"; + DecoderBlockWeights block; + block.snake = load_snake(store, source, prefix + ".0", decoder_channels); + block.conv = binding::conv_transpose1d_from_named_source( + store, + source, + prefix + ".1.conv.weight", + prefix + ".1.conv.bias", + conv_storage_type); + decoder_channels /= 2; + block.residual1 = load_residual_unit(store, source, prefix + ".2", conv_storage_type, decoder_channels); + block.residual3 = load_residual_unit(store, source, prefix + ".3", conv_storage_type, decoder_channels); + block.residual9 = load_residual_unit(store, source, prefix + ".4", conv_storage_type, decoder_channels); + weights->decoder_blocks.push_back(std::move(block)); + } + weights->decoder_final_snake = load_snake(store, source, "decoder.model.5", 96); + weights->decoder_final = binding::conv1d_from_named_source( + store, + source, + "decoder.model.6.conv.weight", + "decoder.model.6.conv.bias", + conv_storage_type); + + store.upload(); + return weights; +} + +struct DecodeGraph { + DecodeGraph( + std::shared_ptr assets, + std::shared_ptr weights, + core::ExecutionContext & execution_context, + size_t graph_arena_bytes, + int64_t frames) + : assets_(std::move(assets)), + weights_(std::move(weights)), + backend_(execution_context.backend()), + backend_type_(execution_context.backend_type()), + threads_(std::max(1, execution_context.config().threads)), + frame_capacity_(frames), + constants_(backend_, threads_, "Fish Audio codec decode constants") { + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio codec decode graph context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "fish_audio.codec.decode", backend_type_}; + constants_.begin_graph(); + for (int64_t codebook = 0; codebook < assets_->config.codec.total_codebooks; ++codebook) { + auto ids = core::make_tensor(ctx, GGML_TYPE_I32, core::TensorShape::from_dims({1, frame_capacity_})); + ggml_set_input(ids.tensor); + code_inputs_.push_back(ids); + } + auto latent = build_decode_quantizer(ctx, constants_, code_inputs_, *weights_); + auto waveform = build_decoder(ctx, latent, *weights_); + output_ = waveform.tensor; + ggml_set_output(output_); + graph_ = ggml_new_graph_custom(ctx_.get(), 1048576, false); + ggml_build_forward_expand(graph_, output_); + constants_.finish_graph(); + constants_.ensure_uploaded(); + gallocr_.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_))); + if (gallocr_ == nullptr || !ggml_gallocr_alloc_graph(gallocr_.get(), graph_)) { + throw std::runtime_error("failed to allocate Fish Audio codec decode graph"); + } + } + + ~DecodeGraph() { + engine::core::release_backend_graph_resources(backend_, graph_); + } + + bool matches(int64_t frames, ggml_backend_t backend, int threads) const { + return frame_capacity_ >= frames && backend_ == backend && threads_ == std::max(1, threads); + } + + runtime::AudioBuffer run(const FishAudioCodes & codes) { + const int64_t codebooks = assets_->config.codec.total_codebooks; + if (codes.codebooks != codebooks || codes.frames <= 0 || + static_cast(codes.codes.size()) != codebooks * codes.frames) { + std::ostringstream oss; + oss << "Fish Audio codec decode code shape mismatch: expected_codebooks=" << codebooks + << " actual_codebooks=" << codes.codebooks + << " frames=" << codes.frames + << " values=" << codes.codes.size() + << " expected_values=" << (codebooks * codes.frames); + throw std::runtime_error(oss.str()); + } + if (codes.frames > frame_capacity_) { + throw std::runtime_error("Fish Audio codec decode request exceeds graph capacity"); + } + for (int64_t codebook = 0; codebook < codebooks; ++codebook) { + std::vector padded(static_cast(frame_capacity_), 0); + for (int64_t frame = 0; frame < codes.frames; ++frame) { + int32_t value = codes.codes[static_cast(codebook * codes.frames + frame)]; + if (codebook == 0) { + value = std::clamp(value, 0, 4095); + } else { + value = std::clamp(value, 0, 1023); + } + padded[static_cast(frame)] = value; + } + core::write_tensor_i32(code_inputs_[static_cast(codebook)], padded); + } + core::set_backend_threads(backend_, threads_); + const ggml_status status = engine::core::compute_backend_graph(backend_, graph_); + ggml_backend_synchronize(backend_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Fish Audio codec decode graph compute failed"); + } + auto values = core::read_tensor_f32(output_); + const int64_t expected_samples = codes.frames * assets_->config.codec.frame_length; + if (static_cast(values.size()) > expected_samples) { + values.resize(static_cast(expected_samples)); + } + return runtime::AudioBuffer{assets_->config.codec.sample_rate, 1, std::move(values)}; + } + +private: + std::shared_ptr assets_; + std::shared_ptr weights_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + int64_t frame_capacity_ = 0; + std::unique_ptr ctx_; + std::vector code_inputs_; + ggml_tensor * output_ = nullptr; + ggml_cgraph * graph_ = nullptr; + std::unique_ptr, GgmlGallocrDeleter> gallocr_; + common::ConstantTensorCache constants_; +}; + +struct EncodeGraph { + EncodeGraph( + std::shared_ptr assets, + std::shared_ptr weights, + core::ExecutionContext & execution_context, + size_t graph_arena_bytes, + int64_t samples, + int64_t frames) + : assets_(std::move(assets)), + weights_(std::move(weights)), + backend_(execution_context.backend()), + backend_type_(execution_context.backend_type()), + threads_(std::max(1, execution_context.config().threads)), + sample_capacity_(samples), + frame_capacity_(frames), + constants_(backend_, threads_, "Fish Audio codec encode constants") { + ggml_init_params params{graph_arena_bytes, nullptr, true}; + ctx_.reset(ggml_init(params)); + if (ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio codec encode graph context"); + } + core::ModuleBuildContext ctx{ctx_.get(), "fish_audio.codec.encode", backend_type_}; + constants_.begin_graph(); + input_ = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, sample_capacity_})); + ggml_set_input(input_.tensor); + auto encoded = build_encoder(ctx, constants_, input_, *weights_); + trace_outputs_.push_back({"fish_audio.codec.encoder_latent", encoded}); + build_encode_quantizer(ctx, constants_, encoded, *weights_, code_outputs_, trace_outputs_); + graph_ = ggml_new_graph_custom(ctx_.get(), 1048576, false); + for (const auto & trace_output : trace_outputs_) { + ggml_set_output(trace_output.second.tensor); + ggml_build_forward_expand(graph_, trace_output.second.tensor); + } + for (ggml_tensor * code_output : code_outputs_) { + ggml_build_forward_expand(graph_, code_output); + } + constants_.finish_graph(); + constants_.ensure_uploaded(); + gallocr_.reset(ggml_gallocr_new(ggml_backend_get_default_buffer_type(backend_))); + if (gallocr_ == nullptr || !ggml_gallocr_alloc_graph(gallocr_.get(), graph_)) { + throw std::runtime_error("failed to allocate Fish Audio codec encode graph"); + } + } + + ~EncodeGraph() { + engine::core::release_backend_graph_resources(backend_, graph_); + } + + bool matches(int64_t samples, int64_t frames, ggml_backend_t backend, int threads) const { + return sample_capacity_ >= samples && + frame_capacity_ >= frames && + backend_ == backend && + threads_ == std::max(1, threads); + } + + FishAudioCodes run(const runtime::AudioBuffer & audio) { + auto mono = prepare_codec_mono(audio, assets_->config.codec.sample_rate); + const int64_t original_samples = static_cast(mono.size()); + const int64_t padded_samples = ceil_div(original_samples, assets_->config.codec.frame_length) * assets_->config.codec.frame_length; + const int64_t frames = ceil_div(original_samples, assets_->config.codec.frame_length); + if (padded_samples > sample_capacity_ || frames > frame_capacity_) { + throw std::runtime_error("Fish Audio codec encode request exceeds graph capacity"); + } + mono.resize(static_cast(sample_capacity_), 0.0F); + core::write_tensor_f32(input_, mono); + core::set_backend_threads(backend_, threads_); + const ggml_status status = engine::core::compute_backend_graph(backend_, graph_); + ggml_backend_synchronize(backend_); + if (status != GGML_STATUS_SUCCESS) { + throw std::runtime_error("Fish Audio codec encode graph compute failed"); + } + if (engine::debug::trace_log_enabled()) { + for (const auto & trace_output : trace_outputs_) { + engine::debug::trace_log_f32( + trace_output.first, + dims_vector(trace_output.second.shape), + core::read_tensor_f32(trace_output.second.tensor)); + } + } + FishAudioCodes out; + out.codebooks = static_cast(code_outputs_.size()); + out.frames = frames; + out.codes.resize(static_cast(out.codebooks * out.frames)); + for (int64_t codebook = 0; codebook < out.codebooks; ++codebook) { + auto values = core::read_tensor_i32(code_outputs_[static_cast(codebook)]); + for (int64_t frame = 0; frame < out.frames; ++frame) { + out.codes[static_cast(codebook * out.frames + frame)] = values[static_cast(frame)]; + } + } + engine::debug::trace_log_i32( + "fish_audio.codec.reference_codes", + {out.codebooks, out.frames}, + out.codes); + return out; + } + +private: + std::shared_ptr assets_; + std::shared_ptr weights_; + ggml_backend_t backend_ = nullptr; + core::BackendType backend_type_ = core::BackendType::Cpu; + int threads_ = 1; + int64_t sample_capacity_ = 0; + int64_t frame_capacity_ = 0; + std::unique_ptr ctx_; + core::TensorValue input_; + std::vector code_outputs_; + std::vector> trace_outputs_; + ggml_cgraph * graph_ = nullptr; + std::unique_ptr, GgmlGallocrDeleter> gallocr_; + common::ConstantTensorCache constants_; +}; + +} // namespace + +class FishAudioCodecRuntime::Impl { +public: + Impl( + std::shared_ptr assets, + core::BackendConfig backend, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType matmul_weight_storage_type, + assets::TensorStorageType conv_weight_storage_type) + : assets_(std::move(assets)), + execution_(std::move(backend)), + threads_(std::max(1, threads)), + graph_arena_bytes_(graph_arena_bytes) { + weights_ = load_weights( + *assets_, + execution_.backend(), + execution_.backend_type(), + weight_context_bytes, + matmul_weight_storage_type, + conv_weight_storage_type); + } + + FishAudioCodes encode_reference(const runtime::AudioBuffer & audio) { + auto mono = prepare_codec_mono(audio, assets_->config.codec.sample_rate); + const int64_t samples = ceil_div(static_cast(mono.size()), assets_->config.codec.frame_length) * + assets_->config.codec.frame_length; + const int64_t frames = ceil_div(static_cast(mono.size()), assets_->config.codec.frame_length); + if (encode_graph_ == nullptr || !encode_graph_->matches(samples, frames, execution_.backend(), threads_)) { + encode_graph_ = std::make_unique(assets_, weights_, execution_, graph_arena_bytes_, samples, frames); + } + return encode_graph_->run(audio); + } + + runtime::AudioBuffer decode(const FishAudioCodes & codes) { + if (decode_graph_ == nullptr || !decode_graph_->matches(codes.frames, execution_.backend(), threads_)) { + decode_graph_ = std::make_unique(assets_, weights_, execution_, graph_arena_bytes_, codes.frames); + } + return decode_graph_->run(codes); + } + + void release_runtime_graphs() { + encode_graph_.reset(); + decode_graph_.reset(); + } + +private: + std::shared_ptr assets_; + core::ExecutionContext execution_; + int threads_ = 1; + size_t graph_arena_bytes_ = 0; + std::shared_ptr weights_; + std::unique_ptr encode_graph_; + std::unique_ptr decode_graph_; +}; + +FishAudioCodecRuntime::FishAudioCodecRuntime( + std::shared_ptr assets, + core::BackendConfig backend, + int threads, + size_t graph_arena_bytes, + size_t weight_context_bytes, + assets::TensorStorageType matmul_weight_storage_type, + assets::TensorStorageType conv_weight_storage_type) + : impl_(std::make_unique( + std::move(assets), + std::move(backend), + threads, + graph_arena_bytes, + weight_context_bytes, + matmul_weight_storage_type, + conv_weight_storage_type)) {} + +FishAudioCodecRuntime::~FishAudioCodecRuntime() = default; + +FishAudioCodes FishAudioCodecRuntime::encode_reference(const runtime::AudioBuffer & audio) { + return impl_->encode_reference(audio); +} + +runtime::AudioBuffer FishAudioCodecRuntime::decode(const FishAudioCodes & codes) { + return impl_->decode(codes); +} + +void FishAudioCodecRuntime::release_runtime_graphs() { + impl_->release_runtime_graphs(); +} + +} // namespace engine::models::fish_audio diff --git a/src/models/fish_audio/generator.cpp b/src/models/fish_audio/generator.cpp new file mode 100644 index 00000000..9813c44a --- /dev/null +++ b/src/models/fish_audio/generator.cpp @@ -0,0 +1,69 @@ +#include "engine/models/fish_audio/generator.h" + +#include "engine/framework/debug/profiler.h" + +#include +#include +#include + +namespace engine::models::fish_audio { +namespace { + +using Clock = std::chrono::steady_clock; + +} // namespace + +FishAudioGenerator::FishAudioGenerator( + std::shared_ptr assets, + std::unique_ptr ar, + std::unique_ptr codec) + : assets_(std::move(assets)), + tokenizer_(assets_), + prompt_builder_(assets_, tokenizer_), + ar_(std::move(ar)), + codec_(std::move(codec)) { + if (assets_ == nullptr || ar_ == nullptr || codec_ == nullptr) { + throw std::runtime_error("Fish Audio generator requires assets, AR runtime, and codec runtime"); + } +} + +FishAudioGenerator::~FishAudioGenerator() = default; + +FishAudioCodes FishAudioGenerator::encode_reference(const runtime::AudioBuffer & audio) { + return codec_->encode_reference(audio); +} + +FishAudioGenerationResult FishAudioGenerator::generate( + const FishAudioRequest & request, + const std::optional & reference_codes, + bool mem_saver) { + engine::debug::trace_log_scalar("fish_audio.request.has_reference", request.reference.has_value()); + engine::debug::trace_log_scalar("fish_audio.request.text_chars", static_cast(request.text.size())); + const auto prompt_start = Clock::now(); + const auto prompt = prompt_builder_.build(request, reference_codes); + engine::debug::timing_log_scalar( + "fish_audio.prompt_build_ms", + engine::debug::elapsed_ms(prompt_start, Clock::now())); + + const auto ar_start = Clock::now(); + FishAudioGenerationResult result; + result.codes = ar_->generate(prompt, request.generation); + engine::debug::trace_log_scalar("fish_audio.generated.frames", result.codes.frames); + engine::debug::trace_log_scalar("fish_audio.generated.codebooks", result.codes.codebooks); + engine::debug::timing_log_scalar( + "fish_audio.ar_generate_ms", + engine::debug::elapsed_ms(ar_start, Clock::now())); + + const auto decode_start = Clock::now(); + result.audio = codec_->decode(result.codes); + engine::debug::timing_log_scalar( + "fish_audio.codec_decode_ms", + engine::debug::elapsed_ms(decode_start, Clock::now())); + if (mem_saver) { + ar_->release_runtime_graphs(); + codec_->release_runtime_graphs(); + } + return result; +} + +} // namespace engine::models::fish_audio diff --git a/src/models/fish_audio/loader.cpp b/src/models/fish_audio/loader.cpp new file mode 100644 index 00000000..e7baf889 --- /dev/null +++ b/src/models/fish_audio/loader.cpp @@ -0,0 +1,135 @@ +#include "engine/models/fish_audio/loader.h" + +#include "engine/framework/assets/model_package.h" +#include "engine/models/fish_audio/session.h" + +#include +#include + +namespace engine::models::fish_audio { +namespace { + +runtime::ModelMetadata metadata(const FishAudioAssets &) { + runtime::ModelMetadata out; + out.family = "fish_audio"; + out.variant = "s2-pro"; + out.description = "Fish Audio S2-Pro loaded from prepared local assets."; + out.config_candidates = {"config.json", "tokenizer_config.json", "tokenizer.json"}; + out.weight_candidates = {"model_audio_cpp.safetensors.index.json", "codec.safetensors", "model.gguf"}; + return out; +} + +runtime::CapabilitySet capabilities(const FishAudioAssets &) { + runtime::CapabilitySet out; + out.supported_tasks = { + {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, + }; + out.languages = {"en", "zh", "auto"}; + out.supports_speaker_reference = true; + out.supports_style_condition = true; + return out; +} + +runtime::ModelCliInterface cli(const FishAudioAssets &) { + runtime::ModelCliInterface out; + out.request_options = { + {"reference_text", "TEXT", "Reference transcript used with speaker reference audio."}, + {"max_new_tokens", "N", "Maximum Fish Audio semantic tokens to generate."}, + {"chunk_length", "N", "Maximum UTF-8 bytes per Fish Speech text batch."}, + {"top_p", "FLOAT", "Top-p sampling value."}, + {"top_k", "N", "Top-k sampling value."}, + {"temperature", "FLOAT", "Sampling temperature."}, + {"seed", "N", "Sampling seed."}, + }; + out.session_options = { + {"fish_audio.mem_saver", "true|false", "Release cached runtime graphs after each request; default false."}, + {"fish_audio.reference_cache_slots", "n", "Prepared reference-audio cache slots; default 1."}, + {"fish_audio.weight_type", "native|f32|f16|bf16|q8_0", "AR matmul weight storage type; default native."}, + {"fish_audio.codec_weight_type", "native|f32|f16|q8_0", "Codec conv/matmul weight storage type; default native."}, + }; + return out; +} + +class FishAudioLoader final : public runtime::IVoiceModelLoader { +public: + std::string family() const override { + return "fish_audio"; + } + + bool can_load(const runtime::ModelLoadRequest & request) const override { + if (request.family_hint.has_value() && *request.family_hint != family()) { + return false; + } + try { + const auto package_spec = engine::assets::default_model_package_spec_path(family()); + (void) engine::assets::load_resource_bundle_from_package_spec(request.model_path, package_spec); + return true; + } catch (...) { + return false; + } + } + + runtime::ModelInspection inspect(const runtime::ModelLoadRequest & request) const override { + const auto assets = load_fish_audio_assets(request.model_path); + runtime::ModelInspection inspection; + inspection.model_root = assets->resources.model_root(); + inspection.metadata = metadata(*assets); + inspection.capabilities = capabilities(*assets); + inspection.cli = cli(*assets); + const auto package_spec = engine::assets::default_model_package_spec_path(family()); + inspection.discovered_configs = runtime::discover_named_assets_from_package_spec( + request.model_path, + package_spec, + engine::assets::ModelPackageResourceKind::Files); + inspection.discovered_weights = runtime::discover_named_assets_from_package_spec( + request.model_path, + package_spec, + engine::assets::ModelPackageResourceKind::Tensors); + return inspection; + } + + std::unique_ptr load(const runtime::ModelLoadRequest & request) const override { + return load_fish_audio_model(request.model_path); + } +}; + +} // namespace + +FishAudioLoadedModel::FishAudioLoadedModel( + runtime::ModelMetadata metadata, + runtime::CapabilitySet capabilities, + std::shared_ptr assets) + : metadata_(std::move(metadata)), + capabilities_(std::move(capabilities)), + assets_(std::move(assets)) {} + +const runtime::ModelMetadata & FishAudioLoadedModel::metadata() const noexcept { + return metadata_; +} + +const runtime::CapabilitySet & FishAudioLoadedModel::capabilities() const noexcept { + return capabilities_; +} + +std::unique_ptr FishAudioLoadedModel::create_task_session( + const runtime::TaskSpec & task, + const runtime::SessionOptions & options) const { + if (task.task != runtime::VoiceTaskKind::Tts || task.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Fish Audio S2-Pro supports offline TTS sessions"); + } + return std::make_unique(task, options, assets_); +} + +std::unique_ptr load_fish_audio_model(const std::filesystem::path & model_path) { + auto assets = load_fish_audio_assets(model_path); + return std::make_unique( + metadata(*assets), + capabilities(*assets), + std::move(assets)); +} + +std::shared_ptr make_fish_audio_loader() { + return std::make_shared(); +} + +} // namespace engine::models::fish_audio diff --git a/src/models/fish_audio/prompt_builder.cpp b/src/models/fish_audio/prompt_builder.cpp new file mode 100644 index 00000000..10b9f393 --- /dev/null +++ b/src/models/fish_audio/prompt_builder.cpp @@ -0,0 +1,101 @@ +#include "engine/models/fish_audio/prompt_builder.h" + +#include +#include +#include + +namespace engine::models::fish_audio { +namespace { + +void append_tokens(std::vector & out, const std::vector & tokens) { + out.insert(out.end(), tokens.begin(), tokens.end()); +} + +std::string reference_text_with_speakers(const std::string & text) { + static const std::regex speaker_re(R"(<\|speaker:\d+\|>)"); + if (std::regex_search(text, speaker_re)) { + return text; + } + return "<|speaker:0|>" + text; +} + +} // namespace + +FishAudioPromptBuilder::FishAudioPromptBuilder( + std::shared_ptr assets, + FishAudioTextTokenizer tokenizer) + : assets_(std::move(assets)), + tokenizer_(std::move(tokenizer)) { + if (assets_ == nullptr) { + throw std::runtime_error("Fish Audio prompt builder requires assets"); + } +} + +FishAudioPrompt FishAudioPromptBuilder::build( + const FishAudioRequest & request, + const std::optional & reference_codes) const { + if (request.text.empty()) { + throw std::runtime_error("Fish Audio request text must not be empty"); + } + const int64_t rows = assets_->config.fast.num_codebooks + 1; + if (rows <= 1) { + throw std::runtime_error("Fish Audio prompt rows are invalid"); + } + + std::vector row0; + if (request.reference.has_value()) { + if (!reference_codes.has_value()) { + throw std::runtime_error("Fish Audio reference request requires encoded reference codes"); + } + if (reference_codes->codebooks != assets_->config.fast.num_codebooks) { + throw std::runtime_error("Fish Audio reference codebook count mismatch"); + } + append_tokens(row0, tokenizer_.encode("<|im_start|>system\n")); + append_tokens(row0, tokenizer_.encode("convert the provided text to speech reference to the following:\n\nText:\n")); + append_tokens(row0, tokenizer_.encode(reference_text_with_speakers(request.reference->text))); + append_tokens(row0, tokenizer_.encode("\n\nSpeech:\n")); + const int32_t semantic_begin = tokenizer_.semantic_begin_id(); + for (int64_t frame = 0; frame < reference_codes->frames; ++frame) { + const int32_t code = reference_codes->codes[static_cast(frame)]; + row0.push_back(semantic_begin + code); + } + append_tokens(row0, tokenizer_.encode("<|im_end|>\n")); + } else { + append_tokens(row0, tokenizer_.encode("<|im_start|>system\n")); + append_tokens(row0, tokenizer_.encode("convert the provided text to speech")); + append_tokens(row0, tokenizer_.encode("<|im_end|>\n")); + } + append_tokens(row0, tokenizer_.encode("<|im_start|>user\n")); + append_tokens(row0, tokenizer_.encode(request.text)); + append_tokens(row0, tokenizer_.encode("<|im_end|>\n")); + append_tokens(row0, tokenizer_.encode("<|im_start|>assistant\n<|voice|>")); + + FishAudioPrompt prompt; + prompt.codebook_rows = rows; + prompt.steps = static_cast(row0.size()); + prompt.text = request.text; + prompt.matrix.assign(static_cast(rows * prompt.steps), 0); + for (int64_t step = 0; step < prompt.steps; ++step) { + prompt.matrix[static_cast(step)] = row0[static_cast(step)]; + } + if (reference_codes.has_value()) { + int64_t semantic_index = 0; + for (int64_t step = 0; step < prompt.steps; ++step) { + const int32_t token = prompt.matrix[static_cast(step)]; + if (token < tokenizer_.semantic_begin_id() || token > tokenizer_.semantic_end_id()) { + continue; + } + if (semantic_index >= reference_codes->frames) { + break; + } + for (int64_t codebook = 0; codebook < reference_codes->codebooks; ++codebook) { + prompt.matrix[static_cast((codebook + 1) * prompt.steps + step)] = + reference_codes->codes[static_cast(codebook * reference_codes->frames + semantic_index)]; + } + ++semantic_index; + } + } + return prompt; +} + +} // namespace engine::models::fish_audio diff --git a/src/models/fish_audio/session.cpp b/src/models/fish_audio/session.cpp new file mode 100644 index 00000000..f362664c --- /dev/null +++ b/src/models/fish_audio/session.cpp @@ -0,0 +1,432 @@ +#include "engine/models/fish_audio/session.h" + +#include "engine/framework/audio/wav_reader.h" +#include "engine/framework/debug/profiler.h" +#include "engine/framework/io/filesystem.h" +#include "engine/framework/runtime/options.h" +#include "engine/framework/runtime/session.h" +#include "engine/models/fish_audio/ar.h" +#include "engine/models/fish_audio/codec.h" +#include "engine/models/fish_audio/generator.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace engine::models::fish_audio { +namespace { + +using Clock = std::chrono::steady_clock; +namespace fs = std::filesystem; + +constexpr size_t kDefaultArGraphArenaBytes = 512ull * 1024ull * 1024ull; +constexpr size_t kDefaultCodecGraphArenaBytes = 512ull * 1024ull * 1024ull; +constexpr size_t kDefaultArWeightContextBytes = 512ull * 1024ull * 1024ull; +constexpr size_t kDefaultCodecWeightContextBytes = 512ull * 1024ull * 1024ull; +constexpr int64_t kDefaultReferenceCacheSlots = 1; +constexpr const char * kReferenceTextOption = "reference_text"; + +std::shared_ptr require_assets(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("Fish Audio session requires assets"); + } + return assets; +} + +assets::TensorStorageType option_weight_type( + const runtime::SessionOptions & options, + const char * key, + assets::TensorStorageType fallback) { + const auto it = options.options.find(key); + if (it == options.options.end()) { + return fallback; + } + return assets::parse_tensor_storage_type(it->second); +} + +void validate_ar_weight_storage(assets::TensorStorageType type, const char * option_name) { + if (type == assets::TensorStorageType::Native || + type == assets::TensorStorageType::F32 || + type == assets::TensorStorageType::F16 || + type == assets::TensorStorageType::BF16 || + type == assets::TensorStorageType::Q8_0) { + return; + } + throw std::runtime_error(std::string(option_name) + " supports native/f32/f16/bf16/q8_0"); +} + +void validate_codec_weight_storage(assets::TensorStorageType type, const char * option_name) { + if (type == assets::TensorStorageType::Native || + type == assets::TensorStorageType::F32 || + type == assets::TensorStorageType::F16 || + type == assets::TensorStorageType::Q8_0) { + return; + } + throw std::runtime_error(std::string(option_name) + " supports native/f32/f16/q8_0"); +} + +bool mem_saver_from_options(const runtime::SessionOptions & options) { + if (const auto value = runtime::find_option(options.options, {"fish_audio.mem_saver", "mem_saver"})) { + return runtime::parse_bool_option(*value, "fish_audio.mem_saver"); + } + return false; +} + +std::size_t resolve_reference_cache_slots(const runtime::SessionOptions & options) { + const int64_t slots = runtime::parse_i64_option( + options.options, + {"fish_audio.reference_cache_slots", "reference_cache_slots"}) + .value_or(kDefaultReferenceCacheSlots); + if (slots < 0) { + throw std::runtime_error("fish_audio.reference_cache_slots must be non-negative"); + } + if (static_cast(slots) > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("fish_audio.reference_cache_slots is too large"); + } + return static_cast(slots); +} + +uint64_t mix_reference_key(uint64_t key, uint64_t value) { + key ^= value; + key *= 1099511628211ull; + return key; +} + +uint64_t hash_audio_samples(const runtime::AudioBuffer & audio) { + uint64_t key = 1469598103934665603ull; + for (const float sample : audio.samples) { + uint32_t bits = 0; + std::memcpy(&bits, &sample, sizeof(bits)); + key = mix_reference_key(key, static_cast(bits)); + } + return key; +} + +FishAudioGenerationOptions generation_options_from_request(const runtime::TaskRequest & request) { + FishAudioGenerationOptions options; + options.max_new_tokens = runtime::parse_i64_option(request.options, {"max_new_tokens", "max_tokens"}) + .value_or(options.max_new_tokens); + options.chunk_length = runtime::parse_i64_option(request.options, {"chunk_length"}) + .value_or(options.chunk_length); + options.top_p = runtime::parse_float_option(request.options, {"top_p"}).value_or(options.top_p); + options.top_k = runtime::parse_int_option(request.options, {"top_k"}).value_or(options.top_k); + options.temperature = runtime::parse_float_option(request.options, {"temperature"}).value_or(options.temperature); + options.seed = runtime::parse_u32_option(request.options, {"seed"}).value_or(options.seed); + if (options.max_new_tokens <= 0) { + throw std::runtime_error("Fish Audio max_new_tokens must be positive"); + } + if (options.chunk_length <= 0) { + throw std::runtime_error("Fish Audio chunk_length must be positive"); + } + if (!(options.top_p > 0.0F && options.top_p <= 1.0F)) { + throw std::runtime_error("Fish Audio top_p must be in (0, 1]"); + } + if (options.top_k <= 0) { + throw std::runtime_error("Fish Audio top_k must be positive"); + } + if (!(options.temperature > 0.0F && options.temperature < 2.0F)) { + throw std::runtime_error("Fish Audio temperature must be in (0, 2)"); + } + return options; +} + +std::string lower_ascii(std::string value) { + std::transform( + value.begin(), + value.end(), + value.begin(), + [](unsigned char ch) { return static_cast(std::tolower(ch)); }); + return value; +} + +bool valid_reference_id_char(unsigned char ch) { + return std::isalnum(ch) != 0 || ch == '-' || ch == '_' || ch == ' '; +} + +void validate_reference_id(const std::string & id) { + if (id.empty() || id.size() > 255) { + throw std::runtime_error( + "Fish Audio cached_voice_id must be 1-255 characters"); + } + for (const unsigned char ch : id) { + if (!valid_reference_id_char(ch)) { + throw std::runtime_error( + "Fish Audio cached_voice_id may only contain alphanumeric characters, hyphens, underscores, and spaces"); + } + } +} + +bool is_supported_saved_reference_audio(const fs::path & path) { + return lower_ascii(path.extension().string()) == ".wav"; +} + +std::vector collect_saved_reference_audio_files(const fs::path & directory) { + std::vector files; + for (const auto & entry : fs::recursive_directory_iterator(directory)) { + if (!entry.is_regular_file() || !is_supported_saved_reference_audio(entry.path())) { + continue; + } + auto lab_path = entry.path(); + lab_path.replace_extension(".lab"); + if (engine::io::is_existing_file(lab_path)) { + files.push_back(entry.path()); + } + } + std::sort(files.begin(), files.end()); + return files; +} + +runtime::AudioBuffer read_saved_reference_audio(const fs::path & path) { + auto wav = engine::audio::read_wav_f32(path); + return runtime::AudioBuffer{wav.sample_rate, wav.channels, std::move(wav.samples)}; +} + +FishAudioReference load_saved_reference( + const FishAudioAssets & assets, + const std::string & reference_id) { + validate_reference_id(reference_id); + const auto reference_dir = engine::io::require_directory( + assets.resources.model_root() / "references" / reference_id, + "Fish Audio cached voice reference"); + const auto audio_files = collect_saved_reference_audio_files(reference_dir); + if (audio_files.empty()) { + throw std::runtime_error( + "Fish Audio cached_voice_id '" + reference_id + + "' requires one WAV reference with a matching .lab file under " + + reference_dir.string()); + } + if (audio_files.size() > 1) { + throw std::runtime_error( + "Fish Audio cached_voice_id '" + reference_id + + "' has multiple WAV references with .lab files; the C++ session expects exactly one reference pair"); + } + auto lab_path = audio_files.front(); + lab_path.replace_extension(".lab"); + return FishAudioReference{ + read_saved_reference_audio(audio_files.front()), + engine::io::read_text_file(lab_path), + reference_id}; +} + +std::string reference_cache_id_from_voice(const std::optional & voice) { + if (voice.has_value() && + voice->speaker.has_value() && + voice->speaker->cached_voice_id.has_value() && + !voice->speaker->cached_voice_id->empty()) { + return *voice->speaker->cached_voice_id; + } + return {}; +} + +bool has_reference_selector(const std::optional & voice) { + if (!voice.has_value() || !voice->speaker.has_value()) { + return false; + } + const auto & speaker = *voice->speaker; + return speaker.audio.has_value() || + (speaker.cached_voice_id.has_value() && !speaker.cached_voice_id->empty()); +} + +std::optional reference_from_voice( + const FishAudioAssets & assets, + const std::optional & voice, + const std::unordered_map & options, + const char * role) { + if (!has_reference_selector(voice)) { + return std::nullopt; + } + const auto & speaker = *voice->speaker; + if (speaker.audio.has_value()) { + auto reference_text = runtime::find_option(options, {kReferenceTextOption}); + if (!reference_text.has_value()) { + throw std::runtime_error( + std::string(role) + " with inline reference audio requires reference_text option"); + } + return FishAudioReference{ + speaker.audio, + *reference_text, + reference_cache_id_from_voice(voice)}; + } + return load_saved_reference(assets, *speaker.cached_voice_id); +} + +} // namespace + +FishAudioSession::FishAudioSession( + runtime::TaskSpec task, + runtime::SessionOptions options, + std::shared_ptr assets) + : RuntimeSessionBase(options), + task_(task), + assets_(require_assets(std::move(assets))), + reference_cache_(resolve_reference_cache_slots(this->options())) { + if (task_.task != runtime::VoiceTaskKind::Tts || task_.mode != runtime::RunMode::Offline) { + throw std::runtime_error("Fish Audio only supports offline TTS sessions"); + } + const auto ar_weight_type = + option_weight_type(options, "fish_audio.weight_type", assets::TensorStorageType::Native); + const auto codec_weight_type = + option_weight_type(options, "fish_audio.codec_weight_type", assets::TensorStorageType::Native); + validate_ar_weight_storage(ar_weight_type, "fish_audio.weight_type"); + validate_codec_weight_storage(codec_weight_type, "fish_audio.codec_weight_type"); + const int threads = options.backend.threads > 0 ? options.backend.threads : 1; + auto ar = std::make_unique( + assets_, + options.backend, + threads, + runtime::parse_size_mb_option(options.options, {"fish_audio.ar_graph_arena_mb"}, kDefaultArGraphArenaBytes), + runtime::parse_size_mb_option(options.options, {"fish_audio.ar_weight_context_mb"}, kDefaultArWeightContextBytes), + ar_weight_type); + auto codec = std::make_unique( + assets_, + options.backend, + threads, + runtime::parse_size_mb_option(options.options, {"fish_audio.codec_graph_arena_mb"}, kDefaultCodecGraphArenaBytes), + runtime::parse_size_mb_option(options.options, {"fish_audio.codec_weight_context_mb"}, kDefaultCodecWeightContextBytes), + codec_weight_type, + codec_weight_type); + generator_ = std::make_unique( + assets_, + std::move(ar), + std::move(codec)); + assets_->model_weights->release_storage(); + assets_->codec_weights->release_storage(); +} + +FishAudioSession::~FishAudioSession() = default; + +std::string FishAudioSession::family() const { + return "fish_audio"; +} + +runtime::VoiceTaskKind FishAudioSession::task_kind() const { + return task_.task; +} + +runtime::RunMode FishAudioSession::run_mode() const { + return task_.mode; +} + +bool FishAudioSession::ReferenceCacheKeyEqual::operator()( + const ReferenceCacheKey & lhs, + const ReferenceCacheKey & rhs) const { + return lhs.source_id == rhs.source_id && + lhs.sample_rate == rhs.sample_rate && + lhs.channels == rhs.channels && + lhs.sample_count == rhs.sample_count && + lhs.sample_hash == rhs.sample_hash; +} + +void FishAudioSession::prepare(const runtime::SessionPreparationRequest & request) { + defaults_.reset(); + FishAudioRequest defaults; + bool has_defaults = false; + if (request.text.has_value()) { + defaults.text = request.text->text; + has_defaults = true; + } + defaults.generation.max_new_tokens = runtime::parse_i64_option(request.options, {"max_new_tokens", "max_tokens"}) + .value_or(defaults.generation.max_new_tokens); + if (auto reference = reference_from_voice(*assets_, request.voice, request.options, "Fish Audio prepare"); + reference.has_value()) { + defaults.reference = std::move(*reference); + if (defaults.reference->audio.has_value()) { + (void) resolve_reference_codes(*defaults.reference); + } + has_defaults = true; + } + if (has_defaults) { + defaults_ = std::move(defaults); + } + mark_prepared(); +} + +FishAudioRequest FishAudioSession::make_request(const runtime::TaskRequest & request) const { + FishAudioRequest out = defaults_.value_or(FishAudioRequest{}); + if (request.text_input.has_value()) { + out.text = request.text_input->text; + } + out.generation = generation_options_from_request(request); + if (auto reference = reference_from_voice(*assets_, request.voice, request.options, "Fish Audio request"); + reference.has_value()) { + out.reference = std::move(*reference); + } else if (request.text_input.has_value()) { + out.reference = std::nullopt; + } + if (out.text.empty()) { + throw std::runtime_error("Fish Audio request text must not be empty"); + } + return out; +} + +const FishAudioCodes & FishAudioSession::resolve_reference_codes(const FishAudioReference & reference) { + ReferenceCacheKey key; + key.source_id = reference.cache_id; + if (reference.cache_id.empty() && !reference.audio.has_value()) { + throw std::runtime_error("Fish Audio cached reference requires reference audio or a reference id"); + } + if (reference.audio.has_value() && reference.cache_id.empty()) { + key.sample_rate = reference.audio->sample_rate; + key.channels = reference.audio->channels; + key.sample_count = static_cast(reference.audio->samples.size()); + key.sample_hash = hash_audio_samples(*reference.audio); + } + if (const auto * cached = reference_cache_.find(key)) { + engine::debug::trace_log_scalar("fish_audio.reference_cache.hit", 1); + engine::debug::trace_log_scalar("fish_audio.reference_cache.slots", static_cast(reference_cache_.capacity())); + engine::debug::trace_log_scalar("fish_audio.reference_cache.entries", static_cast(reference_cache_.size())); + engine::debug::trace_log_scalar("fish_audio.reference_cache.evicted", 0); + return cached->codes; + } + if (!reference.audio.has_value()) { + throw std::runtime_error("Fish Audio reference id is not cached and no reference audio was provided"); + } + const bool will_evict = reference_cache_.capacity() > 0 && reference_cache_.size() >= reference_cache_.capacity(); + const auto start = Clock::now(); + ReferenceCacheEntry entry; + entry.codes = generator_->encode_reference(*reference.audio); + engine::debug::trace_log_scalar("fish_audio.reference.frames", entry.codes.frames); + engine::debug::trace_log_scalar("fish_audio.reference.codebooks", entry.codes.codebooks); + if (reference_cache_.capacity() == 0) { + uncached_reference_ = std::move(entry); + } else { + reference_cache_.put(key, std::move(entry)); + } + engine::debug::trace_log_scalar("fish_audio.reference_cache.hit", 0); + engine::debug::trace_log_scalar("fish_audio.reference_cache.slots", static_cast(reference_cache_.capacity())); + engine::debug::trace_log_scalar("fish_audio.reference_cache.entries", static_cast(reference_cache_.size())); + engine::debug::trace_log_scalar("fish_audio.reference_cache.evicted", will_evict ? 1 : 0); + engine::debug::timing_log_scalar("fish_audio.reference_encode_ms", engine::debug::elapsed_ms(start, Clock::now())); + if (reference_cache_.capacity() == 0) { + return uncached_reference_->codes; + } + const auto * cached = reference_cache_.find(key); + if (cached == nullptr) { + throw std::runtime_error("Fish Audio reference cache insert failed"); + } + return cached->codes; +} + +runtime::TaskResult FishAudioSession::run(const runtime::TaskRequest & request) { + require_prepared("Fish Audio run()"); + const auto wall_start = Clock::now(); + const bool mem_saver = mem_saver_from_options(options()); + auto fish_request = make_request(request); + std::optional reference_codes = std::nullopt; + if (fish_request.reference.has_value()) { + reference_codes = resolve_reference_codes(*fish_request.reference); + } + auto generated = generator_->generate(fish_request, reference_codes, mem_saver); + runtime::TaskResult result; + result.audio_output = std::move(generated.audio); + engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start, Clock::now())); + return result; +} + +} // namespace engine::models::fish_audio diff --git a/src/models/fish_audio/tokenizer_text.cpp b/src/models/fish_audio/tokenizer_text.cpp new file mode 100644 index 00000000..43839f3f --- /dev/null +++ b/src/models/fish_audio/tokenizer_text.cpp @@ -0,0 +1,69 @@ +#include "engine/models/fish_audio/tokenizer_text.h" + +#include "engine/framework/tokenizers/llama_bpe.h" + +#include +#include + +namespace engine::models::fish_audio { +namespace { + +int32_t require_token_id(const engine::tokenizers::LlamaBpeTokenizer & tokenizer, const std::string & token) { + const auto id = tokenizer.find_token_id(token); + if (!id.has_value()) { + throw std::runtime_error("Fish Audio tokenizer missing token: " + token); + } + return *id; +} + +} // namespace + +struct FishAudioTextTokenizer::Impl { + explicit Impl(std::shared_ptr input_assets) + : assets(std::move(input_assets)), + tokenizer(engine::tokenizers::LlamaBpeTokenizerSpec{ + {}, + {}, + assets->resources.require_file("tokenizer_config"), + assets->resources.require_file("tokenizer_json"), + engine::tokenizers::LlamaBpePreTokenizer::Qwen2, + }), + im_end(require_token_id(tokenizer, "<|im_end|>")), + semantic_begin(static_cast(assets->config.semantic_start_token_id)), + semantic_end(static_cast(assets->config.semantic_end_token_id)) {} + + std::shared_ptr assets; + engine::tokenizers::LlamaBpeTokenizer tokenizer; + int32_t im_end = 0; + int32_t semantic_begin = 0; + int32_t semantic_end = 0; +}; + +FishAudioTextTokenizer::FishAudioTextTokenizer(std::shared_ptr assets) { + if (assets == nullptr) { + throw std::runtime_error("Fish Audio text tokenizer requires assets"); + } + impl_ = std::make_shared(std::move(assets)); +} + +std::vector FishAudioTextTokenizer::encode(const std::string & text) const { + return impl_->tokenizer.encode(text, true); +} + +int32_t FishAudioTextTokenizer::token_id(const std::string & token) const { + return require_token_id(impl_->tokenizer, token); +} + +int32_t FishAudioTextTokenizer::im_end_id() const noexcept { + return impl_->im_end; +} + +int32_t FishAudioTextTokenizer::semantic_begin_id() const noexcept { + return impl_->semantic_begin; +} + +int32_t FishAudioTextTokenizer::semantic_end_id() const noexcept { + return impl_->semantic_end; +} + +} // namespace engine::models::fish_audio diff --git a/src/models/higgs_tts/ar.cpp b/src/models/higgs_tts/ar.cpp index 922d4830..ba442bf6 100644 --- a/src/models/higgs_tts/ar.cpp +++ b/src/models/higgs_tts/ar.cpp @@ -56,7 +56,6 @@ modules::QwenDecoderStackConfig make_higgs_qwen_stack_config(const HiggsTextConf out.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; out.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; out.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; - out.runtime.static_cache.transpose_context = false; return out; } diff --git a/tools/audiocpp_cli/audiocpp_cli_path_cases.json b/tools/audiocpp_cli/audiocpp_cli_path_cases.json index 70429432..7054381d 100644 --- a/tools/audiocpp_cli/audiocpp_cli_path_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_path_cases.json @@ -611,6 +611,53 @@ } ] }, + { + "id": "fish_audio_s2_pro_paths", + "coverage": "Fish Audio S2-Pro path coverage for default voice, reference voice clone, inline control tag, AR generation, and codec decode in one offline session", + "family": "fish_audio", + "model": "models/s2-pro", + "task": "tts", + "mode": "offline", + "session_options": { + "fish_audio.weight_type": "native", + "fish_audio.codec_weight_type": "native", + "fish_audio.reference_cache_slots": "1" + }, + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "official_auto_voice_english", + "text": "The field recorder captured a clean reference take, and the operator confirmed that every timestamp matched the written production notes.", + "chunk_length": 200, + "top_p": 0.8, + "repetition_penalty": 1.1, + "temperature": 0.8, + "seed": 1234 + }, + { + "id": "official_reference_voice_clone", + "text": "The studio engineer checked the short voice prompt, confirmed the take was clear, and started the final render.", + "voice_ref": "resources/sample.wav", + "reference_text": "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you.", + "chunk_length": 200, + "top_p": 0.8, + "repetition_penalty": 1.1, + "temperature": 0.8, + "seed": 2234 + }, + { + "id": "official_inline_control_tag", + "text": "[whisper in small voice] The prototype actually worked after the last reset, and the control room stayed quiet until every green light appeared.", + "chunk_length": 200, + "top_p": 0.8, + "repetition_penalty": 1.1, + "temperature": 0.8, + "seed": 3234 + } + ] + }, { "id": "heartmula_music_generation", "coverage": "HeartMuLa text-to-music generation using native model weight types and the warmbench reference request", From 6aadbc0bf0b6249d54108c9b74634998135cd94e Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Mon, 20 Jul 2026 16:07:51 -0400 Subject: [PATCH 07/27] Scope Higgs Qwen decoder optimizations --- .../modules/attention/qwen_decoder.h | 23 +++++++ .../modules/optimizations/fast_kv_modules.h | 14 ++++ .../modules/attention/qwen_decoder.cpp | 17 +++-- .../modules/optimizations/fast_kv_modules.cpp | 67 +++++++++++-------- src/models/higgs_tts/ar.cpp | 3 + .../test_qwen_decoder_packed_projections.cpp | 14 +++- 6 files changed, 105 insertions(+), 33 deletions(-) diff --git a/include/engine/framework/modules/attention/qwen_decoder.h b/include/engine/framework/modules/attention/qwen_decoder.h index ee40a2d2..56b85603 100644 --- a/include/engine/framework/modules/attention/qwen_decoder.h +++ b/include/engine/framework/modules/attention/qwen_decoder.h @@ -24,11 +24,27 @@ enum class QwenDecoderStaticCacheUpdateMode { DirectSetRows, }; +enum class QwenDecoderStaticCacheSetRowsMode { + Exact, + BackendViewOptimized, +}; + enum class QwenDecoderQKVLayout { Separate, PackedQKV, }; +enum class QwenDecoderMLPMode { + Exact, + FusedSwiGLU, + PackedGateUp, +}; + +enum class QwenDecoderPrefixAttentionMode { + Exact, + FlashWithPrefix, +}; + enum class QwenDecoderPositionEncoding { Rotary, None, @@ -55,16 +71,23 @@ struct QwenDecoderActivationCastPolicy { struct QwenDecoderAttentionPolicy { QwenDecoderAttentionMode prefill_mode = QwenDecoderAttentionMode::ManualRepeat; QwenDecoderAttentionMode static_mode = QwenDecoderAttentionMode::FlashGrouped; + QwenDecoderPrefixAttentionMode prefix_mode = QwenDecoderPrefixAttentionMode::Exact; int64_t grouped_query_min_steps = 0; }; struct QwenDecoderStaticCachePolicy { QwenDecoderStaticCacheUpdateMode update_mode = QwenDecoderStaticCacheUpdateMode::ScratchTail; + QwenDecoderStaticCacheSetRowsMode set_rows_mode = QwenDecoderStaticCacheSetRowsMode::Exact; +}; + +struct QwenDecoderMLPPolicy { + QwenDecoderMLPMode mode = QwenDecoderMLPMode::Exact; }; struct QwenDecoderRuntimePolicy { QwenDecoderAttentionPolicy attention; QwenDecoderStaticCachePolicy static_cache; + QwenDecoderMLPPolicy mlp; }; struct QwenDecoderLayerConfig { diff --git a/include/engine/framework/modules/optimizations/fast_kv_modules.h b/include/engine/framework/modules/optimizations/fast_kv_modules.h index 75175b2a..247a8cc2 100644 --- a/include/engine/framework/modules/optimizations/fast_kv_modules.h +++ b/include/engine/framework/modules/optimizations/fast_kv_modules.h @@ -4,8 +4,19 @@ namespace engine::modules { +enum class FastKVSetRowsMode { + Exact, + BackendViewOptimized, +}; + +struct FastKVSetRowsConfig { + FastKVSetRowsMode mode = FastKVSetRowsMode::Exact; +}; + class FastKVSetRowsModule { public: + explicit FastKVSetRowsModule(FastKVSetRowsConfig config = {}); + const core::ModuleSchema & schema() const noexcept; core::TensorValue build( @@ -15,6 +26,9 @@ class FastKVSetRowsModule { const core::TensorValue & row_index) const; static const core::ModuleSchema & static_schema() noexcept; + +private: + FastKVSetRowsConfig config_; }; } // namespace engine::modules diff --git a/src/framework/modules/attention/qwen_decoder.cpp b/src/framework/modules/attention/qwen_decoder.cpp index 833c1d35..6007eb5f 100644 --- a/src/framework/modules/attention/qwen_decoder.cpp +++ b/src/framework/modules/attention/qwen_decoder.cpp @@ -337,7 +337,11 @@ core::TensorValue build_mlp( core::TensorValue gate; core::TensorValue up; std::optional packed_gate_up; - if (weights.gate_up_proj.has_value()) { + const auto mlp_mode = config.runtime.mlp.mode; + if (mlp_mode == QwenDecoderMLPMode::PackedGateUp) { + if (!weights.gate_up_proj.has_value()) { + throw std::runtime_error("QwenMLPWeights.gate_up_proj is required for packed gate/up mode"); + } auto gate_up = LinearModule( { config.hidden_size, @@ -376,7 +380,7 @@ core::TensorValue build_mlp( !config.activation_cast.after_mlp_silu && !config.activation_cast.after_mlp_mul); core::TensorValue gated; - if (can_use_fused_swiglu && packed_gate_up.has_value()) { + if (can_use_fused_swiglu && mlp_mode == QwenDecoderMLPMode::PackedGateUp) { gated = core::wrap_tensor( ggml_swiglu(ctx.ggml, packed_gate_up->tensor), core::TensorShape::from_dims({ @@ -385,7 +389,7 @@ core::TensorValue build_mlp( config.intermediate_size, }), packed_gate_up->type); - } else if (can_use_fused_swiglu) { + } else if (can_use_fused_swiglu && mlp_mode == QwenDecoderMLPMode::FusedSwiGLU) { gated = core::wrap_tensor( ggml_swiglu_split(ctx.ggml, gate.tensor, up.tensor), gate.shape, @@ -518,6 +522,7 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build( } else if (attention_mask.has_value() && (config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGrouped || (prefix_key.has_value() && + config_.runtime.attention.prefix_mode == QwenDecoderPrefixAttentionMode::FlashWithPrefix && config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGroupedViewKV))) { q_heads = core::wrap_tensor(ggml_cont(ctx.ggml, q_heads.tensor), q_heads.shape, q_heads.type); auto k_heads = TransposeModule({{0, 2, 1, 3}, all_k.shape.rank}).build(ctx, all_k); @@ -644,7 +649,11 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail( if (!cache_slot.has_value()) { throw std::runtime_error("Qwen decoder direct static-cache update requires cache_slot"); } - const FastKVSetRowsModule set_rows; + const FastKVSetRowsModule set_rows({ + config_.runtime.static_cache.set_rows_mode == QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized + ? FastKVSetRowsMode::BackendViewOptimized + : FastKVSetRowsMode::Exact, + }); attention_key_cache = set_rows.build(ctx, cache_key, k, *cache_slot); attention_value_cache = set_rows.build(ctx, cache_value, v, *cache_slot); if (config_.activation_cast.enabled && config_.activation_cast.after_static_cache_update) { diff --git a/src/framework/modules/optimizations/fast_kv_modules.cpp b/src/framework/modules/optimizations/fast_kv_modules.cpp index 64389763..eff1e166 100644 --- a/src/framework/modules/optimizations/fast_kv_modules.cpp +++ b/src/framework/modules/optimizations/fast_kv_modules.cpp @@ -30,6 +30,8 @@ const core::ModuleSchema kFastKVSetRowsSchema = { } // namespace +FastKVSetRowsModule::FastKVSetRowsModule(FastKVSetRowsConfig config) : config_(config) {} + const core::ModuleSchema & FastKVSetRowsModule::schema() const noexcept { return static_schema(); } @@ -51,8 +53,14 @@ core::TensorValue FastKVSetRowsModule::build( if (row_index.shape.rank != 1 || (row_index.shape.dims[0] != 1 && row_index.shape.dims[0] != batch)) { throw std::runtime_error("FastKVSetRowsModule row_index must have shape {1} or {batch}"); } - if ((cache.type != GGML_TYPE_F32 && cache.type != GGML_TYPE_F16) || row.type != GGML_TYPE_F32) { - throw std::runtime_error("FastKVSetRowsModule requires an f32/f16 cache and an f32 row tensor"); + const bool optimized = config_.mode == FastKVSetRowsMode::BackendViewOptimized; + if (((!optimized && cache.type != GGML_TYPE_F32) || + (optimized && cache.type != GGML_TYPE_F32 && cache.type != GGML_TYPE_F16)) || + row.type != GGML_TYPE_F32) { + throw std::runtime_error( + optimized + ? "FastKVSetRowsModule requires an f32/f16 cache and an f32 row tensor" + : "FastKVSetRowsModule requires f32 cache and row tensors"); } if (row_index.type != GGML_TYPE_I32 && row_index.type != GGML_TYPE_I64) { throw std::runtime_error("FastKVSetRowsModule requires i32 or i64 row_index tensor"); @@ -69,39 +77,44 @@ core::TensorValue FastKVSetRowsModule::build( } auto flat_cache = core::reshape_tensor(ctx, cache, core::TensorShape::from_dims({steps, row_elems})); auto contiguous_row = tensor_layout::ensure_contiguous_layout_if_needed(ctx, row); - auto flat_row = core::wrap_tensor( - ggml_view_2d( - ctx.ggml, - contiguous_row.tensor, - row_elems, - 1, - contiguous_row.tensor->nb[2], - 0), - core::TensorShape::from_dims({1, row_elems}), - row.type); + auto flat_row = optimized + ? core::wrap_tensor( + ggml_view_2d( + ctx.ggml, + contiguous_row.tensor, + row_elems, + 1, + contiguous_row.tensor->nb[2], + 0), + core::TensorShape::from_dims({1, row_elems}), + row.type) + : core::reshape_tensor(ctx, contiguous_row, core::TensorShape::from_dims({1, row_elems})); ggml_tensor * updated = ggml_set_rows(ctx.ggml, flat_cache.tensor, flat_row.tensor, row_index.tensor); - // ggml_set_rows src[2] is only a legacy dependency anchor for the - // destination. Point it at the underlying cache so the metadata-only - // flatten does not interrupt CUDA's ROPE -> VIEW -> SET_ROWS fusion. - updated->src[2] = cache.tensor; + if (optimized) { + updated->src[2] = cache.tensor; + } auto flat_updated = core::wrap_tensor(updated, flat_cache.shape, cache.type); return core::reshape_tensor(ctx, flat_updated, cache.shape); } auto flat_cache = core::reshape_tensor(ctx, cache, core::TensorShape::from_dims({batch * steps, row_elems})); auto contiguous_row = tensor_layout::ensure_contiguous_layout_if_needed(ctx, row); - auto flat_row = core::wrap_tensor( - ggml_view_2d( - ctx.ggml, - contiguous_row.tensor, - row_elems, - batch, - contiguous_row.tensor->nb[3], - 0), - core::TensorShape::from_dims({batch, row_elems}), - row.type); + auto flat_row = optimized + ? core::wrap_tensor( + ggml_view_2d( + ctx.ggml, + contiguous_row.tensor, + row_elems, + batch, + contiguous_row.tensor->nb[3], + 0), + core::TensorShape::from_dims({batch, row_elems}), + row.type) + : core::reshape_tensor(ctx, contiguous_row, core::TensorShape::from_dims({batch, row_elems})); ggml_tensor * updated = ggml_set_rows(ctx.ggml, flat_cache.tensor, flat_row.tensor, row_index.tensor); - updated->src[2] = cache.tensor; + if (optimized) { + updated->src[2] = cache.tensor; + } auto flat_updated = core::wrap_tensor(updated, flat_cache.shape, cache.type); return core::reshape_tensor(ctx, flat_updated, cache.shape); } diff --git a/src/models/higgs_tts/ar.cpp b/src/models/higgs_tts/ar.cpp index ba442bf6..5749ffef 100644 --- a/src/models/higgs_tts/ar.cpp +++ b/src/models/higgs_tts/ar.cpp @@ -55,7 +55,10 @@ modules::QwenDecoderStackConfig make_higgs_qwen_stack_config(const HiggsTextConf out.use_qk_norm = true; out.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; out.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.runtime.attention.prefix_mode = modules::QwenDecoderPrefixAttentionMode::FlashWithPrefix; out.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; + out.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; return out; } diff --git a/tests/unittests/test_qwen_decoder_packed_projections.cpp b/tests/unittests/test_qwen_decoder_packed_projections.cpp index c34c4535..f66353ff 100644 --- a/tests/unittests/test_qwen_decoder_packed_projections.cpp +++ b/tests/unittests/test_qwen_decoder_packed_projections.cpp @@ -124,6 +124,9 @@ LayerResult run_layer(bool packed) { config.qkv_layout = packed ? engine::modules::QwenDecoderQKVLayout::PackedQKV : engine::modules::QwenDecoderQKVLayout::Separate; + config.runtime.mlp.mode = packed + ? engine::modules::QwenDecoderMLPMode::PackedGateUp + : engine::modules::QwenDecoderMLPMode::Exact; config.use_qk_norm = false; config.runtime.attention.prefill_mode = engine::modules::QwenDecoderAttentionMode::ManualRepeat; @@ -250,7 +253,9 @@ void test_f16_kv_set_rows() { ctx, GGML_TYPE_I64, engine::core::TensorShape::from_dims({1})); - const auto output = engine::modules::FastKVSetRowsModule{}.build(ctx, cache, row, row_index); + const auto output = engine::modules::FastKVSetRowsModule({ + engine::modules::FastKVSetRowsMode::BackendViewOptimized, + }).build(ctx, cache, row, row_index); ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); ggml_build_forward_expand(graph, output.tensor); @@ -313,7 +318,9 @@ void test_f16_kv_set_rows_batched() { ctx, GGML_TYPE_I64, engine::core::TensorShape::from_dims({2})); - const auto output = engine::modules::FastKVSetRowsModule{}.build(ctx, cache, rows, row_indices); + const auto output = engine::modules::FastKVSetRowsModule({ + engine::modules::FastKVSetRowsMode::BackendViewOptimized, + }).build(ctx, cache, rows, row_indices); ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); ggml_build_forward_expand(graph, output.tensor); @@ -447,6 +454,9 @@ void test_higgs_decode_graph_exposes_cuda_fast_paths() { engine::modules::QwenDecoderAttentionMode::FlashGroupedViewKV; config.runtime.static_cache.update_mode = engine::modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + config.runtime.static_cache.set_rows_mode = + engine::modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; + config.runtime.mlp.mode = engine::modules::QwenDecoderMLPMode::PackedGateUp; ggml_cgraph * graph = ggml_new_graph_custom(ggml, kGraphNodes, false); const auto outputs = engine::modules::QwenDecoderLayerModule(config).build_with_static_cache_tail( From 51770d7859b41465a850de41f4dbe814c0a95a45 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Mon, 20 Jul 2026 16:22:44 -0400 Subject: [PATCH 08/27] Scope Higgs cache optimizations --- CMakeLists.txt | 7 - include/engine/framework/runtime/kv_cache.h | 14 +- .../modules/attention/qwen_decoder.cpp | 92 +- src/framework/runtime/kv_cache.cpp | 64 +- src/models/higgs_tts/ar.cpp | 5 +- tests/unittests/test_conv_lowering_matrix.cpp | 834 ------------------ 6 files changed, 137 insertions(+), 879 deletions(-) delete mode 100644 tests/unittests/test_conv_lowering_matrix.cpp diff --git a/CMakeLists.txt b/CMakeLists.txt index fd235a43..df275571 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -921,13 +921,6 @@ if (ENGINE_BUILD_TESTS) COMMAND depthwise_conv1d_lowering_test ) - add_engine_unittest(conv_lowering_matrix_test tests/unittests/test_conv_lowering_matrix.cpp) - - add_test( - NAME conv_lowering_matrix_test - COMMAND conv_lowering_matrix_test - ) - add_engine_unittest(gguf_tensor_source_test tests/unittests/test_gguf_tensor_source.cpp) target_include_directories(gguf_tensor_source_test PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/tests/unittests) diff --git a/include/engine/framework/runtime/kv_cache.h b/include/engine/framework/runtime/kv_cache.h index e38a00e7..e79fc9d1 100644 --- a/include/engine/framework/runtime/kv_cache.h +++ b/include/engine/framework/runtime/kv_cache.h @@ -19,6 +19,10 @@ struct TransformerKVState { std::vector layers; }; +struct TransformerKVCacheOptions { + bool allow_f16_storage = false; +}; + class TransformerKVCache { public: TransformerKVCache() = default; @@ -27,6 +31,12 @@ class TransformerKVCache { int64_t step_elems, std::vector keys, std::vector values); + TransformerKVCache( + int64_t cache_steps, + int64_t step_elems, + std::vector keys, + std::vector values, + TransformerKVCacheOptions options); void import_state(const TransformerKVState & state); TransformerKVState export_state() const; @@ -55,6 +65,7 @@ class TransformerKVCache { int64_t step_elems_ = 0; int64_t valid_steps_ = 0; int64_t current_end_ = 0; + TransformerKVCacheOptions options_; std::vector layers_; }; @@ -65,6 +76,7 @@ core::TensorValue view_transformer_kv_cache_steps( int64_t steps, int64_t heads, int64_t head_dim, - const char * label); + const char * label, + ggml_type view_type = GGML_TYPE_F32); } // namespace engine::runtime diff --git a/src/framework/modules/attention/qwen_decoder.cpp b/src/framework/modules/attention/qwen_decoder.cpp index 6007eb5f..f473e738 100644 --- a/src/framework/modules/attention/qwen_decoder.cpp +++ b/src/framework/modules/attention/qwen_decoder.cpp @@ -334,6 +334,48 @@ core::TensorValue build_mlp( const core::TensorValue & input, const QwenDecoderLayerConfig & config, const QwenMLPWeights & weights) { + if (config.runtime.mlp.mode == QwenDecoderMLPMode::Exact) { + auto gate = LinearModule( + { + config.hidden_size, + config.intermediate_size, + weights.gate_proj.bias.has_value(), + config.projection_precision, + }) + .build(ctx, input, require_linear(weights.gate_proj, false, "QwenMLPWeights.gate_proj")); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { + gate = activation_cast(ctx, gate, config.activation_cast); + } + gate = SiluModule{}.build(ctx, gate); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_silu) { + gate = activation_cast(ctx, gate, config.activation_cast); + } + auto up = LinearModule( + { + config.hidden_size, + config.intermediate_size, + weights.up_proj.bias.has_value(), + config.projection_precision, + }) + .build(ctx, input, require_linear(weights.up_proj, false, "QwenMLPWeights.up_proj")); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { + up = activation_cast(ctx, up, config.activation_cast); + } + auto gated = MulModule{}.build(ctx, gate, up); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_mul) { + gated = activation_cast(ctx, gated, config.activation_cast); + } + auto down = LinearModule( + { + config.intermediate_size, + config.hidden_size, + weights.down_proj.bias.has_value(), + config.projection_precision, + }) + .build(ctx, gated, require_linear(weights.down_proj, false, "QwenMLPWeights.down_proj")); + return down; + } + core::TensorValue gate; core::TensorValue up; std::optional packed_gate_up; @@ -489,22 +531,33 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build( v = core::ensure_backend_addressable_layout(ctx, v); auto q_heads = TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q); - auto attention_prefix_key = prefix_key; - auto attention_prefix_value = prefix_value; - if (attention_prefix_key.has_value() && attention_prefix_key->type != k.type) { - attention_prefix_key = core::wrap_tensor( - ggml_cast(ctx.ggml, attention_prefix_key->tensor, k.type), - attention_prefix_key->shape, - k.type); - } - if (attention_prefix_value.has_value() && attention_prefix_value->type != v.type) { - attention_prefix_value = core::wrap_tensor( - ggml_cast(ctx.ggml, attention_prefix_value->tensor, v.type), - attention_prefix_value->shape, - v.type); - } - auto all_k = attention_prefix_key.has_value() ? ConcatModule({1}).build(ctx, *attention_prefix_key, k) : k; - auto all_v = attention_prefix_value.has_value() ? ConcatModule({1}).build(ctx, *attention_prefix_value, v) : v; + const bool use_prefix_flash = + prefix_key.has_value() && + config_.runtime.attention.prefix_mode == QwenDecoderPrefixAttentionMode::FlashWithPrefix && + config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGroupedViewKV; + core::TensorValue all_k = k; + core::TensorValue all_v = v; + if (use_prefix_flash) { + auto attention_prefix_key = prefix_key; + auto attention_prefix_value = prefix_value; + if (attention_prefix_key->type != k.type) { + attention_prefix_key = core::wrap_tensor( + ggml_cast(ctx.ggml, attention_prefix_key->tensor, k.type), + attention_prefix_key->shape, + k.type); + } + if (attention_prefix_value->type != v.type) { + attention_prefix_value = core::wrap_tensor( + ggml_cast(ctx.ggml, attention_prefix_value->tensor, v.type), + attention_prefix_value->shape, + v.type); + } + all_k = ConcatModule({1}).build(ctx, *attention_prefix_key, k); + all_v = ConcatModule({1}).build(ctx, *attention_prefix_value, v); + } else if (prefix_key.has_value()) { + all_k = ConcatModule({1}).build(ctx, *prefix_key, k); + all_v = ConcatModule({1}).build(ctx, *prefix_value, v); + } core::TensorValue context; if (!prefix_key.has_value() && attention_mask.has_value() && config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGroupedViewKV) { @@ -520,10 +573,9 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build( *attention_mask, config_.attention_precision); } else if (attention_mask.has_value() && - (config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGrouped || - (prefix_key.has_value() && - config_.runtime.attention.prefix_mode == QwenDecoderPrefixAttentionMode::FlashWithPrefix && - config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGroupedViewKV))) { + ((!prefix_key.has_value() && + config_.runtime.attention.prefill_mode == QwenDecoderAttentionMode::FlashGrouped) || + use_prefix_flash)) { q_heads = core::wrap_tensor(ggml_cont(ctx.ggml, q_heads.tensor), q_heads.shape, q_heads.type); auto k_heads = TransposeModule({{0, 2, 1, 3}, all_k.shape.rank}).build(ctx, all_k); auto v_heads = TransposeModule({{0, 2, 1, 3}, all_v.shape.rank}).build(ctx, all_v); diff --git a/src/framework/runtime/kv_cache.cpp b/src/framework/runtime/kv_cache.cpp index 17e236fa..0d7f9208 100644 --- a/src/framework/runtime/kv_cache.cpp +++ b/src/framework/runtime/kv_cache.cpp @@ -11,24 +11,44 @@ namespace engine::runtime { namespace { -void write_cache_tensor(const core::TensorValue & tensor, const std::vector & values) { +void validate_cache_tensor(const core::TensorValue & tensor, const TransformerKVCacheOptions & options) { + if (tensor.type == GGML_TYPE_F32) { + return; + } + if (options.allow_f16_storage && tensor.type == GGML_TYPE_F16) { + return; + } + throw std::runtime_error( + options.allow_f16_storage + ? "TransformerKVCache supports only f32 and f16 cache tensors" + : "TransformerKVCache requires f32 cache tensors"); +} + +void write_cache_tensor( + const core::TensorValue & tensor, + const std::vector & values, + const TransformerKVCacheOptions & options) { + validate_cache_tensor(tensor, options); if (tensor.type == GGML_TYPE_F32) { core::write_tensor_f32(tensor, values); - } else if (tensor.type == GGML_TYPE_F16) { + return; + } + if (options.allow_f16_storage && tensor.type == GGML_TYPE_F16) { core::write_tensor_f16(tensor, values); - } else { - throw std::runtime_error("TransformerKVCache supports only f32 and f16 cache tensors"); + return; } + throw std::runtime_error("TransformerKVCache requires f32 cache tensors"); } -std::vector read_cache_tensor(const core::TensorValue & tensor) { +std::vector read_cache_tensor(const core::TensorValue & tensor, const TransformerKVCacheOptions & options) { + validate_cache_tensor(tensor, options); if (tensor.type == GGML_TYPE_F32) { return core::read_tensor_f32(tensor.tensor); } - if (tensor.type == GGML_TYPE_F16) { + if (options.allow_f16_storage && tensor.type == GGML_TYPE_F16) { return core::read_tensor_f16(tensor.tensor); } - throw std::runtime_error("TransformerKVCache supports only f32 and f16 cache tensors"); + throw std::runtime_error("TransformerKVCache requires f32 cache tensors"); } } // namespace @@ -38,8 +58,17 @@ TransformerKVCache::TransformerKVCache( int64_t step_elems, std::vector keys, std::vector values) + : TransformerKVCache(cache_steps, step_elems, std::move(keys), std::move(values), {}) {} + +TransformerKVCache::TransformerKVCache( + int64_t cache_steps, + int64_t step_elems, + std::vector keys, + std::vector values, + TransformerKVCacheOptions options) : cache_steps_(std::max(0, cache_steps)), - step_elems_(std::max(0, step_elems)) { + step_elems_(std::max(0, step_elems)), + options_(options) { if (step_elems_ <= 0) { throw std::runtime_error("TransformerKVCache requires positive step_elems"); } @@ -49,6 +78,8 @@ TransformerKVCache::TransformerKVCache( const size_t cache_elems = static_cast(cache_steps_ * step_elems_); layers_.reserve(keys.size()); for (size_t layer = 0; layer < keys.size(); ++layer) { + validate_cache_tensor(keys[layer], options_); + validate_cache_tensor(values[layer], options_); layers_.push_back(LayerCache{ std::move(keys[layer]), std::move(values[layer]), @@ -92,8 +123,8 @@ void TransformerKVCache::import_state(const TransformerKVState & state) { std::copy(source.key.begin(), source.key.end(), cache.import_key_scratch.begin()); std::copy(source.value.begin(), source.value.end(), cache.import_value_scratch.begin()); } - write_cache_tensor(cache.key_tensor, cache.import_key_scratch); - write_cache_tensor(cache.value_tensor, cache.import_value_scratch); + write_cache_tensor(cache.key_tensor, cache.import_key_scratch, options_); + write_cache_tensor(cache.value_tensor, cache.import_value_scratch, options_); } } } @@ -109,8 +140,8 @@ TransformerKVState TransformerKVCache::export_state() const { if (keep_elems == 0) { continue; } - const auto key_values = read_cache_tensor(layers_[layer].key_tensor); - const auto value_values = read_cache_tensor(layers_[layer].value_tensor); + const auto key_values = read_cache_tensor(layers_[layer].key_tensor, options_); + const auto value_values = read_cache_tensor(layers_[layer].value_tensor, options_); out.key.assign(key_values.begin(), key_values.begin() + static_cast(keep_elems)); out.value.assign(value_values.begin(), value_values.begin() + static_cast(keep_elems)); } @@ -165,11 +196,11 @@ void TransformerKVCache::trace_log_state(const std::string & name, int64_t num_h return; } const size_t keep_elems = static_cast(valid_steps_ * step_elems_); - const auto first_key = read_cache_tensor(layers_.front().key_tensor); + const auto first_key = read_cache_tensor(layers_.front().key_tensor, options_); std::vector first_key_keep(first_key.begin(), first_key.begin() + static_cast(keep_elems)); debug::trace_log_f32(name + ".layer0.key", {1, valid_steps_, num_heads, head_dim}, first_key_keep); if (layers_.size() > 1) { - const auto last_key = read_cache_tensor(layers_.back().key_tensor); + const auto last_key = read_cache_tensor(layers_.back().key_tensor, options_); std::vector last_key_keep(last_key.begin(), last_key.begin() + static_cast(keep_elems)); debug::trace_log_f32(name + ".layer_last.key", {1, valid_steps_, num_heads, head_dim}, last_key_keep); } @@ -182,7 +213,8 @@ core::TensorValue view_transformer_kv_cache_steps( int64_t steps, int64_t heads, int64_t head_dim, - const char * label) { + const char * label, + ggml_type view_type) { if (start < 0 || steps <= 0 || start + steps > cache.shape.dims[1]) { throw std::runtime_error(std::string(label) + " cache view range is invalid"); } @@ -199,7 +231,7 @@ core::TensorValue view_transformer_kv_cache_steps( cache.tensor->nb[3], static_cast(start) * cache.tensor->nb[2]), core::TensorShape::from_dims({1, steps, heads, head_dim}), - cache.type); + view_type); } } // namespace engine::runtime diff --git a/src/models/higgs_tts/ar.cpp b/src/models/higgs_tts/ar.cpp index 5749ffef..21161241 100644 --- a/src/models/higgs_tts/ar.cpp +++ b/src/models/higgs_tts/ar.cpp @@ -450,11 +450,14 @@ struct HiggsARKVCache::Impl { GGML_TYPE_F16, core::TensorShape::from_dims({1, cache_steps, config.text.num_key_value_heads, dim}))); } + runtime::TransformerKVCacheOptions cache_options; + cache_options.allow_f16_storage = true; cache = runtime::TransformerKVCache( cache_steps, config.text.num_key_value_heads * dim, std::move(key_tensors), - std::move(value_tensors)); + std::move(value_tensors), + cache_options); buffer = ggml_backend_alloc_ctx_tensors(ctx.get(), runtime->backend()); if (buffer == nullptr) { throw std::runtime_error("failed to allocate Higgs TTS AR KV cache"); diff --git a/tests/unittests/test_conv_lowering_matrix.cpp b/tests/unittests/test_conv_lowering_matrix.cpp deleted file mode 100644 index 0344132e..00000000 --- a/tests/unittests/test_conv_lowering_matrix.cpp +++ /dev/null @@ -1,834 +0,0 @@ -#include "engine/framework/core/backend.h" -#include "engine/framework/modules/conv_modules.h" -#include "engine/framework/modules/streaming_conv_modules.h" -#include "engine/framework/modules/structural_modules.h" - -#include -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace { - -constexpr size_t kGraphBytes = 512 * 1024 * 1024; -constexpr size_t kGraphNodes = 16384; -constexpr int kWarmupRounds = 1; -constexpr int kMeasureRounds = 5; - -struct DiffStats { - float max_abs = 0.0f; - double mean_abs = 0.0; - double cosine = 1.0; -}; - -struct RunResult { - bool supported = false; - std::string error; - engine::core::TensorShape shape = {}; - std::vector values; - double avg_ms = 0.0; -}; - -std::vector make_patterned_f32(size_t count, float phase, float scale) { - std::vector values(count, 0.0f); - for (size_t i = 0; i < count; ++i) { - const float x = static_cast(i); - values[i] = scale * ( - std::sin(phase + 0.113f * x) + - 0.5f * std::cos(phase * 0.7f + 0.071f * x)); - } - return values; -} - -int64_t conv_out(int64_t input, int64_t kernel, int stride, int padding, int dilation) { - return (input + 2 * padding - dilation * (kernel - 1) - 1) / stride + 1; -} - -int64_t conv_transpose_out(int64_t input, int64_t kernel, int stride, int padding, int dilation) { - return (input - 1) * stride - 2 * padding + dilation * (kernel - 1) + 1; -} - -const char * backend_name(engine::core::BackendType backend_type) { - switch (backend_type) { - case engine::core::BackendType::Cpu: return "cpu"; - case engine::core::BackendType::Cuda: return "cuda"; - case engine::core::BackendType::Vulkan: return "vulkan"; - case engine::core::BackendType::Metal: return "metal"; - default: return "unknown"; - } -} - -bool same_shape(const engine::core::TensorShape & lhs, const engine::core::TensorShape & rhs) { - if (lhs.rank != rhs.rank) { - return false; - } - for (size_t i = 0; i < lhs.rank; ++i) { - if (lhs.dims[i] != rhs.dims[i]) { - return false; - } - } - return true; -} - -DiffStats diff_values(const std::vector & reference, const std::vector & actual) { - if (reference.size() != actual.size()) { - throw std::runtime_error("value count mismatch"); - } - DiffStats stats; - double dot = 0.0; - double ref_norm = 0.0; - double actual_norm = 0.0; - for (size_t i = 0; i < reference.size(); ++i) { - const float diff = std::fabs(reference[i] - actual[i]); - stats.max_abs = std::max(stats.max_abs, diff); - stats.mean_abs += diff; - dot += static_cast(reference[i]) * static_cast(actual[i]); - ref_norm += static_cast(reference[i]) * static_cast(reference[i]); - actual_norm += static_cast(actual[i]) * static_cast(actual[i]); - } - stats.mean_abs /= static_cast(reference.size()); - if (ref_norm > 0.0 && actual_norm > 0.0) { - stats.cosine = dot / (std::sqrt(ref_norm) * std::sqrt(actual_norm)); - } - return stats; -} - -engine::core::TensorValue add_bias_3d( - engine::core::ModuleBuildContext & ctx, - const engine::core::TensorValue & output, - int64_t channels, - const std::optional & bias) { - if (!bias.has_value()) { - return output; - } - auto output_contiguous = engine::core::ensure_backend_addressable_layout(ctx, output); - auto bias_view = engine::core::reshape_tensor(ctx, *bias, engine::core::TensorShape::from_dims({1, channels, 1})); - auto repeated = engine::core::wrap_tensor( - ggml_repeat(ctx.ggml, bias_view.tensor, output_contiguous.tensor), - output.shape, - GGML_TYPE_F32); - return engine::core::wrap_tensor(ggml_add(ctx.ggml, output_contiguous.tensor, repeated.tensor), output.shape, GGML_TYPE_F32); -} - -engine::core::TensorValue add_bias_4d( - engine::core::ModuleBuildContext & ctx, - const engine::core::TensorValue & output, - int64_t channels, - const std::optional & bias) { - if (!bias.has_value()) { - return output; - } - auto output_contiguous = engine::core::ensure_backend_addressable_layout(ctx, output); - auto bias_view = engine::core::reshape_tensor(ctx, *bias, engine::core::TensorShape::from_dims({1, channels, 1, 1})); - auto repeated = engine::core::wrap_tensor( - ggml_repeat(ctx.ggml, bias_view.tensor, output_contiguous.tensor), - output.shape, - GGML_TYPE_F32); - return engine::core::wrap_tensor(ggml_add(ctx.ggml, output_contiguous.tensor, repeated.tensor), output.shape, GGML_TYPE_F32); -} - -engine::core::TensorValue view_batch_matrix( - engine::core::ModuleBuildContext & ctx, - const engine::core::TensorValue & input, - int64_t batch_index, - int64_t channels, - int64_t frames) { - auto * view = ggml_view_2d( - ctx.ggml, - input.tensor, - frames, - channels, - input.tensor->nb[1], - static_cast(batch_index) * input.tensor->nb[2]); - return engine::core::wrap_tensor(view, engine::core::TensorShape::from_dims({channels, frames}), input.type); -} - -class GraphRunner { -public: - GraphRunner(const char * name, engine::core::BackendType backend_type) : backend_type_(backend_type) { - backend_ = engine::core::init_backend({backend_type, 0, 8}); - engine::core::set_backend_threads(backend_, 8); - ggml_init_params params{}; - params.mem_size = kGraphBytes; - params.mem_buffer = nullptr; - params.no_alloc = true; - ggml_ = ggml_init(params); - if (ggml_ == nullptr) { - throw std::runtime_error("failed to initialize ggml test context"); - } - ctx_.ggml = ggml_; - ctx_.module_instance_name = name; - ctx_.backend_type = backend_type; - } - - ~GraphRunner() { - if (buffer_ != nullptr) { - ggml_backend_buffer_free(buffer_); - } - if (ggml_ != nullptr) { - ggml_free(ggml_); - } - if (backend_ != nullptr) { - ggml_backend_free(backend_); - } - } - - engine::core::TensorValue make_f32(const engine::core::TensorShape & shape) { - return engine::core::make_tensor(ctx_, GGML_TYPE_F32, shape); - } - - engine::core::ModuleBuildContext & ctx() noexcept { return ctx_; } - - RunResult run( - const engine::core::TensorValue & output, - const std::vector>> & writes) { - ggml_cgraph * graph = ggml_new_graph_custom(ggml_, kGraphNodes, false); - ggml_build_forward_expand(graph, output.tensor); - engine::core::validate_backend_graph_supported(backend_, graph, "conv_lowering_matrix"); - buffer_ = ggml_backend_alloc_ctx_tensors(ggml_, backend_); - if (buffer_ == nullptr) { - throw std::runtime_error("failed to allocate backend tensors"); - } - for (const auto & write : writes) { - engine::core::write_tensor_f32(write.first, write.second); - } - for (int i = 0; i < kWarmupRounds; ++i) { - if (ggml_backend_graph_compute(backend_, graph) != GGML_STATUS_SUCCESS) { - throw std::runtime_error("warmup graph compute failed"); - } - } - double total_ms = 0.0; - for (int i = 0; i < kMeasureRounds; ++i) { - const auto start = std::chrono::steady_clock::now(); - if (ggml_backend_graph_compute(backend_, graph) != GGML_STATUS_SUCCESS) { - throw std::runtime_error("graph compute failed"); - } - const auto end = std::chrono::steady_clock::now(); - total_ms += std::chrono::duration(end - start).count(); - } - RunResult result; - result.supported = true; - result.shape = output.shape; - result.avg_ms = total_ms / static_cast(kMeasureRounds); - engine::core::read_tensor_f32_into(output.tensor, result.values); - return result; - } - -private: - engine::core::BackendType backend_type_; - ggml_backend_t backend_ = nullptr; - ggml_backend_buffer_t buffer_ = nullptr; - ggml_context * ggml_ = nullptr; - engine::core::ModuleBuildContext ctx_{}; -}; - -template -RunResult run_guarded( - const char * label, - engine::core::BackendType backend_type, - Fn && fn) { - try { - GraphRunner runner(label, backend_type); - return fn(runner); - } catch (const std::exception & ex) { - RunResult result; - result.supported = false; - result.error = ex.what(); - return result; - } -} - -struct Conv1dCase { - const char * name; - int64_t batch; - int64_t in_channels; - int64_t out_channels; - int64_t frames; - int64_t kernel; - int stride; - int padding; - int dilation; - bool bias; -}; - -RunResult run_conv1d(const Conv1dCase & c, const char * candidate, engine::core::BackendType backend_type) { - return run_guarded(candidate, backend_type, [&](GraphRunner & runner) { - const auto input_shape = engine::core::TensorShape::from_dims({c.batch, c.in_channels, c.frames}); - const auto weight_shape = engine::core::TensorShape::from_dims({c.out_channels, c.in_channels, c.kernel}); - const auto bias_shape = engine::core::TensorShape::from_dims({c.out_channels}); - auto input = runner.make_f32(input_shape); - auto weight = runner.make_f32(weight_shape); - std::optional bias = c.bias ? std::optional(runner.make_f32(bias_shape)) : std::nullopt; - - engine::core::TensorValue output; - if (std::string(candidate) == "native") { - const auto output_shape = engine::core::TensorShape::from_dims( - {c.batch, c.out_channels, conv_out(c.frames, c.kernel, c.stride, c.padding, c.dilation)}); - if (c.batch == 1) { - output = engine::core::wrap_tensor( - ggml_conv_1d(runner.ctx().ggml, weight.tensor, input.tensor, c.stride, c.padding, c.dilation), - output_shape, - GGML_TYPE_F32); - } else { - for (int64_t batch = 0; batch < c.batch; ++batch) { - auto batch_input = view_batch_matrix(runner.ctx(), input, batch, c.in_channels, c.frames); - auto batch_output = engine::core::wrap_tensor( - ggml_conv_1d(runner.ctx().ggml, weight.tensor, batch_input.tensor, c.stride, c.padding, c.dilation), - engine::core::TensorShape::from_dims({1, c.out_channels, output_shape.dims[2]}), - GGML_TYPE_F32); - output = output.valid() ? engine::modules::ConcatModule({0}).build(runner.ctx(), output, batch_output) : batch_output; - } - } - output = add_bias_3d(runner.ctx(), output, c.out_channels, bias); - } else if (std::string(candidate) == "conv2d_normal") { - auto x4 = engine::core::reshape_tensor(runner.ctx(), input, engine::core::TensorShape::from_dims({c.batch, c.in_channels, 1, c.frames})); - auto w4 = engine::core::reshape_tensor(runner.ctx(), weight, engine::core::TensorShape::from_dims({c.out_channels, c.in_channels, 1, c.kernel})); - auto y4 = engine::core::wrap_tensor( - ggml_conv_2d(runner.ctx().ggml, w4.tensor, x4.tensor, c.stride, 1, c.padding, 0, c.dilation, 1), - engine::core::TensorShape::from_dims({c.batch, c.out_channels, 1, conv_out(c.frames, c.kernel, c.stride, c.padding, c.dilation)}), - GGML_TYPE_F32); - y4 = add_bias_4d(runner.ctx(), y4, c.out_channels, bias); - output = engine::core::reshape_tensor(runner.ctx(), y4, engine::core::TensorShape::from_dims({c.batch, c.out_channels, y4.shape.dims[3]})); - } else if (std::string(candidate) == "conv2d_direct") { - auto x4 = engine::core::reshape_tensor(runner.ctx(), input, engine::core::TensorShape::from_dims({c.batch, c.in_channels, 1, c.frames})); - auto w4 = engine::core::reshape_tensor(runner.ctx(), weight, engine::core::TensorShape::from_dims({c.out_channels, c.in_channels, 1, c.kernel})); - auto y4 = engine::core::wrap_tensor( - ggml_conv_2d_direct(runner.ctx().ggml, w4.tensor, x4.tensor, c.stride, 1, c.padding, 0, c.dilation, 1), - engine::core::TensorShape::from_dims({c.batch, c.out_channels, 1, conv_out(c.frames, c.kernel, c.stride, c.padding, c.dilation)}), - GGML_TYPE_F32); - y4 = add_bias_4d(runner.ctx(), y4, c.out_channels, bias); - output = engine::core::reshape_tensor(runner.ctx(), y4, engine::core::TensorShape::from_dims({c.batch, c.out_channels, y4.shape.dims[3]})); - } else { - throw std::runtime_error("unknown conv1d candidate"); - } - - std::vector>> writes; - writes.push_back({input, make_patterned_f32(static_cast(input_shape.num_elements()), 0.19f, 0.031f)}); - writes.push_back({weight, make_patterned_f32(static_cast(weight_shape.num_elements()), 0.47f, 0.017f)}); - if (bias) { - writes.push_back({*bias, make_patterned_f32(static_cast(bias_shape.num_elements()), 0.83f, 0.011f)}); - } - return runner.run(output, writes); - }); -} - -struct Conv2dCase { - const char * name; - int64_t batch; - int64_t in_channels; - int64_t out_channels; - int64_t height; - int64_t width; - int64_t kernel_h; - int64_t kernel_w; - int stride_h; - int stride_w; - int padding_h; - int padding_w; - int dilation_h; - int dilation_w; - bool bias; -}; - -RunResult run_conv2d(const Conv2dCase & c, const char * candidate, engine::core::BackendType backend_type) { - return run_guarded(candidate, backend_type, [&](GraphRunner & runner) { - const auto input_shape = engine::core::TensorShape::from_dims({c.batch, c.in_channels, c.height, c.width}); - const auto weight_shape = engine::core::TensorShape::from_dims({c.out_channels, c.in_channels, c.kernel_h, c.kernel_w}); - const auto bias_shape = engine::core::TensorShape::from_dims({c.out_channels}); - auto input = runner.make_f32(input_shape); - auto weight = runner.make_f32(weight_shape); - std::optional bias = c.bias ? std::optional(runner.make_f32(bias_shape)) : std::nullopt; - - engine::core::TensorValue output; - const auto output_shape = engine::core::TensorShape::from_dims({ - c.batch, - c.out_channels, - conv_out(c.height, c.kernel_h, c.stride_h, c.padding_h, c.dilation_h), - conv_out(c.width, c.kernel_w, c.stride_w, c.padding_w, c.dilation_w), - }); - if (std::string(candidate) == "im2col_matmul") { - output = engine::core::wrap_tensor( - ggml_conv_2d(runner.ctx().ggml, weight.tensor, input.tensor, c.stride_w, c.stride_h, c.padding_w, c.padding_h, c.dilation_w, c.dilation_h), - output_shape, - GGML_TYPE_F32); - output = add_bias_4d(runner.ctx(), output, c.out_channels, bias); - } else if (std::string(candidate) == "direct") { - output = engine::core::wrap_tensor( - ggml_conv_2d_direct(runner.ctx().ggml, weight.tensor, input.tensor, c.stride_w, c.stride_h, c.padding_w, c.padding_h, c.dilation_w, c.dilation_h), - output_shape, - GGML_TYPE_F32); - output = add_bias_4d(runner.ctx(), output, c.out_channels, bias); - } else { - throw std::runtime_error("unknown conv2d candidate"); - } - std::vector>> writes; - writes.push_back({input, make_patterned_f32(static_cast(input_shape.num_elements()), 0.21f, 0.021f)}); - writes.push_back({weight, make_patterned_f32(static_cast(weight_shape.num_elements()), 0.51f, 0.013f)}); - if (bias) { - writes.push_back({*bias, make_patterned_f32(static_cast(bias_shape.num_elements()), 0.91f, 0.009f)}); - } - return runner.run(output, writes); - }); -} - -struct Depthwise1dCase { - const char * name; - int64_t batch; - int64_t channels; - int64_t frames; - int64_t kernel; - int stride; - int padding; - int dilation; - bool bias; -}; - -RunResult run_depthwise1d(const Depthwise1dCase & c, const char * candidate, engine::core::BackendType backend_type) { - return run_guarded(candidate, backend_type, [&](GraphRunner & runner) { - const auto input_shape = engine::core::TensorShape::from_dims({c.batch, c.channels, c.frames}); - const auto weight_shape = engine::core::TensorShape::from_dims({c.channels, 1, c.kernel}); - const auto bias_shape = engine::core::TensorShape::from_dims({c.channels}); - auto input = runner.make_f32(input_shape); - auto weight = runner.make_f32(weight_shape); - std::optional bias = c.bias ? std::optional(runner.make_f32(bias_shape)) : std::nullopt; - engine::core::TensorValue output; - if (std::string(candidate) == "dw2d_direct") { - auto x4 = engine::core::reshape_tensor(runner.ctx(), input, engine::core::TensorShape::from_dims({c.batch, c.channels, 1, c.frames})); - auto w4 = engine::core::reshape_tensor(runner.ctx(), weight, engine::core::TensorShape::from_dims({c.channels, 1, 1, c.kernel})); - auto y4 = engine::core::wrap_tensor( - ggml_conv_2d_dw_direct(runner.ctx().ggml, w4.tensor, x4.tensor, c.stride, 1, c.padding, 0, c.dilation, 1), - engine::core::TensorShape::from_dims({c.batch, c.channels, 1, conv_out(c.frames, c.kernel, c.stride, c.padding, c.dilation)}), - GGML_TYPE_F32); - y4 = add_bias_4d(runner.ctx(), y4, c.channels, bias); - output = engine::core::reshape_tensor(runner.ctx(), y4, engine::core::TensorShape::from_dims({c.batch, c.channels, y4.shape.dims[3]})); - } else if (std::string(candidate) == "native_1d_dw") { - if (c.batch != 1) { - throw std::runtime_error("native ggml_conv_1d_dw asserts for batched rank-3 input; slice batch first"); - } - output = engine::core::wrap_tensor( - ggml_conv_1d_dw(runner.ctx().ggml, weight.tensor, input.tensor, c.stride, c.padding, c.dilation), - engine::core::TensorShape::from_dims({c.batch, c.channels, conv_out(c.frames, c.kernel, c.stride, c.padding, c.dilation)}), - GGML_TYPE_F32); - output = add_bias_3d(runner.ctx(), output, c.channels, bias); - } else { - throw std::runtime_error("unknown depthwise1d candidate"); - } - std::vector>> writes; - writes.push_back({input, make_patterned_f32(static_cast(input_shape.num_elements()), 0.23f, 0.025f)}); - writes.push_back({weight, make_patterned_f32(static_cast(weight_shape.num_elements()), 0.53f, 0.015f)}); - if (bias) { - writes.push_back({*bias, make_patterned_f32(static_cast(bias_shape.num_elements()), 0.93f, 0.007f)}); - } - return runner.run(output, writes); - }); -} - -struct Pointwise1dCase { - const char * name; - int64_t batch; - int64_t in_channels; - int64_t out_channels; - int64_t frames; - bool bias; -}; - -RunResult run_pointwise1d(const Pointwise1dCase & c, const char * candidate, engine::core::BackendType backend_type) { - return run_guarded(candidate, backend_type, [&](GraphRunner & runner) { - const auto input_shape = engine::core::TensorShape::from_dims({c.batch, c.in_channels, c.frames}); - const auto weight_shape = engine::core::TensorShape::from_dims({c.out_channels, c.in_channels, 1}); - const auto bias_shape = engine::core::TensorShape::from_dims({c.out_channels}); - auto input = runner.make_f32(input_shape); - auto weight = runner.make_f32(weight_shape); - std::optional bias = c.bias ? std::optional(runner.make_f32(bias_shape)) : std::nullopt; - - engine::core::TensorValue output; - if (std::string(candidate) == "conv1d_kernel1") { - const auto output_shape = engine::core::TensorShape::from_dims({c.batch, c.out_channels, c.frames}); - if (c.batch == 1) { - output = engine::core::wrap_tensor( - ggml_conv_1d(runner.ctx().ggml, weight.tensor, input.tensor, 1, 0, 1), - output_shape, - GGML_TYPE_F32); - } else { - for (int64_t batch = 0; batch < c.batch; ++batch) { - auto batch_input = view_batch_matrix(runner.ctx(), input, batch, c.in_channels, c.frames); - auto batch_output = engine::core::wrap_tensor( - ggml_conv_1d(runner.ctx().ggml, weight.tensor, batch_input.tensor, 1, 0, 1), - engine::core::TensorShape::from_dims({1, c.out_channels, c.frames}), - GGML_TYPE_F32); - output = output.valid() ? engine::modules::ConcatModule({0}).build(runner.ctx(), output, batch_output) : batch_output; - } - } - output = add_bias_3d(runner.ctx(), output, c.out_channels, bias); - } else if (std::string(candidate) == "linear_matmul") { - auto x = engine::modules::TransposeModule({{0, 2, 1, 3}, 3}).build(runner.ctx(), input); - x = engine::core::ensure_backend_addressable_layout(runner.ctx(), x); - auto matrix = engine::core::reshape_tensor(runner.ctx(), x, engine::core::TensorShape::from_dims({c.batch * c.frames, c.in_channels})); - auto w2 = engine::core::reshape_tensor(runner.ctx(), weight, engine::core::TensorShape::from_dims({c.out_channels, c.in_channels})); - auto projected = engine::core::wrap_tensor( - ggml_mul_mat(runner.ctx().ggml, w2.tensor, matrix.tensor), - engine::core::TensorShape::from_dims({c.batch * c.frames, c.out_channels}), - GGML_TYPE_F32); - if (bias) { - projected = engine::core::wrap_tensor(ggml_add(runner.ctx().ggml, projected.tensor, bias->tensor), projected.shape, GGML_TYPE_F32); - } - auto y = engine::core::reshape_tensor(runner.ctx(), projected, engine::core::TensorShape::from_dims({c.batch, c.frames, c.out_channels})); - output = engine::modules::TransposeModule({{0, 2, 1, 3}, 3}).build(runner.ctx(), y); - } else { - throw std::runtime_error("unknown pointwise1d candidate"); - } - std::vector>> writes; - writes.push_back({input, make_patterned_f32(static_cast(input_shape.num_elements()), 0.24f, 0.027f)}); - writes.push_back({weight, make_patterned_f32(static_cast(weight_shape.num_elements()), 0.54f, 0.014f)}); - if (bias) { - writes.push_back({*bias, make_patterned_f32(static_cast(bias_shape.num_elements()), 0.94f, 0.007f)}); - } - return runner.run(output, writes); - }); -} - -RunResult run_depthwise2d(const Conv2dCase & c, const char * candidate, engine::core::BackendType backend_type) { - return run_guarded(candidate, backend_type, [&](GraphRunner & runner) { - const auto input_shape = engine::core::TensorShape::from_dims({c.batch, c.in_channels, c.height, c.width}); - const auto weight_shape = engine::core::TensorShape::from_dims({c.in_channels, 1, c.kernel_h, c.kernel_w}); - const auto bias_shape = engine::core::TensorShape::from_dims({c.in_channels}); - auto input = runner.make_f32(input_shape); - auto weight = runner.make_f32(weight_shape); - std::optional bias = c.bias ? std::optional(runner.make_f32(bias_shape)) : std::nullopt; - const auto output_shape = engine::core::TensorShape::from_dims({ - c.batch, - c.in_channels, - conv_out(c.height, c.kernel_h, c.stride_h, c.padding_h, c.dilation_h), - conv_out(c.width, c.kernel_w, c.stride_w, c.padding_w, c.dilation_w), - }); - engine::core::TensorValue output; - if (std::string(candidate) == "direct") { - output = engine::core::wrap_tensor( - ggml_conv_2d_dw_direct(runner.ctx().ggml, weight.tensor, input.tensor, c.stride_w, c.stride_h, c.padding_w, c.padding_h, c.dilation_w, c.dilation_h), - output_shape, - GGML_TYPE_F32); - output = add_bias_4d(runner.ctx(), output, c.in_channels, bias); - } else if (std::string(candidate) == "im2col_matmul") { - output = engine::core::wrap_tensor( - ggml_conv_2d_dw(runner.ctx().ggml, weight.tensor, input.tensor, c.stride_w, c.stride_h, c.padding_w, c.padding_h, c.dilation_w, c.dilation_h), - output_shape, - GGML_TYPE_F32); - output = add_bias_4d(runner.ctx(), output, c.in_channels, bias); - } else { - throw std::runtime_error("unknown depthwise2d candidate"); - } - std::vector>> writes; - writes.push_back({input, make_patterned_f32(static_cast(input_shape.num_elements()), 0.25f, 0.023f)}); - writes.push_back({weight, make_patterned_f32(static_cast(weight_shape.num_elements()), 0.55f, 0.012f)}); - if (bias) { - writes.push_back({*bias, make_patterned_f32(static_cast(bias_shape.num_elements()), 0.95f, 0.008f)}); - } - return runner.run(output, writes); - }); -} - -struct ConvTranspose1dCase { - const char * name; - int64_t batch; - int64_t in_channels; - int64_t out_channels; - int64_t frames; - int64_t kernel; - int stride; - int padding; - int dilation; - bool bias; -}; - -engine::core::TensorValue build_conv_transpose_native( - engine::core::ModuleBuildContext & ctx, - const ConvTranspose1dCase & c, - const engine::core::TensorValue & input, - const engine::core::TensorValue & weight, - const std::optional & bias) { - if (c.padding != 0 || c.dilation != 1) { - throw std::runtime_error("native ggml_conv_transpose_1d supports only padding=0 and dilation=1"); - } - engine::core::TensorValue output; - for (int64_t batch = 0; batch < c.batch; ++batch) { - auto matrix = view_batch_matrix(ctx, input, batch, c.in_channels, c.frames); - auto batch_out = engine::core::wrap_tensor( - ggml_conv_transpose_1d(ctx.ggml, weight.tensor, matrix.tensor, c.stride, c.padding, c.dilation), - engine::core::TensorShape::from_dims({1, c.out_channels, conv_transpose_out(c.frames, c.kernel, c.stride, c.padding, c.dilation)}), - GGML_TYPE_F32); - output = output.valid() ? engine::modules::ConcatModule({0}).build(ctx, output, batch_out) : batch_out; - } - return add_bias_3d(ctx, output, c.out_channels, bias); -} - -engine::core::TensorValue build_conv_transpose_col2im( - engine::core::ModuleBuildContext & ctx, - const ConvTranspose1dCase & c, - const engine::core::TensorValue & input, - const engine::core::TensorValue & weight, - const std::optional & bias) { - if (c.dilation != 1) { - throw std::runtime_error("col2im lowering currently supports only dilation=1"); - } - auto * weight_perm = ggml_reshape_2d( - ctx.ggml, - ggml_cont(ctx.ggml, ggml_permute(ctx.ggml, weight.tensor, 1, 2, 0, 3)), - c.in_channels, - c.kernel * c.out_channels); - ggml_tensor * bias_matrix = nullptr; - if (c.bias) { - if (!bias.has_value()) { - throw std::runtime_error("missing bias"); - } - bias_matrix = ggml_reshape_2d(ctx.ggml, bias->tensor, 1, c.out_channels); - } - engine::core::TensorValue output; - for (int64_t batch = 0; batch < c.batch; ++batch) { - auto * batch_input = ggml_view_2d( - ctx.ggml, - input.tensor, - input.tensor->ne[0], - input.tensor->ne[1], - input.tensor->nb[1], - static_cast(batch) * input.tensor->nb[2]); - auto * transposed_input = ggml_cont(ctx.ggml, ggml_transpose(ctx.ggml, batch_input)); - auto * columns = ggml_mul_mat(ctx.ggml, weight_perm, transposed_input); - auto * batch_output = ggml_col2im_1d(ctx.ggml, columns, c.stride, static_cast(c.out_channels), c.padding); - if (bias_matrix != nullptr) { - batch_output = ggml_add(ctx.ggml, batch_output, bias_matrix); - } - auto batch_value = engine::core::wrap_tensor( - ggml_reshape_3d(ctx.ggml, batch_output, batch_output->ne[0], batch_output->ne[1], 1), - engine::core::TensorShape::from_dims({1, c.out_channels, batch_output->ne[0]}), - GGML_TYPE_F32); - output = output.valid() ? engine::modules::ConcatModule({0}).build(ctx, output, batch_value) : batch_value; - } - return output; -} - -RunResult run_conv_transpose1d(const ConvTranspose1dCase & c, const char * candidate, engine::core::BackendType backend_type) { - if (backend_type == engine::core::BackendType::Cpu && std::string(candidate) == "matmul_col2im") { - RunResult result; - result.supported = false; - result.error = "current ggml CPU backend aborts for COL2IM_1D"; - return result; - } - return run_guarded(candidate, backend_type, [&](GraphRunner & runner) { - const auto input_shape = engine::core::TensorShape::from_dims({c.batch, c.in_channels, c.frames}); - const auto weight_shape = engine::core::TensorShape::from_dims({c.in_channels, c.out_channels, c.kernel}); - const auto bias_shape = engine::core::TensorShape::from_dims({c.out_channels}); - auto input = runner.make_f32(input_shape); - auto weight = runner.make_f32(weight_shape); - std::optional bias = c.bias ? std::optional(runner.make_f32(bias_shape)) : std::nullopt; - engine::core::TensorValue output; - if (std::string(candidate) == "native_direct") { - output = build_conv_transpose_native(runner.ctx(), c, input, weight, bias); - } else if (std::string(candidate) == "matmul_col2im") { - output = build_conv_transpose_col2im(runner.ctx(), c, input, weight, bias); - } else { - throw std::runtime_error("unknown conv_transpose1d candidate"); - } - std::vector>> writes; - writes.push_back({input, make_patterned_f32(static_cast(input_shape.num_elements()), 0.27f, 0.019f)}); - writes.push_back({weight, make_patterned_f32(static_cast(weight_shape.num_elements()), 0.57f, 0.011f)}); - if (bias) { - writes.push_back({*bias, make_patterned_f32(static_cast(bias_shape.num_elements()), 0.97f, 0.006f)}); - } - return runner.run(output, writes); - }); -} - -struct MatrixRow { - std::string module; - std::string case_name; - std::string candidate; - engine::core::BackendType backend; - RunResult result; - std::optional diff; -}; - -void print_row(const MatrixRow & row) { - std::cout << "| " << row.module - << " | " << row.case_name - << " | " << row.candidate - << " | " << backend_name(row.backend) - << " | "; - if (!row.result.supported) { - std::string error = row.result.error; - std::replace(error.begin(), error.end(), '|', '/'); - std::cout << "unsupported | - | - | - | - | " << error << " |\n"; - return; - } - std::cout << "ok | " << row.result.shape.to_string() - << " | " << std::fixed << std::setprecision(4) << row.result.avg_ms - << " | "; - if (row.diff.has_value()) { - std::cout << std::scientific << std::setprecision(3) << row.diff->max_abs - << " | " << row.diff->mean_abs - << " | " << std::fixed << std::setprecision(9) << row.diff->cosine << " |\n"; - } else { - std::cout << "- | - | - |\n"; - } -} - -void add_result( - std::vector & rows, - const std::string & module, - const std::string & case_name, - const std::string & candidate, - engine::core::BackendType backend, - const RunResult & result, - const RunResult & reference) { - MatrixRow row{module, case_name, candidate, backend, result, std::nullopt}; - if (result.supported && reference.supported) { - if (!same_shape(reference.shape, result.shape)) { - row.result.supported = false; - row.result.error = "shape mismatch vs reference " + reference.shape.to_string(); - } else { - row.diff = diff_values(reference.values, result.values); - } - } - rows.push_back(std::move(row)); -} - -bool backend_available(engine::core::BackendType backend_type) { - try { - GraphRunner runner("conv_lowering_matrix.probe", backend_type); - return true; - } catch (...) { - return false; - } -} - -} // namespace - -int main() { - try { - std::vector backends = {engine::core::BackendType::Cpu}; - if (backend_available(engine::core::BackendType::Cuda)) { - backends.push_back(engine::core::BackendType::Cuda); - } else { - std::cout << "[SKIP] cuda backend unavailable\n"; - } - if (backend_available(engine::core::BackendType::Vulkan)) { - backends.push_back(engine::core::BackendType::Vulkan); - } else { - std::cout << "[SKIP] vulkan backend unavailable\n"; - } - - std::vector rows; - - const std::vector conv1d_cases = { - {"citrinet_like_large_regular", 1, 80, 256, 256, 11, 1, 5, 1, true}, - {"bigvgan_like_resblock", 1, 192, 192, 384, 7, 1, 3, 1, true}, - {"batched_stride_regular", 2, 64, 128, 160, 5, 2, 2, 1, true}, - {"dilated_regular", 1, 128, 128, 192, 3, 1, 2, 2, false}, - }; - const std::vector conv1d_candidates = {"native", "conv2d_normal", "conv2d_direct"}; - for (const auto & c : conv1d_cases) { - const auto reference = run_conv1d(c, "native", engine::core::BackendType::Cpu); - for (const auto backend : backends) { - for (const auto & candidate : conv1d_candidates) { - add_result(rows, "Conv1dModule", c.name, candidate, backend, run_conv1d(c, candidate.c_str(), backend), reference); - } - } - } - - const std::vector conv2d_cases = { - {"spectrogram_small_kernel", 1, 64, 128, 20, 160, 3, 3, 1, 1, 1, 1, 1, 1, true}, - {"conv1d_lowered_shape", 1, 256, 256, 1, 384, 1, 7, 1, 1, 0, 3, 1, 1, true}, - {"batched_feature_map", 2, 32, 64, 12, 96, 3, 5, 1, 2, 1, 2, 1, 1, false}, - }; - const std::vector conv2d_candidates = {"im2col_matmul", "direct"}; - for (const auto & c : conv2d_cases) { - const auto reference = run_conv2d(c, "im2col_matmul", engine::core::BackendType::Cpu); - for (const auto backend : backends) { - for (const auto & candidate : conv2d_candidates) { - add_result(rows, "Conv2dModule", c.name, candidate, backend, run_conv2d(c, candidate.c_str(), backend), reference); - } - } - } - - const std::vector depthwise1d_cases = { - {"conformer_like_depthwise", 1, 256, 192, 31, 1, 15, 1, true}, - {"tokenizer_stride_depthwise", 2, 96, 256, 7, 2, 3, 1, true}, - {"dilated_depthwise", 1, 128, 160, 5, 1, 4, 2, false}, - }; - const std::vector depthwise1d_candidates = {"dw2d_direct", "native_1d_dw"}; - for (const auto & c : depthwise1d_cases) { - const auto reference = run_depthwise1d(c, "dw2d_direct", engine::core::BackendType::Cpu); - for (const auto backend : backends) { - for (const auto & candidate : depthwise1d_candidates) { - add_result(rows, "DepthwiseConv1dModule", c.name, candidate, backend, run_depthwise1d(c, candidate.c_str(), backend), reference); - } - } - } - - const std::vector pointwise1d_cases = { - {"conformer_projection", 1, 256, 512, 192, true}, - {"batched_token_projection", 2, 192, 384, 160, true}, - {"vocoder_channel_mix", 1, 192, 192, 384, false}, - }; - const std::vector pointwise1d_candidates = {"conv1d_kernel1", "linear_matmul"}; - for (const auto & c : pointwise1d_cases) { - const auto reference = run_pointwise1d(c, "conv1d_kernel1", engine::core::BackendType::Cpu); - for (const auto backend : backends) { - for (const auto & candidate : pointwise1d_candidates) { - add_result(rows, "PointwiseConv1dModule", c.name, candidate, backend, run_pointwise1d(c, candidate.c_str(), backend), reference); - } - } - } - - const std::vector depthwise2d_cases = { - {"depthwise_1d_lowered_shape", 1, 192, 192, 1, 384, 1, 7, 1, 1, 0, 3, 1, 1, true}, - {"image_depthwise_small", 1, 64, 64, 24, 80, 3, 3, 1, 1, 1, 1, 1, 1, true}, - }; - const std::vector depthwise2d_candidates = {"direct", "im2col_matmul"}; - for (const auto & c : depthwise2d_cases) { - const auto reference = run_depthwise2d(c, "direct", engine::core::BackendType::Cpu); - for (const auto backend : backends) { - for (const auto & candidate : depthwise2d_candidates) { - add_result(rows, "DepthwiseConv2dModule", c.name, candidate, backend, run_depthwise2d(c, candidate.c_str(), backend), reference); - } - } - } - - const std::vector conv_transpose_cases = { - {"qwen3_like_stride5_padding0", 1, 256, 128, 96, 10, 5, 0, 1, true}, - {"vocoder_stride2_padding1", 1, 192, 96, 192, 4, 2, 1, 1, true}, - {"batched_stride2_no_bias", 2, 128, 128, 96, 2, 2, 0, 1, false}, - {"dilated_unsupported_probe", 1, 64, 64, 96, 3, 2, 0, 2, true}, - }; - const std::vector conv_transpose_candidates = {"native_direct", "matmul_col2im"}; - for (const auto & c : conv_transpose_cases) { - const auto reference = run_conv_transpose1d(c, "native_direct", engine::core::BackendType::Cpu); - for (const auto backend : backends) { - for (const auto & candidate : conv_transpose_candidates) { - add_result(rows, "ConvTranspose1dModule", c.name, candidate, backend, run_conv_transpose1d(c, candidate.c_str(), backend), reference); - } - } - } - - std::cout << "| module | case | candidate | backend | status | shape | avg_ms | max_abs_vs_cpu_ref | mean_abs_vs_cpu_ref | cosine_vs_cpu_ref |\n"; - std::cout << "| --- | --- | --- | --- | --- | --- | ---: | ---: | ---: | ---: |\n"; - for (const auto & row : rows) { - print_row(row); - } - } catch (const std::exception & ex) { - std::cerr << "[FAIL] " << ex.what() << '\n'; - return 1; - } - return 0; -} From d81583c246a8b978a3dd913c1c4d6e1672238743 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Mon, 20 Jul 2026 19:42:07 -0400 Subject: [PATCH 09/27] Restore Qwen decoder exact MLP parity --- src/framework/modules/attention/qwen_decoder.cpp | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/framework/modules/attention/qwen_decoder.cpp b/src/framework/modules/attention/qwen_decoder.cpp index f473e738..188fbc11 100644 --- a/src/framework/modules/attention/qwen_decoder.cpp +++ b/src/framework/modules/attention/qwen_decoder.cpp @@ -373,6 +373,9 @@ core::TensorValue build_mlp( config.projection_precision, }) .build(ctx, gated, require_linear(weights.down_proj, false, "QwenMLPWeights.down_proj")); + if (config.activation_cast.enabled && config.activation_cast.after_mlp_projection) { + down = activation_cast(ctx, down, config.activation_cast); + } return down; } @@ -746,8 +749,6 @@ QwenDecoderLayerOutputs QwenDecoderLayerModule::build_with_static_cache_tail( config_.runtime.attention.static_mode == QwenDecoderAttentionMode::ManualRepeatThenGroupedQuery) { k_heads = repeat_kv_heads(ctx, k_heads, kv_repeats); v_heads = repeat_kv_heads(ctx, v_heads, kv_repeats); - k_heads = core::wrap_tensor(ggml_cont(ctx.ggml, k_heads.tensor), k_heads.shape, k_heads.type); - v_heads = core::wrap_tensor(ggml_cont(ctx.ggml, v_heads.tensor), v_heads.shape, v_heads.type); context = attention_from_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask); } else if (config_.runtime.attention.static_mode == QwenDecoderAttentionMode::FlashGroupedViewKV) { context = flash_attention_from_grouped_heads_view_kv( From a0350926f7c88f28246eed2cef74e957e8853ee6 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Mon, 20 Jul 2026 20:25:34 -0400 Subject: [PATCH 10/27] Align Higgs and Fish token cap defaults --- include/engine/models/fish_audio/types.h | 4 +- include/engine/models/higgs_tts/generator.h | 2 +- src/models/fish_audio/loader.cpp | 5 +- src/models/fish_audio/session.cpp | 55 ++++++++++++++----- src/models/higgs_tts/generator.cpp | 2 +- src/models/higgs_tts/loader.cpp | 5 +- src/models/higgs_tts/session.cpp | 17 ++++-- ...audiocpp_cli_longform_tts_clone_cases.json | 3 +- .../audiocpp_cli/audiocpp_cli_path_cases.json | 6 +- 9 files changed, 66 insertions(+), 33 deletions(-) diff --git a/include/engine/models/fish_audio/types.h b/include/engine/models/fish_audio/types.h index 3c05e94a..68f7bffb 100644 --- a/include/engine/models/fish_audio/types.h +++ b/include/engine/models/fish_audio/types.h @@ -10,8 +10,8 @@ namespace engine::models::fish_audio { struct FishAudioGenerationOptions { - int64_t max_new_tokens = 1024; - int64_t chunk_length = 200; + int64_t max_new_tokens = 2048; + int64_t text_chunk_size = 200; float top_p = 0.8F; int top_k = 30; float temperature = 0.8F; diff --git a/include/engine/models/higgs_tts/generator.h b/include/engine/models/higgs_tts/generator.h index 4eb3e761..d943383a 100644 --- a/include/engine/models/higgs_tts/generator.h +++ b/include/engine/models/higgs_tts/generator.h @@ -14,7 +14,7 @@ namespace engine::models::higgs_tts { struct HiggsGenerationOptions { - int64_t max_tokens = 1024; + int64_t max_tokens = 2048; float temperature = 1.0F; std::optional top_p; std::optional top_k; diff --git a/src/models/fish_audio/loader.cpp b/src/models/fish_audio/loader.cpp index e7baf889..2f3a025b 100644 --- a/src/models/fish_audio/loader.cpp +++ b/src/models/fish_audio/loader.cpp @@ -34,8 +34,9 @@ runtime::ModelCliInterface cli(const FishAudioAssets &) { runtime::ModelCliInterface out; out.request_options = { {"reference_text", "TEXT", "Reference transcript used with speaker reference audio."}, - {"max_new_tokens", "N", "Maximum Fish Audio semantic tokens to generate."}, - {"chunk_length", "N", "Maximum UTF-8 bytes per Fish Speech text batch."}, + {"max_new_tokens", "N", "Maximum Fish Audio semantic tokens to generate; default 2048, 0 uses the default."}, + {"text_chunk_size", "N", "Long-form text chunk size; default 200."}, + {"text_chunk_mode", "default|tag_aware|japanese|endline", "Framework text chunking mode."}, {"top_p", "FLOAT", "Top-p sampling value."}, {"top_k", "N", "Top-k sampling value."}, {"temperature", "FLOAT", "Sampling temperature."}, diff --git a/src/models/fish_audio/session.cpp b/src/models/fish_audio/session.cpp index f362664c..7074a364 100644 --- a/src/models/fish_audio/session.cpp +++ b/src/models/fish_audio/session.cpp @@ -5,6 +5,7 @@ #include "engine/framework/io/filesystem.h" #include "engine/framework/runtime/options.h" #include "engine/framework/runtime/session.h" +#include "engine/framework/text/chunking.h" #include "engine/models/fish_audio/ar.h" #include "engine/models/fish_audio/codec.h" #include "engine/models/fish_audio/generator.h" @@ -110,19 +111,25 @@ uint64_t hash_audio_samples(const runtime::AudioBuffer & audio) { FishAudioGenerationOptions generation_options_from_request(const runtime::TaskRequest & request) { FishAudioGenerationOptions options; - options.max_new_tokens = runtime::parse_i64_option(request.options, {"max_new_tokens", "max_tokens"}) - .value_or(options.max_new_tokens); - options.chunk_length = runtime::parse_i64_option(request.options, {"chunk_length"}) - .value_or(options.chunk_length); + if (const auto value = runtime::parse_i64_option(request.options, {"max_new_tokens", "max_tokens"})) { + if (*value < 0) { + throw std::runtime_error("Fish Audio max_new_tokens must be non-negative"); + } + if (*value > 0) { + options.max_new_tokens = *value; + } + } + options.text_chunk_size = + engine::text::parse_text_chunk_size_override(request.options).value_or(options.text_chunk_size); options.top_p = runtime::parse_float_option(request.options, {"top_p"}).value_or(options.top_p); options.top_k = runtime::parse_int_option(request.options, {"top_k"}).value_or(options.top_k); options.temperature = runtime::parse_float_option(request.options, {"temperature"}).value_or(options.temperature); options.seed = runtime::parse_u32_option(request.options, {"seed"}).value_or(options.seed); if (options.max_new_tokens <= 0) { - throw std::runtime_error("Fish Audio max_new_tokens must be positive"); + throw std::runtime_error("Fish Audio max_new_tokens must be positive after default resolution"); } - if (options.chunk_length <= 0) { - throw std::runtime_error("Fish Audio chunk_length must be positive"); + if (options.text_chunk_size <= 0) { + throw std::runtime_error("Fish Audio text_chunk_size must be positive"); } if (!(options.top_p > 0.0F && options.top_p <= 1.0F)) { throw std::runtime_error("Fish Audio top_p must be in (0, 1]"); @@ -331,8 +338,14 @@ void FishAudioSession::prepare(const runtime::SessionPreparationRequest & reques defaults.text = request.text->text; has_defaults = true; } - defaults.generation.max_new_tokens = runtime::parse_i64_option(request.options, {"max_new_tokens", "max_tokens"}) - .value_or(defaults.generation.max_new_tokens); + if (const auto value = runtime::parse_i64_option(request.options, {"max_new_tokens", "max_tokens"})) { + if (*value < 0) { + throw std::runtime_error("Fish Audio max_new_tokens must be non-negative"); + } + if (*value > 0) { + defaults.generation.max_new_tokens = *value; + } + } if (auto reference = reference_from_voice(*assets_, request.voice, request.options, "Fish Audio prepare"); reference.has_value()) { defaults.reference = std::move(*reference); @@ -417,14 +430,26 @@ runtime::TaskResult FishAudioSession::run(const runtime::TaskRequest & request) require_prepared("Fish Audio run()"); const auto wall_start = Clock::now(); const bool mem_saver = mem_saver_from_options(options()); - auto fish_request = make_request(request); - std::optional reference_codes = std::nullopt; - if (fish_request.reference.has_value()) { - reference_codes = resolve_reference_codes(*fish_request.reference); + const auto request_options = generation_options_from_request(request); + const auto text_chunk_mode = + engine::text::parse_text_chunk_mode_override(request.options).value_or(engine::text::TextChunkMode::Default); + const auto chunk_requests = runtime::chunk_text_request(request, request_options.text_chunk_size, text_chunk_mode); + engine::debug::trace_log_scalar("fish_audio.text_chunk_size", request_options.text_chunk_size); + engine::debug::trace_log_scalar("fish_audio.text_chunk_mode", engine::text::text_chunk_mode_name(text_chunk_mode)); + engine::debug::trace_log_scalar("fish_audio.text_chunk_count", static_cast(chunk_requests.size())); + + runtime::AudioBuffer merged_audio; + for (const auto & chunk_request : chunk_requests) { + auto fish_request = make_request(chunk_request); + std::optional reference_codes = std::nullopt; + if (fish_request.reference.has_value()) { + reference_codes = resolve_reference_codes(*fish_request.reference); + } + auto generated = generator_->generate(fish_request, reference_codes, mem_saver); + runtime::append_audio_buffer(merged_audio, generated.audio); } - auto generated = generator_->generate(fish_request, reference_codes, mem_saver); runtime::TaskResult result; - result.audio_output = std::move(generated.audio); + result.audio_output = std::move(merged_audio); engine::debug::timing_log_scalar("session.wall_ms", engine::debug::elapsed_ms(wall_start, Clock::now())); return result; } diff --git a/src/models/higgs_tts/generator.cpp b/src/models/higgs_tts/generator.cpp index 058fa34d..b635d60b 100644 --- a/src/models/higgs_tts/generator.cpp +++ b/src/models/higgs_tts/generator.cpp @@ -38,7 +38,7 @@ int64_t bucketed_initial_cache_steps(int64_t prompt_steps, int64_t max_tokens) { void validate_generation_options(const HiggsGenerationOptions & options) { if (options.max_tokens <= 0) { - throw std::runtime_error("Higgs TTS max_tokens must be positive"); + throw std::runtime_error("Higgs TTS max_tokens must be positive after default resolution"); } if (!(options.temperature > 0.0F)) { throw std::runtime_error("Higgs TTS temperature must be positive"); diff --git a/src/models/higgs_tts/loader.cpp b/src/models/higgs_tts/loader.cpp index a39dc9a8..1ec4a4b4 100644 --- a/src/models/higgs_tts/loader.cpp +++ b/src/models/higgs_tts/loader.cpp @@ -37,13 +37,14 @@ runtime::CapabilitySet capabilities(const HiggsAssets &) { runtime::ModelCliInterface cli(const HiggsAssets &) { runtime::ModelCliInterface out; out.request_options = { - {"max_tokens", "n", "Maximum generated AR tokens."}, + {"max_tokens", "n", "Maximum generated AR tokens; default 2048, 0 uses the default."}, {"temperature", "float", "AR sampling temperature."}, {"top_k", "n", "AR top-k sampling limit."}, {"top_p", "float", "AR nucleus sampling probability."}, {"repetition_penalty", "float", "Accepted for Python API compatibility; Higgs audio sampling does not consume it."}, {"seed", "n", "Torch RNG seed."}, - {"text_chunk_size", "n", "Long-form text chunk size."}, + {"text_chunk_size", "n", "Long-form text chunk size; default 1024."}, + {"text_chunk_mode", "default|tag_aware|japanese|endline", "Framework text chunking mode."}, }; out.session_options = { {"higgs_tts.weight_type", "native|f32|f16|bf16|q8_0", "AR and codec weight storage type."}, diff --git a/src/models/higgs_tts/session.cpp b/src/models/higgs_tts/session.cpp index 7f67cf7c..931ae10c 100644 --- a/src/models/higgs_tts/session.cpp +++ b/src/models/higgs_tts/session.cpp @@ -16,7 +16,7 @@ namespace { using Clock = std::chrono::steady_clock; -constexpr int64_t kDefaultTextChunkSize = 4096; +constexpr int64_t kDefaultTextChunkSize = 1024; void validate_matmul_weight_storage(assets::TensorStorageType storage_type, const char * option_name) { if (storage_type == assets::TensorStorageType::Native || @@ -64,16 +64,18 @@ HiggsGenerationOptions generation_options_from_request( const runtime::TaskRequest & request, const HiggsConfig & config) { HiggsGenerationOptions options; - options.max_tokens = 1024; + options.max_tokens = 2048; options.temperature = 0.8F; options.top_p = 0.8F; options.top_k = 30; options.repetition_penalty = 1.1F; if (const auto value = runtime::parse_int_option(request.options, {"max_tokens"})) { - if (*value <= 0) { - throw std::runtime_error("Higgs TTS max_tokens must be positive"); + if (*value < 0) { + throw std::runtime_error("Higgs TTS max_tokens must be non-negative"); + } + if (*value > 0) { + options.max_tokens = *value; } - options.max_tokens = *value; } if (const auto value = runtime::parse_float_option(request.options, {"temperature"})) { options.temperature = *value; @@ -208,12 +210,15 @@ runtime::TaskResult HiggsTTSSession::run(const runtime::TaskRequest & request) { const auto wall_start = Clock::now(); const int64_t text_chunk_size = engine::text::parse_text_chunk_size_override(request.options).value_or(kDefaultTextChunkSize); - const auto chunk_requests = runtime::chunk_text_request(request, text_chunk_size); + const auto text_chunk_mode = + engine::text::parse_text_chunk_mode_override(request.options).value_or(engine::text::TextChunkMode::Default); + const auto chunk_requests = runtime::chunk_text_request(request, text_chunk_size, text_chunk_mode); const std::string reference_text = runtime::find_option(request.options, {"reference_text"}).value_or(""); const auto * reference_audio = find_reference_audio(request); const HiggsCodecEncodeOutput * reference_codes = reference_audio != nullptr ? &resolve_reference_codes(*reference_audio, reference_text) : nullptr; debug::trace_log_scalar("higgs_tts.text_chunk_size", text_chunk_size); + debug::trace_log_scalar("higgs_tts.text_chunk_mode", engine::text::text_chunk_mode_name(text_chunk_mode)); debug::trace_log_scalar("higgs_tts.text_chunk_count", static_cast(chunk_requests.size())); runtime::AudioBuffer merged_audio; diff --git a/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json b/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json index f1b9230a..575e2b28 100644 --- a/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json @@ -239,7 +239,8 @@ "id": "clone_longform", "text": "At dawn the harbor station opens its tall windows and the first clerk begins a careful report for the day. She notes the weather above the river, the slow cargo boats beyond the bridge, and the market voices arriving from the eastern road. A brass clock marks each quarter hour while porters stack wooden crates, bakers carry warm bread across the square, and a violinist practices the same bright phrase under the stone archway. By midmorning the keeper of the lighthouse sends a message about shifting currents, the museum guide unlocks a cabinet of maps, and a teacher leads a quiet line of students toward the ferry. In the afternoon a painter describes the silver color of the water, a mechanic jokes with the tram driver, and the station master reads an announcement that asks every traveler to keep close watch over letters, tickets, and parcels. After sunset the same clerk continues the report because new visitors keep arriving from the inland road. She explains that a florist carries pale roses past the fountain, two carpenters compare measurements beside the warehouse door, and the watchman checks each lock before the tide reaches its highest mark. A child laughs when the tram bell rings, a cook lowers a basket of fruit to the cellar, and three sailors unfold a chart that shows old channels, sandbars, and safe turning points for the morning crossing. Near midnight the lamps still glow on wet stone, the last cart rattles toward the market gate, and the report ends by saying that the harbor remains orderly, the wind has softened, the ferries are secure, and the town can rest until the next sunrise returns over the water. On the following morning the clerk resumes the record with even greater care because a week of inspections is about to begin. She writes that a ferry captain checks the mooring ropes one by one, a bookseller arranges travel guides beside the station cafe, and a pair of gardeners lift wet soil into bright clay pots near the west entrance. The bakery sends out trays of seed bread, the telegraph operator copies three official notices, and a tailor unfolds navy cloth across a polished wooden counter while customers wait in a line that bends toward the fountain. Before noon a surveyor compares bridge numbers against an old ledger, two cousins argue cheerfully about the best route to the fish market, and a choir director rehearses a patient scale that echoes against the warehouse wall. The lighthouse keeper reports that the northern channel is calmer than expected, the harbor pilot recommends a slower turn near the sandbar, and the customs officer stamps a packet of forms before waving a cart through the side gate. Later the schoolteacher returns with another group of students, asking them to observe the colors of rope, paint, stone, and water so they can write more exact descriptions in the classroom. A photographer kneels beside a rain barrel to capture the reflection of the clock tower, a mechanic tightens a brass hinge on the tram door, and an elderly traveler asks the clerk whether the evening ferry still stops at the orchard village beyond the marsh. As dusk arrives, lamps are trimmed again, shutters are tested against the wind, and the station kitchen sends bowls of soup to workers who remain on the late shift. The report continues with notes about a carpenter measuring floorboards in the east hall, a florist tying silver ribbon around the last stems of the day, and a violin case resting open on a bench beside the ticket window while its owner copies melody marks into a notebook. Long after the market gate closes, the clerk still writes that the harbor road stays busy, the river glints beneath scattered lamps, and the town maintains its patient rhythm of signals, footsteps, voices, bells, and distant engines. On the third day the clerk decides the record should be more precise, so she marks each event by the quarter hour and notes which sounds carry farthest through the station concourse. At first light she hears broom bristles on the stone steps, kettle lids in the cafe kitchen, and the slow scrape of crates being nudged across a loading cart beside the river wall. A messenger in a green coat delivers two canvas pouches, the ticket agent counts rolled coins into a brass tray, and a mother reads directions aloud while her son traces the painted ferry schedule with one curious finger. Midmorning brings a burst of sunlight across the waiting hall, making every brass handle shine while the museum guide escorts visitors toward the gallery of maps and navigational instruments. A porter pauses to describe the oldest compass in the display, a student sketches the harbor outline in graphite, and an apprentice clockmaker compares the station bell to a pocket watch that once belonged to his grandfather. By noon the fish market sends salt and seaweed scents through the open doors, tram wheels hiss at the curb, and the baker from the square exchanges a laugh with the florist who is carrying fresh lilies to the hotel veranda. The clerk writes that a cooper rolls three narrow barrels toward the cellar ramp, a translator copies weather bulletins for inland travelers, and a painter in a blue scarf studies the changing color of the tide as if each small wave might explain a different part of the sky. In the late afternoon the station master reviews freight tags, the customs officer checks a parcel of glassware, and a choir of children crosses the square singing a phrase so soft that the watchman removes his cap to listen. Evening settles slowly; lamps brighten in sequence, a cook inventories apples and onions in the pantry, and two sailors spread a faded chart on a crate so they can debate whether the shoals have shifted since the previous autumn. Before sleep the clerk closes the day with a final note that every vessel is accounted for, every platform has been swept, every lock has been tested twice, and the harbor seems ready to welcome another tide, another market, and another patient stream of voices at sunrise.", "voice_ref": "resources/a.wav", - "reference_text": "This little work was finished in the year eighteen o three, and intended for immediate publication." + "reference_text": "This little work was finished in the year eighteen o three, and intended for immediate publication.", + "text_chunk_size": 512 } ] }, diff --git a/tools/audiocpp_cli/audiocpp_cli_path_cases.json b/tools/audiocpp_cli/audiocpp_cli_path_cases.json index 7054381d..992f24e5 100644 --- a/tools/audiocpp_cli/audiocpp_cli_path_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_path_cases.json @@ -630,7 +630,7 @@ { "id": "official_auto_voice_english", "text": "The field recorder captured a clean reference take, and the operator confirmed that every timestamp matched the written production notes.", - "chunk_length": 200, + "text_chunk_size": 200, "top_p": 0.8, "repetition_penalty": 1.1, "temperature": 0.8, @@ -641,7 +641,7 @@ "text": "The studio engineer checked the short voice prompt, confirmed the take was clear, and started the final render.", "voice_ref": "resources/sample.wav", "reference_text": "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you.", - "chunk_length": 200, + "text_chunk_size": 200, "top_p": 0.8, "repetition_penalty": 1.1, "temperature": 0.8, @@ -650,7 +650,7 @@ { "id": "official_inline_control_tag", "text": "[whisper in small voice] The prototype actually worked after the last reset, and the control room stayed quiet until every green light appeared.", - "chunk_length": 200, + "text_chunk_size": 200, "top_p": 0.8, "repetition_penalty": 1.1, "temperature": 0.8, From 0fc8f9c75dc9c4743fdbf02a20bb4c31add6e9df Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Mon, 20 Jul 2026 22:28:15 -0400 Subject: [PATCH 11/27] Fix Fish Audio longform generation lifecycle --- include/engine/models/fish_audio/generator.h | 1 + .../engine/models/fish_audio/prompt_builder.h | 5 +- include/engine/models/fish_audio/types.h | 5 + src/models/fish_audio/ar.cpp | 224 +++++++++++++----- src/models/fish_audio/generator.cpp | 4 +- src/models/fish_audio/prompt_builder.cpp | 65 +++-- src/models/fish_audio/session.cpp | 13 +- .../compare_audiocpp_cli_path_results.py | 37 +-- 8 files changed, 240 insertions(+), 114 deletions(-) diff --git a/include/engine/models/fish_audio/generator.h b/include/engine/models/fish_audio/generator.h index 374c9aac..4f9c7800 100644 --- a/include/engine/models/fish_audio/generator.h +++ b/include/engine/models/fish_audio/generator.h @@ -27,6 +27,7 @@ class FishAudioGenerator { FishAudioGenerationResult generate( const FishAudioRequest & request, const std::optional & reference_codes, + const std::optional & previous_turn, bool mem_saver); private: diff --git a/include/engine/models/fish_audio/prompt_builder.h b/include/engine/models/fish_audio/prompt_builder.h index 825e42bf..6ed69e49 100644 --- a/include/engine/models/fish_audio/prompt_builder.h +++ b/include/engine/models/fish_audio/prompt_builder.h @@ -9,7 +9,10 @@ class FishAudioPromptBuilder { public: FishAudioPromptBuilder(std::shared_ptr assets, FishAudioTextTokenizer tokenizer); - FishAudioPrompt build(const FishAudioRequest & request, const std::optional & reference_codes) const; + FishAudioPrompt build( + const FishAudioRequest & request, + const std::optional & reference_codes, + const std::optional & previous_turn) const; private: std::shared_ptr assets_; diff --git a/include/engine/models/fish_audio/types.h b/include/engine/models/fish_audio/types.h index 68f7bffb..268bc2ea 100644 --- a/include/engine/models/fish_audio/types.h +++ b/include/engine/models/fish_audio/types.h @@ -36,6 +36,11 @@ struct FishAudioCodes { int64_t frames = 0; }; +struct FishAudioConversationTurn { + std::string text; + FishAudioCodes codes; +}; + struct FishAudioPrompt { std::vector matrix; int64_t codebook_rows = 0; diff --git a/src/models/fish_audio/ar.cpp b/src/models/fish_audio/ar.cpp index 91704cd9..90a1ffdf 100644 --- a/src/models/fish_audio/ar.cpp +++ b/src/models/fish_audio/ar.cpp @@ -613,21 +613,43 @@ FishStaticDecoderOutputs build_fish_static_decoder( int64_t cache_steps, const core::TensorValue & attention_mask, const core::TensorValue & cache_slot, + std::vector cache_keys, + std::vector cache_values, bool norm_fastlayer_input) { - auto decoder = modules::QwenCausalDecoderModule(config).build_static_cache_tail( - ctx, - graph, - input, - positions, - weights, - cache_steps, - attention_mask, - cache_slot); - auto fast_hidden = norm_fastlayer_input ? decoder.hidden : decoder.sequence; + if (cache_keys.size() != weights.stack.layers.size() || cache_values.size() != weights.stack.layers.size()) { + throw std::runtime_error("Fish Audio static decoder cache layer count mismatch"); + } + const int64_t step_elems = config.stack.num_key_value_heads * config.stack.head_dim; + auto x = input; + const auto layer_config = modules::qwen_decoder_layer_config_from_stack(config.stack); + const modules::QwenDecoderLayerModule layer_module(layer_config); + for (size_t layer_index = 0; layer_index < weights.stack.layers.size(); ++layer_index) { + auto out = layer_module.build_with_static_cache_tail( + ctx, + graph, + x, + positions, + weights.stack.layers[layer_index], + cache_keys[layer_index], + cache_values[layer_index], + cache_slot, + attention_mask); + x = out.output; + } + auto hidden = modules::RMSNormModule({config.stack.hidden_size, config.stack.rms_norm_eps, true, false}) + .build(ctx, x, weights.final_norm); + const auto logits = modules::LinearModule({ + config.stack.hidden_size, + config.logits_size, + config.use_lm_head_bias, + config.lm_head_precision, + }) + .build(ctx, hidden, weights.lm_head); + auto fast_hidden = norm_fastlayer_input ? hidden : x; return { fast_hidden, - decoder.logits, - std::move(decoder.cache), + logits, + runtime::TransformerKVCache(cache_steps, step_elems, std::move(cache_keys), std::move(cache_values)), }; } @@ -960,23 +982,60 @@ class FishAudioARRuntime::Impl { StepGraph(std::shared_ptr runtime, int64_t cache_steps) : runtime_(std::move(runtime)), cache_steps_(cache_steps) { - ggml_init_params params{runtime_->graph_arena_bytes(), nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { + ggml_init_params state_params{8ull * 1024ull * 1024ull, nullptr, true}; + state_ctx_.reset(ggml_init(state_params)); + if (state_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio AR step state context"); + } + ggml_init_params graph_params{runtime_->graph_arena_bytes(), nullptr, true}; + graph_ctx_.reset(ggml_init(graph_params)); + if (graph_ctx_ == nullptr) { throw std::runtime_error("failed to initialize Fish Audio AR step context"); } const auto & assets = runtime_->assets(); const auto & config = assets.config.text; - core::ModuleBuildContext ctx{ctx_.get(), "fish_audio.ar.step", runtime_->backend_type()}; + input_ = ggml_new_tensor_3d(state_ctx_.get(), GGML_TYPE_F32, config.dim, 1, 1); + position_ = ggml_new_tensor_1d(state_ctx_.get(), GGML_TYPE_I32, 1); + cache_slot_ = ggml_new_tensor_1d(state_ctx_.get(), GGML_TYPE_I32, 1); + mask_ = ggml_new_tensor_4d(state_ctx_.get(), GGML_TYPE_F16, cache_steps_, 1, 1, 1); + std::vector cache_keys; + std::vector cache_values; + cache_keys.reserve(runtime_->weights().slow_layers.size()); + cache_values.reserve(runtime_->weights().slow_layers.size()); + for (size_t layer = 0; layer < runtime_->weights().slow_layers.size(); ++layer) { + cache_keys.push_back(core::wrap_tensor( + ggml_new_tensor_4d( + state_ctx_.get(), + GGML_TYPE_F32, + config.head_dim, + config.n_local_heads, + cache_steps_, + 1), + core::TensorShape::from_dims({1, cache_steps_, config.n_local_heads, config.head_dim}), + GGML_TYPE_F32)); + cache_values.push_back(core::wrap_tensor( + ggml_new_tensor_4d( + state_ctx_.get(), + GGML_TYPE_F32, + config.head_dim, + config.n_local_heads, + cache_steps_, + 1), + core::TensorShape::from_dims({1, cache_steps_, config.n_local_heads, config.head_dim}), + GGML_TYPE_F32)); + } + state_buffer_ = ggml_backend_alloc_ctx_tensors(state_ctx_.get(), runtime_->backend()); + if (state_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate Fish Audio AR step state tensors"); + } + + core::ModuleBuildContext ctx{graph_ctx_.get(), "fish_audio.ar.step", runtime_->backend_type()}; auto input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, config.dim})); - input_ = input.tensor; - position_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); - cache_slot_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); - mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, cache_steps_, 1, 1, 1); + input = core::wrap_tensor(ggml_cpy(ctx.ggml, input_, input.tensor), input.shape, input.type); auto position_value = core::wrap_tensor(position_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); auto cache_slot_value = core::wrap_tensor(cache_slot_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); auto mask_value = core::wrap_tensor(mask_, core::TensorShape::from_dims({1, 1, 1, cache_steps_}), GGML_TYPE_F16); - graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + graph_ = ggml_new_graph_custom(graph_ctx_.get(), 65536, false); auto & constants = runtime_->slow_step_constants(); constants.begin_graph(); auto decoder = build_fish_static_decoder( @@ -989,6 +1048,8 @@ class FishAudioARRuntime::Impl { cache_steps_, mask_value, cache_slot_value, + std::move(cache_keys), + std::move(cache_values), assets.config.norm_fastlayer_input); cache_ = std::move(decoder.cache); hidden_ = decoder.hidden.tensor; @@ -999,8 +1060,10 @@ class FishAudioARRuntime::Impl { ggml_build_forward_expand(graph_, hidden_); constants.finish_graph(); constants.ensure_uploaded(); - buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); - if (buffer_ == nullptr) { + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { throw std::runtime_error("failed to allocate Fish Audio AR step tensors"); } mask_scratch_.assign(static_cast(cache_steps_), ggml_fp32_to_fp16(-INFINITY)); @@ -1008,8 +1071,11 @@ class FishAudioARRuntime::Impl { ~StepGraph() { core::release_backend_graph_resources(runtime_->backend(), graph_); - if (buffer_ != nullptr) { - ggml_backend_buffer_free(buffer_); + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + if (state_buffer_ != nullptr) { + ggml_backend_buffer_free(state_buffer_); } } @@ -1073,7 +1139,8 @@ class FishAudioARRuntime::Impl { private: std::shared_ptr runtime_; int64_t cache_steps_ = 0; - std::unique_ptr ctx_; + std::unique_ptr state_ctx_; + std::unique_ptr graph_ctx_; ggml_tensor * input_ = nullptr; ggml_tensor * position_ = nullptr; ggml_tensor * cache_slot_ = nullptr; @@ -1083,28 +1150,66 @@ class FishAudioARRuntime::Impl { runtime::TransformerKVCache cache_; std::vector mask_scratch_; ggml_cgraph * graph_ = nullptr; - ggml_backend_buffer_t buffer_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t state_buffer_ = nullptr; }; class FastGraph { public: explicit FastGraph(std::shared_ptr runtime) : runtime_(std::move(runtime)) { - ggml_init_params params{runtime_->graph_arena_bytes(), nullptr, true}; - ctx_.reset(ggml_init(params)); - if (ctx_ == nullptr) { + ggml_init_params state_params{8ull * 1024ull * 1024ull, nullptr, true}; + state_ctx_.reset(ggml_init(state_params)); + if (state_ctx_ == nullptr) { + throw std::runtime_error("failed to initialize Fish Audio fast AR state context"); + } + ggml_init_params graph_params{runtime_->graph_arena_bytes(), nullptr, true}; + graph_ctx_.reset(ggml_init(graph_params)); + if (graph_ctx_ == nullptr) { throw std::runtime_error("failed to initialize Fish Audio fast AR context"); } const auto & config = runtime_->assets().config.fast; const auto & weights = runtime_->weights(); - core::ModuleBuildContext ctx{ctx_.get(), "fish_audio.ar.fast", runtime_->backend_type()}; + input_ = ggml_new_tensor_3d(state_ctx_.get(), GGML_TYPE_F32, config.dim, 1, 1); + position_ = ggml_new_tensor_1d(state_ctx_.get(), GGML_TYPE_I32, 1); + mask_ = ggml_new_tensor_4d(state_ctx_.get(), GGML_TYPE_F16, config.num_codebooks, 1, 1, 1); + std::vector cache_keys; + std::vector cache_values; + cache_keys.reserve(weights.fast_layers.size()); + cache_values.reserve(weights.fast_layers.size()); + for (size_t layer = 0; layer < weights.fast_layers.size(); ++layer) { + cache_keys.push_back(core::wrap_tensor( + ggml_new_tensor_4d( + state_ctx_.get(), + GGML_TYPE_F32, + config.head_dim, + config.n_local_heads, + config.num_codebooks, + 1), + core::TensorShape::from_dims({1, config.num_codebooks, config.n_local_heads, config.head_dim}), + GGML_TYPE_F32)); + cache_values.push_back(core::wrap_tensor( + ggml_new_tensor_4d( + state_ctx_.get(), + GGML_TYPE_F32, + config.head_dim, + config.n_local_heads, + config.num_codebooks, + 1), + core::TensorShape::from_dims({1, config.num_codebooks, config.n_local_heads, config.head_dim}), + GGML_TYPE_F32)); + } + state_buffer_ = ggml_backend_alloc_ctx_tensors(state_ctx_.get(), runtime_->backend()); + if (state_buffer_ == nullptr) { + throw std::runtime_error("failed to allocate Fish Audio fast AR state tensors"); + } + + core::ModuleBuildContext ctx{graph_ctx_.get(), "fish_audio.ar.fast", runtime_->backend_type()}; auto input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, config.dim})); - input_ = input.tensor; - position_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, 1); - mask_ = ggml_new_tensor_4d(ctx_.get(), GGML_TYPE_F16, config.num_codebooks, 1, 1, 1); + input = core::wrap_tensor(ggml_cpy(ctx.ggml, input_, input.tensor), input.shape, input.type); auto position_value = core::wrap_tensor(position_, core::TensorShape::from_dims({1}), GGML_TYPE_I32); auto mask_value = core::wrap_tensor(mask_, core::TensorShape::from_dims({1, 1, 1, config.num_codebooks}), GGML_TYPE_F16); - graph_ = ggml_new_graph_custom(ctx_.get(), 32768, false); + graph_ = ggml_new_graph_custom(graph_ctx_.get(), 32768, false); auto & constants = runtime_->fast_constants(); constants.begin_graph(); modules::QwenCausalDecoderWeights decoder_weights; @@ -1114,27 +1219,28 @@ class FishAudioARRuntime::Impl { } decoder_weights.final_norm = binding::norm_data(constants, weights.fast_norm); decoder_weights.lm_head = binding::linear_data(constants, weights.fast_output); - auto decoder = modules::QwenCausalDecoderModule(make_fast_decoder_config(config)) - .build_static_cache_tail( - ctx, - graph_, - input, - position_value, - decoder_weights, - config.num_codebooks, - mask_value, - position_value); - for (size_t layer = 0; layer < weights.fast_layers.size(); ++layer) { - cache_keys_.push_back(decoder.cache.key_tensor(layer).tensor); - cache_values_.push_back(decoder.cache.value_tensor(layer).tensor); - } + auto decoder = build_fish_static_decoder( + ctx, + graph_, + input, + position_value, + decoder_weights, + make_fast_decoder_config(config), + config.num_codebooks, + mask_value, + position_value, + std::move(cache_keys), + std::move(cache_values), + true); logits_ = decoder.logits.tensor; ggml_set_output(logits_); ggml_build_forward_expand(graph_, logits_); constants.finish_graph(); constants.ensure_uploaded(); - buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); - if (buffer_ == nullptr) { + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { throw std::runtime_error("failed to allocate Fish Audio fast AR graph"); } mask_scratch_.assign(static_cast(config.num_codebooks), ggml_fp32_to_fp16(-INFINITY)); @@ -1142,8 +1248,11 @@ class FishAudioARRuntime::Impl { ~FastGraph() { core::release_backend_graph_resources(runtime_->backend(), graph_); - if (buffer_ != nullptr) { - ggml_backend_buffer_free(buffer_); + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + } + if (state_buffer_ != nullptr) { + ggml_backend_buffer_free(state_buffer_); } } @@ -1190,16 +1299,16 @@ class FishAudioARRuntime::Impl { private: std::shared_ptr runtime_; - std::unique_ptr ctx_; + std::unique_ptr state_ctx_; + std::unique_ptr graph_ctx_; ggml_tensor * input_ = nullptr; ggml_tensor * position_ = nullptr; ggml_tensor * mask_ = nullptr; ggml_tensor * logits_ = nullptr; - std::vector cache_keys_; - std::vector cache_values_; std::vector mask_scratch_; ggml_cgraph * graph_ = nullptr; - ggml_backend_buffer_t buffer_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; + ggml_backend_buffer_t state_buffer_ = nullptr; }; void ensure_prefill_graph(int64_t steps, FishARProfile & profile) { @@ -1250,7 +1359,6 @@ class FishAudioARRuntime::Impl { const auto biased = apply_semantic_bias(config, im_end_id(), slow_logits); profile.sample_bias_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); timing_start = Clock::now(); - // Upstream Python accepts repetition_penalty on the request but does not apply it in this generation path. int32_t main_token = sample_from_logits( biased, options.temperature, diff --git a/src/models/fish_audio/generator.cpp b/src/models/fish_audio/generator.cpp index 9813c44a..90b45f33 100644 --- a/src/models/fish_audio/generator.cpp +++ b/src/models/fish_audio/generator.cpp @@ -36,11 +36,13 @@ FishAudioCodes FishAudioGenerator::encode_reference(const runtime::AudioBuffer & FishAudioGenerationResult FishAudioGenerator::generate( const FishAudioRequest & request, const std::optional & reference_codes, + const std::optional & previous_turn, bool mem_saver) { engine::debug::trace_log_scalar("fish_audio.request.has_reference", request.reference.has_value()); engine::debug::trace_log_scalar("fish_audio.request.text_chars", static_cast(request.text.size())); + engine::debug::trace_log_scalar("fish_audio.request.has_previous_turn", previous_turn.has_value()); const auto prompt_start = Clock::now(); - const auto prompt = prompt_builder_.build(request, reference_codes); + const auto prompt = prompt_builder_.build(request, reference_codes, previous_turn); engine::debug::timing_log_scalar( "fish_audio.prompt_build_ms", engine::debug::elapsed_ms(prompt_start, Clock::now())); diff --git a/src/models/fish_audio/prompt_builder.cpp b/src/models/fish_audio/prompt_builder.cpp index 10b9f393..c4da2492 100644 --- a/src/models/fish_audio/prompt_builder.cpp +++ b/src/models/fish_audio/prompt_builder.cpp @@ -11,6 +11,11 @@ void append_tokens(std::vector & out, const std::vector & toke out.insert(out.end(), tokens.begin(), tokens.end()); } +struct CodeSpan { + int64_t start = 0; + const FishAudioCodes * codes = nullptr; +}; + std::string reference_text_with_speakers(const std::string & text) { static const std::regex speaker_re(R"(<\|speaker:\d+\|>)"); if (std::regex_search(text, speaker_re)) { @@ -19,6 +24,23 @@ std::string reference_text_with_speakers(const std::string & text) { return "<|speaker:0|>" + text; } +void append_code_span( + std::vector & row0, + std::vector & spans, + const FishAudioTextTokenizer & tokenizer, + const FishAudioCodes & codes, + int64_t expected_codebooks) { + if (codes.codebooks != expected_codebooks) { + throw std::runtime_error("Fish Audio prompt codebook count mismatch"); + } + const int64_t start = static_cast(row0.size()); + const int32_t semantic_begin = tokenizer.semantic_begin_id(); + for (int64_t frame = 0; frame < codes.frames; ++frame) { + row0.push_back(semantic_begin + codes.codes[static_cast(frame)]); + } + spans.push_back({start, &codes}); +} + } // namespace FishAudioPromptBuilder::FishAudioPromptBuilder( @@ -33,7 +55,8 @@ FishAudioPromptBuilder::FishAudioPromptBuilder( FishAudioPrompt FishAudioPromptBuilder::build( const FishAudioRequest & request, - const std::optional & reference_codes) const { + const std::optional & reference_codes, + const std::optional & previous_turn) const { if (request.text.empty()) { throw std::runtime_error("Fish Audio request text must not be empty"); } @@ -43,28 +66,30 @@ FishAudioPrompt FishAudioPromptBuilder::build( } std::vector row0; + std::vector code_spans; if (request.reference.has_value()) { if (!reference_codes.has_value()) { throw std::runtime_error("Fish Audio reference request requires encoded reference codes"); } - if (reference_codes->codebooks != assets_->config.fast.num_codebooks) { - throw std::runtime_error("Fish Audio reference codebook count mismatch"); - } append_tokens(row0, tokenizer_.encode("<|im_start|>system\n")); append_tokens(row0, tokenizer_.encode("convert the provided text to speech reference to the following:\n\nText:\n")); append_tokens(row0, tokenizer_.encode(reference_text_with_speakers(request.reference->text))); append_tokens(row0, tokenizer_.encode("\n\nSpeech:\n")); - const int32_t semantic_begin = tokenizer_.semantic_begin_id(); - for (int64_t frame = 0; frame < reference_codes->frames; ++frame) { - const int32_t code = reference_codes->codes[static_cast(frame)]; - row0.push_back(semantic_begin + code); - } + append_code_span(row0, code_spans, tokenizer_, *reference_codes, assets_->config.fast.num_codebooks); append_tokens(row0, tokenizer_.encode("<|im_end|>\n")); } else { append_tokens(row0, tokenizer_.encode("<|im_start|>system\n")); append_tokens(row0, tokenizer_.encode("convert the provided text to speech")); append_tokens(row0, tokenizer_.encode("<|im_end|>\n")); } + if (previous_turn.has_value()) { + append_tokens(row0, tokenizer_.encode("<|im_start|>user\n")); + append_tokens(row0, tokenizer_.encode(previous_turn->text)); + append_tokens(row0, tokenizer_.encode("<|im_end|>\n")); + append_tokens(row0, tokenizer_.encode("<|im_start|>assistant\n<|voice|>")); + append_code_span(row0, code_spans, tokenizer_, previous_turn->codes, assets_->config.fast.num_codebooks); + append_tokens(row0, tokenizer_.encode("<|im_end|>\n")); + } append_tokens(row0, tokenizer_.encode("<|im_start|>user\n")); append_tokens(row0, tokenizer_.encode(request.text)); append_tokens(row0, tokenizer_.encode("<|im_end|>\n")); @@ -78,21 +103,19 @@ FishAudioPrompt FishAudioPromptBuilder::build( for (int64_t step = 0; step < prompt.steps; ++step) { prompt.matrix[static_cast(step)] = row0[static_cast(step)]; } - if (reference_codes.has_value()) { - int64_t semantic_index = 0; - for (int64_t step = 0; step < prompt.steps; ++step) { - const int32_t token = prompt.matrix[static_cast(step)]; - if (token < tokenizer_.semantic_begin_id() || token > tokenizer_.semantic_end_id()) { - continue; - } - if (semantic_index >= reference_codes->frames) { - break; + for (const auto & span : code_spans) { + if (span.codes == nullptr) { + throw std::runtime_error("Fish Audio prompt code span is missing codes"); + } + for (int64_t frame = 0; frame < span.codes->frames; ++frame) { + const int64_t step = span.start + frame; + if (step < 0 || step >= prompt.steps) { + throw std::runtime_error("Fish Audio prompt code span exceeds prompt length"); } - for (int64_t codebook = 0; codebook < reference_codes->codebooks; ++codebook) { + for (int64_t codebook = 0; codebook < span.codes->codebooks; ++codebook) { prompt.matrix[static_cast((codebook + 1) * prompt.steps + step)] = - reference_codes->codes[static_cast(codebook * reference_codes->frames + semantic_index)]; + span.codes->codes[static_cast(codebook * span.codes->frames + frame)]; } - ++semantic_index; } } return prompt; diff --git a/src/models/fish_audio/session.cpp b/src/models/fish_audio/session.cpp index 7074a364..15218fe2 100644 --- a/src/models/fish_audio/session.cpp +++ b/src/models/fish_audio/session.cpp @@ -439,14 +439,19 @@ runtime::TaskResult FishAudioSession::run(const runtime::TaskRequest & request) engine::debug::trace_log_scalar("fish_audio.text_chunk_count", static_cast(chunk_requests.size())); runtime::AudioBuffer merged_audio; - for (const auto & chunk_request : chunk_requests) { + std::optional reference_codes = std::nullopt; + std::optional previous_turn = std::nullopt; + for (size_t chunk_index = 0; chunk_index < chunk_requests.size(); ++chunk_index) { + const auto & chunk_request = chunk_requests[chunk_index]; auto fish_request = make_request(chunk_request); - std::optional reference_codes = std::nullopt; - if (fish_request.reference.has_value()) { + if (fish_request.reference.has_value() && !reference_codes.has_value()) { reference_codes = resolve_reference_codes(*fish_request.reference); } - auto generated = generator_->generate(fish_request, reference_codes, mem_saver); + auto generated = generator_->generate(fish_request, reference_codes, previous_turn, mem_saver); runtime::append_audio_buffer(merged_audio, generated.audio); + if (chunk_requests.size() > 1) { + previous_turn = FishAudioConversationTurn{fish_request.text, std::move(generated.codes)}; + } } runtime::TaskResult result; result.audio_output = std::move(merged_audio); diff --git a/tools/audiocpp_cli/compare_audiocpp_cli_path_results.py b/tools/audiocpp_cli/compare_audiocpp_cli_path_results.py index 19864d8d..9f3f90b0 100644 --- a/tools/audiocpp_cli/compare_audiocpp_cli_path_results.py +++ b/tools/audiocpp_cli/compare_audiocpp_cli_path_results.py @@ -11,6 +11,7 @@ import wave from pathlib import Path +import librosa import numpy as np @@ -122,29 +123,6 @@ def stft_magnitude(samples: np.ndarray, n_fft: int = 1024, hop: int = 256) -> np return spec -def hz_to_mel(freq: np.ndarray | float) -> np.ndarray | float: - return 2595.0 * np.log10(1.0 + np.asarray(freq) / 700.0) - - -def mel_to_hz(mel: np.ndarray) -> np.ndarray: - return 700.0 * (np.power(10.0, mel / 2595.0) - 1.0) - - -def mel_filterbank(sample_rate: int, n_fft: int, n_mels: int = 80) -> np.ndarray: - freq_bins = n_fft // 2 + 1 - mel_points = np.linspace(float(hz_to_mel(0.0)), float(hz_to_mel(sample_rate / 2.0)), n_mels + 2) - hz_points = mel_to_hz(mel_points) - bins = np.floor((n_fft + 1) * hz_points / sample_rate).astype(np.int64) - filters = np.zeros((n_mels, freq_bins), dtype=np.float32) - for mel_index in range(n_mels): - left, center, right = bins[mel_index : mel_index + 3] - if center > left: - filters[mel_index, left:center] = (np.arange(left, center) - left) / (center - left) - if right > center: - filters[mel_index, center:right] = (right - np.arange(center, right)) / (right - center) - return filters - - def wav_similarity_detail(src_path: Path, baseline_path: Path) -> tuple[float, str]: src_rate, src_audio = read_wav_f32(src_path) baseline_rate, baseline_audio = read_wav_f32(baseline_path) @@ -163,12 +141,13 @@ def wav_similarity_detail(src_path: Path, baseline_path: Path) -> tuple[float, s stft_cos = cosine_similarity(src_mag, baseline_mag) log_stft_cos = cosine_similarity(np.log1p(src_mag), np.log1p(baseline_mag)) - if src_rate == baseline_rate and freq_bins > 0 and frames > 0: - filters = mel_filterbank(src_rate, 1024) - filters = filters[:, :freq_bins] - src_mel = np.log1p(filters @ src_mag) - baseline_mel = np.log1p(filters @ baseline_mag) - log_mel_cos = cosine_similarity(src_mel, baseline_mel) + if src_rate == baseline_rate and src_mono.size > 0 and baseline_mono.size > 0: + src_mel = librosa.feature.melspectrogram(y=mono(src_audio), sr=src_rate) + baseline_mel = librosa.feature.melspectrogram(y=mono(baseline_audio), sr=baseline_rate) + mel_frames = min(src_mel.shape[1], baseline_mel.shape[1]) + src_log_mel = librosa.power_to_db(src_mel[:, :mel_frames], ref=1.0) + baseline_log_mel = librosa.power_to_db(baseline_mel[:, :mel_frames], ref=1.0) + log_mel_cos = cosine_similarity(src_log_mel, baseline_log_mel) log_mel_text = f"{log_mel_cos:.9f}" else: log_mel_text = "n/a" From 1f67472df3c7930f87f4e3fa7baf5e604b3eabef Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Mon, 20 Jul 2026 23:30:07 -0400 Subject: [PATCH 12/27] Reduce Higgs TTS runtime graph memory --- include/engine/models/higgs_tts/codec.h | 2 + include/engine/models/higgs_tts/session.h | 21 ++++++- src/models/higgs_tts/codec.cpp | 41 +++++++++---- src/models/higgs_tts/generator.cpp | 1 + src/models/higgs_tts/loader.cpp | 1 + src/models/higgs_tts/session.cpp | 71 +++++++++++++++-------- 6 files changed, 101 insertions(+), 36 deletions(-) diff --git a/include/engine/models/higgs_tts/codec.h b/include/engine/models/higgs_tts/codec.h index 8e160e91..fa646843 100644 --- a/include/engine/models/higgs_tts/codec.h +++ b/include/engine/models/higgs_tts/codec.h @@ -113,6 +113,8 @@ class HiggsCodecRuntime { const std::vector & codes, int64_t frames, int64_t codebooks) const; + void release_encode_graph(); + void release_runtime_graphs(); private: std::shared_ptr assets_; diff --git a/include/engine/models/higgs_tts/session.h b/include/engine/models/higgs_tts/session.h index 0898f452..c0ef18a5 100644 --- a/include/engine/models/higgs_tts/session.h +++ b/include/engine/models/higgs_tts/session.h @@ -1,5 +1,6 @@ #pragma once +#include "engine/framework/runtime/cache_slots.h" #include "engine/framework/runtime/session_base.h" #include "engine/models/higgs_tts/assets.h" #include "engine/models/higgs_tts/ar.h" @@ -31,12 +32,25 @@ class HiggsTTSSession final private: struct ReferenceCacheEntry { - std::string reference_text; + HiggsCodecEncodeOutput codes; + }; + + struct ReferenceCacheKey { int sample_rate = 0; int channels = 0; uint64_t sample_count = 0; uint64_t sample_hash = 0; - HiggsCodecEncodeOutput codes; + std::string reference_text; + }; + + struct ReferenceCacheKeyEqual { + bool operator()(const ReferenceCacheKey & lhs, const ReferenceCacheKey & rhs) const noexcept { + return lhs.sample_rate == rhs.sample_rate && + lhs.channels == rhs.channels && + lhs.sample_count == rhs.sample_count && + lhs.sample_hash == rhs.sample_hash && + lhs.reference_text == rhs.reference_text; + } }; HiggsGenerationRequest make_generation_request( @@ -58,7 +72,8 @@ class HiggsTTSSession final std::shared_ptr ar_; std::shared_ptr codec_; std::unique_ptr generator_; - std::optional reference_cache_; + runtime::CacheSlots reference_cache_; + std::optional uncached_reference_; }; } // namespace engine::models::higgs_tts diff --git a/src/models/higgs_tts/codec.cpp b/src/models/higgs_tts/codec.cpp index 82d94cc8..32a16820 100644 --- a/src/models/higgs_tts/codec.cpp +++ b/src/models/higgs_tts/codec.cpp @@ -1137,8 +1137,14 @@ class HiggsCodecEncodeGraph { ggml_set_output(outputs_[codebook]); ggml_build_forward_expand(graph_, outputs_[codebook]); } - buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); - if (buffer_ == nullptr) { + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } throw std::runtime_error("failed to allocate Higgs TTS codec encode graph"); } engine::debug::timing_log_scalar("higgs_tts.codec.encode.graph.build_ms", @@ -1147,8 +1153,8 @@ class HiggsCodecEncodeGraph { ~HiggsCodecEncodeGraph() { engine::core::release_backend_graph_resources(runtime_->backend(), graph_); - if (buffer_ != nullptr) { - ggml_backend_buffer_free(buffer_); + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); } } @@ -1212,7 +1218,7 @@ class HiggsCodecEncodeGraph { ggml_tensor * semantic_input_ = nullptr; std::array(kCodecCodebooks)> outputs_ = {}; ggml_cgraph * graph_ = nullptr; - ggml_backend_buffer_t buffer_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; }; class HiggsCodecDecodeGraph { @@ -1251,8 +1257,14 @@ class HiggsCodecDecodeGraph { ggml_set_output(output_); graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); ggml_build_forward_expand(graph_, output_); - buffer_ = ggml_backend_alloc_ctx_tensors(ctx_.get(), runtime_->backend()); - if (buffer_ == nullptr) { + gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend())); + if (gallocr_ == nullptr || + !ggml_gallocr_reserve(gallocr_, graph_) || + !ggml_gallocr_alloc_graph(gallocr_, graph_)) { + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); + gallocr_ = nullptr; + } throw std::runtime_error("failed to allocate Higgs TTS codec decode graph"); } code_scratch_.assign(static_cast(capacity_frames_ * kCodecCodebooks), 0); @@ -1263,8 +1275,8 @@ class HiggsCodecDecodeGraph { ~HiggsCodecDecodeGraph() { engine::core::release_backend_graph_resources(runtime_->backend(), graph_); - if (buffer_ != nullptr) { - ggml_backend_buffer_free(buffer_); + if (gallocr_ != nullptr) { + ggml_gallocr_free(gallocr_); } } @@ -1337,7 +1349,7 @@ class HiggsCodecDecodeGraph { std::vector code_scratch_; std::vector frame_mask_values_; ggml_cgraph * graph_ = nullptr; - ggml_backend_buffer_t buffer_ = nullptr; + ggml_gallocr_t gallocr_ = nullptr; }; HiggsCodecWeights load_higgs_codec_decode_weights(const HiggsAssets & assets, @@ -1611,4 +1623,13 @@ HiggsCodecDecodeOutput HiggsCodecRuntime::decode_codes(const std::vectorrelease_runtime_graphs(); return result; } diff --git a/src/models/higgs_tts/loader.cpp b/src/models/higgs_tts/loader.cpp index 1ec4a4b4..9a4a0872 100644 --- a/src/models/higgs_tts/loader.cpp +++ b/src/models/higgs_tts/loader.cpp @@ -55,6 +55,7 @@ runtime::ModelCliInterface cli(const HiggsAssets &) { {"higgs_tts.ar_decode_graph_arena_mb", "n", "AR decode graph arena size."}, {"higgs_tts.codec_decode_graph_arena_mb", "n", "Codec decode graph arena size."}, {"higgs_tts.codec_encode_graph_arena_mb", "n", "Codec encode graph arena size."}, + {"higgs_tts.reference_cache_slots", "n", "Encoded reference-audio cache slots; default 1."}, }; return out; } diff --git a/src/models/higgs_tts/session.cpp b/src/models/higgs_tts/session.cpp index 931ae10c..8c4f3bfe 100644 --- a/src/models/higgs_tts/session.cpp +++ b/src/models/higgs_tts/session.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -17,6 +18,7 @@ namespace { using Clock = std::chrono::steady_clock; constexpr int64_t kDefaultTextChunkSize = 1024; +constexpr int64_t kDefaultReferenceCacheSlots = 1; void validate_matmul_weight_storage(assets::TensorStorageType storage_type, const char * option_name) { if (storage_type == assets::TensorStorageType::Native || @@ -48,6 +50,20 @@ uint64_t hash_audio_samples(const runtime::AudioBuffer & audio) { return hash; } +std::size_t resolve_reference_cache_slots(const runtime::SessionOptions & options) { + const int64_t slots = runtime::parse_i64_option( + options.options, + {"higgs_tts.reference_cache_slots", "reference_cache_slots"}) + .value_or(kDefaultReferenceCacheSlots); + if (slots < 0) { + throw std::runtime_error("higgs_tts.reference_cache_slots must be non-negative"); + } + if (static_cast(slots) > static_cast(std::numeric_limits::max())) { + throw std::runtime_error("higgs_tts.reference_cache_slots is too large"); + } + return static_cast(slots); +} + const runtime::AudioBuffer * find_reference_audio(const runtime::TaskRequest & request) { if (request.voice.has_value() && request.voice->speaker.has_value() @@ -109,7 +125,8 @@ HiggsTTSSession::HiggsTTSSession( std::shared_ptr assets) : RuntimeSessionBase(options), task_(task), - assets_(std::move(assets)) { + assets_(std::move(assets)), + reference_cache_(resolve_reference_cache_slots(this->options())) { if (assets_ == nullptr) { throw std::runtime_error("Higgs TTS session requires assets"); } @@ -155,6 +172,7 @@ HiggsTTSSession::HiggsTTSSession( key != "higgs_tts.ar_decode_graph_arena_mb" && key != "higgs_tts.codec_decode_graph_arena_mb" && key != "higgs_tts.codec_encode_graph_arena_mb" && + key != "higgs_tts.reference_cache_slots" && key != "higgs_tts.weight_type" && key != "higgs_tts.ar_weight_type" && key != "higgs_tts.codec_weight_type") { @@ -243,35 +261,42 @@ const HiggsCodecEncodeOutput & HiggsTTSSession::resolve_reference_codes( const std::string & reference_text) { const uint64_t sample_count = static_cast(audio.samples.size()); const uint64_t sample_hash = hash_audio_samples(audio); + ReferenceCacheKey key; + key.reference_text = reference_text; + key.sample_rate = audio.sample_rate; + key.channels = audio.channels; + key.sample_count = sample_count; + key.sample_hash = sample_hash; debug::trace_log_scalar("higgs_tts.reference_audio.sample_rate", audio.sample_rate); debug::trace_log_scalar("higgs_tts.reference_audio.channels", audio.channels); debug::trace_log_f32("higgs_tts.reference_audio.samples", {static_cast(audio.samples.size())}, audio.samples); - const bool cache_hit = reference_cache_.has_value() - && reference_cache_->reference_text == reference_text - && reference_cache_->sample_rate == audio.sample_rate - && reference_cache_->channels == audio.channels - && reference_cache_->sample_count == sample_count - && reference_cache_->sample_hash == sample_hash; - if (!cache_hit) { - const auto encode_start = Clock::now(); - ReferenceCacheEntry entry; - entry.reference_text = reference_text; - entry.sample_rate = audio.sample_rate; - entry.channels = audio.channels; - entry.sample_count = sample_count; - entry.sample_hash = sample_hash; - entry.codes = codec_->encode_reference(audio); - debug::trace_log_scalar("higgs_tts.reference_codes.frames", entry.codes.frames); - debug::trace_log_scalar("higgs_tts.reference_codes.codebooks", entry.codes.codebooks); - debug::trace_log_i32("higgs_tts.reference_codes.values", - {entry.codes.frames, entry.codes.codebooks}, - entry.codes.codes); - reference_cache_ = std::move(entry); + debug::trace_log_scalar("higgs_tts.reference_cache.capacity", static_cast(reference_cache_.capacity())); + debug::trace_log_scalar("higgs_tts.reference_cache.size", static_cast(reference_cache_.size())); + if (const auto * cached = reference_cache_.find(key)) { + debug::trace_log_scalar("higgs_tts.reference_cache.hit", 1); + return cached->codes; + } + debug::trace_log_scalar("higgs_tts.reference_cache.hit", 0); + + const auto encode_start = Clock::now(); + ReferenceCacheEntry entry; + entry.codes = codec_->encode_reference(audio); + codec_->release_encode_graph(); + debug::trace_log_scalar("higgs_tts.reference_codes.frames", entry.codes.frames); + debug::trace_log_scalar("higgs_tts.reference_codes.codebooks", entry.codes.codebooks); + debug::trace_log_i32("higgs_tts.reference_codes.values", + {entry.codes.frames, entry.codes.codebooks}, + entry.codes.codes); + if (reference_cache_.capacity() == 0) { + uncached_reference_ = std::move(entry); debug::timing_log_scalar("higgs_tts.codec.encode_reference_ms", engine::debug::elapsed_ms(encode_start)); + return uncached_reference_->codes; } - return reference_cache_->codes; + reference_cache_.put(key, std::move(entry)); + debug::timing_log_scalar("higgs_tts.codec.encode_reference_ms", engine::debug::elapsed_ms(encode_start)); + return reference_cache_.find(key)->codes; } HiggsGenerationRequest HiggsTTSSession::make_generation_request( From f297cae9139c9dcff0acc9ae8cf2c3907737ea10 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Tue, 21 Jul 2026 00:13:45 -0400 Subject: [PATCH 13/27] Rename Higgs Audio TTS family --- CMakeLists.txt | 20 +-- README.md | 3 +- docs/tts.md | 10 +- .../{higgs_tts => higgs_audio_tts}/ar.h | 6 +- .../{higgs_tts => higgs_audio_tts}/assets.h | 4 +- .../codebooks.h | 4 +- .../{higgs_tts => higgs_audio_tts}/codec.h | 6 +- .../generator.h | 12 +- .../{higgs_tts => higgs_audio_tts}/loader.h | 10 +- .../{higgs_tts => higgs_audio_tts}/sampler.h | 6 +- .../{higgs_tts => higgs_audio_tts}/session.h | 12 +- .../tokenizer_text.h | 6 +- .../{higgs_tts.json => higgs_audio_tts.json} | 2 +- src/framework/runtime/registry.cpp | 4 +- .../{higgs_tts => higgs_audio_tts}/ar.cpp | 34 ++--- .../{higgs_tts => higgs_audio_tts}/assets.cpp | 8 +- .../codebooks.cpp | 6 +- .../{higgs_tts => higgs_audio_tts}/codec.cpp | 40 +++--- .../generator.cpp | 118 +++++++++--------- .../{higgs_tts => higgs_audio_tts}/loader.cpp | 45 ++++--- .../sampler.cpp | 6 +- .../session.cpp | 86 ++++++------- .../tokenizer_text.cpp | 6 +- .../{higgs_tts => higgs_audio_tts}/.gitignore | 0 .../{higgs_tts => higgs_audio_tts}/README.md | 10 +- .../compare_warmbench_results.py | 4 +- .../higgs_audio_tts_cuda_bench_cases.json} | 0 .../higgs_audio_tts_cuda_mixed_cases.json} | 0 .../higgs_audio_tts_cuda_perf_cases.json} | 0 .../higgs_audio_tts_python_warm_bench.py} | 16 +-- .../higgs_audio_tts_sampler_logits.bin} | Bin .../higgs_audio_tts_warm_bench.cpp} | 16 +-- .../higgs_audio_tts_warm_bench_cases.json} | 0 .../run_cuda_performance.ps1 | 4 +- tests/warmbench.py | 22 ++-- ...audiocpp_cli_longform_tts_clone_cases.json | 4 +- .../audiocpp_cli/audiocpp_cli_path_cases.json | 8 +- 37 files changed, 274 insertions(+), 264 deletions(-) rename include/engine/models/{higgs_tts => higgs_audio_tts}/ar.h (97%) rename include/engine/models/{higgs_tts => higgs_audio_tts}/assets.h (94%) rename include/engine/models/{higgs_tts => higgs_audio_tts}/codebooks.h (85%) rename include/engine/models/{higgs_tts => higgs_audio_tts}/codec.h (96%) rename include/engine/models/{higgs_tts => higgs_audio_tts}/generator.h (87%) rename include/engine/models/{higgs_tts => higgs_audio_tts}/loader.h (71%) rename include/engine/models/{higgs_tts => higgs_audio_tts}/sampler.h (90%) rename include/engine/models/{higgs_tts => higgs_audio_tts}/session.h (90%) rename include/engine/models/{higgs_tts => higgs_audio_tts}/tokenizer_text.h (83%) rename model_specs/{higgs_tts.json => higgs_audio_tts.json} (96%) rename src/models/{higgs_tts => higgs_audio_tts}/ar.cpp (97%) rename src/models/{higgs_tts => higgs_audio_tts}/assets.cpp (97%) rename src/models/{higgs_tts => higgs_audio_tts}/codebooks.cpp (95%) rename src/models/{higgs_tts => higgs_audio_tts}/codec.cpp (98%) rename src/models/{higgs_tts => higgs_audio_tts}/generator.cpp (83%) rename src/models/{higgs_tts => higgs_audio_tts}/loader.cpp (73%) rename src/models/{higgs_tts => higgs_audio_tts}/sampler.cpp (99%) rename src/models/{higgs_tts => higgs_audio_tts}/session.cpp (76%) rename src/models/{higgs_tts => higgs_audio_tts}/tokenizer_text.cpp (95%) rename tests/{higgs_tts => higgs_audio_tts}/.gitignore (100%) rename tests/{higgs_tts => higgs_audio_tts}/README.md (82%) rename tests/{higgs_tts => higgs_audio_tts}/compare_warmbench_results.py (98%) rename tests/{higgs_tts/higgs_tts_cuda_bench_cases.json => higgs_audio_tts/higgs_audio_tts_cuda_bench_cases.json} (100%) rename tests/{higgs_tts/higgs_tts_cuda_mixed_cases.json => higgs_audio_tts/higgs_audio_tts_cuda_mixed_cases.json} (100%) rename tests/{higgs_tts/higgs_tts_cuda_perf_cases.json => higgs_audio_tts/higgs_audio_tts_cuda_perf_cases.json} (100%) rename tests/{higgs_tts/higgs_tts_python_warm_bench.py => higgs_audio_tts/higgs_audio_tts_python_warm_bench.py} (96%) rename tests/{higgs_tts/higgs_tts_sampler_logits.bin => higgs_audio_tts/higgs_audio_tts_sampler_logits.bin} (100%) rename tests/{higgs_tts/higgs_tts_warm_bench.cpp => higgs_audio_tts/higgs_audio_tts_warm_bench.cpp} (95%) rename tests/{higgs_tts/higgs_tts_warm_bench_cases.json => higgs_audio_tts/higgs_audio_tts_warm_bench_cases.json} (100%) rename tests/{higgs_tts => higgs_audio_tts}/run_cuda_performance.ps1 (93%) diff --git a/CMakeLists.txt b/CMakeLists.txt index df275571..ab2d05f1 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -379,15 +379,15 @@ add_library(engine_runtime STATIC src/models/higgs_audio_stt/postprocess.cpp src/models/higgs_audio_stt/session.cpp src/models/higgs_audio_stt/loader.cpp - src/models/higgs_tts/ar.cpp - src/models/higgs_tts/assets.cpp - src/models/higgs_tts/codec.cpp - src/models/higgs_tts/codebooks.cpp - src/models/higgs_tts/generator.cpp - src/models/higgs_tts/loader.cpp - src/models/higgs_tts/sampler.cpp - src/models/higgs_tts/session.cpp - src/models/higgs_tts/tokenizer_text.cpp + src/models/higgs_audio_tts/ar.cpp + src/models/higgs_audio_tts/assets.cpp + src/models/higgs_audio_tts/codec.cpp + src/models/higgs_audio_tts/codebooks.cpp + src/models/higgs_audio_tts/generator.cpp + src/models/higgs_audio_tts/loader.cpp + src/models/higgs_audio_tts/sampler.cpp + src/models/higgs_audio_tts/session.cpp + src/models/higgs_audio_tts/tokenizer_text.cpp src/models/irodori_tts/assets.cpp src/models/irodori_tts/codec.cpp src/models/irodori_tts/condition_encoder.cpp @@ -731,7 +731,7 @@ if (ENGINE_BUILD_WARMBENCH) add_engine_warmbench(chatterbox_warm_bench tests/chatterbox/chatterbox_warm_bench.cpp) add_engine_warmbench(citrinet_asr_warm_bench tests/citrinet_asr/citrinet_asr_warm_bench.cpp) add_engine_warmbench(higgs_audio_stt_warm_bench tests/higgs_audio_stt/higgs_audio_stt_warm_bench.cpp) - add_engine_warmbench(higgs_tts_warm_bench tests/higgs_tts/higgs_tts_warm_bench.cpp) + add_engine_warmbench(higgs_audio_tts_warm_bench tests/higgs_audio_tts/higgs_audio_tts_warm_bench.cpp) add_engine_warmbench(hviske_asr_warm_bench tests/hviske_asr/hviske_asr_warm_bench.cpp) add_engine_warmbench(index_tts2_warm_bench tests/index_tts2/index_tts2_warm_bench.cpp) add_engine_warmbench(irodori_tts_warm_bench tests/irodori_tts/irodori_tts_warm_bench.cpp) diff --git a/README.md b/README.md index f084cbc0..7af8c41a 100644 --- a/README.md +++ b/README.md @@ -53,6 +53,7 @@ audio.cpp would not be moving this quickly without generous contributors bringin | **citrinet_asr** | ASR | en | Citrinet-256 | | **heartmula** | music generation | zh, en, ja, ko, es | HeartMuLa-oss-3B with HeartCodec-oss | | **higgs_audio_stt** | ASR | en | Higgs Audio v3 STT | +| **higgs_audio_tts** | TTS, voice cloning | auto | Higgs Audio v3 TTS 4B | | **htdemucs** | source separation | lang agnostic | HTDemucs, HTDemucs_ft | | **hviske_asr** | ASR | da | Hviske v5.3 | | **marblenet_vad** | VAD | lang agnostic | MarbleNet VAD | @@ -89,7 +90,7 @@ Community model ports live under `community_models` to make the ownership bounda |---|---|---|---|---| | **outetts** | TTS, voice cloning | en, ar, zh, nl, fr, de, it, ja, ko, lt, ru, es, pt, be, bn, ka, hu, lv, fa, pl, sw, ta, uk | Mirek [@mirek190](https://github.com/mirek190) | Llama-OuteTTS-1.0-1B TTS and voice cloning support | -WIP: Higgs Audio v3 TTS 4B, Fish Audio S2 Pro. +WIP: Fish Audio S2 Pro. PocketTTS language selection is a model-load option. When the model path points at the PocketTTS root, the loader uses `english` unless you pass `--load-option language=`. Kyutai's normal non-English PocketTTS releases are smaller distilled language models intended for the fast PocketTTS path. The `_24l` variants are larger 24-layer, undistilled preview models that can sound better but are slower. Kyutai currently publishes French only as `french_24l`, not as a normal distilled `french` language directory, so French is not listed as a normal PocketTTS language here. diff --git a/docs/tts.md b/docs/tts.md index 1813d0e2..14e06c65 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -10,7 +10,7 @@ | OmniVoice | `omnivoice` | `tts` | [OmniVoice](#omnivoice) | | PocketTTS | `pocket_tts` | `tts` | [PocketTTS](#pockettts) | | VoxCPM2 | `voxcpm2` | `tts`, `vdes` | [VoxCPM2](#voxcpm2) | -| Higgs Audio v3 TTS | `higgs_tts` | `tts` | [Higgs Audio v3 TTS](#higgs-audio-v3-tts) | +| Higgs Audio v3 TTS | `higgs_audio_tts` | `tts` | [Higgs Audio v3 TTS](#higgs-audio-v3-tts) | | IndexTTS2 | `index_tts2` | `tts` | [IndexTTS2](#indextts2) | | Irodori-TTS | `irodori_tts` | `tts`, `vdes` | [Irodori-TTS](#irodori-tts) | | OuteTTS | `outetts` | `tts`, `clon` | [OuteTTS](#outetts) | @@ -304,7 +304,7 @@ Higgs Audio v3 TTS is a voice-clone TTS model. The current integration uses the | Field | Value | |---|---| -| Family | `higgs_tts` | +| Family | `higgs_audio_tts` | | Model directory | `models/higgs-audio-v3-tts-4b` | | Task | `tts` | | Modes | `offline` | @@ -313,15 +313,15 @@ Higgs Audio v3 TTS is a voice-clone TTS model. The current integration uses the | Built-in voices | Not exposed | ```bash -audiocpp_cli --task tts --family higgs_tts --model models/higgs-audio-v3-tts-4b --backend cuda --text "Hello from Higgs Audio." --voice-ref assets/resources/b.wav --reference-text "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you." --out out.wav +audiocpp_cli --task tts --family higgs_audio_tts --model models/higgs-audio-v3-tts-4b --backend cuda --text "Hello from Higgs Audio." --voice-ref assets/resources/b.wav --reference-text "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you." --out out.wav ``` | Option | Values | Default | Meaning | |---|---|---:|---| | `--voice-ref` | WAV path | required | Reference speaker audio. | | `--reference-text` | text | empty string | Transcript for reference audio. | -| `--text-chunk-size` | integer chars | `512` | Long-form chunk size. | -| `--max-tokens` | integer | `1024` | Maximum generated AR tokens per chunk. | +| `--text-chunk-size` | integer chars | `1024` | Long-form chunk size. | +| `--max-tokens` | integer | `2048` | Maximum generated AR tokens per chunk. | | `--temperature` | float | `0.8` | AR sampling temperature. | | `--top-k` | integer | `30` | AR top-k sampling limit. The narrower default is less prone to premature EOC than the Python client's `50`. | | `--top-p` | float | `0.8` | AR nucleus sampling limit. The Python client's unfiltered equivalent is `1.0`. | diff --git a/include/engine/models/higgs_tts/ar.h b/include/engine/models/higgs_audio_tts/ar.h similarity index 97% rename from include/engine/models/higgs_tts/ar.h rename to include/engine/models/higgs_audio_tts/ar.h index fd47ba3f..195bfdf3 100644 --- a/include/engine/models/higgs_tts/ar.h +++ b/include/engine/models/higgs_audio_tts/ar.h @@ -5,7 +5,7 @@ #include "engine/framework/core/module.h" #include "engine/framework/modules/attention/qwen_decoder.h" #include "engine/framework/runtime/kv_cache.h" -#include "engine/models/higgs_tts/assets.h" +#include "engine/models/higgs_audio_tts/assets.h" #include #include @@ -16,7 +16,7 @@ namespace engine::core { class BackendWeightStore; } -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { struct HiggsQwenDecoderStackWeights { std::vector layers; @@ -165,4 +165,4 @@ class HiggsARDecodeGraph { std::unique_ptr impl_; }; -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_tts/assets.h b/include/engine/models/higgs_audio_tts/assets.h similarity index 94% rename from include/engine/models/higgs_tts/assets.h rename to include/engine/models/higgs_audio_tts/assets.h index a6899778..aadc195e 100644 --- a/include/engine/models/higgs_tts/assets.h +++ b/include/engine/models/higgs_audio_tts/assets.h @@ -11,7 +11,7 @@ namespace engine::assets { class TensorSource; } -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { struct HiggsTextConfig { std::string model_type; @@ -62,4 +62,4 @@ struct HiggsAssets { std::shared_ptr load_higgs_assets(const std::filesystem::path & model_path); -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_tts/codebooks.h b/include/engine/models/higgs_audio_tts/codebooks.h similarity index 85% rename from include/engine/models/higgs_tts/codebooks.h rename to include/engine/models/higgs_audio_tts/codebooks.h index 5209b8a1..06024ab9 100644 --- a/include/engine/models/higgs_tts/codebooks.h +++ b/include/engine/models/higgs_audio_tts/codebooks.h @@ -3,7 +3,7 @@ #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { constexpr int32_t kHiggsBocId = 1024; constexpr int32_t kHiggsEocId = 1025; @@ -21,4 +21,4 @@ std::vector reverse_higgs_delay_pattern( int64_t delayed_frames, int64_t codebooks); -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_tts/codec.h b/include/engine/models/higgs_audio_tts/codec.h similarity index 96% rename from include/engine/models/higgs_tts/codec.h rename to include/engine/models/higgs_audio_tts/codec.h index fa646843..28516eea 100644 --- a/include/engine/models/higgs_tts/codec.h +++ b/include/engine/models/higgs_audio_tts/codec.h @@ -7,7 +7,7 @@ #include "engine/framework/modules/conv_modules.h" #include "engine/framework/modules/linear_module.h" #include "engine/framework/runtime/session.h" -#include "engine/models/higgs_tts/assets.h" +#include "engine/models/higgs_audio_tts/assets.h" #include #include @@ -19,7 +19,7 @@ namespace engine::core { class BackendWeightStore; } -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { class HiggsCodecDecodeGraph; class HiggsCodecEncodeGraph; @@ -135,4 +135,4 @@ HiggsCodecWeights load_higgs_codec_decode_weights( size_t weight_context_bytes, assets::TensorStorageType weight_storage_type); -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_tts/generator.h b/include/engine/models/higgs_audio_tts/generator.h similarity index 87% rename from include/engine/models/higgs_tts/generator.h rename to include/engine/models/higgs_audio_tts/generator.h index d943383a..aeb15f9b 100644 --- a/include/engine/models/higgs_tts/generator.h +++ b/include/engine/models/higgs_audio_tts/generator.h @@ -1,9 +1,9 @@ #pragma once -#include "engine/models/higgs_tts/ar.h" -#include "engine/models/higgs_tts/codec.h" -#include "engine/models/higgs_tts/sampler.h" -#include "engine/models/higgs_tts/tokenizer_text.h" +#include "engine/models/higgs_audio_tts/ar.h" +#include "engine/models/higgs_audio_tts/codec.h" +#include "engine/models/higgs_audio_tts/sampler.h" +#include "engine/models/higgs_audio_tts/tokenizer_text.h" #include #include @@ -11,7 +11,7 @@ #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { struct HiggsGenerationOptions { int64_t max_tokens = 2048; @@ -74,4 +74,4 @@ class HiggsGenerator { std::unique_ptr decode_graph_; }; -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_tts/loader.h b/include/engine/models/higgs_audio_tts/loader.h similarity index 71% rename from include/engine/models/higgs_tts/loader.h rename to include/engine/models/higgs_audio_tts/loader.h index 6f17c809..b5fb9d58 100644 --- a/include/engine/models/higgs_tts/loader.h +++ b/include/engine/models/higgs_audio_tts/loader.h @@ -1,12 +1,12 @@ #pragma once #include "engine/framework/runtime/model.h" -#include "engine/models/higgs_tts/assets.h" +#include "engine/models/higgs_audio_tts/assets.h" #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { class HiggsTTSLoadedModel final : public runtime::ILoadedVoiceModel { public: @@ -27,7 +27,7 @@ class HiggsTTSLoadedModel final : public runtime::ILoadedVoiceModel { std::shared_ptr assets_; }; -std::unique_ptr load_higgs_tts_model(const std::filesystem::path & model_path); -std::shared_ptr make_higgs_tts_loader(); +std::unique_ptr load_higgs_audio_tts_model(const std::filesystem::path & model_path); +std::shared_ptr make_higgs_audio_tts_loader(); -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_tts/sampler.h b/include/engine/models/higgs_audio_tts/sampler.h similarity index 90% rename from include/engine/models/higgs_tts/sampler.h rename to include/engine/models/higgs_audio_tts/sampler.h index 5d40a2a2..0ed74c1d 100644 --- a/include/engine/models/higgs_tts/sampler.h +++ b/include/engine/models/higgs_audio_tts/sampler.h @@ -1,13 +1,13 @@ #pragma once #include "engine/framework/sampling/torch_random.h" -#include "engine/models/higgs_tts/codebooks.h" +#include "engine/models/higgs_audio_tts/codebooks.h" #include #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { constexpr int64_t kHiggsMaxTopK = 1026; @@ -51,4 +51,4 @@ class HiggsCodebookSampler { std::vector scratch_codes_; }; -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_tts/session.h b/include/engine/models/higgs_audio_tts/session.h similarity index 90% rename from include/engine/models/higgs_tts/session.h rename to include/engine/models/higgs_audio_tts/session.h index c0ef18a5..6b19ff2a 100644 --- a/include/engine/models/higgs_tts/session.h +++ b/include/engine/models/higgs_audio_tts/session.h @@ -2,10 +2,10 @@ #include "engine/framework/runtime/cache_slots.h" #include "engine/framework/runtime/session_base.h" -#include "engine/models/higgs_tts/assets.h" -#include "engine/models/higgs_tts/ar.h" -#include "engine/models/higgs_tts/codec.h" -#include "engine/models/higgs_tts/generator.h" +#include "engine/models/higgs_audio_tts/assets.h" +#include "engine/models/higgs_audio_tts/ar.h" +#include "engine/models/higgs_audio_tts/codec.h" +#include "engine/models/higgs_audio_tts/generator.h" #include #include @@ -13,7 +13,7 @@ #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { class HiggsTTSSession final : public runtime::RuntimeSessionBase @@ -76,4 +76,4 @@ class HiggsTTSSession final std::optional uncached_reference_; }; -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/include/engine/models/higgs_tts/tokenizer_text.h b/include/engine/models/higgs_audio_tts/tokenizer_text.h similarity index 83% rename from include/engine/models/higgs_tts/tokenizer_text.h rename to include/engine/models/higgs_audio_tts/tokenizer_text.h index b43f99f1..33cb41a1 100644 --- a/include/engine/models/higgs_tts/tokenizer_text.h +++ b/include/engine/models/higgs_audio_tts/tokenizer_text.h @@ -1,13 +1,13 @@ #pragma once -#include "engine/models/higgs_tts/assets.h" +#include "engine/models/higgs_audio_tts/assets.h" #include #include #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { struct HiggsPromptRequest { std::string text; @@ -34,4 +34,4 @@ class HiggsTextTokenizer { std::shared_ptr impl_; }; -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/model_specs/higgs_tts.json b/model_specs/higgs_audio_tts.json similarity index 96% rename from model_specs/higgs_tts.json rename to model_specs/higgs_audio_tts.json index a28cf9cd..798d4d6d 100644 --- a/model_specs/higgs_tts.json +++ b/model_specs/higgs_audio_tts.json @@ -1,5 +1,5 @@ { - "family": "higgs_tts", + "family": "higgs_audio_tts", "sources": [ { "format": "gguf", diff --git a/src/framework/runtime/registry.cpp b/src/framework/runtime/registry.cpp index 78bdcce2..6e3be298 100644 --- a/src/framework/runtime/registry.cpp +++ b/src/framework/runtime/registry.cpp @@ -14,7 +14,7 @@ #include "engine/models/fish_audio/loader.h" #include "engine/models/heartmula/loader.h" #include "engine/models/higgs_audio_stt/loader.h" -#include "engine/models/higgs_tts/loader.h" +#include "engine/models/higgs_audio_tts/loader.h" #include "engine/models/hviske_asr/loader.h" #include "engine/models/index_tts2/loader.h" #include "engine/models/irodori_tts/loader.h" @@ -261,7 +261,7 @@ ModelRegistry make_default_registry(const std::optional & engine::models::fish_audio::make_fish_audio_loader(), engine::models::heartmula::make_heartmula_loader(), engine::models::higgs_audio_stt::make_higgs_audio_stt_loader(), - engine::models::higgs_tts::make_higgs_tts_loader(), + engine::models::higgs_audio_tts::make_higgs_audio_tts_loader(), engine::models::hviske_asr::make_hviske_asr_loader(), engine::models::irodori_tts::make_irodori_tts_loader(), engine::models::nemotron_asr::make_nemotron_asr_loader(), diff --git a/src/models/higgs_tts/ar.cpp b/src/models/higgs_audio_tts/ar.cpp similarity index 97% rename from src/models/higgs_tts/ar.cpp rename to src/models/higgs_audio_tts/ar.cpp index 21161241..c113cb04 100644 --- a/src/models/higgs_tts/ar.cpp +++ b/src/models/higgs_audio_tts/ar.cpp @@ -1,4 +1,4 @@ -#include "engine/models/higgs_tts/ar.h" +#include "engine/models/higgs_audio_tts/ar.h" #include "engine/framework/core/backend.h" #include "engine/framework/core/backend_weight_store.h" @@ -24,7 +24,7 @@ #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { namespace { namespace modules = engine::modules; @@ -345,7 +345,7 @@ HiggsARWeights load_higgs_ar_weights( weights.store = std::make_shared( backend, backend_type, - "higgs_tts.ar.weights", + "higgs_audio_tts.ar.weights", weight_context_bytes); weights.text_embedding = weights.store->load_tensor( source, @@ -435,7 +435,7 @@ struct HiggsARKVCache::Impl { if (ctx == nullptr) { throw std::runtime_error("failed to initialize Higgs TTS AR KV cache context"); } - core::ModuleBuildContext build_ctx{ctx.get(), "higgs_tts.ar.kv_cache", runtime->backend_type()}; + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_audio_tts.ar.kv_cache", runtime->backend_type()}; std::vector key_tensors; std::vector value_tensors; key_tensors.reserve(tensor_weights.decoder.layers.size()); @@ -580,7 +580,7 @@ struct HiggsARDecodeGraph::Impl { } const auto & config = runtime->assets().config; const auto & tensor_weights = runtime->weights(); - core::ModuleBuildContext build_ctx{ctx.get(), "higgs_tts.ar.decode", runtime->backend_type()}; + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_audio_tts.ar.decode", runtime->backend_type()}; fused_code_ids = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, config.audio.num_codebooks); auto x = build_higgs_decode_code_embedding( @@ -630,7 +630,7 @@ struct HiggsARDecodeGraph::Impl { fused_code_ids_values.assign(static_cast(config.audio.num_codebooks), 0); attention_mask_values.assign(static_cast(cache_steps), ggml_fp32_to_fp16(-INFINITY)); engine::debug::timing_log_scalar( - "higgs_tts.ar.decode.graph.build_ms", + "higgs_audio_tts.ar.decode.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); } @@ -714,7 +714,7 @@ struct HiggsARDecodeGraph::Impl { const double input_upload_delta_ms = engine::debug::elapsed_ms(timing_start, Clock::now()); input_upload_ms += input_upload_delta_ms; if (log_timing) { - engine::debug::timing_log_scalar("higgs_tts.ar.decode.step0.input_upload_ms", input_upload_delta_ms); + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.step0.input_upload_ms", input_upload_delta_ms); } timing_start = Clock::now(); attention_mask_values[static_cast(cache_slot_value)] = ggml_fp32_to_fp16(0.0F); @@ -726,7 +726,7 @@ struct HiggsARDecodeGraph::Impl { const double mask_upload_delta_ms = engine::debug::elapsed_ms(timing_start, Clock::now()); mask_upload_ms += mask_upload_delta_ms; if (log_timing) { - engine::debug::timing_log_scalar("higgs_tts.ar.decode.step0.mask_upload_ms", mask_upload_delta_ms); + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.step0.mask_upload_ms", mask_upload_delta_ms); } timing_start = Clock::now(); @@ -737,7 +737,7 @@ struct HiggsARDecodeGraph::Impl { const double graph_compute_delta_ms = engine::debug::elapsed_ms(timing_start, Clock::now()); graph_compute_ms += graph_compute_delta_ms; if (log_timing) { - engine::debug::timing_log_scalar("higgs_tts.ar.decode.step0.graph.compute_ms", graph_compute_delta_ms); + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.step0.graph.compute_ms", graph_compute_delta_ms); } if (status != GGML_STATUS_SUCCESS) { throw std::runtime_error("Higgs TTS AR decode graph compute failed"); @@ -753,7 +753,7 @@ struct HiggsARDecodeGraph::Impl { const double output_read_delta_ms = engine::debug::elapsed_ms(timing_start, Clock::now()); output_read_ms += output_read_delta_ms; if (log_timing) { - engine::debug::timing_log_scalar("higgs_tts.ar.decode.step0.output_read_ms", output_read_delta_ms); + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.step0.output_read_ms", output_read_delta_ms); } cache->advance_after_direct_append(1); @@ -808,7 +808,7 @@ struct HiggsARPrefillGraph::Impl { throw std::runtime_error("Higgs TTS AR suffix prefill requires a target KV cache"); } if (layerwise) { - engine::debug::timing_log_scalar("higgs_tts.ar.prefill.graph.build_ms", 0.0); + engine::debug::timing_log_scalar("higgs_audio_tts.ar.prefill.graph.build_ms", 0.0); return; } const auto build_start = Clock::now(); @@ -819,7 +819,7 @@ struct HiggsARPrefillGraph::Impl { } const auto & config = runtime->assets().config; const auto & tensor_weights = runtime->weights(); - core::ModuleBuildContext build_ctx{ctx.get(), "higgs_tts.ar.prefill", runtime->backend_type()}; + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_audio_tts.ar.prefill", runtime->backend_type()}; text_tokens = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, run_steps); fused_code_ids = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_I32, config.audio.num_codebooks, run_steps); @@ -915,7 +915,7 @@ struct HiggsARPrefillGraph::Impl { positions_values = modules::qwen_position_ids(run_steps, start_step); attention_mask_values = modules::qwen_causal_suffix_mask_values(1, run_steps, start_step); engine::debug::timing_log_scalar( - "higgs_tts.ar.prefill.graph.build_ms", + "higgs_audio_tts.ar.prefill.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); } @@ -944,7 +944,7 @@ struct HiggsARPrefillGraph::Impl { if (ctx == nullptr) { throw std::runtime_error("failed to initialize Higgs TTS AR embedding graph context"); } - core::ModuleBuildContext build_ctx{ctx.get(), "higgs_tts.ar.prefill.embedding", runtime.backend_type()}; + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_audio_tts.ar.prefill.embedding", runtime.backend_type()}; text_tokens = ggml_new_tensor_1d(ctx.get(), GGML_TYPE_I32, steps); fused_code_ids = ggml_new_tensor_2d(ctx.get(), GGML_TYPE_I32, config.audio.num_codebooks, steps); text_gate = ggml_new_tensor_3d(ctx.get(), GGML_TYPE_F32, 1, steps, 1); @@ -1020,7 +1020,7 @@ struct HiggsARPrefillGraph::Impl { if (ctx == nullptr) { throw std::runtime_error("failed to initialize Higgs TTS AR layer prefill graph context"); } - core::ModuleBuildContext build_ctx{ctx.get(), "higgs_tts.ar.prefill.layer", runtime.backend_type()}; + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_audio_tts.ar.prefill.layer", runtime.backend_type()}; auto x = core::make_tensor( build_ctx, GGML_TYPE_F32, @@ -1116,7 +1116,7 @@ struct HiggsARPrefillGraph::Impl { if (ctx == nullptr) { throw std::runtime_error("failed to initialize Higgs TTS AR final prefill graph context"); } - core::ModuleBuildContext build_ctx{ctx.get(), "higgs_tts.ar.prefill.final", runtime.backend_type()}; + core::ModuleBuildContext build_ctx{ctx.get(), "higgs_audio_tts.ar.prefill.final", runtime.backend_type()}; auto x = core::make_tensor( build_ctx, GGML_TYPE_F32, @@ -1364,4 +1364,4 @@ void HiggsARDecodeGraph::run_step_into( impl_->run_step_into(input, output, log_timing); } -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_tts/assets.cpp b/src/models/higgs_audio_tts/assets.cpp similarity index 97% rename from src/models/higgs_tts/assets.cpp rename to src/models/higgs_audio_tts/assets.cpp index 350be213..7a163d45 100644 --- a/src/models/higgs_tts/assets.cpp +++ b/src/models/higgs_audio_tts/assets.cpp @@ -1,4 +1,4 @@ -#include "engine/models/higgs_tts/assets.h" +#include "engine/models/higgs_audio_tts/assets.h" #include "engine/framework/assets/model_package.h" #include "engine/framework/io/config.h" @@ -7,7 +7,7 @@ #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { namespace json = engine::io::json; namespace { @@ -161,11 +161,11 @@ std::shared_ptr load_higgs_assets(const std::filesystem::path HiggsAssets assets; assets.resources = assets::load_resource_bundle_from_package_spec( model_path, - assets::default_model_package_spec_path("higgs_tts")); + assets::default_model_package_spec_path("higgs_audio_tts")); assets.config = parse_config(assets.resources); assets.weights = assets.resources.open_tensor_source("weights"); validate_weight_anchors(assets); return std::make_shared(std::move(assets)); } -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_tts/codebooks.cpp b/src/models/higgs_audio_tts/codebooks.cpp similarity index 95% rename from src/models/higgs_tts/codebooks.cpp rename to src/models/higgs_audio_tts/codebooks.cpp index 9cf0b0b6..54c15d38 100644 --- a/src/models/higgs_tts/codebooks.cpp +++ b/src/models/higgs_audio_tts/codebooks.cpp @@ -1,9 +1,9 @@ -#include "engine/models/higgs_tts/codebooks.h" +#include "engine/models/higgs_audio_tts/codebooks.h" #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { namespace { void require_codebook_matrix( @@ -76,4 +76,4 @@ std::vector reverse_higgs_delay_pattern( return raw; } -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_tts/codec.cpp b/src/models/higgs_audio_tts/codec.cpp similarity index 98% rename from src/models/higgs_tts/codec.cpp rename to src/models/higgs_audio_tts/codec.cpp index 32a16820..658919d3 100644 --- a/src/models/higgs_tts/codec.cpp +++ b/src/models/higgs_audio_tts/codec.cpp @@ -1,4 +1,4 @@ -#include "engine/models/higgs_tts/codec.h" +#include "engine/models/higgs_audio_tts/codec.h" #include "engine/framework/audio/conversion.h" #include "engine/framework/audio/resampling.h" @@ -25,7 +25,7 @@ #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { namespace { using Clock = std::chrono::steady_clock; @@ -1120,7 +1120,7 @@ class HiggsCodecEncodeGraph { throw std::runtime_error("failed to initialize Higgs TTS codec encode graph context"); } core::ModuleBuildContext build_ctx{ - ctx_.get(), "higgs_tts.codec.encode", runtime_->backend_type()}; + ctx_.get(), "higgs_audio_tts.codec.encode", runtime_->backend_type()}; acoustic_input_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_F32, acoustic_samples_); semantic_input_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_F32, semantic_samples_); ggml_set_input(acoustic_input_); @@ -1147,7 +1147,7 @@ class HiggsCodecEncodeGraph { } throw std::runtime_error("failed to allocate Higgs TTS codec encode graph"); } - engine::debug::timing_log_scalar("higgs_tts.codec.encode.graph.build_ms", + engine::debug::timing_log_scalar("higgs_audio_tts.codec.encode.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); } @@ -1177,12 +1177,12 @@ class HiggsCodecEncodeGraph { acoustic_input_, acoustic.data(), 0, acoustic.size() * sizeof(float)); ggml_backend_tensor_set( semantic_input_, semantic.data(), 0, semantic.size() * sizeof(float)); - engine::debug::timing_log_scalar("higgs_tts.codec.encode_input_upload_ms", + engine::debug::timing_log_scalar("higgs_audio_tts.codec.encode_input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); core::set_backend_threads(runtime_->backend(), runtime_->threads()); timing_start = Clock::now(); const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); - engine::debug::timing_log_scalar("higgs_tts.codec.encode.graph.compute_ms", + engine::debug::timing_log_scalar("higgs_audio_tts.codec.encode.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); if (status != GGML_STATUS_SUCCESS) { throw std::runtime_error("Higgs TTS codec encode graph compute failed"); @@ -1203,7 +1203,7 @@ class HiggsCodecEncodeGraph { codebook_codes[static_cast(frame)]; } } - engine::debug::timing_log_scalar("higgs_tts.codec.encode_output_read_ms", + engine::debug::timing_log_scalar("higgs_audio_tts.codec.encode_output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); return out; } @@ -1238,7 +1238,7 @@ class HiggsCodecDecodeGraph { throw std::runtime_error("failed to initialize Higgs TTS codec decode graph context"); } core::ModuleBuildContext build_ctx{ - ctx_.get(), "higgs_tts.codec.decode", runtime_->backend_type()}; + ctx_.get(), "higgs_audio_tts.codec.decode", runtime_->backend_type()}; codes_ = ggml_new_tensor_2d(ctx_.get(), GGML_TYPE_I32, kCodecCodebooks, capacity_frames_); frame_mask_ = ggml_new_tensor_2d(ctx_.get(), GGML_TYPE_F32, 1, capacity_frames_); ggml_set_input(codes_); @@ -1269,7 +1269,7 @@ class HiggsCodecDecodeGraph { } code_scratch_.assign(static_cast(capacity_frames_ * kCodecCodebooks), 0); frame_mask_values_.assign(static_cast(capacity_frames_), 0.0F); - engine::debug::timing_log_scalar("higgs_tts.codec.decode.graph.build_ms", + engine::debug::timing_log_scalar("higgs_audio_tts.codec.decode.graph.build_ms", engine::debug::elapsed_ms(build_start, Clock::now())); } @@ -1314,12 +1314,12 @@ class HiggsCodecDecodeGraph { codes_, code_scratch_.data(), 0, code_scratch_.size() * sizeof(int32_t)); ggml_backend_tensor_set( frame_mask_, frame_mask_values_.data(), 0, frame_mask_values_.size() * sizeof(float)); - engine::debug::timing_log_scalar("higgs_tts.codec.decode_input_upload_ms", + engine::debug::timing_log_scalar("higgs_audio_tts.codec.decode_input_upload_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); core::set_backend_threads(runtime_->backend(), runtime_->threads()); timing_start = Clock::now(); const ggml_status status = engine::core::compute_backend_graph(runtime_->backend(), graph_); - engine::debug::timing_log_scalar("higgs_tts.codec.decode.graph.compute_ms", + engine::debug::timing_log_scalar("higgs_audio_tts.codec.decode.graph.compute_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); if (status != GGML_STATUS_SUCCESS) { throw std::runtime_error("Higgs TTS codec decode graph compute failed"); @@ -1334,7 +1334,7 @@ class HiggsCodecDecodeGraph { out.values.resize(static_cast(out.samples)); timing_start = Clock::now(); ggml_backend_tensor_get(output_, out.values.data(), 0, out.values.size() * sizeof(float)); - engine::debug::timing_log_scalar("higgs_tts.codec.decode_output_read_ms", + engine::debug::timing_log_scalar("higgs_audio_tts.codec.decode_output_read_ms", engine::debug::elapsed_ms(timing_start, Clock::now())); return out; } @@ -1365,7 +1365,7 @@ HiggsCodecWeights load_higgs_codec_decode_weights(const HiggsAssets & assets, } HiggsCodecWeights weights; weights.store = std::make_shared( - backend, backend_type, "higgs_tts.codec.weights", weight_context_bytes); + backend, backend_type, "higgs_audio_tts.codec.weights", weight_context_bytes); const auto & source = *assets.weights; load_hubert_semantic_model_weights(weights, source, weight_storage_type); weights.quantizers.reserve(kCodecCodebooks); @@ -1527,11 +1527,11 @@ HiggsCodecRuntime::encode_reference(const runtime::AudioBuffer & audio) const { static_cast(semantic_16k.size()), frames); } - engine::debug::trace_log_scalar("higgs_tts.codec.encode.input_frames", frames); - engine::debug::trace_log_f32("higgs_tts.codec.encode.input_acoustic_24k", + engine::debug::trace_log_scalar("higgs_audio_tts.codec.encode.input_frames", frames); + engine::debug::trace_log_f32("higgs_audio_tts.codec.encode.input_acoustic_24k", {static_cast(acoustic_24k.size())}, acoustic_24k); - engine::debug::trace_log_f32("higgs_tts.codec.encode.input_semantic_16k", + engine::debug::trace_log_f32("higgs_audio_tts.codec.encode.input_semantic_16k", {static_cast(semantic_16k.size())}, semantic_16k); return encode_graph_->run(acoustic_24k, semantic_16k, frames); @@ -1549,9 +1549,9 @@ HiggsCodecDecodeOutput HiggsCodecRuntime::decode_codes(const std::vector(codes.size()) != frames * codebooks) { throw std::runtime_error("Higgs TTS codec decode code count mismatch"); } - engine::debug::trace_log_scalar("higgs_tts.codec.decode.input_frames", frames); - engine::debug::trace_log_scalar("higgs_tts.codec.decode.input_codebooks", codebooks); - engine::debug::trace_log_i32("higgs_tts.codec.decode.input_codes", + engine::debug::trace_log_scalar("higgs_audio_tts.codec.decode.input_frames", frames); + engine::debug::trace_log_scalar("higgs_audio_tts.codec.decode.input_codebooks", codebooks); + engine::debug::trace_log_i32("higgs_audio_tts.codec.decode.input_codes", {frames, codebooks}, codes); @@ -1632,4 +1632,4 @@ void HiggsCodecRuntime::release_runtime_graphs() { decode_graph_.reset(); } -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_tts/generator.cpp b/src/models/higgs_audio_tts/generator.cpp similarity index 83% rename from src/models/higgs_tts/generator.cpp rename to src/models/higgs_audio_tts/generator.cpp index 328d79bc..3cfc0dac 100644 --- a/src/models/higgs_tts/generator.cpp +++ b/src/models/higgs_audio_tts/generator.cpp @@ -1,10 +1,10 @@ -#include "engine/models/higgs_tts/generator.h" +#include "engine/models/higgs_audio_tts/generator.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/debug/trace.h" #include "engine/framework/runtime/options.h" #include "engine/framework/sampling/torch_random.h" -#include "engine/models/higgs_tts/codebooks.h" +#include "engine/models/higgs_audio_tts/codebooks.h" #include #include @@ -15,7 +15,7 @@ #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { namespace { using Clock = std::chrono::steady_clock; @@ -255,28 +255,28 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re "Higgs TTS generation got reference codebooks without reference codes"); } validate_generation_options(request.options); - engine::debug::trace_log_scalar("higgs_tts.request.text", request.text); - engine::debug::trace_log_scalar("higgs_tts.request.reference_text", request.reference_text); - engine::debug::trace_log_scalar("higgs_tts.request.text_chars", request.text.size()); - engine::debug::trace_log_scalar("higgs_tts.request.reference_text_chars", + engine::debug::trace_log_scalar("higgs_audio_tts.request.text", request.text); + engine::debug::trace_log_scalar("higgs_audio_tts.request.reference_text", request.reference_text); + engine::debug::trace_log_scalar("higgs_audio_tts.request.text_chars", request.text.size()); + engine::debug::trace_log_scalar("higgs_audio_tts.request.reference_text_chars", request.reference_text.size()); - engine::debug::trace_log_scalar("higgs_tts.request.max_tokens", + engine::debug::trace_log_scalar("higgs_audio_tts.request.max_tokens", request.options.max_tokens); - engine::debug::trace_log_scalar("higgs_tts.request.temperature", request.options.temperature); - engine::debug::trace_log_scalar("higgs_tts.request.top_p", + engine::debug::trace_log_scalar("higgs_audio_tts.request.temperature", request.options.temperature); + engine::debug::trace_log_scalar("higgs_audio_tts.request.top_p", request.options.top_p.has_value() ? std::to_string(*request.options.top_p) : "none"); - engine::debug::trace_log_scalar("higgs_tts.request.top_k", + engine::debug::trace_log_scalar("higgs_audio_tts.request.top_k", request.options.top_k.has_value() ? std::to_string(*request.options.top_k) : "none"); - engine::debug::trace_log_scalar("higgs_tts.request.repetition_penalty", + engine::debug::trace_log_scalar("higgs_audio_tts.request.repetition_penalty", request.options.repetition_penalty); - engine::debug::trace_log_scalar("higgs_tts.request.has_seed", request.options.seed.has_value()); - engine::debug::trace_log_scalar("higgs_tts.request.seed", + engine::debug::trace_log_scalar("higgs_audio_tts.request.has_seed", request.options.seed.has_value()); + engine::debug::trace_log_scalar("higgs_audio_tts.request.seed", request.options.seed.has_value() ? std::to_string(*request.options.seed) : "none"); if (has_reference) { - engine::debug::trace_log_i32("higgs_tts.request.reference_codes", + engine::debug::trace_log_i32("higgs_audio_tts.request.reference_codes", {request.reference_frames, request.reference_codebooks}, request.reference_codes); } @@ -306,36 +306,36 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re request.reference_text, delayed_reference_frames, }); - engine::debug::trace_log_scalar("higgs_tts.prompt.text_tokens", prompt.text_ids.size()); - engine::debug::trace_log_i32("higgs_tts.prompt.text_ids", + engine::debug::trace_log_scalar("higgs_audio_tts.prompt.text_tokens", prompt.text_ids.size()); + engine::debug::trace_log_i32("higgs_audio_tts.prompt.text_ids", {static_cast(prompt.text_ids.size())}, prompt.text_ids); - engine::debug::trace_log_scalar("higgs_tts.prompt.reference_text_tokens", + engine::debug::trace_log_scalar("higgs_audio_tts.prompt.reference_text_tokens", prompt.reference_text_ids.size()); - engine::debug::trace_log_i32("higgs_tts.prompt.reference_text_ids", + engine::debug::trace_log_i32("higgs_audio_tts.prompt.reference_text_ids", {static_cast(prompt.reference_text_ids.size())}, prompt.reference_text_ids); const auto prepared = make_prepared_prompt(prompt, *delayed_reference_codes, delayed_reference_frames, config); - engine::debug::trace_log_scalar("higgs_tts.prompt.tokens", prepared.prompt.token_ids.size()); - engine::debug::trace_log_i32("higgs_tts.prompt.token_ids", + engine::debug::trace_log_scalar("higgs_audio_tts.prompt.tokens", prepared.prompt.token_ids.size()); + engine::debug::trace_log_i32("higgs_audio_tts.prompt.token_ids", {static_cast(prepared.prompt.token_ids.size())}, prepared.prompt.token_ids); - engine::debug::trace_log_scalar("higgs_tts.prompt.delayed_reference_rows", + engine::debug::trace_log_scalar("higgs_audio_tts.prompt.delayed_reference_rows", delayed_reference_frames); - engine::debug::trace_log_i32("higgs_tts.prompt.delayed_reference_codes", + engine::debug::trace_log_i32("higgs_audio_tts.prompt.delayed_reference_codes", {delayed_reference_frames, config.audio.num_codebooks}, *delayed_reference_codes); - engine::debug::trace_log_i32("higgs_tts.ar.prefill.text_tokens", + engine::debug::trace_log_i32("higgs_audio_tts.ar.prefill.text_tokens", {prepared.ar_input.steps}, prepared.ar_input.text_tokens); - engine::debug::trace_log_i32("higgs_tts.ar.prefill.fused_code_ids", + engine::debug::trace_log_i32("higgs_audio_tts.ar.prefill.fused_code_ids", {prepared.ar_input.steps, config.audio.num_codebooks}, prepared.ar_input.fused_code_ids); - engine::debug::trace_log_f32("higgs_tts.ar.prefill.text_gate", + engine::debug::trace_log_f32("higgs_audio_tts.ar.prefill.text_gate", {prepared.ar_input.steps}, prepared.ar_input.text_gate); - engine::debug::trace_log_f32("higgs_tts.ar.prefill.code_gate", + engine::debug::trace_log_f32("higgs_audio_tts.ar.prefill.code_gate", {prepared.ar_input.steps}, prepared.ar_input.code_gate); const int64_t prompt_steps = prepared.ar_input.steps; @@ -351,8 +351,8 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re std::equal(matching_reference_cache->prefix_tokens.begin(), matching_reference_cache->prefix_tokens.end(), prepared.prompt.token_ids.begin()); - engine::debug::trace_log_scalar("higgs_tts.generator.reference_prefix_cache_hit", reference_cache_hit); - engine::debug::trace_log_scalar("higgs_tts.generator.reference_prefix_steps", prepared.prefix_steps); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.reference_prefix_cache_hit", reference_cache_hit); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.reference_prefix_steps", prepared.prefix_steps); const int64_t max_cache_steps = prompt_steps + request.options.max_tokens; const int64_t initial_cache_steps = bucketed_initial_cache_steps(prompt_steps, request.options.max_tokens); const bool cache_rebuild = @@ -372,11 +372,11 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re } else { ar_kv_cache_->reset(); } - engine::debug::trace_log_scalar("higgs_tts.generator.reference_kv_cache_hit", reference_kv_cache_hit); - engine::debug::trace_log_scalar("higgs_tts.generator.prefill_start_step", prefill_start_step); - engine::debug::trace_log_scalar("higgs_tts.generator.prefill_run_steps", prompt_steps - prefill_start_step); - engine::debug::trace_log_scalar("higgs_tts.generator.kv_cache_steps", initial_cache_steps); - engine::debug::trace_log_scalar("higgs_tts.generator.kv_cache_rebuild", cache_rebuild); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.reference_kv_cache_hit", reference_kv_cache_hit); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.prefill_start_step", prefill_start_step); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.prefill_run_steps", prompt_steps - prefill_start_step); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.kv_cache_steps", initial_cache_steps); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.kv_cache_rebuild", cache_rebuild); if (prefill_graph_ == nullptr || !prefill_graph_->matches(*ar_, prompt_steps, prefill_start_step)) { prefill_graph_.reset(); @@ -395,9 +395,9 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re decode_graph_->import_prefill_state(prefill_output.kv_state); } HiggsARDecodeOutput prefill = std::move(prefill_output.output); - engine::debug::timing_log_scalar("higgs_tts.generator.prefill_ms", + engine::debug::timing_log_scalar("higgs_audio_tts.generator.prefill_ms", engine::debug::elapsed_ms(prefill_start, Clock::now())); - engine::debug::trace_log_f32("higgs_tts.sampler.prefill_logits", + engine::debug::trace_log_f32("higgs_audio_tts.sampler.prefill_logits", {config.audio.num_codebooks, config.audio.vocab_size}, prefill.codebook_logits); @@ -415,18 +415,18 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re cuda_sampling_policy_ = engine::sampling::resolve_torch_cuda_sampling_policy( ar_->backend_type(), ar_->device(), - "higgs_tts.cuda_sampling_policy", + "higgs_audio_tts.cuda_sampling_policy", "Higgs TTS", engine::sampling::TorchCudaSamplingPolicyFailureMode::StrictCuda); } sampling.cuda_policy = *cuda_sampling_policy_; - engine::debug::trace_log_scalar("higgs_tts.sampler.temperature", sampling.temperature); - engine::debug::trace_log_scalar("higgs_tts.sampler.has_seed", sampling.has_seed); - engine::debug::trace_log_scalar("higgs_tts.sampler.seed", sampling.seed); - engine::debug::trace_log_scalar("higgs_tts.sampler.top_p", + engine::debug::trace_log_scalar("higgs_audio_tts.sampler.temperature", sampling.temperature); + engine::debug::trace_log_scalar("higgs_audio_tts.sampler.has_seed", sampling.has_seed); + engine::debug::trace_log_scalar("higgs_audio_tts.sampler.seed", sampling.seed); + engine::debug::trace_log_scalar("higgs_audio_tts.sampler.top_p", sampling.top_p.has_value() ? std::to_string(*sampling.top_p) : "none"); - engine::debug::trace_log_scalar("higgs_tts.sampler.top_k", + engine::debug::trace_log_scalar("higgs_audio_tts.sampler.top_k", sampling.top_k.has_value() ? std::to_string(*sampling.top_k) : "none"); HiggsGenerationResult result; @@ -436,7 +436,7 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re static_cast(prefill.codebook_logits.size()), state, sampling); - engine::debug::trace_log_i32("higgs_tts.sampler.first_output_codes", + engine::debug::trace_log_i32("higgs_audio_tts.sampler.first_output_codes", {static_cast(first_sampled.size())}, first_sampled); result.delayed_codes.insert( @@ -464,7 +464,7 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re decode_graph_.reset(); ar_kv_cache_ = std::make_unique(ar_, grown_cache_steps); ar_kv_cache_->import_state(kv_state); - engine::debug::trace_log_scalar("higgs_tts.generator.kv_cache_grown_steps", grown_cache_steps); + engine::debug::trace_log_scalar("higgs_audio_tts.generator.kv_cache_grown_steps", grown_cache_steps); decode_graph_ = std::make_unique( ar_, ar_kv_cache_->cache_steps(), *ar_kv_cache_, ar_decode_graph_arena_bytes_); decode_graph_->begin_decode_run(); @@ -475,7 +475,7 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re const auto step_start = Clock::now(); decode_graph_->run_step_into(input, decoded, !logged_decode_step_timing); if (!logged_decode_step_timing) { - engine::debug::timing_log_scalar("higgs_tts.generator.decode.step0.ar_ms", + engine::debug::timing_log_scalar("higgs_audio_tts.generator.decode.step0.ar_ms", engine::debug::elapsed_ms(step_start, Clock::now())); } const auto sample_start = Clock::now(); @@ -485,7 +485,7 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re sampling); sampler_total_ms += engine::debug::elapsed_ms(sample_start, Clock::now()); if (!logged_decode_step_timing) { - engine::debug::timing_log_scalar("higgs_tts.generator.decode.step0.sampler_ms", + engine::debug::timing_log_scalar("higgs_audio_tts.generator.decode.step0.sampler_ms", sampler_total_ms); logged_decode_step_timing = true; } @@ -495,13 +495,13 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re } } decode_timing_total.add(decode_graph_->timing()); - engine::debug::timing_log_scalar("higgs_tts.ar.decode.steps", decode_timing_total.steps); - engine::debug::timing_log_scalar("higgs_tts.ar.decode.input_upload_ms", decode_timing_total.input_upload_ms); - engine::debug::timing_log_scalar("higgs_tts.ar.decode.mask_upload_ms", decode_timing_total.mask_upload_ms); - engine::debug::timing_log_scalar("higgs_tts.ar.decode.graph.compute_ms", decode_timing_total.graph_compute_ms); - engine::debug::timing_log_scalar("higgs_tts.ar.decode.output_read_ms", decode_timing_total.output_read_ms); - engine::debug::timing_log_scalar("higgs_tts.generator.decode.sampler_ms", sampler_total_ms); - engine::debug::timing_log_scalar("higgs_tts.generator.decode_ms", + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.steps", decode_timing_total.steps); + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.input_upload_ms", decode_timing_total.input_upload_ms); + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.mask_upload_ms", decode_timing_total.mask_upload_ms); + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.graph.compute_ms", decode_timing_total.graph_compute_ms); + engine::debug::timing_log_scalar("higgs_audio_tts.ar.decode.output_read_ms", decode_timing_total.output_read_ms); + engine::debug::timing_log_scalar("higgs_audio_tts.generator.decode.sampler_ms", sampler_total_ms); + engine::debug::timing_log_scalar("higgs_audio_tts.generator.decode_ms", engine::debug::elapsed_ms(decode_start, Clock::now())); if (!state.generation_done) { throw std::runtime_error("Higgs TTS generation reached max_tokens before EOC"); @@ -510,11 +510,11 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re result.raw_codes = reverse_higgs_delay_pattern( result.delayed_codes, result.delayed_frames, config.audio.num_codebooks); result.raw_frames = result.delayed_frames - (config.audio.num_codebooks - 1); - engine::debug::trace_log_i32("higgs_tts.generator.delayed_codes", + engine::debug::trace_log_i32("higgs_audio_tts.generator.delayed_codes", {result.delayed_frames, config.audio.num_codebooks}, result.delayed_codes); const int64_t delayed_head_rows = std::min(result.delayed_frames, 8); - engine::debug::trace_log_i32("higgs_tts.generator.delayed_codes_head8", + engine::debug::trace_log_i32("higgs_audio_tts.generator.delayed_codes_head8", {delayed_head_rows, config.audio.num_codebooks}, std::vector( result.delayed_codes.begin(), @@ -531,19 +531,19 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re code = 0; } } - engine::debug::trace_log_i32("higgs_tts.generator.raw_codes_for_codec", + engine::debug::trace_log_i32("higgs_audio_tts.generator.raw_codes_for_codec", {result.raw_frames, config.audio.num_codebooks}, result.raw_codes); const auto codec_start = Clock::now(); result.audio = codec_->decode_codes(result.raw_codes, result.raw_frames, config.audio.num_codebooks); - engine::debug::trace_log_f32("higgs_tts.codec.decode.output_audio", + engine::debug::trace_log_f32("higgs_audio_tts.codec.decode.output_audio", {result.audio.samples}, result.audio.values); - engine::debug::timing_log_scalar("higgs_tts.generator.codec_decode_ms", + engine::debug::timing_log_scalar("higgs_audio_tts.generator.codec_decode_ms", engine::debug::elapsed_ms(codec_start, Clock::now())); codec_->release_runtime_graphs(); return result; } -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_tts/loader.cpp b/src/models/higgs_audio_tts/loader.cpp similarity index 73% rename from src/models/higgs_tts/loader.cpp rename to src/models/higgs_audio_tts/loader.cpp index 9a4a0872..d14b1db9 100644 --- a/src/models/higgs_tts/loader.cpp +++ b/src/models/higgs_audio_tts/loader.cpp @@ -1,17 +1,17 @@ -#include "engine/models/higgs_tts/loader.h" +#include "engine/models/higgs_audio_tts/loader.h" #include "engine/framework/assets/model_package.h" -#include "engine/models/higgs_tts/session.h" +#include "engine/models/higgs_audio_tts/session.h" #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { namespace { runtime::ModelMetadata metadata(const HiggsAssets &) { runtime::ModelMetadata out; - out.family = "higgs_tts"; + out.family = "higgs_audio_tts"; out.variant = "v3-4b"; out.description = "Higgs Audio v3 TTS loaded from local SGLang-Omni compatible assets."; out.config_candidates = { @@ -47,15 +47,15 @@ runtime::ModelCliInterface cli(const HiggsAssets &) { {"text_chunk_mode", "default|tag_aware|japanese|endline", "Framework text chunking mode."}, }; out.session_options = { - {"higgs_tts.weight_type", "native|f32|f16|bf16|q8_0", "AR and codec weight storage type."}, - {"higgs_tts.ar_weight_type", "native|f32|f16|bf16|q8_0", "Autoregressive decoder weight storage type."}, - {"higgs_tts.codec_weight_type", "native|f32|f16|bf16|q8_0", "Codec weight storage type."}, - {"higgs_tts.ar_weight_context_mb", "n", "AR weight context size."}, - {"higgs_tts.codec_weight_context_mb", "n", "Codec weight context size."}, - {"higgs_tts.ar_decode_graph_arena_mb", "n", "AR decode graph arena size."}, - {"higgs_tts.codec_decode_graph_arena_mb", "n", "Codec decode graph arena size."}, - {"higgs_tts.codec_encode_graph_arena_mb", "n", "Codec encode graph arena size."}, - {"higgs_tts.reference_cache_slots", "n", "Encoded reference-audio cache slots; default 1."}, + {"higgs_audio_tts.weight_type", "native|f32|f16|bf16|q8_0", "AR and codec weight storage type."}, + {"higgs_audio_tts.ar_weight_type", "native|f32|f16|bf16|q8_0", "Autoregressive decoder weight storage type."}, + {"higgs_audio_tts.codec_weight_type", "native|f32|f16|bf16|q8_0", "Codec weight storage type."}, + {"higgs_audio_tts.ar_weight_context_mb", "n", "AR weight context size."}, + {"higgs_audio_tts.codec_weight_context_mb", "n", "Codec weight context size."}, + {"higgs_audio_tts.ar_decode_graph_arena_mb", "n", "AR decode graph arena size."}, + {"higgs_audio_tts.codec_decode_graph_arena_mb", "n", "Codec decode graph arena size."}, + {"higgs_audio_tts.codec_encode_graph_arena_mb", "n", "Codec encode graph arena size."}, + {"higgs_audio_tts.reference_cache_slots", "n", "Encoded reference-audio cache slots; default 1."}, }; return out; } @@ -63,7 +63,16 @@ runtime::ModelCliInterface cli(const HiggsAssets &) { class HiggsTTSLoader final : public runtime::IVoiceModelLoader { public: std::string family() const override { - return "higgs_tts"; + return "higgs_audio_tts"; + } + + runtime::CapabilitySet advertised_capabilities() const override { + runtime::CapabilitySet out; + out.supported_tasks = { + {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, + }; + out.supports_speaker_reference = true; + return out; } bool can_load(const runtime::ModelLoadRequest & request) const override { @@ -99,7 +108,7 @@ class HiggsTTSLoader final : public runtime::IVoiceModelLoader { } std::unique_ptr load(const runtime::ModelLoadRequest & request) const override { - return load_higgs_tts_model(request.model_path); + return load_higgs_audio_tts_model(request.model_path); } }; @@ -133,7 +142,7 @@ std::unique_ptr HiggsTTSLoadedModel::create_task_ses return std::make_unique(task, options, assets_); } -std::unique_ptr load_higgs_tts_model(const std::filesystem::path & model_path) { +std::unique_ptr load_higgs_audio_tts_model(const std::filesystem::path & model_path) { auto assets = load_higgs_assets(model_path); return std::make_unique( metadata(*assets), @@ -141,8 +150,8 @@ std::unique_ptr load_higgs_tts_model(const std::filesystem: std::move(assets)); } -std::shared_ptr make_higgs_tts_loader() { +std::shared_ptr make_higgs_audio_tts_loader() { return std::make_shared(); } -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_tts/sampler.cpp b/src/models/higgs_audio_tts/sampler.cpp similarity index 99% rename from src/models/higgs_tts/sampler.cpp rename to src/models/higgs_audio_tts/sampler.cpp index 9c6496c3..ccebd9b7 100644 --- a/src/models/higgs_tts/sampler.cpp +++ b/src/models/higgs_audio_tts/sampler.cpp @@ -1,4 +1,4 @@ -#include "engine/models/higgs_tts/sampler.h" +#include "engine/models/higgs_audio_tts/sampler.h" #include "engine/framework/sampling/torch_random.h" @@ -9,7 +9,7 @@ #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { namespace { constexpr float kGreedyTemperatureThreshold = 1.0e-5F; @@ -455,4 +455,4 @@ const std::vector & HiggsCodebookSampler::step(const float * logits, return scratch_codes_; } -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_tts/session.cpp b/src/models/higgs_audio_tts/session.cpp similarity index 76% rename from src/models/higgs_tts/session.cpp rename to src/models/higgs_audio_tts/session.cpp index 8c4f3bfe..8c3d5328 100644 --- a/src/models/higgs_tts/session.cpp +++ b/src/models/higgs_audio_tts/session.cpp @@ -1,4 +1,4 @@ -#include "engine/models/higgs_tts/session.h" +#include "engine/models/higgs_audio_tts/session.h" #include "engine/framework/debug/profiler.h" #include "engine/framework/debug/trace.h" @@ -12,7 +12,7 @@ #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { namespace { using Clock = std::chrono::steady_clock; @@ -53,13 +53,13 @@ uint64_t hash_audio_samples(const runtime::AudioBuffer & audio) { std::size_t resolve_reference_cache_slots(const runtime::SessionOptions & options) { const int64_t slots = runtime::parse_i64_option( options.options, - {"higgs_tts.reference_cache_slots", "reference_cache_slots"}) + {"higgs_audio_tts.reference_cache_slots", "reference_cache_slots"}) .value_or(kDefaultReferenceCacheSlots); if (slots < 0) { - throw std::runtime_error("higgs_tts.reference_cache_slots must be non-negative"); + throw std::runtime_error("higgs_audio_tts.reference_cache_slots must be non-negative"); } if (static_cast(slots) > static_cast(std::numeric_limits::max())) { - throw std::runtime_error("higgs_tts.reference_cache_slots is too large"); + throw std::runtime_error("higgs_audio_tts.reference_cache_slots is too large"); } return static_cast(slots); } @@ -141,41 +141,41 @@ HiggsTTSSession::HiggsTTSSession( } ar_weight_context_bytes_ = runtime::parse_size_mb_option( - options.options, {"higgs_tts.ar_weight_context_mb"}, ar_weight_context_bytes_); + options.options, {"higgs_audio_tts.ar_weight_context_mb"}, ar_weight_context_bytes_); codec_weight_context_bytes_ = runtime::parse_size_mb_option( - options.options, {"higgs_tts.codec_weight_context_mb"}, codec_weight_context_bytes_); + options.options, {"higgs_audio_tts.codec_weight_context_mb"}, codec_weight_context_bytes_); ar_decode_graph_arena_bytes_ = runtime::parse_size_mb_option( - options.options, {"higgs_tts.ar_decode_graph_arena_mb"}, ar_decode_graph_arena_bytes_); + options.options, {"higgs_audio_tts.ar_decode_graph_arena_mb"}, ar_decode_graph_arena_bytes_); codec_decode_graph_arena_bytes_ = runtime::parse_size_mb_option( - options.options, {"higgs_tts.codec_decode_graph_arena_mb"}, codec_decode_graph_arena_bytes_); + options.options, {"higgs_audio_tts.codec_decode_graph_arena_mb"}, codec_decode_graph_arena_bytes_); codec_encode_graph_arena_bytes_ = runtime::parse_size_mb_option( - options.options, {"higgs_tts.codec_encode_graph_arena_mb"}, codec_encode_graph_arena_bytes_); + options.options, {"higgs_audio_tts.codec_encode_graph_arena_mb"}, codec_encode_graph_arena_bytes_); - if (const auto it = options.options.find("higgs_tts.weight_type"); it != options.options.end()) { + if (const auto it = options.options.find("higgs_audio_tts.weight_type"); it != options.options.end()) { const auto storage_type = assets::parse_tensor_storage_type(it->second); - validate_matmul_weight_storage(storage_type, "higgs_tts.weight_type"); + validate_matmul_weight_storage(storage_type, "higgs_audio_tts.weight_type"); ar_weight_storage_type_ = storage_type; codec_weight_storage_type_ = storage_type; } - if (const auto it = options.options.find("higgs_tts.ar_weight_type"); it != options.options.end()) { + if (const auto it = options.options.find("higgs_audio_tts.ar_weight_type"); it != options.options.end()) { ar_weight_storage_type_ = assets::parse_tensor_storage_type(it->second); - validate_matmul_weight_storage(ar_weight_storage_type_, "higgs_tts.ar_weight_type"); + validate_matmul_weight_storage(ar_weight_storage_type_, "higgs_audio_tts.ar_weight_type"); } - if (const auto it = options.options.find("higgs_tts.codec_weight_type"); it != options.options.end()) { + if (const auto it = options.options.find("higgs_audio_tts.codec_weight_type"); it != options.options.end()) { codec_weight_storage_type_ = assets::parse_tensor_storage_type(it->second); - validate_matmul_weight_storage(codec_weight_storage_type_, "higgs_tts.codec_weight_type"); + validate_matmul_weight_storage(codec_weight_storage_type_, "higgs_audio_tts.codec_weight_type"); } for (const auto & [key, _] : options.options) { - if (key.rfind("higgs_tts.", 0) == 0 && - key != "higgs_tts.ar_weight_context_mb" && - key != "higgs_tts.codec_weight_context_mb" && - key != "higgs_tts.ar_decode_graph_arena_mb" && - key != "higgs_tts.codec_decode_graph_arena_mb" && - key != "higgs_tts.codec_encode_graph_arena_mb" && - key != "higgs_tts.reference_cache_slots" && - key != "higgs_tts.weight_type" && - key != "higgs_tts.ar_weight_type" && - key != "higgs_tts.codec_weight_type") { + if (key.rfind("higgs_audio_tts.", 0) == 0 && + key != "higgs_audio_tts.ar_weight_context_mb" && + key != "higgs_audio_tts.codec_weight_context_mb" && + key != "higgs_audio_tts.ar_decode_graph_arena_mb" && + key != "higgs_audio_tts.codec_decode_graph_arena_mb" && + key != "higgs_audio_tts.codec_encode_graph_arena_mb" && + key != "higgs_audio_tts.reference_cache_slots" && + key != "higgs_audio_tts.weight_type" && + key != "higgs_audio_tts.ar_weight_type" && + key != "higgs_audio_tts.codec_weight_type") { throw std::runtime_error("unknown Higgs TTS session option: " + key); } } @@ -200,7 +200,7 @@ HiggsTTSSession::HiggsTTSSession( } std::string HiggsTTSSession::family() const { - return "higgs_tts"; + return "higgs_audio_tts"; } runtime::VoiceTaskKind HiggsTTSSession::task_kind() const { @@ -235,9 +235,9 @@ runtime::TaskResult HiggsTTSSession::run(const runtime::TaskRequest & request) { const auto * reference_audio = find_reference_audio(request); const HiggsCodecEncodeOutput * reference_codes = reference_audio != nullptr ? &resolve_reference_codes(*reference_audio, reference_text) : nullptr; - debug::trace_log_scalar("higgs_tts.text_chunk_size", text_chunk_size); - debug::trace_log_scalar("higgs_tts.text_chunk_mode", engine::text::text_chunk_mode_name(text_chunk_mode)); - debug::trace_log_scalar("higgs_tts.text_chunk_count", static_cast(chunk_requests.size())); + debug::trace_log_scalar("higgs_audio_tts.text_chunk_size", text_chunk_size); + debug::trace_log_scalar("higgs_audio_tts.text_chunk_mode", engine::text::text_chunk_mode_name(text_chunk_mode)); + debug::trace_log_scalar("higgs_audio_tts.text_chunk_count", static_cast(chunk_requests.size())); runtime::AudioBuffer merged_audio; for (const auto & chunk_request : chunk_requests) { @@ -267,35 +267,35 @@ const HiggsCodecEncodeOutput & HiggsTTSSession::resolve_reference_codes( key.channels = audio.channels; key.sample_count = sample_count; key.sample_hash = sample_hash; - debug::trace_log_scalar("higgs_tts.reference_audio.sample_rate", audio.sample_rate); - debug::trace_log_scalar("higgs_tts.reference_audio.channels", audio.channels); - debug::trace_log_f32("higgs_tts.reference_audio.samples", + debug::trace_log_scalar("higgs_audio_tts.reference_audio.sample_rate", audio.sample_rate); + debug::trace_log_scalar("higgs_audio_tts.reference_audio.channels", audio.channels); + debug::trace_log_f32("higgs_audio_tts.reference_audio.samples", {static_cast(audio.samples.size())}, audio.samples); - debug::trace_log_scalar("higgs_tts.reference_cache.capacity", static_cast(reference_cache_.capacity())); - debug::trace_log_scalar("higgs_tts.reference_cache.size", static_cast(reference_cache_.size())); + debug::trace_log_scalar("higgs_audio_tts.reference_cache.capacity", static_cast(reference_cache_.capacity())); + debug::trace_log_scalar("higgs_audio_tts.reference_cache.size", static_cast(reference_cache_.size())); if (const auto * cached = reference_cache_.find(key)) { - debug::trace_log_scalar("higgs_tts.reference_cache.hit", 1); + debug::trace_log_scalar("higgs_audio_tts.reference_cache.hit", 1); return cached->codes; } - debug::trace_log_scalar("higgs_tts.reference_cache.hit", 0); + debug::trace_log_scalar("higgs_audio_tts.reference_cache.hit", 0); const auto encode_start = Clock::now(); ReferenceCacheEntry entry; entry.codes = codec_->encode_reference(audio); codec_->release_encode_graph(); - debug::trace_log_scalar("higgs_tts.reference_codes.frames", entry.codes.frames); - debug::trace_log_scalar("higgs_tts.reference_codes.codebooks", entry.codes.codebooks); - debug::trace_log_i32("higgs_tts.reference_codes.values", + debug::trace_log_scalar("higgs_audio_tts.reference_codes.frames", entry.codes.frames); + debug::trace_log_scalar("higgs_audio_tts.reference_codes.codebooks", entry.codes.codebooks); + debug::trace_log_i32("higgs_audio_tts.reference_codes.values", {entry.codes.frames, entry.codes.codebooks}, entry.codes.codes); if (reference_cache_.capacity() == 0) { uncached_reference_ = std::move(entry); - debug::timing_log_scalar("higgs_tts.codec.encode_reference_ms", engine::debug::elapsed_ms(encode_start)); + debug::timing_log_scalar("higgs_audio_tts.codec.encode_reference_ms", engine::debug::elapsed_ms(encode_start)); return uncached_reference_->codes; } reference_cache_.put(key, std::move(entry)); - debug::timing_log_scalar("higgs_tts.codec.encode_reference_ms", engine::debug::elapsed_ms(encode_start)); + debug::timing_log_scalar("higgs_audio_tts.codec.encode_reference_ms", engine::debug::elapsed_ms(encode_start)); return reference_cache_.find(key)->codes; } @@ -324,4 +324,4 @@ HiggsGenerationRequest HiggsTTSSession::make_generation_request( return out; } -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/src/models/higgs_tts/tokenizer_text.cpp b/src/models/higgs_audio_tts/tokenizer_text.cpp similarity index 95% rename from src/models/higgs_tts/tokenizer_text.cpp rename to src/models/higgs_audio_tts/tokenizer_text.cpp index 6c94a403..9cfd9739 100644 --- a/src/models/higgs_tts/tokenizer_text.cpp +++ b/src/models/higgs_audio_tts/tokenizer_text.cpp @@ -1,4 +1,4 @@ -#include "engine/models/higgs_tts/tokenizer_text.h" +#include "engine/models/higgs_audio_tts/tokenizer_text.h" #include "engine/framework/tokenizers/llama_bpe.h" @@ -7,7 +7,7 @@ #include #include -namespace engine::models::higgs_tts { +namespace engine::models::higgs_audio_tts { namespace { int32_t require_token_id(const engine::tokenizers::LlamaBpeTokenizer & tokenizer, const std::string & token) { @@ -91,4 +91,4 @@ HiggsPromptEncoding HiggsTextTokenizer::encode_prompt(const HiggsPromptRequest & return encoding; } -} // namespace engine::models::higgs_tts +} // namespace engine::models::higgs_audio_tts diff --git a/tests/higgs_tts/.gitignore b/tests/higgs_audio_tts/.gitignore similarity index 100% rename from tests/higgs_tts/.gitignore rename to tests/higgs_audio_tts/.gitignore diff --git a/tests/higgs_tts/README.md b/tests/higgs_audio_tts/README.md similarity index 82% rename from tests/higgs_tts/README.md rename to tests/higgs_audio_tts/README.md index 393a1e63..dbf8e476 100644 --- a/tests/higgs_tts/README.md +++ b/tests/higgs_audio_tts/README.md @@ -7,14 +7,14 @@ grouped FlashAttention, packed SwiGLU, direct KV updates, and the `ROPE -> VIEW -> SET_ROWS` fusion pattern. ```powershell -cmake --build build/windows-cuda-release --config Release --target qwen_decoder_packed_projection_test higgs_tts_warm_bench -j 8 +cmake --build build/windows-cuda-release --config Release --target qwen_decoder_packed_projection_test higgs_audio_tts_warm_bench -j 8 ctest --test-dir build/windows-cuda-release -C Release -R qwen_decoder_packed_projection_test --output-on-failure ``` Run the fixed-seed, five-request CUDA benchmark and save every generated WAV: ```powershell -tests/higgs_tts/run_cuda_performance.ps1 ` +tests/higgs_audio_tts/run_cuda_performance.ps1 ` -Model ../models/higgs-audio-v3-tts-4b_Q8/higgs-audio-v3-tts-4b_Q8.gguf ` -Label candidate ``` @@ -22,15 +22,15 @@ tests/higgs_tts/run_cuda_performance.ps1 ` Compare a candidate run with a prior result directory request by request: ```powershell -tests/higgs_tts/run_cuda_performance.ps1 ` +tests/higgs_audio_tts/run_cuda_performance.ps1 ` -Model ../models/higgs-audio-v3-tts-4b_Q8/higgs-audio-v3-tts-4b_Q8.gguf ` -Label candidate ` - -Baseline tests/higgs_tts/results/baseline + -Baseline tests/higgs_audio_tts/results/baseline ``` The comparison reports frame counts, wall time, RTF, speedup, waveform cosine, and 80-band log-mel cosine. Result WAVs, logs, and JSON reports are written below -`tests/higgs_tts/results/`, which is intentionally ignored by Git. +`tests/higgs_audio_tts/results/`, which is intentionally ignored by Git. The comparison helper requires Python 3 with NumPy. Add `-RequireSameFrames` when comparing paths that are expected to be diff --git a/tests/higgs_tts/compare_warmbench_results.py b/tests/higgs_audio_tts/compare_warmbench_results.py similarity index 98% rename from tests/higgs_tts/compare_warmbench_results.py rename to tests/higgs_audio_tts/compare_warmbench_results.py index 3facf3a5..1ebc8a43 100644 --- a/tests/higgs_tts/compare_warmbench_results.py +++ b/tests/higgs_audio_tts/compare_warmbench_results.py @@ -2,7 +2,7 @@ """Compare Higgs warmbench runs request by request. Each result directory is expected to contain timing.log and audio/audio_N.wav, -as emitted by higgs_tts_warm_bench. The report includes exact frame counts, +as emitted by higgs_audio_tts_warm_bench. The report includes exact frame counts, wall time, RTF, speedup, waveform cosine, and log-mel cosine per request. """ @@ -72,7 +72,7 @@ def log_mel(samples: np.ndarray, sample_rate: int, n_fft: int = 1024, hop: int = def timings(path: Path) -> dict[int, float]: result: dict[int, float] = {} for line in (path / "timing.log").read_text(encoding="utf-8").splitlines(): - prefix = "higgs_tts.cpp.request_" + prefix = "higgs_audio_tts.cpp.request_" if not line.startswith(prefix) or ".wall_ms=" not in line: continue index_text, value = line[len(prefix) :].split(".wall_ms=", 1) diff --git a/tests/higgs_tts/higgs_tts_cuda_bench_cases.json b/tests/higgs_audio_tts/higgs_audio_tts_cuda_bench_cases.json similarity index 100% rename from tests/higgs_tts/higgs_tts_cuda_bench_cases.json rename to tests/higgs_audio_tts/higgs_audio_tts_cuda_bench_cases.json diff --git a/tests/higgs_tts/higgs_tts_cuda_mixed_cases.json b/tests/higgs_audio_tts/higgs_audio_tts_cuda_mixed_cases.json similarity index 100% rename from tests/higgs_tts/higgs_tts_cuda_mixed_cases.json rename to tests/higgs_audio_tts/higgs_audio_tts_cuda_mixed_cases.json diff --git a/tests/higgs_tts/higgs_tts_cuda_perf_cases.json b/tests/higgs_audio_tts/higgs_audio_tts_cuda_perf_cases.json similarity index 100% rename from tests/higgs_tts/higgs_tts_cuda_perf_cases.json rename to tests/higgs_audio_tts/higgs_audio_tts_cuda_perf_cases.json diff --git a/tests/higgs_tts/higgs_tts_python_warm_bench.py b/tests/higgs_audio_tts/higgs_audio_tts_python_warm_bench.py similarity index 96% rename from tests/higgs_tts/higgs_tts_python_warm_bench.py rename to tests/higgs_audio_tts/higgs_audio_tts_python_warm_bench.py index cd4c3b80..7eeba322 100644 --- a/tests/higgs_tts/higgs_tts_python_warm_bench.py +++ b/tests/higgs_audio_tts/higgs_audio_tts_python_warm_bench.py @@ -35,7 +35,7 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description="Python reference Higgs Audio v3 TTS warmbench.") - parser.add_argument("--family", default="higgs_tts") + parser.add_argument("--family", default="higgs_audio_tts") parser.add_argument("--model", type=Path, default=DEFAULT_MODEL) parser.add_argument("--reference-root", type=Path, default=REFERENCE_ROOT) parser.add_argument("--backend", choices=("cuda",), default="cuda") @@ -57,8 +57,8 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--port", type=int, default=18180) parser.add_argument("--server-timeout-sec", type=float, default=300.0) parser.add_argument("--output-dir", type=Path, default=None) - parser.add_argument("--audio-out", type=Path, default=Path("higgs_tts_python_audio.wav")) - parser.add_argument("--timing-file", type=Path, default=Path("higgs_tts_python_timing.log")) + parser.add_argument("--audio-out", type=Path, default=Path("higgs_audio_tts_python_audio.wav")) + parser.add_argument("--timing-file", type=Path, default=Path("higgs_audio_tts_python_timing.log")) parser.add_argument("--summary-file", type=Path, default=None) return parser.parse_args() @@ -322,9 +322,9 @@ def main() -> int: raise RuntimeError("Higgs TTS warmbench request sequence is empty") output_dir = args.output_dir or args.audio_out.parent output_dir.mkdir(parents=True, exist_ok=True) - timing_lines = ["higgs_tts.python.model_load_excluded=1"] + timing_lines = ["higgs_audio_tts.python.model_load_excluded=1"] server = start_server(args) - timing_lines.append(f"higgs_tts.python.server_pid={server.pid}") + timing_lines.append(f"higgs_audio_tts.python.server_pid={server.pid}") try: wait_for_server(args, server) for _ in range(max(0, args.warmup)): @@ -338,8 +338,8 @@ def main() -> int: summary, wall_ms = run_request(args, request, audio_path) total_ms += wall_ms average_ms = total_ms / float(max(1, args.iterations)) - timing_lines.append(f"higgs_tts.python.request_{request_index}.wall_ms={average_ms:.6f}") - print(f"higgs_tts.python.wall_ms={average_ms}") + timing_lines.append(f"higgs_audio_tts.python.request_{request_index}.wall_ms={average_ms:.6f}") + print(f"higgs_audio_tts.python.wall_ms={average_ms}") steps.append( { "request_index": request_index, @@ -353,7 +353,7 @@ def main() -> int: "metrics": {"wall_ms": average_ms}, } ) - summary_payload = {"family": "higgs_tts", "backend": args.backend, "sequence_steps": steps} + summary_payload = {"family": "higgs_audio_tts", "backend": args.backend, "sequence_steps": steps} if args.summary_file: args.summary_file.parent.mkdir(parents=True, exist_ok=True) args.summary_file.write_text(json.dumps(summary_payload, ensure_ascii=False) + "\n", encoding="utf-8") diff --git a/tests/higgs_tts/higgs_tts_sampler_logits.bin b/tests/higgs_audio_tts/higgs_audio_tts_sampler_logits.bin similarity index 100% rename from tests/higgs_tts/higgs_tts_sampler_logits.bin rename to tests/higgs_audio_tts/higgs_audio_tts_sampler_logits.bin diff --git a/tests/higgs_tts/higgs_tts_warm_bench.cpp b/tests/higgs_audio_tts/higgs_audio_tts_warm_bench.cpp similarity index 95% rename from tests/higgs_tts/higgs_tts_warm_bench.cpp rename to tests/higgs_audio_tts/higgs_audio_tts_warm_bench.cpp index 73c400b9..2ddb68f6 100644 --- a/tests/higgs_tts/higgs_tts_warm_bench.cpp +++ b/tests/higgs_audio_tts/higgs_audio_tts_warm_bench.cpp @@ -265,7 +265,7 @@ int main(int argc, char ** argv) { } const std::filesystem::path output_dir = arg_value(argc, argv, "--output-dir", ""); const std::filesystem::path timing_path = - arg_value(argc, argv, "--timing-file", "/tmp/higgs_tts_warm_bench_timing.log"); + arg_value(argc, argv, "--timing-file", "/tmp/higgs_audio_tts_warm_bench_timing.log"); const std::filesystem::path log_path = arg_value(argc, argv, "--log-file", ""); if (has_flag(argc, argv, "--enable-trace")) { if (log_path.empty()) { @@ -276,7 +276,7 @@ int main(int argc, char ** argv) { engine::runtime::ModelLoadRequest load_request; load_request.model_path = model_path; - load_request.family_hint = "higgs_tts"; + load_request.family_hint = "higgs_audio_tts"; auto registry = engine::runtime::make_default_registry(); auto model = registry.load(load_request); @@ -305,7 +305,7 @@ int main(int argc, char ** argv) { std::filesystem::create_directories(output_dir); } - std::vector timing_lines{"higgs_tts.cpp.model_load_excluded=1"}; + std::vector timing_lines{"higgs_audio_tts.cpp.model_load_excluded=1"}; engine::io::json::Value::Array steps; steps.reserve(requests.size()); for (size_t request_index = 0; request_index < requests.size(); ++request_index) { @@ -331,15 +331,15 @@ int main(int argc, char ** argv) { last_result.audio_output->samples); } timing_lines.push_back( - "higgs_tts.cpp.request_" + std::to_string(request_index) + ".wall_ms=" + std::to_string(wall_ms)); + "higgs_audio_tts.cpp.request_" + std::to_string(request_index) + ".wall_ms=" + std::to_string(wall_ms)); const auto & audio = *last_result.audio_output; const double frames = static_cast( audio.samples.size() / static_cast(std::max(1, audio.channels))); const double duration_sec = audio.sample_rate > 0 ? frames / audio.sample_rate : 0.0; const double rtf = duration_sec > 0.0 ? wall_ms / 1000.0 / duration_sec : 0.0; timing_lines.push_back( - "higgs_tts.cpp.request_" + std::to_string(request_index) + ".rtf=" + std::to_string(rtf)); - std::cout << "higgs_tts.cpp.request=" << request_index + "higgs_audio_tts.cpp.request_" + std::to_string(request_index) + ".rtf=" + std::to_string(rtf)); + std::cout << "higgs_audio_tts.cpp.request=" << request_index << " wall_ms=" << wall_ms << " rtf=" << rtf << "\n"; steps.push_back(step_json(last_result, static_cast(request_index), wall_ms, audio_path)); @@ -347,14 +347,14 @@ int main(int argc, char ** argv) { write_timing(timing_path, timing_lines); const auto summary = engine::io::json::Value::make_object({ - {"family", string("higgs_tts")}, + {"family", string("higgs_audio_tts")}, {"backend", string(backend_name)}, {"sequence_steps", engine::io::json::Value::make_array(std::move(steps))}, }); std::cout << "summary_json=" << engine::io::json::stringify(summary) << "\n"; return 0; } catch (const std::exception & ex) { - std::cerr << "higgs_tts_warm_bench failed: " << ex.what() << "\n"; + std::cerr << "higgs_audio_tts_warm_bench failed: " << ex.what() << "\n"; return 1; } } diff --git a/tests/higgs_tts/higgs_tts_warm_bench_cases.json b/tests/higgs_audio_tts/higgs_audio_tts_warm_bench_cases.json similarity index 100% rename from tests/higgs_tts/higgs_tts_warm_bench_cases.json rename to tests/higgs_audio_tts/higgs_audio_tts_warm_bench_cases.json diff --git a/tests/higgs_tts/run_cuda_performance.ps1 b/tests/higgs_audio_tts/run_cuda_performance.ps1 similarity index 93% rename from tests/higgs_tts/run_cuda_performance.ps1 rename to tests/higgs_audio_tts/run_cuda_performance.ps1 index 456a3549..cf15c446 100644 --- a/tests/higgs_tts/run_cuda_performance.ps1 +++ b/tests/higgs_audio_tts/run_cuda_performance.ps1 @@ -14,8 +14,8 @@ param( $ErrorActionPreference = "Stop" $RepoRoot = (Resolve-Path (Join-Path $PSScriptRoot "../..")).Path $ModelPath = (Resolve-Path $Model).Path -$Bench = Join-Path $RepoRoot "$BuildDir/bin/higgs_tts_warm_bench.exe" -$Cases = Join-Path $PSScriptRoot "higgs_tts_cuda_perf_cases.json" +$Bench = Join-Path $RepoRoot "$BuildDir/bin/higgs_audio_tts_warm_bench.exe" +$Cases = Join-Path $PSScriptRoot "higgs_audio_tts_cuda_perf_cases.json" $ResultDir = Join-Path $PSScriptRoot "results/$Label" if (-not (Test-Path -LiteralPath $Bench)) { diff --git a/tests/warmbench.py b/tests/warmbench.py index f1180f36..5067e124 100644 --- a/tests/warmbench.py +++ b/tests/warmbench.py @@ -460,21 +460,21 @@ "log_mel_cosine_min": 0.90, "cpp_session_options": ["heartmula.weight_type=f32"], }, - "higgs_tts": { - "kind": "higgs_tts", + "higgs_audio_tts": { + "kind": "higgs_audio_tts", "modes": ["offline"], - "cpp_bin": "build/debug/bin/higgs_tts_warm_bench", - "python_script": "tests/higgs_tts/higgs_tts_python_warm_bench.py", + "cpp_bin": "build/debug/bin/higgs_audio_tts_warm_bench", + "python_script": "tests/higgs_audio_tts/higgs_audio_tts_python_warm_bench.py", "python_conda_env": "qwen3-tts", "model": "models/higgs-audio-v3-tts-4b", - "case_catalog": "tests/higgs_tts/higgs_tts_warm_bench_cases.json", + "case_catalog": "tests/higgs_audio_tts/higgs_audio_tts_warm_bench_cases.json", "default_case_name": "default", "default_requests_per_session": 1, "default_warmup": 0, "wav_cosine_min": 0.90, "log_mel_cosine_min": 0.90, "length_ratio_min": 0.98, - "cpp_session_options": ["higgs_tts.codec_weight_type=f32"], + "cpp_session_options": ["higgs_audio_tts.codec_weight_type=f32"], }, "index_tts2": { "kind": "index_tts2", @@ -4215,7 +4215,7 @@ def build_heartmula_commands( return python_command, cpp_command -def build_higgs_tts_commands( +def build_higgs_audio_tts_commands( config: dict[str, Any], backend: str, args: argparse.Namespace, @@ -4370,7 +4370,7 @@ def validate_sequence_result(summary: dict[str, Any], request_count: int, kind: and len(step.get("stems", [])) > 0 and isinstance(step.get("metrics", {}), dict) for step in steps) - elif kind in {"vevo2", "seed_vc", "miocodec", "voxcpm2", "supertonic", "vibevoice", "irodori_tts", "heartmula", "higgs_tts", "index_tts2"}: + elif kind in {"vevo2", "seed_vc", "miocodec", "voxcpm2", "supertonic", "vibevoice", "irodori_tts", "heartmula", "higgs_audio_tts", "index_tts2"}: payload_valid = all( isinstance(step.get("stems", []), list) and len(step.get("stems", [])) > 0 @@ -4514,10 +4514,10 @@ def run_scenario( irodori_requests, request_manifest = resolve_vevo2_case(config, args) args.requests_per_session = len(irodori_requests) python_command, cpp_command = build_irodori_tts_commands(scenario_config, backend, args, scenario_dir, irodori_requests) - elif scenario_config["kind"] == "higgs_tts": + elif scenario_config["kind"] == "higgs_audio_tts": higgs_requests, request_manifest = resolve_vevo2_case(config, args) args.requests_per_session = len(higgs_requests) - python_command, cpp_command = build_higgs_tts_commands(scenario_config, backend, args, scenario_dir, higgs_requests) + python_command, cpp_command = build_higgs_audio_tts_commands(scenario_config, backend, args, scenario_dir, higgs_requests) elif scenario_config["kind"] == "index_tts2": index_tts2_requests, request_manifest = resolve_vevo2_case(config, args) args.requests_per_session = len(index_tts2_requests) @@ -4865,7 +4865,7 @@ def run_scenario( cpp_step_path = cpp_step_paths[request_index] if request_index < len(cpp_step_paths) else "" append_log(master_log, f"PYTHON OUTPUT family={family} mode={mode} backend={backend} request={request_index} path={python_step_path} valid={int(file_is_nonempty(python_step_path))}") append_log(master_log, f"CPP OUTPUT family={family} mode={mode} backend={backend} request={request_index} path={cpp_step_path} valid={int(file_is_nonempty(cpp_step_path))}") - elif scenario_config["kind"] in {"vevo2", "seed_vc", "miocodec", "voxcpm2", "supertonic", "vibevoice", "irodori_tts", "heartmula", "higgs_tts", "index_tts2"}: + elif scenario_config["kind"] in {"vevo2", "seed_vc", "miocodec", "voxcpm2", "supertonic", "vibevoice", "irodori_tts", "heartmula", "higgs_audio_tts", "index_tts2"}: python_valid = validate_sequence_result(python_summary, args.requests_per_session, scenario_config["kind"]) cpp_valid = validate_sequence_result(cpp_summary, args.requests_per_session, scenario_config["kind"]) python_step_paths = write_sequence_step_artifacts(python_summary.get("sequence_steps", []), scenario_dir / "python_json", "python") diff --git a/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json b/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json index 575e2b28..c34b69fe 100644 --- a/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json @@ -225,9 +225,9 @@ ] }, { - "id": "higgs_tts_voice_clone_longform", + "id": "higgs_audio_tts_voice_clone_longform", "coverage": "Higgs Audio v3 voice clone with framework long-form text chunking, bounded AR generation, and codec decode", - "family": "higgs_tts", + "family": "higgs_audio_tts", "model": "models/higgs-audio-v3-tts-4b", "task": "tts", "mode": "offline", diff --git a/tools/audiocpp_cli/audiocpp_cli_path_cases.json b/tools/audiocpp_cli/audiocpp_cli_path_cases.json index 992f24e5..eddc048a 100644 --- a/tools/audiocpp_cli/audiocpp_cli_path_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_path_cases.json @@ -548,9 +548,9 @@ ] }, { - "id": "higgs_tts_voice_clone_chunked", + "id": "higgs_audio_tts_voice_clone_chunked", "coverage": "Higgs Audio v3 voice clone path with framework text chunking, AR generation, and codec decode", - "family": "higgs_tts", + "family": "higgs_audio_tts", "model": "models/higgs-audio-v3-tts-4b", "task": "tts", "mode": "offline", @@ -571,9 +571,9 @@ ] }, { - "id": "higgs_tts_voice_clone_cache_pollution", + "id": "higgs_audio_tts_voice_clone_cache_pollution", "coverage": "Higgs Audio v3 long-lived session cache pollution check with short, long, medium, and long requests", - "family": "higgs_tts", + "family": "higgs_audio_tts", "model": "models/higgs-audio-v3-tts-4b", "task": "tts", "mode": "offline", From 7bb62552506a040c74c23231833d7a0c83b5d6f8 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Tue, 21 Jul 2026 00:26:09 -0400 Subject: [PATCH 14/27] Add path test model override --- tools/audiocpp_cli/run_audiocpp_cli_path_tests.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tools/audiocpp_cli/run_audiocpp_cli_path_tests.py b/tools/audiocpp_cli/run_audiocpp_cli_path_tests.py index ba66b1c5..8c45d9e0 100644 --- a/tools/audiocpp_cli/run_audiocpp_cli_path_tests.py +++ b/tools/audiocpp_cli/run_audiocpp_cli_path_tests.py @@ -349,7 +349,16 @@ def resolve_model_path(models_root: Path, value: str) -> Path: return models_root / path +def resolve_model_override(value: Path | None) -> Path | None: + if value is None: + return None + if value.is_absolute(): + return value + return REPO_ROOT / value + + def build_command(args: argparse.Namespace, case: dict[str, Any], case_dir: Path) -> list[str]: + model_path = resolve_model_override(args.model_path) or resolve_model_path(args.models_root, case["model"]) command = [ str(args.audiocpp_cli_bin), "--task", @@ -357,7 +366,7 @@ def build_command(args: argparse.Namespace, case: dict[str, Any], case_dir: Path "--family", case["family"], "--model", - str(resolve_model_path(args.models_root, case["model"])), + str(model_path), "--backend", case.get("backend", args.backend), "--mode", @@ -498,6 +507,7 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--cases", type=Path, default=DEFAULT_CASES) parser.add_argument("--audiocpp-cli-bin", type=Path, default=DEFAULT_AUDIOCPP_CLI_BIN) parser.add_argument("--models-root", type=Path, default=DEFAULT_MODELS_ROOT) + parser.add_argument("--model-path", type=Path, help="Override the model path for every selected case") parser.add_argument("--backend", default="cuda", choices=["cpu", "cuda", "vulkan", "metal", "best"]) parser.add_argument("--device", type=int, default=0) parser.add_argument("--threads", type=int, default=DEFAULT_THREADS) From 9c447a66a1ef6cf78222cc30a9f38233b8ff71a8 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Tue, 21 Jul 2026 00:44:54 -0400 Subject: [PATCH 15/27] Optimize Fish Audio AR decoder path --- src/models/fish_audio/ar.cpp | 69 +++++++++++++++++++++++++----------- 1 file changed, 48 insertions(+), 21 deletions(-) diff --git a/src/models/fish_audio/ar.cpp b/src/models/fish_audio/ar.cpp index 90a1ffdf..f33e608c 100644 --- a/src/models/fish_audio/ar.cpp +++ b/src/models/fish_audio/ar.cpp @@ -88,15 +88,12 @@ struct GgmlContextDeleter { struct FishLayerWeights { assets::TensorDataF32 input_norm; - core::TensorValue q_proj; - core::TensorValue k_proj; - core::TensorValue v_proj; + core::TensorValue qkv_proj; core::TensorValue o_proj; std::optional q_norm; std::optional k_norm; assets::TensorDataF32 post_norm; - core::TensorValue gate_proj; - core::TensorValue up_proj; + core::TensorValue gate_up_proj; core::TensorValue down_proj; }; @@ -155,12 +152,14 @@ modules::QwenCausalDecoderConfig make_slow_decoder_config(const FishAudioTextCon out.stack.rope_theta = config.rope_base; out.stack.rope_type = GGML_ROPE_TYPE_NORMAL; out.stack.attention_precision = GGML_PREC_F32; - out.stack.qkv_layout = modules::QwenDecoderQKVLayout::Separate; + out.stack.qkv_layout = modules::QwenDecoderQKVLayout::PackedQKV; out.stack.use_qk_norm = config.attention_qk_norm; out.stack.activation_cast = fish_activation_cast_policy(); - out.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::ManualRepeat; - out.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + out.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; out.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.stack.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; + out.stack.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; out.logits_size = config.vocab_size; out.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; out.lm_head_precision = GGML_PREC_F32; @@ -179,12 +178,14 @@ modules::QwenCausalDecoderConfig make_fast_decoder_config(const FishAudioFastCon out.stack.rope_theta = config.rope_base; out.stack.rope_type = GGML_ROPE_TYPE_NORMAL; out.stack.attention_precision = GGML_PREC_F32; - out.stack.qkv_layout = modules::QwenDecoderQKVLayout::Separate; + out.stack.qkv_layout = modules::QwenDecoderQKVLayout::PackedQKV; out.stack.use_qk_norm = config.attention_qk_norm; out.stack.activation_cast = fish_activation_cast_policy(); - out.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::ManualRepeat; - out.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::ManualRepeat; + out.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; + out.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; out.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; + out.stack.runtime.static_cache.set_rows_mode = modules::QwenDecoderStaticCacheSetRowsMode::BackendViewOptimized; + out.stack.runtime.mlp.mode = modules::QwenDecoderMLPMode::PackedGateUp; out.logits_size = config.vocab_size; out.logits_mode = modules::QwenCausalDecoderLogitsMode::LastStep; out.lm_head_precision = GGML_PREC_F32; @@ -197,9 +198,7 @@ modules::QwenDecoderLayerWeights bind_layer( bool use_qk_norm) { modules::QwenDecoderLayerWeights out; out.input_norm = binding::norm_data(constants, weights.input_norm); - out.self_attention.q_weight = weights.q_proj; - out.self_attention.k_weight = weights.k_proj; - out.self_attention.v_weight = weights.v_proj; + out.self_attention.qkv_weight = weights.qkv_proj; out.self_attention.out_weight = weights.o_proj; if (use_qk_norm) { if (!weights.q_norm.has_value() || !weights.k_norm.has_value()) { @@ -209,8 +208,7 @@ modules::QwenDecoderLayerWeights bind_layer( out.k_norm = binding::norm_data(constants, *weights.k_norm); } out.post_norm = binding::norm_data(constants, weights.post_norm); - out.mlp.gate_proj = binding::linear_data(constants, weights.gate_proj); - out.mlp.up_proj = binding::linear_data(constants, weights.up_proj); + out.mlp.gate_up_proj = binding::linear_data(constants, weights.gate_up_proj); out.mlp.down_proj = binding::linear_data(constants, weights.down_proj); return out; } @@ -336,18 +334,47 @@ FishLayerWeights load_layer( assets::TensorStorageType storage_type) { FishLayerWeights w; w.input_norm = source.require_f32_tensor(prefix + ".attention_norm.weight", {hidden}); - w.q_proj = store.load_tensor(source, prefix + ".attention.q_proj.weight", storage_type, {heads * head_dim, hidden}); - w.k_proj = store.load_tensor(source, prefix + ".attention.k_proj.weight", storage_type, {kv_heads * head_dim, hidden}); - w.v_proj = store.load_tensor(source, prefix + ".attention.v_proj.weight", storage_type, {kv_heads * head_dim, hidden}); + { + const auto q = source.require_tensor(prefix + ".attention.q_proj.weight", storage_type, {heads * head_dim, hidden}); + const auto k = source.require_tensor(prefix + ".attention.k_proj.weight", storage_type, {kv_heads * head_dim, hidden}); + const auto v = source.require_tensor(prefix + ".attention.v_proj.weight", storage_type, {kv_heads * head_dim, hidden}); + if (q.type != k.type || q.type != v.type) { + throw std::runtime_error("Fish Audio packed QKV weights require matching storage types"); + } + std::vector packed; + packed.reserve(q.bytes.size() + k.bytes.size() + v.bytes.size()); + packed.insert(packed.end(), q.bytes.begin(), q.bytes.end()); + packed.insert(packed.end(), k.bytes.begin(), k.bytes.end()); + packed.insert(packed.end(), v.bytes.begin(), v.bytes.end()); + w.qkv_proj = store.make_tensor( + core::TensorShape::from_dims({(heads + 2 * kv_heads) * head_dim, hidden}), + q.type, + packed.data(), + packed.size()); + } w.o_proj = store.load_tensor(source, prefix + ".attention.wo.weight", storage_type, {hidden, heads * head_dim}); if (qk_norm) { w.q_norm = source.require_f32_tensor(prefix + ".attention.q_norm.weight", {head_dim}); w.k_norm = source.require_f32_tensor(prefix + ".attention.k_norm.weight", {head_dim}); } w.post_norm = source.require_f32_tensor(prefix + ".ffn_norm.weight", {hidden}); - w.gate_proj = store.load_tensor(source, prefix + ".feed_forward.w1.weight", storage_type, {intermediate, hidden}); w.down_proj = store.load_tensor(source, prefix + ".feed_forward.w2.weight", storage_type, {hidden, intermediate}); - w.up_proj = store.load_tensor(source, prefix + ".feed_forward.w3.weight", storage_type, {intermediate, hidden}); + { + const auto gate = source.require_tensor(prefix + ".feed_forward.w1.weight", storage_type, {intermediate, hidden}); + const auto up = source.require_tensor(prefix + ".feed_forward.w3.weight", storage_type, {intermediate, hidden}); + if (gate.type != up.type) { + throw std::runtime_error("Fish Audio packed gate/up weights require matching storage types"); + } + std::vector packed; + packed.reserve(gate.bytes.size() + up.bytes.size()); + packed.insert(packed.end(), gate.bytes.begin(), gate.bytes.end()); + packed.insert(packed.end(), up.bytes.begin(), up.bytes.end()); + w.gate_up_proj = store.make_tensor( + core::TensorShape::from_dims({intermediate * 2, hidden}), + gate.type, + packed.data(), + packed.size()); + } return w; } From 804cd9cf870a555f5619011453a7c15fc5b4a113 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Tue, 21 Jul 2026 01:27:17 -0400 Subject: [PATCH 16/27] Optimize Fish Audio KV cache lifecycle --- include/engine/models/fish_audio/types.h | 2 +- src/models/fish_audio/ar.cpp | 108 +++++++++++++---------- src/models/fish_audio/loader.cpp | 2 +- 3 files changed, 64 insertions(+), 48 deletions(-) diff --git a/include/engine/models/fish_audio/types.h b/include/engine/models/fish_audio/types.h index 268bc2ea..d76523ea 100644 --- a/include/engine/models/fish_audio/types.h +++ b/include/engine/models/fish_audio/types.h @@ -10,7 +10,7 @@ namespace engine::models::fish_audio { struct FishAudioGenerationOptions { - int64_t max_new_tokens = 2048; + int64_t max_new_tokens = 1024; int64_t text_chunk_size = 200; float top_p = 0.8F; int top_k = 30; diff --git a/src/models/fish_audio/ar.cpp b/src/models/fish_audio/ar.cpp index f33e608c..720382dc 100644 --- a/src/models/fish_audio/ar.cpp +++ b/src/models/fish_audio/ar.cpp @@ -48,7 +48,6 @@ struct FishARProfile { double prefill_input_upload_ms = 0.0; double prefill_graph_ms = 0.0; double prefill_output_read_ms = 0.0; - double prefill_state_read_ms = 0.0; double step_input_upload_ms = 0.0; double step_mask_upload_ms = 0.0; double step_graph_ms = 0.0; @@ -57,7 +56,6 @@ struct FishARProfile { double fast_mask_upload_ms = 0.0; double fast_graph_ms = 0.0; double fast_output_read_ms = 0.0; - double import_prefill_state_ms = 0.0; double sample_bias_ms = 0.0; double sample_main_ms = 0.0; double sample_high_ms = 0.0; @@ -117,7 +115,11 @@ struct SlowForwardOutput { struct SlowPrefillOutput { SlowForwardOutput forward; - runtime::TransformerKVState state; +}; + +struct FishPrefillCacheTarget { + std::vector keys; + std::vector values; }; modules::QwenDecoderActivationCastPolicy fish_activation_cast_policy() { @@ -811,8 +813,8 @@ class FishAudioARRuntime::Impl { if (max_new_tokens <= 0) { throw std::runtime_error("Fish Audio prompt leaves no room for generated tokens"); } - ensure_prefill_graph(prompt.steps, profile); ensure_step_graph(prompt.steps + max_new_tokens, profile); + ensure_prefill_graph(prompt.steps, profile); ensure_fast_graph(profile); SampleState sample; sample.seed = options.seed; @@ -830,9 +832,7 @@ class FishAudioARRuntime::Impl { } append_frame(generated_frame_major, frame); ++profile.generated_frames; - timing_start = Clock::now(); - step_graph_->import_state(prefill.state); - profile.import_prefill_state_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); + step_graph_->finish_prefill(prompt.steps); bool ended_by_im_end = false; for (int64_t step = 1; step < max_new_tokens; ++step) { timing_start = Clock::now(); @@ -874,17 +874,25 @@ class FishAudioARRuntime::Impl { private: class PrefillGraph { public: - PrefillGraph(std::shared_ptr runtime, int64_t steps) + PrefillGraph( + std::shared_ptr runtime, + int64_t steps, + FishPrefillCacheTarget target_cache) : runtime_(std::move(runtime)), - steps_(steps) { + steps_(steps), + target_cache_(std::move(target_cache)) { + const auto & assets = runtime_->assets(); + const auto & config = assets.config.text; + if (target_cache_.keys.size() != runtime_->weights().slow_layers.size() || + target_cache_.values.size() != runtime_->weights().slow_layers.size()) { + throw std::runtime_error("Fish Audio prefill target cache layer count mismatch"); + } ggml_init_params params{runtime_->graph_arena_bytes(), nullptr, true}; ctx_.reset(ggml_init(params)); if (ctx_ == nullptr) { throw std::runtime_error("failed to initialize Fish Audio AR prefill context"); } core::ModuleBuildContext ctx{ctx_.get(), "fish_audio.ar.prefill", runtime_->backend_type()}; - const auto & assets = runtime_->assets(); - const auto & config = assets.config.text; auto input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, steps_, config.dim})); input_ = input.tensor; positions_ = ggml_new_tensor_1d(ctx_.get(), GGML_TYPE_I32, steps_); @@ -903,32 +911,39 @@ class FishAudioARRuntime::Impl { bind_slow_weights(*constants_, runtime_->weights(), config), make_slow_decoder_config(config), assets.config.norm_fastlayer_input); - for (const auto & layer : decoder.state.layers) { + graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); + for (size_t layer_index = 0; layer_index < decoder.state.layers.size(); ++layer_index) { + const auto & layer = decoder.state.layers[layer_index]; if (!layer.key.has_value() || !layer.value.has_value()) { throw std::runtime_error("Fish Audio prefill decoder did not produce K/V state"); } - keys_.push_back(layer.key->tensor); - values_.push_back(layer.value->tensor); + auto key_dest = runtime::view_transformer_kv_cache_steps( + ctx, + target_cache_.keys[layer_index], + 0, + steps_, + config.n_local_heads, + config.head_dim, + "Fish Audio prefill key cache", + target_cache_.keys[layer_index].type); + auto value_dest = runtime::view_transformer_kv_cache_steps( + ctx, + target_cache_.values[layer_index], + 0, + steps_, + config.n_local_heads, + config.head_dim, + "Fish Audio prefill value cache", + target_cache_.values[layer_index].type); + ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), layer.key->tensor, key_dest.tensor)); + ggml_build_forward_expand(graph_, ggml_cpy(ctx_.get(), layer.value->tensor, value_dest.tensor)); } hidden_ = decoder.hidden.tensor; logits_ = decoder.logits.tensor; ggml_set_output(hidden_); - for (ggml_tensor * key : keys_) { - ggml_set_output(key); - } - for (ggml_tensor * value : values_) { - ggml_set_output(value); - } ggml_set_output(logits_); - graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); ggml_build_forward_expand(graph_, logits_); ggml_build_forward_expand(graph_, hidden_); - for (ggml_tensor * key : keys_) { - ggml_build_forward_expand(graph_, key); - } - for (ggml_tensor * value : values_) { - ggml_build_forward_expand(graph_, value); - } constants_->finish_graph(); constants_->ensure_uploaded(); gallocr_ = ggml_gallocr_new(ggml_backend_get_default_buffer_type(runtime_->backend())); @@ -972,18 +987,6 @@ class FishAudioARRuntime::Impl { ggml_backend_tensor_get(logits_, out.forward.logits.data(), 0, out.forward.logits.size() * sizeof(float)); ggml_backend_tensor_get(hidden_, out.forward.hidden.data(), 0, out.forward.hidden.size() * sizeof(float)); profile.prefill_output_read_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); - out.state.current_end = steps_; - out.state.layers.resize(keys_.size()); - const size_t values_per_layer = static_cast(steps_ * config.n_local_heads * config.head_dim); - timing_start = Clock::now(); - for (size_t i = 0; i < keys_.size(); ++i) { - out.state.layers[i].valid_steps = steps_; - out.state.layers[i].key.resize(values_per_layer); - out.state.layers[i].value.resize(values_per_layer); - ggml_backend_tensor_get(keys_[i], out.state.layers[i].key.data(), 0, values_per_layer * sizeof(float)); - ggml_backend_tensor_get(values_[i], out.state.layers[i].value.data(), 0, values_per_layer * sizeof(float)); - } - profile.prefill_state_read_ms += engine::debug::elapsed_ms(timing_start, Clock::now()); return out; } @@ -997,11 +1000,10 @@ class FishAudioARRuntime::Impl { ggml_tensor * positions_ = nullptr; ggml_tensor * hidden_ = nullptr; ggml_tensor * logits_ = nullptr; - std::vector keys_; - std::vector values_; ggml_cgraph * graph_ = nullptr; ggml_gallocr_t gallocr_ = nullptr; std::unique_ptr constants_; + FishPrefillCacheTarget target_cache_; }; class StepGraph { @@ -1108,8 +1110,20 @@ class FishAudioARRuntime::Impl { int64_t cache_steps() const noexcept { return cache_steps_; } - void import_state(const runtime::TransformerKVState & state) { - cache_.import_state(state); + FishPrefillCacheTarget prefill_target_cache() const { + FishPrefillCacheTarget out; + out.keys.reserve(runtime_->weights().slow_layers.size()); + out.values.reserve(runtime_->weights().slow_layers.size()); + for (size_t layer = 0; layer < runtime_->weights().slow_layers.size(); ++layer) { + out.keys.push_back(cache_.key_tensor(layer)); + out.values.push_back(cache_.value_tensor(layer)); + } + return out; + } + + void finish_prefill(int64_t steps) { + cache_.retain_prefix(0); + cache_.advance_after_direct_append(steps); const auto masked = ggml_fp32_to_fp16(-INFINITY); const auto visible = ggml_fp32_to_fp16(0.0F); std::fill(mask_scratch_.begin(), mask_scratch_.end(), masked); @@ -1341,7 +1355,10 @@ class FishAudioARRuntime::Impl { void ensure_prefill_graph(int64_t steps, FishARProfile & profile) { if (!prefill_graph_ || prefill_graph_->steps() != steps) { const auto build_start = Clock::now(); - prefill_graph_ = std::make_unique(runtime_, steps); + if (!step_graph_) { + throw std::runtime_error("Fish Audio AR prefill requires a step graph"); + } + prefill_graph_ = std::make_unique(runtime_, steps, step_graph_->prefill_target_cache()); profile.graph_build_prefill_ms += engine::debug::elapsed_ms(build_start, Clock::now()); } } @@ -1350,6 +1367,7 @@ class FishAudioARRuntime::Impl { if (!step_graph_ || step_graph_->cache_steps() < cache_steps) { const auto build_start = Clock::now(); step_graph_ = std::make_unique(runtime_, cache_steps); + prefill_graph_.reset(); profile.graph_build_step_ms += engine::debug::elapsed_ms(build_start, Clock::now()); } } @@ -1449,7 +1467,6 @@ class FishAudioARRuntime::Impl { engine::debug::timing_log_scalar("fish_audio.ar.profile.prefill_input_upload_ms", profile.prefill_input_upload_ms); engine::debug::timing_log_scalar("fish_audio.ar.profile.prefill_graph_ms", profile.prefill_graph_ms); engine::debug::timing_log_scalar("fish_audio.ar.profile.prefill_output_read_ms", profile.prefill_output_read_ms); - engine::debug::timing_log_scalar("fish_audio.ar.profile.prefill_state_read_ms", profile.prefill_state_read_ms); engine::debug::timing_log_scalar("fish_audio.ar.profile.step_input_upload_ms", profile.step_input_upload_ms); engine::debug::timing_log_scalar("fish_audio.ar.profile.step_mask_upload_ms", profile.step_mask_upload_ms); engine::debug::timing_log_scalar("fish_audio.ar.profile.step_graph_ms", profile.step_graph_ms); @@ -1458,7 +1475,6 @@ class FishAudioARRuntime::Impl { engine::debug::timing_log_scalar("fish_audio.ar.profile.fast_mask_upload_ms", profile.fast_mask_upload_ms); engine::debug::timing_log_scalar("fish_audio.ar.profile.fast_graph_ms", profile.fast_graph_ms); engine::debug::timing_log_scalar("fish_audio.ar.profile.fast_output_read_ms", profile.fast_output_read_ms); - engine::debug::timing_log_scalar("fish_audio.ar.profile.import_prefill_state_ms", profile.import_prefill_state_ms); engine::debug::timing_log_scalar("fish_audio.ar.profile.sample_bias_ms", profile.sample_bias_ms); engine::debug::timing_log_scalar("fish_audio.ar.profile.sample_main_ms", profile.sample_main_ms); engine::debug::timing_log_scalar("fish_audio.ar.profile.sample_high_ms", profile.sample_high_ms); diff --git a/src/models/fish_audio/loader.cpp b/src/models/fish_audio/loader.cpp index 2f3a025b..c172a69d 100644 --- a/src/models/fish_audio/loader.cpp +++ b/src/models/fish_audio/loader.cpp @@ -34,7 +34,7 @@ runtime::ModelCliInterface cli(const FishAudioAssets &) { runtime::ModelCliInterface out; out.request_options = { {"reference_text", "TEXT", "Reference transcript used with speaker reference audio."}, - {"max_new_tokens", "N", "Maximum Fish Audio semantic tokens to generate; default 2048, 0 uses the default."}, + {"max_new_tokens", "N", "Maximum Fish Audio semantic tokens to generate; default 1024, 0 uses the default."}, {"text_chunk_size", "N", "Long-form text chunk size; default 200."}, {"text_chunk_mode", "default|tag_aware|japanese|endline", "Framework text chunking mode."}, {"top_p", "FLOAT", "Top-p sampling value."}, From 32a0ff34b7d3bc29adf7c90eb374b2067c14b903 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Tue, 21 Jul 2026 01:58:09 -0400 Subject: [PATCH 17/27] Add BF16 KV cache support for Fish Audio --- include/engine/framework/core/backend.h | 4 ++ include/engine/framework/runtime/kv_cache.h | 1 + src/framework/core/backend.cpp | 38 +++++++++++++++++-- .../modules/optimizations/fast_kv_modules.cpp | 4 +- src/framework/runtime/kv_cache.cpp | 14 ++++++- src/models/fish_audio/ar.cpp | 25 +++++++----- 6 files changed, 70 insertions(+), 16 deletions(-) diff --git a/include/engine/framework/core/backend.h b/include/engine/framework/core/backend.h index 4597e357..8c186826 100644 --- a/include/engine/framework/core/backend.h +++ b/include/engine/framework/core/backend.h @@ -70,12 +70,16 @@ void write_tensor_f32_slice(const TensorValue & tensor, size_t element_offset, c void write_tensor_f32(const TensorValue & tensor, const std::vector & values); void write_tensor_f16(const TensorValue & tensor, const float * values, size_t count); void write_tensor_f16(const TensorValue & tensor, const std::vector & values); +void write_tensor_bf16(const TensorValue & tensor, const float * values, size_t count); +void write_tensor_bf16(const TensorValue & tensor, const std::vector & values); void write_tensor_i32(const TensorValue & tensor, const int32_t * values, size_t count); void write_tensor_i32(const TensorValue & tensor, const std::vector & values); void read_tensor_f32_into(const ggml_tensor * tensor, std::vector & values); std::vector read_tensor_f32(const ggml_tensor * tensor); void read_tensor_f16_into(const ggml_tensor * tensor, std::vector & values); std::vector read_tensor_f16(const ggml_tensor * tensor); +void read_tensor_bf16_into(const ggml_tensor * tensor, std::vector & values); +std::vector read_tensor_bf16(const ggml_tensor * tensor); void read_tensor_i32_into(const ggml_tensor * tensor, std::vector & values); std::vector read_tensor_i32(const ggml_tensor * tensor); diff --git a/include/engine/framework/runtime/kv_cache.h b/include/engine/framework/runtime/kv_cache.h index e79fc9d1..81156423 100644 --- a/include/engine/framework/runtime/kv_cache.h +++ b/include/engine/framework/runtime/kv_cache.h @@ -21,6 +21,7 @@ struct TransformerKVState { struct TransformerKVCacheOptions { bool allow_f16_storage = false; + bool allow_bf16_storage = false; }; class TransformerKVCache { diff --git a/src/framework/core/backend.cpp b/src/framework/core/backend.cpp index 255d59b4..62cf854f 100644 --- a/src/framework/core/backend.cpp +++ b/src/framework/core/backend.cpp @@ -469,6 +469,26 @@ void write_tensor_f16(const TensorValue & tensor, const std::vector & val write_tensor_f16(tensor, values.data(), values.size()); } +void write_tensor_bf16(const TensorValue & tensor, const float * values, size_t count) { + if (tensor.type != GGML_TYPE_BF16) { + throw std::runtime_error("write_tensor_bf16 requires GGML_TYPE_BF16 tensor"); + } + if (tensor.shape.num_elements() != static_cast(count)) { + throw std::runtime_error( + "write_tensor_bf16 value count does not match tensor shape for tensor '" + + std::string(tensor.tensor != nullptr ? tensor.tensor->name : "") + + "': expected " + std::to_string(tensor.shape.num_elements()) + + ", got " + std::to_string(count)); + } + std::vector bf16_values(count); + ggml_fp32_to_bf16_row(values, bf16_values.data(), static_cast(count)); + ggml_backend_tensor_set(tensor.tensor, bf16_values.data(), 0, count * sizeof(ggml_bf16_t)); +} + +void write_tensor_bf16(const TensorValue & tensor, const std::vector & values) { + write_tensor_bf16(tensor, values.data(), values.size()); +} + void write_tensor_i32(const TensorValue & tensor, const int32_t * values, size_t count) { if (tensor.type != GGML_TYPE_I32) { throw std::runtime_error("write_tensor_i32 requires GGML_TYPE_I32 tensor"); @@ -543,9 +563,21 @@ std::vector read_tensor_f16(const ggml_tensor * tensor) { return values; } -void read_tensor_i32_into(const ggml_tensor * tensor, std::vector & values) { - read_tensor_typed_into(tensor, GGML_TYPE_I32, values); -} +void read_tensor_bf16_into(const ggml_tensor * tensor, std::vector & values) { + const auto bf16_values = read_tensor_typed(tensor, GGML_TYPE_BF16); + values.resize(bf16_values.size()); + ggml_bf16_to_fp32_row(bf16_values.data(), values.data(), static_cast(values.size())); +} + +std::vector read_tensor_bf16(const ggml_tensor * tensor) { + std::vector values; + read_tensor_bf16_into(tensor, values); + return values; +} + +void read_tensor_i32_into(const ggml_tensor * tensor, std::vector & values) { + read_tensor_typed_into(tensor, GGML_TYPE_I32, values); +} std::vector read_tensor_i32(const ggml_tensor * tensor) { return read_tensor_typed(tensor, GGML_TYPE_I32); diff --git a/src/framework/modules/optimizations/fast_kv_modules.cpp b/src/framework/modules/optimizations/fast_kv_modules.cpp index eff1e166..aa879b0e 100644 --- a/src/framework/modules/optimizations/fast_kv_modules.cpp +++ b/src/framework/modules/optimizations/fast_kv_modules.cpp @@ -55,11 +55,11 @@ core::TensorValue FastKVSetRowsModule::build( } const bool optimized = config_.mode == FastKVSetRowsMode::BackendViewOptimized; if (((!optimized && cache.type != GGML_TYPE_F32) || - (optimized && cache.type != GGML_TYPE_F32 && cache.type != GGML_TYPE_F16)) || + (optimized && cache.type != GGML_TYPE_F32 && cache.type != GGML_TYPE_F16 && cache.type != GGML_TYPE_BF16)) || row.type != GGML_TYPE_F32) { throw std::runtime_error( optimized - ? "FastKVSetRowsModule requires an f32/f16 cache and an f32 row tensor" + ? "FastKVSetRowsModule requires an f32/f16/bf16 cache and an f32 row tensor" : "FastKVSetRowsModule requires f32 cache and row tensors"); } if (row_index.type != GGML_TYPE_I32 && row_index.type != GGML_TYPE_I64) { diff --git a/src/framework/runtime/kv_cache.cpp b/src/framework/runtime/kv_cache.cpp index 0d7f9208..555cfd44 100644 --- a/src/framework/runtime/kv_cache.cpp +++ b/src/framework/runtime/kv_cache.cpp @@ -18,9 +18,12 @@ void validate_cache_tensor(const core::TensorValue & tensor, const TransformerKV if (options.allow_f16_storage && tensor.type == GGML_TYPE_F16) { return; } + if (options.allow_bf16_storage && tensor.type == GGML_TYPE_BF16) { + return; + } throw std::runtime_error( - options.allow_f16_storage - ? "TransformerKVCache supports only f32 and f16 cache tensors" + options.allow_f16_storage || options.allow_bf16_storage + ? "TransformerKVCache supports only f32/f16/bf16 cache tensors when enabled" : "TransformerKVCache requires f32 cache tensors"); } @@ -37,6 +40,10 @@ void write_cache_tensor( core::write_tensor_f16(tensor, values); return; } + if (options.allow_bf16_storage && tensor.type == GGML_TYPE_BF16) { + core::write_tensor_bf16(tensor, values); + return; + } throw std::runtime_error("TransformerKVCache requires f32 cache tensors"); } @@ -48,6 +55,9 @@ std::vector read_cache_tensor(const core::TensorValue & tensor, const Tra if (options.allow_f16_storage && tensor.type == GGML_TYPE_F16) { return core::read_tensor_f16(tensor.tensor); } + if (options.allow_bf16_storage && tensor.type == GGML_TYPE_BF16) { + return core::read_tensor_bf16(tensor.tensor); + } throw std::runtime_error("TransformerKVCache requires f32 cache tensors"); } diff --git a/src/models/fish_audio/ar.cpp b/src/models/fish_audio/ar.cpp index 720382dc..ecb32aea 100644 --- a/src/models/fish_audio/ar.cpp +++ b/src/models/fish_audio/ar.cpp @@ -675,10 +675,17 @@ FishStaticDecoderOutputs build_fish_static_decoder( }) .build(ctx, hidden, weights.lm_head); auto fast_hidden = norm_fastlayer_input ? hidden : x; + runtime::TransformerKVCacheOptions cache_options; + cache_options.allow_bf16_storage = true; return { fast_hidden, logits, - runtime::TransformerKVCache(cache_steps, step_elems, std::move(cache_keys), std::move(cache_values)), + runtime::TransformerKVCache( + cache_steps, + step_elems, + std::move(cache_keys), + std::move(cache_values), + cache_options), }; } @@ -1035,23 +1042,23 @@ class FishAudioARRuntime::Impl { cache_keys.push_back(core::wrap_tensor( ggml_new_tensor_4d( state_ctx_.get(), - GGML_TYPE_F32, + GGML_TYPE_BF16, config.head_dim, config.n_local_heads, cache_steps_, 1), core::TensorShape::from_dims({1, cache_steps_, config.n_local_heads, config.head_dim}), - GGML_TYPE_F32)); + GGML_TYPE_BF16)); cache_values.push_back(core::wrap_tensor( ggml_new_tensor_4d( state_ctx_.get(), - GGML_TYPE_F32, + GGML_TYPE_BF16, config.head_dim, config.n_local_heads, cache_steps_, 1), core::TensorShape::from_dims({1, cache_steps_, config.n_local_heads, config.head_dim}), - GGML_TYPE_F32)); + GGML_TYPE_BF16)); } state_buffer_ = ggml_backend_alloc_ctx_tensors(state_ctx_.get(), runtime_->backend()); if (state_buffer_ == nullptr) { @@ -1222,23 +1229,23 @@ class FishAudioARRuntime::Impl { cache_keys.push_back(core::wrap_tensor( ggml_new_tensor_4d( state_ctx_.get(), - GGML_TYPE_F32, + GGML_TYPE_BF16, config.head_dim, config.n_local_heads, config.num_codebooks, 1), core::TensorShape::from_dims({1, config.num_codebooks, config.n_local_heads, config.head_dim}), - GGML_TYPE_F32)); + GGML_TYPE_BF16)); cache_values.push_back(core::wrap_tensor( ggml_new_tensor_4d( state_ctx_.get(), - GGML_TYPE_F32, + GGML_TYPE_BF16, config.head_dim, config.n_local_heads, config.num_codebooks, 1), core::TensorShape::from_dims({1, config.num_codebooks, config.n_local_heads, config.head_dim}), - GGML_TYPE_F32)); + GGML_TYPE_BF16)); } state_buffer_ = ggml_backend_alloc_ctx_tensors(state_ctx_.get(), runtime_->backend()); if (state_buffer_ == nullptr) { From 16a0511f7f5e9e545daec5246a029d29c7e187c6 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Tue, 21 Jul 2026 02:23:54 -0400 Subject: [PATCH 18/27] Align Fish Audio codec graph lifecycle --- include/engine/models/fish_audio/codec.h | 1 + src/models/fish_audio/codec.cpp | 8 ++++++++ src/models/fish_audio/generator.cpp | 6 ++++-- src/models/fish_audio/loader.cpp | 2 +- src/models/fish_audio/session.cpp | 3 --- 5 files changed, 14 insertions(+), 6 deletions(-) diff --git a/include/engine/models/fish_audio/codec.h b/include/engine/models/fish_audio/codec.h index 5427fef4..2da7556e 100644 --- a/include/engine/models/fish_audio/codec.h +++ b/include/engine/models/fish_audio/codec.h @@ -23,6 +23,7 @@ class FishAudioCodecRuntime { FishAudioCodes encode_reference(const runtime::AudioBuffer & audio); runtime::AudioBuffer decode(const FishAudioCodes & codes); + void release_encode_graph(); void release_runtime_graphs(); private: diff --git a/src/models/fish_audio/codec.cpp b/src/models/fish_audio/codec.cpp index 0c22cf88..f9f4c924 100644 --- a/src/models/fish_audio/codec.cpp +++ b/src/models/fish_audio/codec.cpp @@ -1081,6 +1081,10 @@ class FishAudioCodecRuntime::Impl { return decode_graph_->run(codes); } + void release_encode_graph() { + encode_graph_.reset(); + } + void release_runtime_graphs() { encode_graph_.reset(); decode_graph_.reset(); @@ -1123,6 +1127,10 @@ runtime::AudioBuffer FishAudioCodecRuntime::decode(const FishAudioCodes & codes) return impl_->decode(codes); } +void FishAudioCodecRuntime::release_encode_graph() { + impl_->release_encode_graph(); +} + void FishAudioCodecRuntime::release_runtime_graphs() { impl_->release_runtime_graphs(); } diff --git a/src/models/fish_audio/generator.cpp b/src/models/fish_audio/generator.cpp index 90b45f33..f1666729 100644 --- a/src/models/fish_audio/generator.cpp +++ b/src/models/fish_audio/generator.cpp @@ -30,7 +30,9 @@ FishAudioGenerator::FishAudioGenerator( FishAudioGenerator::~FishAudioGenerator() = default; FishAudioCodes FishAudioGenerator::encode_reference(const runtime::AudioBuffer & audio) { - return codec_->encode_reference(audio); + auto codes = codec_->encode_reference(audio); + codec_->release_encode_graph(); + return codes; } FishAudioGenerationResult FishAudioGenerator::generate( @@ -61,9 +63,9 @@ FishAudioGenerationResult FishAudioGenerator::generate( engine::debug::timing_log_scalar( "fish_audio.codec_decode_ms", engine::debug::elapsed_ms(decode_start, Clock::now())); + codec_->release_runtime_graphs(); if (mem_saver) { ar_->release_runtime_graphs(); - codec_->release_runtime_graphs(); } return result; } diff --git a/src/models/fish_audio/loader.cpp b/src/models/fish_audio/loader.cpp index c172a69d..60e7b951 100644 --- a/src/models/fish_audio/loader.cpp +++ b/src/models/fish_audio/loader.cpp @@ -43,7 +43,7 @@ runtime::ModelCliInterface cli(const FishAudioAssets &) { {"seed", "N", "Sampling seed."}, }; out.session_options = { - {"fish_audio.mem_saver", "true|false", "Release cached runtime graphs after each request; default false."}, + {"fish_audio.mem_saver", "true|false", "Release cached AR runtime graphs after each request; default false."}, {"fish_audio.reference_cache_slots", "n", "Prepared reference-audio cache slots; default 1."}, {"fish_audio.weight_type", "native|f32|f16|bf16|q8_0", "AR matmul weight storage type; default native."}, {"fish_audio.codec_weight_type", "native|f32|f16|q8_0", "Codec conv/matmul weight storage type; default native."}, diff --git a/src/models/fish_audio/session.cpp b/src/models/fish_audio/session.cpp index 15218fe2..12606d81 100644 --- a/src/models/fish_audio/session.cpp +++ b/src/models/fish_audio/session.cpp @@ -349,9 +349,6 @@ void FishAudioSession::prepare(const runtime::SessionPreparationRequest & reques if (auto reference = reference_from_voice(*assets_, request.voice, request.options, "Fish Audio prepare"); reference.has_value()) { defaults.reference = std::move(*reference); - if (defaults.reference->audio.has_value()) { - (void) resolve_reference_codes(*defaults.reference); - } has_defaults = true; } if (has_defaults) { From ee682d0f2d11ce3cf46f7456976ed119a65bf964 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Tue, 21 Jul 2026 09:26:07 -0400 Subject: [PATCH 19/27] Stabilize Fish Audio longform generation --- src/models/fish_audio/ar.cpp | 8 +++++++ ...audiocpp_cli_longform_tts_clone_cases.json | 24 +++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/src/models/fish_audio/ar.cpp b/src/models/fish_audio/ar.cpp index ecb32aea..c052ad24 100644 --- a/src/models/fish_audio/ar.cpp +++ b/src/models/fish_audio/ar.cpp @@ -1251,6 +1251,14 @@ class FishAudioARRuntime::Impl { if (state_buffer_ == nullptr) { throw std::runtime_error("failed to allocate Fish Audio fast AR state tensors"); } + for (const auto & cache : cache_keys) { + std::vector zeros(static_cast(ggml_nbytes(cache.tensor)), 0); + ggml_backend_tensor_set(cache.tensor, zeros.data(), 0, zeros.size()); + } + for (const auto & cache : cache_values) { + std::vector zeros(static_cast(ggml_nbytes(cache.tensor)), 0); + ggml_backend_tensor_set(cache.tensor, zeros.data(), 0, zeros.size()); + } core::ModuleBuildContext ctx{graph_ctx_.get(), "fish_audio.ar.fast", runtime_->backend_type()}; auto input = core::make_tensor(ctx, GGML_TYPE_F32, core::TensorShape::from_dims({1, 1, config.dim})); diff --git a/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json b/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json index c34b69fe..3314e72f 100644 --- a/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json +++ b/tools/audiocpp_cli/audiocpp_cli_longform_tts_clone_cases.json @@ -244,6 +244,30 @@ } ] }, + { + "id": "fish_audio_voice_clone_longform", + "coverage": "Fish Audio S2-Pro voice clone with framework long-form text chunking, bounded AR generation, and codec decode", + "family": "fish_audio", + "model": "models/s2-pro", + "task": "tts", + "mode": "offline", + "outputs": [ + "audio" + ], + "requests": [ + { + "id": "clone_longform", + "text": "At dawn the harbor station opens its tall windows and the first clerk begins a careful report for the day. She notes the weather above the river, the slow cargo boats beyond the bridge, and the market voices arriving from the eastern road. A brass clock marks each quarter hour while porters stack wooden crates, bakers carry warm bread across the square, and a violinist practices the same bright phrase under the stone archway. By midmorning the keeper of the lighthouse sends a message about shifting currents, the museum guide unlocks a cabinet of maps, and a teacher leads a quiet line of students toward the ferry. In the afternoon a painter describes the silver color of the water, a mechanic jokes with the tram driver, and the station master reads an announcement that asks every traveler to keep close watch over letters, tickets, and parcels. After sunset the same clerk continues the report because new visitors keep arriving from the inland road. She explains that a florist carries pale roses past the fountain, two carpenters compare measurements beside the warehouse door, and the watchman checks each lock before the tide reaches its highest mark. A child laughs when the tram bell rings, a cook lowers a basket of fruit to the cellar, and three sailors unfold a chart that shows old channels, sandbars, and safe turning points for the morning crossing. Near midnight the lamps still glow on wet stone, the last cart rattles toward the market gate, and the report ends by saying that the harbor remains orderly, the wind has softened, the ferries are secure, and the town can rest until the next sunrise returns over the water. On the following morning the clerk resumes the record with even greater care because a week of inspections is about to begin. She writes that a ferry captain checks the mooring ropes one by one, a bookseller arranges travel guides beside the station cafe, and a pair of gardeners lift wet soil into bright clay pots near the west entrance. The bakery sends out trays of seed bread, the telegraph operator copies three official notices, and a tailor unfolds navy cloth across a polished wooden counter while customers wait in a line that bends toward the fountain. Before noon a surveyor compares bridge numbers against an old ledger, two cousins argue cheerfully about the best route to the fish market, and a choir director rehearses a patient scale that echoes against the warehouse wall. The lighthouse keeper reports that the northern channel is calmer than expected, the harbor pilot recommends a slower turn near the sandbar, and the customs officer stamps a packet of forms before waving a cart through the side gate. Later the schoolteacher returns with another group of students, asking them to observe the colors of rope, paint, stone, and water so they can write more exact descriptions in the classroom. A photographer kneels beside a rain barrel to capture the reflection of the clock tower, a mechanic tightens a brass hinge on the tram door, and an elderly traveler asks the clerk whether the evening ferry still stops at the orchard village beyond the marsh. As dusk arrives, lamps are trimmed again, shutters are tested against the wind, and the station kitchen sends bowls of soup to workers who remain on the late shift. The report continues with notes about a carpenter measuring floorboards in the east hall, a florist tying silver ribbon around the last stems of the day, and a violin case resting open on a bench beside the ticket window while its owner copies melody marks into a notebook. Long after the market gate closes, the clerk still writes that the harbor road stays busy, the river glints beneath scattered lamps, and the town maintains its patient rhythm of signals, footsteps, voices, bells, and distant engines. On the third day the clerk decides the record should be more precise, so she marks each event by the quarter hour and notes which sounds carry farthest through the station concourse. At first light she hears broom bristles on the stone steps, kettle lids in the cafe kitchen, and the slow scrape of crates being nudged across a loading cart beside the river wall. A messenger in a green coat delivers two canvas pouches, the ticket agent counts rolled coins into a brass tray, and a mother reads directions aloud while her son traces the painted ferry schedule with one curious finger. Midmorning brings a burst of sunlight across the waiting hall, making every brass handle shine while the museum guide escorts visitors toward the gallery of maps and navigational instruments. A porter pauses to describe the oldest compass in the display, a student sketches the harbor outline in graphite, and an apprentice clockmaker compares the station bell to a pocket watch that once belonged to his grandfather. By noon the fish market sends salt and seaweed scents through the open doors, tram wheels hiss at the curb, and the baker from the square exchanges a laugh with the florist who is carrying fresh lilies to the hotel veranda. The clerk writes that a cooper rolls three narrow barrels toward the cellar ramp, a translator copies weather bulletins for inland travelers, and a painter in a blue scarf studies the changing color of the tide as if each small wave might explain a different part of the sky. In the late afternoon the station master reviews freight tags, the customs officer checks a parcel of glassware, and a choir of children crosses the square singing a phrase so soft that the watchman removes his cap to listen. Evening settles slowly; lamps brighten in sequence, a cook inventories apples and onions in the pantry, and two sailors spread a faded chart on a crate so they can debate whether the shoals have shifted since the previous autumn. Before sleep the clerk closes the day with a final note that every vessel is accounted for, every platform has been swept, every lock has been tested twice, and the harbor seems ready to welcome another tide, another market, and another patient stream of voices at sunrise.", + "voice_ref": "resources/a.wav", + "reference_text": "This little work was finished in the year eighteen o three, and intended for immediate publication.", + "text_chunk_size": 200, + "top_p": 0.8, + "repetition_penalty": 1.1, + "temperature": 0.8, + "seed": 1234 + } + ] + }, { "id": "index_tts2_voice_clone_longform", "coverage": "IndexTTS2 voice clone with shared long-form text for chunking and RTF measurement", From 4e78d4dd1f7d230de93d17c4534a077a8ea488c7 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Tue, 21 Jul 2026 10:03:28 -0400 Subject: [PATCH 20/27] Allow CPU sampling fallback for Higgs and Fish TTS --- .../engine/models/higgs_audio_tts/sampler.h | 2 ++ src/models/fish_audio/ar.cpp | 14 +++++++++- src/models/fish_audio/codec.cpp | 6 ++++ src/models/higgs_audio_tts/generator.cpp | 10 +++++-- src/models/higgs_audio_tts/sampler.cpp | 28 +++++++++++++++++-- src/models/higgs_audio_tts/session.cpp | 4 --- 6 files changed, 53 insertions(+), 11 deletions(-) diff --git a/include/engine/models/higgs_audio_tts/sampler.h b/include/engine/models/higgs_audio_tts/sampler.h index 0ed74c1d..76943d75 100644 --- a/include/engine/models/higgs_audio_tts/sampler.h +++ b/include/engine/models/higgs_audio_tts/sampler.h @@ -5,6 +5,7 @@ #include #include +#include #include namespace engine::models::higgs_audio_tts { @@ -20,6 +21,7 @@ struct HiggsSamplingOptions { bool has_seed = false; uint64_t seed = 0; HiggsCudaSamplingPolicy cuda_policy; + std::mt19937 * fallback_rng = nullptr; }; struct HiggsSamplerState { diff --git a/src/models/fish_audio/ar.cpp b/src/models/fish_audio/ar.cpp index c052ad24..208b0de6 100644 --- a/src/models/fish_audio/ar.cpp +++ b/src/models/fish_audio/ar.cpp @@ -24,6 +24,7 @@ #include #include #include +#include #include #include #include @@ -453,6 +454,7 @@ FishARWeights load_ar_weights( struct SampleState { uint64_t seed = 0; uint64_t call_index = 0; + std::mt19937 rng; std::vector previous_main; }; @@ -534,6 +536,15 @@ int32_t sample_from_logits( const sampling::TorchCudaSamplingPolicy & policy) { const auto distribution = logits_to_distribution(logits, temperature, top_p, top_k); const uint64_t call_index = state.call_index++; + if (!policy.cuda_fast_path) { + std::vector weights; + weights.reserve(distribution.candidates.size()); + for (const auto & candidate : distribution.candidates) { + weights.push_back(static_cast(std::max(candidate.probability, 0.0F))); + } + std::discrete_distribution sampler(weights.begin(), weights.end()); + return distribution.candidates[sampler(state.rng)].index; + } int32_t best = 0; double best_score = -std::numeric_limits::infinity(); for (const auto & candidate : distribution.candidates) { @@ -799,7 +810,7 @@ class FishAudioARRuntime::Impl { backend_config.device, "fish_audio.ar.cuda_sampling_policy", "Fish Audio", - sampling::TorchCudaSamplingPolicyFailureMode::StrictCuda)) {} + sampling::TorchCudaSamplingPolicyFailureMode::FallbackToDefault)) {} ~Impl() { step_graph_.reset(); @@ -825,6 +836,7 @@ class FishAudioARRuntime::Impl { ensure_fast_graph(profile); SampleState sample; sample.seed = options.seed; + sample.rng.seed(options.seed); sample.previous_main.assign(static_cast(kRasWindow), 0); auto timing_start = Clock::now(); auto embeddings = build_slow_embeddings(assets.config, weights, prompt.matrix.data(), prompt.steps); diff --git a/src/models/fish_audio/codec.cpp b/src/models/fish_audio/codec.cpp index f9f4c924..b47150cb 100644 --- a/src/models/fish_audio/codec.cpp +++ b/src/models/fish_audio/codec.cpp @@ -199,6 +199,9 @@ core::TensorValue zero_prefix_like(core::ModuleBuildContext & ctx, const core::T throw std::runtime_error("Fish Audio zero_prefix_like requires positive frames"); } auto first = modules::SliceModule({2, 0, 1}).build(ctx, input); + if (ctx.backend_type == core::BackendType::Cpu) { + first = core::ensure_backend_addressable_layout(ctx, first); + } first = core::wrap_tensor(ggml_scale(ctx.ggml, first.tensor, 0.0F), first.shape, GGML_TYPE_F32); return modules::RepeatModule({core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], frames})}) .build(ctx, first); @@ -209,6 +212,9 @@ core::TensorValue zero_suffix_like(core::ModuleBuildContext & ctx, const core::T throw std::runtime_error("Fish Audio zero_suffix_like requires positive frames"); } auto last = modules::SliceModule({2, input.shape.dims[2] - 1, 1}).build(ctx, input); + if (ctx.backend_type == core::BackendType::Cpu) { + last = core::ensure_backend_addressable_layout(ctx, last); + } last = core::wrap_tensor(ggml_scale(ctx.ggml, last.tensor, 0.0F), last.shape, GGML_TYPE_F32); return modules::RepeatModule({core::TensorShape::from_dims({input.shape.dims[0], input.shape.dims[1], frames})}) .build(ctx, last); diff --git a/src/models/higgs_audio_tts/generator.cpp b/src/models/higgs_audio_tts/generator.cpp index 3cfc0dac..2820067e 100644 --- a/src/models/higgs_audio_tts/generator.cpp +++ b/src/models/higgs_audio_tts/generator.cpp @@ -411,15 +411,19 @@ HiggsGenerationResult HiggsGenerator::generate(const HiggsGenerationRequest & re // Higgs audio-codebook sampler does not consume it. sampling.has_seed = request.options.seed.has_value(); sampling.seed = request.options.seed.value_or(runtime::random_u64_seed()); - if (!cuda_sampling_policy_.has_value()) { + std::mt19937 fallback_rng(static_cast(sampling.seed)); + sampling.fallback_rng = &fallback_rng; + if (!sampling.has_seed && !cuda_sampling_policy_.has_value()) { cuda_sampling_policy_ = engine::sampling::resolve_torch_cuda_sampling_policy( ar_->backend_type(), ar_->device(), "higgs_audio_tts.cuda_sampling_policy", "Higgs TTS", - engine::sampling::TorchCudaSamplingPolicyFailureMode::StrictCuda); + engine::sampling::TorchCudaSamplingPolicyFailureMode::FallbackToDefault); + } + if (cuda_sampling_policy_.has_value()) { + sampling.cuda_policy = *cuda_sampling_policy_; } - sampling.cuda_policy = *cuda_sampling_policy_; engine::debug::trace_log_scalar("higgs_audio_tts.sampler.temperature", sampling.temperature); engine::debug::trace_log_scalar("higgs_audio_tts.sampler.has_seed", sampling.has_seed); engine::debug::trace_log_scalar("higgs_audio_tts.sampler.seed", sampling.seed); diff --git a/src/models/higgs_audio_tts/sampler.cpp b/src/models/higgs_audio_tts/sampler.cpp index ccebd9b7..9238fb7c 100644 --- a/src/models/higgs_audio_tts/sampler.cpp +++ b/src/models/higgs_audio_tts/sampler.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include namespace engine::models::higgs_audio_tts { @@ -292,9 +293,27 @@ int32_t sample_unseeded_torch_multinomial(const std::vector & probs, const std::vector & candidates, uint64_t seed, uint64_t call_index, - const HiggsCudaSamplingPolicy & policy) { - require_cuda_sampling_policy(policy); + const HiggsCudaSamplingPolicy & policy, + std::mt19937 & fallback_rng) { const int64_t vocab_size = static_cast(probs.size()); + if (!policy.cuda_fast_path) { + std::vector weights; + if (candidates.empty()) { + weights.reserve(probs.size()); + for (float prob : probs) { + weights.push_back(static_cast(std::max(prob, 0.0F))); + } + std::discrete_distribution distribution(weights.begin(), weights.end()); + return static_cast(distribution(fallback_rng)); + } + weights.reserve(candidates.size()); + for (const int64_t index : candidates) { + weights.push_back(static_cast(std::max(probs[static_cast(index)], 0.0F))); + } + std::discrete_distribution distribution(weights.begin(), weights.end()); + return static_cast(candidates[distribution(fallback_rng)]); + } + require_cuda_sampling_policy(policy); double best_rank = -std::numeric_limits::infinity(); int32_t best = -1; @@ -374,8 +393,11 @@ int32_t sample_codebook_row(const float * logits, return sample_seeded_sglang_gumbel( scratch.probs, scratch.kept, options.seed & 0x7FFFFFFFull, call_index); } + if (options.fallback_rng == nullptr) { + throw std::runtime_error("Higgs TTS sampler fallback RNG is missing"); + } return sample_unseeded_torch_multinomial( - scratch.probs, scratch.kept, options.seed, call_index, options.cuda_policy); + scratch.probs, scratch.kept, options.seed, call_index, options.cuda_policy, *options.fallback_rng); } } // namespace diff --git a/src/models/higgs_audio_tts/session.cpp b/src/models/higgs_audio_tts/session.cpp index 8c3d5328..0a5d4779 100644 --- a/src/models/higgs_audio_tts/session.cpp +++ b/src/models/higgs_audio_tts/session.cpp @@ -136,10 +136,6 @@ HiggsTTSSession::HiggsTTSSession( if (task_.task != runtime::VoiceTaskKind::Tts) { throw std::runtime_error("Higgs TTS only supports the Tts task"); } - if (options.backend.type != core::BackendType::Cuda) { - throw std::runtime_error("Higgs TTS generation requires CUDA backend"); - } - ar_weight_context_bytes_ = runtime::parse_size_mb_option( options.options, {"higgs_audio_tts.ar_weight_context_mb"}, ar_weight_context_bytes_); codec_weight_context_bytes_ = runtime::parse_size_mb_option( From 48fd05a593dadaf06968bd9e2d0c4ea64471a5d1 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Tue, 21 Jul 2026 10:22:15 -0400 Subject: [PATCH 21/27] Align Fish Audio seed defaults --- src/models/fish_audio/generator.cpp | 1 + src/models/fish_audio/loader.cpp | 2 +- src/models/fish_audio/session.cpp | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/models/fish_audio/generator.cpp b/src/models/fish_audio/generator.cpp index f1666729..a5aa4586 100644 --- a/src/models/fish_audio/generator.cpp +++ b/src/models/fish_audio/generator.cpp @@ -43,6 +43,7 @@ FishAudioGenerationResult FishAudioGenerator::generate( engine::debug::trace_log_scalar("fish_audio.request.has_reference", request.reference.has_value()); engine::debug::trace_log_scalar("fish_audio.request.text_chars", static_cast(request.text.size())); engine::debug::trace_log_scalar("fish_audio.request.has_previous_turn", previous_turn.has_value()); + engine::debug::trace_log_scalar("fish_audio.sampler.seed", request.generation.seed); const auto prompt_start = Clock::now(); const auto prompt = prompt_builder_.build(request, reference_codes, previous_turn); engine::debug::timing_log_scalar( diff --git a/src/models/fish_audio/loader.cpp b/src/models/fish_audio/loader.cpp index 60e7b951..a6c67653 100644 --- a/src/models/fish_audio/loader.cpp +++ b/src/models/fish_audio/loader.cpp @@ -40,7 +40,7 @@ runtime::ModelCliInterface cli(const FishAudioAssets &) { {"top_p", "FLOAT", "Top-p sampling value."}, {"top_k", "N", "Top-k sampling value."}, {"temperature", "FLOAT", "Sampling temperature."}, - {"seed", "N", "Sampling seed."}, + {"seed", "N", "Sampling seed for reproducible output; omitted uses a random seed."}, }; out.session_options = { {"fish_audio.mem_saver", "true|false", "Release cached AR runtime graphs after each request; default false."}, diff --git a/src/models/fish_audio/session.cpp b/src/models/fish_audio/session.cpp index 12606d81..f23c6d28 100644 --- a/src/models/fish_audio/session.cpp +++ b/src/models/fish_audio/session.cpp @@ -124,7 +124,7 @@ FishAudioGenerationOptions generation_options_from_request(const runtime::TaskRe options.top_p = runtime::parse_float_option(request.options, {"top_p"}).value_or(options.top_p); options.top_k = runtime::parse_int_option(request.options, {"top_k"}).value_or(options.top_k); options.temperature = runtime::parse_float_option(request.options, {"temperature"}).value_or(options.temperature); - options.seed = runtime::parse_u32_option(request.options, {"seed"}).value_or(options.seed); + options.seed = runtime::parse_u32_option(request.options, {"seed"}).value_or(runtime::random_u32_seed()); if (options.max_new_tokens <= 0) { throw std::runtime_error("Fish Audio max_new_tokens must be positive after default resolution"); } From eacfa3527cd5c9925383b5e5189fbde1bcac060e Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Tue, 21 Jul 2026 11:01:18 -0400 Subject: [PATCH 22/27] Support Fish Audio Vulkan AR cache path --- src/models/fish_audio/ar.cpp | 45 ++++++++++++++++++++++-------------- 1 file changed, 28 insertions(+), 17 deletions(-) diff --git a/src/models/fish_audio/ar.cpp b/src/models/fish_audio/ar.cpp index 208b0de6..945495fc 100644 --- a/src/models/fish_audio/ar.cpp +++ b/src/models/fish_audio/ar.cpp @@ -123,8 +123,11 @@ struct FishPrefillCacheTarget { std::vector values; }; -modules::QwenDecoderActivationCastPolicy fish_activation_cast_policy() { +modules::QwenDecoderActivationCastPolicy fish_activation_cast_policy(core::BackendType backend_type) { modules::QwenDecoderActivationCastPolicy policy; + if (backend_type == core::BackendType::Vulkan) { + return policy; + } policy.enabled = true; policy.type = GGML_TYPE_BF16; policy.after_input_norm = true; @@ -143,7 +146,9 @@ modules::QwenDecoderActivationCastPolicy fish_activation_cast_policy() { return policy; } -modules::QwenCausalDecoderConfig make_slow_decoder_config(const FishAudioTextConfig & config) { +modules::QwenCausalDecoderConfig make_slow_decoder_config( + const FishAudioTextConfig & config, + core::BackendType backend_type) { modules::QwenCausalDecoderConfig out; out.stack.hidden_size = config.dim; out.stack.num_attention_heads = config.n_head; @@ -157,7 +162,7 @@ modules::QwenCausalDecoderConfig make_slow_decoder_config(const FishAudioTextCon out.stack.attention_precision = GGML_PREC_F32; out.stack.qkv_layout = modules::QwenDecoderQKVLayout::PackedQKV; out.stack.use_qk_norm = config.attention_qk_norm; - out.stack.activation_cast = fish_activation_cast_policy(); + out.stack.activation_cast = fish_activation_cast_policy(backend_type); out.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; out.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; out.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; @@ -169,7 +174,9 @@ modules::QwenCausalDecoderConfig make_slow_decoder_config(const FishAudioTextCon return out; } -modules::QwenCausalDecoderConfig make_fast_decoder_config(const FishAudioFastConfig & config) { +modules::QwenCausalDecoderConfig make_fast_decoder_config( + const FishAudioFastConfig & config, + core::BackendType backend_type) { modules::QwenCausalDecoderConfig out; out.stack.hidden_size = config.dim; out.stack.num_attention_heads = config.n_head; @@ -183,7 +190,7 @@ modules::QwenCausalDecoderConfig make_fast_decoder_config(const FishAudioFastCon out.stack.attention_precision = GGML_PREC_F32; out.stack.qkv_layout = modules::QwenDecoderQKVLayout::PackedQKV; out.stack.use_qk_norm = config.attention_qk_norm; - out.stack.activation_cast = fish_activation_cast_policy(); + out.stack.activation_cast = fish_activation_cast_policy(backend_type); out.stack.runtime.attention.prefill_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; out.stack.runtime.attention.static_mode = modules::QwenDecoderAttentionMode::FlashGroupedViewKV; out.stack.runtime.static_cache.update_mode = modules::QwenDecoderStaticCacheUpdateMode::DirectSetRows; @@ -687,7 +694,7 @@ FishStaticDecoderOutputs build_fish_static_decoder( .build(ctx, hidden, weights.lm_head); auto fast_hidden = norm_fastlayer_input ? hidden : x; runtime::TransformerKVCacheOptions cache_options; - cache_options.allow_bf16_storage = true; + cache_options.allow_bf16_storage = !cache_keys.empty() && cache_keys.front().type == GGML_TYPE_BF16; return { fast_hidden, logits, @@ -928,7 +935,7 @@ class FishAudioARRuntime::Impl { input, positions_value, bind_slow_weights(*constants_, runtime_->weights(), config), - make_slow_decoder_config(config), + make_slow_decoder_config(config, runtime_->backend_type()), assets.config.norm_fastlayer_input); graph_ = ggml_new_graph_custom(ctx_.get(), 65536, false); for (size_t layer_index = 0; layer_index < decoder.state.layers.size(); ++layer_index) { @@ -1050,27 +1057,29 @@ class FishAudioARRuntime::Impl { std::vector cache_values; cache_keys.reserve(runtime_->weights().slow_layers.size()); cache_values.reserve(runtime_->weights().slow_layers.size()); + const ggml_type cache_type = + runtime_->backend_type() == core::BackendType::Vulkan ? GGML_TYPE_F32 : GGML_TYPE_BF16; for (size_t layer = 0; layer < runtime_->weights().slow_layers.size(); ++layer) { cache_keys.push_back(core::wrap_tensor( ggml_new_tensor_4d( state_ctx_.get(), - GGML_TYPE_BF16, + cache_type, config.head_dim, config.n_local_heads, cache_steps_, 1), core::TensorShape::from_dims({1, cache_steps_, config.n_local_heads, config.head_dim}), - GGML_TYPE_BF16)); + cache_type)); cache_values.push_back(core::wrap_tensor( ggml_new_tensor_4d( state_ctx_.get(), - GGML_TYPE_BF16, + cache_type, config.head_dim, config.n_local_heads, cache_steps_, 1), core::TensorShape::from_dims({1, cache_steps_, config.n_local_heads, config.head_dim}), - GGML_TYPE_BF16)); + cache_type)); } state_buffer_ = ggml_backend_alloc_ctx_tensors(state_ctx_.get(), runtime_->backend()); if (state_buffer_ == nullptr) { @@ -1092,7 +1101,7 @@ class FishAudioARRuntime::Impl { input, position_value, bind_slow_weights(constants, runtime_->weights(), config), - make_slow_decoder_config(config), + make_slow_decoder_config(config, runtime_->backend_type()), cache_steps_, mask_value, cache_slot_value, @@ -1237,27 +1246,29 @@ class FishAudioARRuntime::Impl { std::vector cache_values; cache_keys.reserve(weights.fast_layers.size()); cache_values.reserve(weights.fast_layers.size()); + const ggml_type cache_type = + runtime_->backend_type() == core::BackendType::Vulkan ? GGML_TYPE_F32 : GGML_TYPE_BF16; for (size_t layer = 0; layer < weights.fast_layers.size(); ++layer) { cache_keys.push_back(core::wrap_tensor( ggml_new_tensor_4d( state_ctx_.get(), - GGML_TYPE_BF16, + cache_type, config.head_dim, config.n_local_heads, config.num_codebooks, 1), core::TensorShape::from_dims({1, config.num_codebooks, config.n_local_heads, config.head_dim}), - GGML_TYPE_BF16)); + cache_type)); cache_values.push_back(core::wrap_tensor( ggml_new_tensor_4d( state_ctx_.get(), - GGML_TYPE_BF16, + cache_type, config.head_dim, config.n_local_heads, config.num_codebooks, 1), core::TensorShape::from_dims({1, config.num_codebooks, config.n_local_heads, config.head_dim}), - GGML_TYPE_BF16)); + cache_type)); } state_buffer_ = ggml_backend_alloc_ctx_tensors(state_ctx_.get(), runtime_->backend()); if (state_buffer_ == nullptr) { @@ -1293,7 +1304,7 @@ class FishAudioARRuntime::Impl { input, position_value, decoder_weights, - make_fast_decoder_config(config), + make_fast_decoder_config(config, runtime_->backend_type()), config.num_codebooks, mask_value, position_value, From 70463388dd4372b44efc537288fa9559817e30f4 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Tue, 21 Jul 2026 12:03:14 -0400 Subject: [PATCH 23/27] Update GGUF model manager packages --- README.md | 12 ++-- docs/gguf.md | 2 + docs/tts.md | 59 ++++++++++++++++++- src/framework/runtime/registry.cpp | 9 --- tools/model_manager.py | 94 ++++++++++++++---------------- 5 files changed, 107 insertions(+), 69 deletions(-) diff --git a/README.md b/README.md index bbe04188..851f37ab 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,7 @@ audio.cpp would not be moving this quickly without generous contributors bringin | **ace_step** | music generation, music editing | 50+ langs | ACE-Step 1.5 Turbo and Base with acestep-5Hz-lm-1.7B | | **chatterbox** | TTS, voice cloning, voice conversion | ar, da, de, el, en, es, fi, fr, hi, it, ko, ms, nl, no, pl, pt, sv, sw, tr | Chatterbox with 0.5B backbone | | **citrinet_asr** | ASR | en | Citrinet-256 | +| **fish_audio** | TTS, voice cloning | auto, en, zh | Fish Audio S2 Pro | | **heartmula** | music generation | zh, en, ja, ko, es | HeartMuLa-oss-3B with HeartCodec-oss | | **higgs_audio_stt** | ASR | en | Higgs Audio v3 STT | | **higgs_audio_tts** | TTS, voice cloning | auto | Higgs Audio v3 TTS 4B | @@ -91,8 +92,6 @@ Community model ports live under `community_models` to make the ownership bounda | **outetts** | TTS, voice cloning | en, ar, zh, nl, fr, de, it, ja, ko, lt, ru, es, pt, be, bn, ka, hu, lv, fa, pl, sw, ta, uk | Mirek [@mirek190](https://github.com/mirek190) | Llama-OuteTTS-1.0-1B TTS and voice cloning support | | **vietneu_tts** | TTS, voice cloning | vi, en | Phuoc [@phuocnguyen90](https://github.com/phuocnguyen90) | [VieNeu-TTS-v3-Turbo](vietneu_tts.md) TTS and voice cloning support | -WIP: Fish Audio S2 Pro. Parked loaders that are not registered in this tree keep their catalog entries as `UnsupportedSource`; see [docs/maintainers/loader_and_catalog.md](docs/maintainers/loader_and_catalog.md). - PocketTTS language selection is a model-load option. When the model path points at the PocketTTS root, the loader uses `english` unless you pass `--load-option language=`. Kyutai's normal non-English PocketTTS releases are smaller distilled language models intended for the fast PocketTTS path. The `_24l` variants are larger 24-layer, undistilled preview models that can sound better but are slower. Kyutai currently publishes French only as `french_24l`, not as a normal distilled `french` language directory, so French is not listed as a normal PocketTTS language here. ## Docker @@ -374,20 +373,22 @@ Recommended top-level install packages: `Yes` means Hugging Face has a ready-to-use repo that the framework can download as-is. `No` means the tool must assemble, convert, or post-process files before the framework can use them. Packages whose loaders are not registered in this release tree are listed as **Unavailable** (see [docs/maintainers/loader_and_catalog.md](docs/maintainers/loader_and_catalog.md)). +For shared audio.cpp GGUF packages, the model manager installs the default Q8_0 GGUF. Other precision variants can be downloaded directly from [audio-cpp/audio.cpp-gguf](https://huggingface.co/audio-cpp/audio.cpp-gguf); see [docs/gguf.md](docs/gguf.md) for GGUF support status. + | Package id | Model | HF ready-to-use repo | |---|---|---| | `ace_step` | ACE-Step 1.5 Turbo/Base | No | | `chatterbox` | Chatterbox | **Yes** | | `citrinet_asr` | Citrinet ASR converted layout | No | +| `fish_audio_s2_pro` | Fish Audio S2 Pro GGUF Q8_0 | **Yes** | | `heartmula` | HeartMuLa | No | | `higgs_audio_stt` | Higgs Audio STT | No | -| `higgs_audio_v3_tts_4b` | Higgs Audio v3 TTS 4B | **Yes** | +| `higgs_audio_v3_tts_4b` | Higgs Audio v3 TTS 4B GGUF Q8_0 | **Yes** | | `htdemucs` | HTDemucs | No | | `hviske_asr` | Hviske ASR | **Yes** | | `irodori_tts_500m_v3` | Irodori-TTS 500M v3 | No | | `irodori_tts_600m_v3_voice_design` | Irodori-TTS 600M v3 VoiceDesign | No | | `index_tts2` | IndexTTS-2 | **Yes** | -| `kokoro_82m_bf16` | Kokoro 82M bf16 | Unavailable (loader not in this tree) | | `marblenet_vad` | MarbleNet VAD converted layout | No | | `mel_band_roformer` | Mel-Band RoFormer MLX | **Yes** | | `miocodec_25hz_44k_v2` | MioCodec 25Hz 44.1kHz v2 | No | @@ -400,7 +401,6 @@ Recommended top-level install packages: | `nemotron_asr` | Nemotron ASR | **Yes** | | `omnivoice` | OmniVoice | **Yes** | | `outetts_1_0_1b` | OuteTTS 1.0 1B with IBM DAC codec and Qwen3-aligned voice cloning | No | -| `parakeet_tdt_0_6b_v3` | Parakeet TDT 0.6B v3 | Unavailable (loader not in this tree) | | `pocket_tts` | PocketTTS | **Yes** | | `qwen3_asr_0_6b` | Qwen3 ASR 0.6B | **Yes** | | `qwen3_asr_1_7b_hf` | Qwen3 ASR 1.7B HF | **Yes** | @@ -613,7 +613,6 @@ For TTS-family models, the measured one-shot RTF is: | model | audio len (s) | wall time (s) | RTF | x faster than real time | |---|---:|---:|---:|---:| | chatterbox | 9.72 | 2.45 | 0.252 | 3.97x | -| kokoro tts | 10.15 | 0.64 | 0.063 | 15.90x | | miotts | 20.40 | 3.30 | 0.162 | 6.18x | | moss_tts_local | 9.60 | 0.97 | 0.101 | 9.91x | | omnivoice | 9.00 | 1.32 | 0.146 | 6.84x | @@ -628,7 +627,6 @@ For long-form TTS tests, each run uses the same 6,026-character, 1,028-word inpu | model | audio len (s) | wall time (s) | RTF | x faster than real time | |---|---:|---:|---:|---:| | chatterbox | 391.24 | 58.57 | 0.150 | 6.68x | -| kokoro tts | 371.17 | 7.19 | 0.019 | 51.60x | | index tts2 | 422.12 | 139.95 | 0.332 | 3.02x | | miotts | 399.16 | 66.59 | 0.167 | 5.99x | | moss_tts_nano | 391.20 | 43.16 | 0.110 | 9.06x | diff --git a/docs/gguf.md b/docs/gguf.md index 68911b34..3b9474c8 100644 --- a/docs/gguf.md +++ b/docs/gguf.md @@ -259,8 +259,10 @@ Status labels: | `ace_step` | No | --- | --- | --- | --- | | `chatterbox` | No | --- | --- | --- | --- | | `citrinet_asr` | Done | Pass | --- | --- | Pass | +| `fish_audio` | Done | Pass | --- | Pass | Pass | | `heartmula` | No | --- | --- | --- | --- | | `higgs_audio_stt` | Done | Pass | --- | Pass | Pass | +| `higgs_audio_tts` | Done | Pass | --- | Pass | Pass | | `htdemucs` | Done | Pass | --- | Pass | Pass (drift) | | `hviske_asr` | Done | Pass | --- | --- | Pass | | `index_tts2` | Done | Pass | Pass | Pass (drift) | No (similarity drift, frame drift, text minor drift) | diff --git a/docs/tts.md b/docs/tts.md index 14e06c65..93d4338e 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -11,6 +11,7 @@ | PocketTTS | `pocket_tts` | `tts` | [PocketTTS](#pockettts) | | VoxCPM2 | `voxcpm2` | `tts`, `vdes` | [VoxCPM2](#voxcpm2) | | Higgs Audio v3 TTS | `higgs_audio_tts` | `tts` | [Higgs Audio v3 TTS](#higgs-audio-v3-tts) | +| Fish Audio S2 Pro | `fish_audio` | `tts` | [Fish Audio S2 Pro](#fish-audio-s2-pro) | | IndexTTS2 | `index_tts2` | `tts` | [IndexTTS2](#indextts2) | | Irodori-TTS | `irodori_tts` | `tts`, `vdes` | [Irodori-TTS](#irodori-tts) | | OuteTTS | `outetts` | `tts`, `clon` | [OuteTTS](#outetts) | @@ -305,7 +306,7 @@ Higgs Audio v3 TTS is a voice-clone TTS model. The current integration uses the | Field | Value | |---|---| | Family | `higgs_audio_tts` | -| Model directory | `models/higgs-audio-v3-tts-4b` | +| Model directory | `models/Higgs-Audio-v3-TTS-4B-GGUF` when installed through the model manager | | Task | `tts` | | Modes | `offline` | | Languages | Model auto-handles supported languages | @@ -313,7 +314,13 @@ Higgs Audio v3 TTS is a voice-clone TTS model. The current integration uses the | Built-in voices | Not exposed | ```bash -audiocpp_cli --task tts --family higgs_audio_tts --model models/higgs-audio-v3-tts-4b --backend cuda --text "Hello from Higgs Audio." --voice-ref assets/resources/b.wav --reference-text "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you." --out out.wav +audiocpp_cli --task tts --family higgs_audio_tts --model models/Higgs-Audio-v3-TTS-4B-GGUF --backend cuda --text "Hello from Higgs Audio." --voice-ref assets/resources/b.wav --reference-text "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you." --out out.wav +``` + +The model manager installs the Q8_0 standalone GGUF package by default: + +```bash +python3 tools/model_manager.py install --models-root models higgs_audio_v3_tts_4b ``` | Option | Values | Default | Meaning | @@ -327,6 +334,54 @@ audiocpp_cli --task tts --family higgs_audio_tts --model models/higgs-audio-v3-t | `--top-p` | float | `0.8` | AR nucleus sampling limit. The Python client's unfiltered equivalent is `1.0`. | | `--repetition-penalty` | float | `1.1` | Accepted for Python API compatibility; Higgs audio-code sampling does not consume it. | +## Fish Audio S2 Pro + +Fish Audio S2 Pro is a TTS and reference voice-clone model. The integration uses the framework text chunker for long-form input, caches prepared reference audio in the session, and supports GGUF loading through the package spec path. + +| Field | Value | +|---|---| +| Family | `fish_audio` | +| Model directory | `models/Fish-Audio-S2-Pro-GGUF` when installed through the model manager | +| Task | `tts` | +| Modes | `offline` | +| Languages | Model auto-handles language; tested paths cover English and Chinese-style prompts | +| Voice input | Optional reference WAV through `--voice-ref`; transcript through `--reference-text` when known | +| Built-in voices | Not exposed | + +Text-to-speech: + +```bash +audiocpp_cli --task tts --family fish_audio --model models/Fish-Audio-S2-Pro-GGUF --backend cuda --text "Hello from Fish Audio." --out out.wav +``` + +Reference voice clone: + +```bash +audiocpp_cli --task tts --family fish_audio --model models/Fish-Audio-S2-Pro-GGUF --backend cuda --text "The final render is ready for review." --voice-ref assets/resources/b.wav --reference-text "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you." --out out.wav +``` + +The model manager installs the Q8_0 standalone GGUF package by default: + +```bash +python3 tools/model_manager.py install --models-root models fish_audio_s2_pro +``` + +| Option | Values | Default | Meaning | +|---|---|---:|---| +| `--voice-ref` | WAV path | not set | Reference speaker audio for voice cloning. | +| `--reference-text` | text | empty string | Transcript for reference audio. | +| `--max-new-tokens` | integer | `1024` | Maximum generated semantic tokens per chunk. `0` uses the default. | +| `--text-chunk-size` | integer chars | `200` | Long-form chunk size. | +| `--text-chunk-mode` | `default`, `tag_aware`, `japanese`, `endline` | `default` | Framework text chunking mode. | +| `--temperature` | float | `0.8` | Sampling temperature. | +| `--top-k` | integer | `30` | Top-k sampling limit. | +| `--top-p` | float | `0.8` | Nucleus sampling limit. | +| `--seed` | integer | random when omitted | Sampling seed for reproducible output. | +| `--session-option fish_audio.mem_saver=true|false` | bool | `false` | Release cached AR runtime graphs after each request. | +| `--session-option fish_audio.reference_cache_slots=` | integer | `1` | Prepared reference-audio cache slots. | +| `--session-option fish_audio.weight_type=` | `native`, `f32`, `f16`, `bf16`, `q8_0` | `native` | AR matmul weight storage type. | +| `--session-option fish_audio.codec_weight_type=` | `native`, `f32`, `f16`, `q8_0` | `native` | Codec conv/matmul weight storage type. | + ## IndexTTS2 IndexTTS2 is a Chinese and English TTS model with voice cloning and expressive emotion controls. It requires a speaker reference through the framework `--voice-ref` path. diff --git a/src/framework/runtime/registry.cpp b/src/framework/runtime/registry.cpp index b6fb1e98..571f3de5 100644 --- a/src/framework/runtime/registry.cpp +++ b/src/framework/runtime/registry.cpp @@ -4,11 +4,6 @@ #include "engine/framework/assets/model_package.h" #include "engine/framework/io/config.h" #include "engine/framework/io/filesystem.h" -// Parked loaders (sources not in this release tree). When commenting these out, -// also mark matching ModelPackage entries UnsupportedSource — see -// docs/maintainers/loader_and_catalog.md and tools/check_loader_catalog_sync.py. -// #include "engine/models/kokoro_tts/loader.h" -// #include "engine/models/parakeet_tdt/loader.h" #include "engine/models/ace_step/loader.h" #include "engine/models/chatterbox/loader.h" #include "engine/models/citrinet_asr/session.h" @@ -245,10 +240,6 @@ ModelRegistry make_registry_from_config( ModelRegistry make_default_registry(const std::optional & config_path) { const std::vector> available_loaders = { - // Parked loaders — keep catalog packages UnsupportedSource while these stay commented. - // See docs/maintainers/loader_and_catalog.md. - // engine::models::kokoro_tts::make_kokoro_tts_loader(), - // engine::models::parakeet_tdt::make_parakeet_tdt_loader(), engine::models::ace_step::make_ace_step_loader(), engine::models::demucs::make_htdemucs_loader(), engine::models::roformer::make_mel_band_roformer_loader(), diff --git a/tools/model_manager.py b/tools/model_manager.py index 1fd998c3..ed1e20b8 100644 --- a/tools/model_manager.py +++ b/tools/model_manager.py @@ -107,6 +107,7 @@ class SnapshotSource: include_prefixes: tuple[str, ...] = () include_suffixes: tuple[str, ...] = () exclude_prefixes: tuple[str, ...] = () + strip_prefix: str = "" @dataclasses.dataclass(frozen=True) @@ -216,22 +217,6 @@ def package_usage_examples(package: ModelPackage) -> list[str]: "vae/diffusion_pytorch_model.safetensors", ), ), - ModelPackage( - id="kokoro_82m_bf16", - display_name="Kokoro 82M bf16", - target_directory="Kokoro-82M-bf16", - source=UnsupportedSource( - reason=( - "kokoro_tts loader is not registered in this release tree yet " - "(commented out in src/framework/runtime/registry.cpp). " - "Re-enable the loader, add model_specs/kokoro_tts.json, then " - "restore a SnapshotSource here." - ), - ), - required_files=("config.json", "kokoro-v1_0.safetensors", "voices/af_heart.safetensors"), - family="kokoro_tts", - tasks=("tts",), - ), ModelPackage( id="moss_tts_nano_100m", display_name="MOSS-TTS-Nano 100M", @@ -421,6 +406,20 @@ def package_usage_examples(package: ModelPackage) -> list[str]: tasks=("asr",), description="Native Hugging Face checkpoint for Voxtral realtime ASR; no conversion is required.", ), + ModelPackage( + id="fish_audio_s2_pro", + display_name="Fish Audio S2 Pro GGUF", + target_directory="Fish-Audio-S2-Pro-GGUF", + source=SnapshotSource( + repo_id="audio-cpp/audio.cpp-gguf", + include_prefixes=("Fish-Audio-S2-Pro-GGUF/fish-audio-s2-pro-q8_0.gguf",), + strip_prefix="Fish-Audio-S2-Pro-GGUF/", + ), + required_files=("fish-audio-s2-pro-q8_0.gguf",), + family="fish_audio", + tasks=("tts",), + description="Standalone audio.cpp Q8_0 GGUF package for Fish Audio S2 Pro.", + ), ModelPackage( id="higgs_audio_stt", display_name="Higgs Audio STT", @@ -605,22 +604,6 @@ def package_usage_examples(package: ModelPackage) -> list[str]: source=SnapshotSource(repo_id="nvidia/diar_sortformer_4spk-v1"), required_files=("config.json", "model.safetensors", "processor_config.json"), ), - ModelPackage( - id="parakeet_tdt_0_6b_v3", - display_name="Parakeet TDT 0.6B v3", - target_directory="parakeet-tdt-0.6b-v3", - source=UnsupportedSource( - reason=( - "parakeet_tdt loader is not registered in this release tree yet " - "(commented out in src/framework/runtime/registry.cpp). " - "Re-enable the loader, add model_specs/parakeet_tdt.json, then " - "restore a SnapshotSource here." - ), - ), - required_files=("config.json", "model.safetensors", "processor_config.json", "tokenizer.json"), - family="parakeet_tdt", - tasks=("asr",), - ), ModelPackage( id="pocket_tts", display_name="PocketTTS", @@ -844,19 +827,17 @@ def package_usage_examples(package: ModelPackage) -> list[str]: ), ModelPackage( id="higgs_audio_v3_tts_4b", - display_name="Higgs Audio v3 TTS 4B", - target_directory="higgs-audio-v3-tts-4b", - source=SnapshotSource(repo_id="bosonai/higgs-audio-v3-tts-4b"), - required_files=( - "chat_template.jinja", - "config.json", - "model.safetensors.index.json", - "model.safetensors", - "tokenizer.json", - "tokenizer_config.json", + display_name="Higgs Audio v3 TTS 4B GGUF", + target_directory="Higgs-Audio-v3-TTS-4B-GGUF", + source=SnapshotSource( + repo_id="audio-cpp/audio.cpp-gguf", + include_prefixes=("Higgs-Audio-v3-TTS-4B-GGUF/higgs-audio-v3-tts-4b-q8_0.gguf",), + strip_prefix="Higgs-Audio-v3-TTS-4B-GGUF/", ), + required_files=("higgs-audio-v3-tts-4b-q8_0.gguf",), family="higgs_audio_tts", tasks=("tts",), + description="Standalone audio.cpp Q8_0 GGUF package for Higgs Audio v3 TTS 4B.", ), ModelPackage( id="heartmula", @@ -1395,7 +1376,7 @@ def _default_tasks_from_family(family: str) -> list[str]: return [] if "forced_aligner" in key or key.endswith("_aligner") or key.endswith("_align"): return ["align"] - if key.endswith("_asr") or key.endswith("_stt") or key in {"parakeet_tdt", "whisper", "voxtral_realtime"}: + if key.endswith("_asr") or key.endswith("_stt") or key in {"whisper", "voxtral_realtime"}: return ["asr"] if "vad" in key: return ["vad"] @@ -1412,8 +1393,6 @@ def _default_tasks_from_family(family: str) -> list[str]: if key.endswith("_asr") or key.endswith("_stt"): return ["asr"] if "tts" in key or key in { - "kokoro", - "kokoro_tts", "chatterbox", "voxcpm2", "omnivoice", @@ -1476,6 +1455,7 @@ def package_payload(package: ModelPackage) -> dict[str, object]: "include_prefixes": list(source.include_prefixes), "include_suffixes": list(source.include_suffixes), "exclude_prefixes": list(source.exclude_prefixes), + "strip_prefix": source.strip_prefix, } installable = True elif isinstance(source, CompositeSnapshotSource): @@ -1490,6 +1470,7 @@ def package_payload(package: ModelPackage) -> dict[str, object]: "include_prefixes": list(placement.source.include_prefixes), "include_suffixes": list(placement.source.include_suffixes), "exclude_prefixes": list(placement.source.exclude_prefixes), + "strip_prefix": placement.source.strip_prefix, } for placement in source.placements ], @@ -1569,11 +1550,22 @@ def http_json(url: str) -> object: return json.load(response) -def list_hf_files(source: SnapshotSource) -> list[tuple[str, int | None]]: +def local_snapshot_path(source: SnapshotSource, remote_path: str) -> str: + if not source.strip_prefix: + return remote_path + if not remote_path.startswith(source.strip_prefix): + raise RuntimeError(f"snapshot path does not start with strip_prefix: {remote_path}") + local_path = remote_path[len(source.strip_prefix):] + if not local_path: + raise RuntimeError(f"snapshot strip_prefix removed full path: {remote_path}") + return local_path + + +def list_hf_files(source: SnapshotSource) -> list[tuple[str, str, int | None]]: payload = http_json(hf_tree_url(source)) if not isinstance(payload, list): raise RuntimeError(f"unexpected HuggingFace tree payload for {source.repo_id}") - files: list[tuple[str, int | None]] = [] + files: list[tuple[str, str, int | None]] = [] for entry in payload: if not isinstance(entry, dict): continue @@ -1588,7 +1580,7 @@ def list_hf_files(source: SnapshotSource) -> list[tuple[str, int | None]]: if any(path.startswith(prefix) for prefix in source.exclude_prefixes): continue size = entry.get("size") - files.append((path, size if isinstance(size, int) else None)) + files.append((path, local_snapshot_path(source, path), size if isinstance(size, int) else None)) if not files: raise RuntimeError(f"no installable files found for {source.repo_id}") return files @@ -1639,11 +1631,11 @@ def install_snapshot_into_dir( validate: bool = True, ) -> None: files = list_hf_files(source) - for relative, expected_size in files: + for remote, relative, expected_size in files: destination = destination_root / relative destination.parent.mkdir(parents=True, exist_ok=True) - print(f"download {relative}") - download_file(hf_resolve_url(source, relative), destination, expected_size) + print(f"download {remote}") + download_file(hf_resolve_url(source, remote), destination, expected_size) if validate: validate_required_files_list(required_files, destination_root, source.repo_id) From 91ce2ea8e5c516ca22cd2f68565709657d938f03 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Tue, 21 Jul 2026 12:14:06 -0400 Subject: [PATCH 24/27] Clarify GGUF directory load errors --- src/framework/assets/model_package.cpp | 38 ++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/src/framework/assets/model_package.cpp b/src/framework/assets/model_package.cpp index d93b5819..f025ff60 100644 --- a/src/framework/assets/model_package.cpp +++ b/src/framework/assets/model_package.cpp @@ -56,6 +56,41 @@ bool is_gguf_file(const std::filesystem::path & path) { return extension == ".gguf"; } +std::vector directory_gguf_files(const std::filesystem::path & path) { + std::vector files; + if (!engine::io::is_existing_directory(path)) { + return files; + } + for (const auto & entry : std::filesystem::directory_iterator(path)) { + const auto candidate = entry.path(); + if (is_gguf_file(candidate)) { + files.push_back(candidate.filename().string()); + } + } + std::sort(files.begin(), files.end()); + return files; +} + +std::string directory_gguf_hint(std::string_view family) { + if (!active_model_path.has_value()) { + return {}; + } + const auto files = directory_gguf_files(*active_model_path); + if (files.empty()) { + return {}; + } + std::string message = "model directory has no default GGUF for family '" + std::string(family) + + "': " + active_model_path->string() + "; found: "; + for (size_t i = 0; i < files.size(); ++i) { + if (i != 0) { + message += ", "; + } + message += files[i]; + } + message += "; pass the GGUF file directly with --model, or rename it to model.gguf"; + return message; +} + std::optional active_gguf_path() { if (!active_model_path.has_value()) return std::nullopt; @@ -378,6 +413,9 @@ std::filesystem::path default_model_package_spec_path(std::string_view family) { if (const auto external = discover_external_model_spec(family)) { return *external; } + if (const auto hint = directory_gguf_hint(family); !hint.empty()) { + throw std::runtime_error(hint); + } throw std::runtime_error("model package spec not found for family '" + std::string(family) + "' (provide --model-spec-override, embed it in the GGUF, enable " "AUDIOCPP_DEPLOYMENT_BUILD, or install model_specs/" + From 0d55861e3de9cc9333d6c36a3939790fd0fb1401 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Tue, 21 Jul 2026 12:31:42 -0400 Subject: [PATCH 25/27] Align new model CLI catalog interfaces --- README.md | 2 +- docs/asr.md | 12 ++++++------ docs/tts.md | 10 +++++----- src/models/fish_audio/loader.cpp | 10 ++++++++++ tools/model_manager.py | 31 ++++++++++--------------------- 5 files changed, 32 insertions(+), 33 deletions(-) diff --git a/README.md b/README.md index 851f37ab..b980d19d 100644 --- a/README.md +++ b/README.md @@ -420,7 +420,7 @@ For shared audio.cpp GGUF packages, the model manager installs the default Q8_0 | `vibevoice_1_5b` | VibeVoice 1.5B | No | | `vibevoice_7b` | VibeVoice 7B | No | | `vibevoice_asr` | VibeVoice ASR | No | -| `voxtral_realtime` | Voxtral Mini 4B Realtime | **Yes** | +| `voxtral_realtime` | Voxtral Mini 4B Realtime GGUF Q8_0 | **Yes** | | `voxcpm2` | VoxCPM2 | No | > [!WARNING] diff --git a/docs/asr.md b/docs/asr.md index e08bcb52..65da80fd 100644 --- a/docs/asr.md +++ b/docs/asr.md @@ -255,12 +255,12 @@ audiocpp_cli --task asr --family vibevoice_asr --model models/VibeVoice-ASR --ba ## Voxtral Realtime -Voxtral Realtime is a Mistral realtime ASR model with offline and streaming sessions. It accepts the native Hugging Face model directory and standalone audio.cpp GGUF packages. +Voxtral Realtime is a Mistral realtime ASR model with offline and streaming sessions. The model manager installs the Q8_0 standalone GGUF package by default; native Hugging Face directories and other standalone GGUF variants can also be used when provided directly. | Field | Value | |---|---| | Family | `voxtral_realtime` | -| Model directory | `models/Voxtral-Mini-4B-Realtime-2602` or a standalone Voxtral GGUF | +| Model path | `models/Voxtral-Mini-4B-Realtime-2602-GGUF/voxtral-mini-4b-realtime-2602-q8_0.gguf` when installed through the model manager | | Task | `asr` | | Modes | `offline`, `streaming` | | Output | Transcription text | @@ -270,19 +270,19 @@ Voxtral Realtime is a Mistral realtime ASR model with offline and streaming sess Offline CLI: ```bash -audiocpp_cli --task asr --family voxtral_realtime --model --backend cuda --threads 8 --audio assets/resources/sample.wav --text-out transcript.txt +audiocpp_cli --task asr --family voxtral_realtime --model models/Voxtral-Mini-4B-Realtime-2602-GGUF/voxtral-mini-4b-realtime-2602-q8_0.gguf --backend cuda --threads 8 --audio assets/resources/sample.wav --text-out transcript.txt ``` Sampling and token-cap options can be passed through request options: ```bash -audiocpp_cli --task asr --family voxtral_realtime --model --backend cuda --threads 8 --audio assets/resources/sample.wav --text-out transcript.txt --request-option max_new_tokens=256 --do-sample false --temperature 1.0 --top-p 1.0 --top-k 50 --seed 1234 +audiocpp_cli --task asr --family voxtral_realtime --model models/Voxtral-Mini-4B-Realtime-2602-GGUF/voxtral-mini-4b-realtime-2602-q8_0.gguf --backend cuda --threads 8 --audio assets/resources/sample.wav --text-out transcript.txt --request-option max_new_tokens=256 --do-sample false --temperature 1.0 --top-p 1.0 --top-k 50 --seed 1234 ``` Streaming CLI: ```bash -audiocpp_cli --task asr --family voxtral_realtime --model --backend cuda --threads 8 --mode streaming --audio assets/resources/sample.wav --text-out transcript.txt +audiocpp_cli --task asr --family voxtral_realtime --model models/Voxtral-Mini-4B-Realtime-2602-GGUF/voxtral-mini-4b-realtime-2602-q8_0.gguf --backend cuda --threads 8 --mode streaming --audio assets/resources/sample.wav --text-out transcript.txt ``` Streaming server config: @@ -299,7 +299,7 @@ Streaming server config: { "id": "voxtral-stream", "family": "voxtral_realtime", - "path": "/path/to/Voxtral-Mini-4B-Realtime-2602", + "path": "/path/to/voxtral-mini-4b-realtime-2602-q8_0.gguf", "task": "asr", "mode": "streaming" } diff --git a/docs/tts.md b/docs/tts.md index 93d4338e..8f350f45 100644 --- a/docs/tts.md +++ b/docs/tts.md @@ -306,7 +306,7 @@ Higgs Audio v3 TTS is a voice-clone TTS model. The current integration uses the | Field | Value | |---|---| | Family | `higgs_audio_tts` | -| Model directory | `models/Higgs-Audio-v3-TTS-4B-GGUF` when installed through the model manager | +| Model path | `models/Higgs-Audio-v3-TTS-4B-GGUF/higgs-audio-v3-tts-4b-q8_0.gguf` when installed through the model manager | | Task | `tts` | | Modes | `offline` | | Languages | Model auto-handles supported languages | @@ -314,7 +314,7 @@ Higgs Audio v3 TTS is a voice-clone TTS model. The current integration uses the | Built-in voices | Not exposed | ```bash -audiocpp_cli --task tts --family higgs_audio_tts --model models/Higgs-Audio-v3-TTS-4B-GGUF --backend cuda --text "Hello from Higgs Audio." --voice-ref assets/resources/b.wav --reference-text "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you." --out out.wav +audiocpp_cli --task tts --family higgs_audio_tts --model models/Higgs-Audio-v3-TTS-4B-GGUF/higgs-audio-v3-tts-4b-q8_0.gguf --backend cuda --text "Hello from Higgs Audio." --voice-ref assets/resources/b.wav --reference-text "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you." --out out.wav ``` The model manager installs the Q8_0 standalone GGUF package by default: @@ -341,7 +341,7 @@ Fish Audio S2 Pro is a TTS and reference voice-clone model. The integration uses | Field | Value | |---|---| | Family | `fish_audio` | -| Model directory | `models/Fish-Audio-S2-Pro-GGUF` when installed through the model manager | +| Model path | `models/Fish-Audio-S2-Pro-GGUF/fish-audio-s2-pro-q8_0.gguf` when installed through the model manager | | Task | `tts` | | Modes | `offline` | | Languages | Model auto-handles language; tested paths cover English and Chinese-style prompts | @@ -351,13 +351,13 @@ Fish Audio S2 Pro is a TTS and reference voice-clone model. The integration uses Text-to-speech: ```bash -audiocpp_cli --task tts --family fish_audio --model models/Fish-Audio-S2-Pro-GGUF --backend cuda --text "Hello from Fish Audio." --out out.wav +audiocpp_cli --task tts --family fish_audio --model models/Fish-Audio-S2-Pro-GGUF/fish-audio-s2-pro-q8_0.gguf --backend cuda --text "Hello from Fish Audio." --out out.wav ``` Reference voice clone: ```bash -audiocpp_cli --task tts --family fish_audio --model models/Fish-Audio-S2-Pro-GGUF --backend cuda --text "The final render is ready for review." --voice-ref assets/resources/b.wav --reference-text "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you." --out out.wav +audiocpp_cli --task tts --family fish_audio --model models/Fish-Audio-S2-Pro-GGUF/fish-audio-s2-pro-q8_0.gguf --backend cuda --text "The final render is ready for review." --voice-ref assets/resources/b.wav --reference-text "Some call me nature. Others call me Mother Nature. I've been here for over 4.5 billion years. 22,500 times longer than you." --out out.wav ``` The model manager installs the Q8_0 standalone GGUF package by default: diff --git a/src/models/fish_audio/loader.cpp b/src/models/fish_audio/loader.cpp index a6c67653..f8f6a3d7 100644 --- a/src/models/fish_audio/loader.cpp +++ b/src/models/fish_audio/loader.cpp @@ -57,6 +57,16 @@ class FishAudioLoader final : public runtime::IVoiceModelLoader { return "fish_audio"; } + runtime::CapabilitySet advertised_capabilities() const override { + runtime::CapabilitySet out; + out.supported_tasks = { + {runtime::VoiceTaskKind::Tts, {runtime::RunMode::Offline}}, + }; + out.supports_speaker_reference = true; + out.supports_style_condition = true; + return out; + } + bool can_load(const runtime::ModelLoadRequest & request) const override { if (request.family_hint.has_value() && *request.family_hint != family()) { return false; diff --git a/tools/model_manager.py b/tools/model_manager.py index ed1e20b8..8afd7514 100644 --- a/tools/model_manager.py +++ b/tools/model_manager.py @@ -150,6 +150,7 @@ class ModelPackage: standalone: bool | None = None parent_package_id: str | None = None tasks: tuple[str, ...] = () + modes: tuple[str, ...] = () gated: bool | None = None @@ -381,30 +382,18 @@ def package_usage_examples(package: ModelPackage) -> list[str]: ), ModelPackage( id="voxtral_realtime", - display_name="Voxtral Mini 4B Realtime", - target_directory="Voxtral-Mini-4B-Realtime-2602", + display_name="Voxtral Mini 4B Realtime GGUF", + target_directory="Voxtral-Mini-4B-Realtime-2602-GGUF", source=SnapshotSource( - repo_id="mistralai/Voxtral-Mini-4B-Realtime-2602", - include_prefixes=( - "config.json", - "generation_config.json", - "model.safetensors", - "params.json", - "processor_config.json", - "tekken.json", - ), - ), - required_files=( - "config.json", - "generation_config.json", - "model.safetensors", - "params.json", - "processor_config.json", - "tekken.json", + repo_id="audio-cpp/audio.cpp-gguf", + include_prefixes=("Voxtral-Mini-4B-Realtime-2602-GGUF/voxtral-mini-4b-realtime-2602-q8_0.gguf",), + strip_prefix="Voxtral-Mini-4B-Realtime-2602-GGUF/", ), + required_files=("voxtral-mini-4b-realtime-2602-q8_0.gguf",), family="voxtral_realtime", tasks=("asr",), - description="Native Hugging Face checkpoint for Voxtral realtime ASR; no conversion is required.", + modes=("offline", "streaming"), + description="Standalone audio.cpp Q8_0 GGUF package for Voxtral realtime ASR.", ), ModelPackage( id="fish_audio_s2_pro", @@ -1512,7 +1501,7 @@ def package_payload(package: ModelPackage) -> dict[str, object]: "source": source_payload, "family": family, "tasks": tasks, - "modes": ["offline"] if tasks else [], + "modes": list(package.modes) if package.modes else (["offline"] if tasks else []), "standalone": standalone, "parent_package_id": parent_package_id, "gated": _package_is_gated(package), From 2f6aa9b6c4ee3bae7a985c29b21e3781faf9b437 Mon Sep 17 00:00:00 2001 From: 0xShug0 <231717474+0xShug0@users.noreply.github.com> Date: Tue, 21 Jul 2026 13:19:20 -0400 Subject: [PATCH 26/27] Support Fish Audio codec normalization on Metal --- src/models/fish_audio/codec.cpp | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/models/fish_audio/codec.cpp b/src/models/fish_audio/codec.cpp index b47150cb..f63c6978 100644 --- a/src/models/fish_audio/codec.cpp +++ b/src/models/fish_audio/codec.cpp @@ -320,12 +320,16 @@ core::TensorValue causal_conv_transpose1d( } core::TensorValue l2_normalize_last(core::ModuleBuildContext & ctx, const core::TensorValue & input) { - auto squared = modules::MulModule{}.build(ctx, input, input); + const bool materialize_input = ctx.backend_type == core::BackendType::Metal; + const auto normalized_input = materialize_input + ? core::ensure_backend_addressable_layout(ctx, input) + : input; + auto squared = modules::MulModule{}.build(ctx, normalized_input, normalized_input); auto sum = modules::ReduceSumModule({static_cast(input.shape.rank - 1)}).build(ctx, squared); auto shifted = core::wrap_tensor(ggml_scale_bias(ctx.ggml, sum.tensor, 1.0F, 1.0e-12F), sum.shape, GGML_TYPE_F32); auto denom = modules::SqrtModule{}.build(ctx, shifted); - auto repeated = modules::RepeatModule({input.shape}).build(ctx, denom); - return core::wrap_tensor(ggml_div(ctx.ggml, input.tensor, repeated.tensor), input.shape, GGML_TYPE_F32); + auto repeated = modules::RepeatModule({normalized_input.shape}).build(ctx, denom); + return core::wrap_tensor(ggml_div(ctx.ggml, normalized_input.tensor, repeated.tensor), normalized_input.shape, GGML_TYPE_F32); } core::TensorValue build_mlp( From 3b5345e3f467da742e3e98f294744a6746b8e406 Mon Sep 17 00:00:00 2001 From: RANGROO Date: Tue, 28 Jul 2026 22:27:22 -0700 Subject: [PATCH 27/27] Fix PocketTTS prompt graph overflow --- src/models/pocket_tts/flow_lm.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/models/pocket_tts/flow_lm.cpp b/src/models/pocket_tts/flow_lm.cpp index e5b5e5f1..21398831 100644 --- a/src/models/pocket_tts/flow_lm.cpp +++ b/src/models/pocket_tts/flow_lm.cpp @@ -14,6 +14,8 @@ namespace engine::models::pocket_tts { namespace { +constexpr size_t kPromptGraphNodeCapacity = 262144; + modules::TransformerEncoderBlockWeights make_transformer_layer_weights( core::ModuleBuildContext & ctx, const models::pocket_tts::PocketTTSBackendWeights & weights, @@ -383,7 +385,8 @@ class FlowLMStepRuntime { core::write_tensor_f32(attention_mask_, attention_mask_buffer_); core::set_backend_threads(backend_, threads_); if (prompt_steps_ > 0) { - prompt_graph_ = ggml_new_graph_custom(ggml_ctx_, 32768, false); + // Long or dense prompts add per-step KV transfer nodes beyond the default graph capacity. + prompt_graph_ = ggml_new_graph_custom(ggml_ctx_, kPromptGraphNodeCapacity, false); ggml_build_forward_expand(prompt_graph_, prompt_output_.tensor); for (size_t step = 0; step < prompt_step_key_sources_.size(); ++step) { for (size_t layer = 0; layer < prompt_step_key_sources_[step].size(); ++layer) {