Promote shared MOSS runtime pieces

CMakeLists.txt

@@ -394,6 +394,7 @@ add_library(engine_core OBJECT
     src/framework/modules/activation_modules.cpp
     src/framework/modules/norm_modules.cpp
     src/framework/modules/lookup_modules.cpp
+    src/framework/modules/multi_codebook_embedding.cpp
     src/framework/modules/structural_modules.cpp
     src/framework/modules/asr_helpers.cpp
     src/framework/modules/conv_modules.cpp
@@ -472,6 +473,7 @@ add_library(engine_core OBJECT
     src/framework/codecs/fish_dac_codec_runtime.cpp
     src/framework/codecs/mel_latent_vae44k_runtime.cpp
     src/framework/codecs/mimi_codec_runtime.cpp
+    src/framework/codecs/moss_audio_tokenizer_codec_runtime.cpp
     src/framework/codecs/neural_audio.cpp
     src/framework/codecs/redae_codec_runtime.cpp
     src/framework/conditioners/clap_audio_conditioner_runtime.cpp
@@ -799,12 +801,6 @@ audiocpp_add_model(moss
         src/models/moss/moss_tts_nano/prompt_builder.cpp
         src/models/moss/moss_tts_nano/session.cpp
         src/models/moss/moss_tts_nano/tokenizer_text.cpp
-        src/models/moss/shared/audio_tokenizer_decoder.cpp
-        src/models/moss/shared/audio_tokenizer_encoder.cpp
-        src/models/moss/shared/audio_tokenizer_config.cpp
-        src/models/moss/shared/audio_tokenizer_quantizer.cpp
-        src/models/moss/shared/sampling.cpp
-        src/models/moss/shared/token_rows.cpp
         src/models/moss/moss_tts_local/depth_transformer.cpp
         src/models/moss/moss_tts_local/generator.cpp
         src/models/moss/moss_tts_local/loader.cpp
@@ -1195,8 +1191,6 @@ audiocpp_add_model(vietneu_tts
         engine/community_models/vietneu_tts/loader.h
     LOADERS
         engine::models::vietneu_tts::make_vietneu_tts_loader
-    DEPENDS
-        moss
 )
 
 audiocpp_add_model(qwen3_asr

include/engine/community_models/moss_voicegen/delay_decoder.h

@@ -2,6 +2,7 @@
 
 #include "engine/community_models/moss_voicegen/assets.h"
 #include "engine/community_models/moss_voicegen/heads.h"
+#include "engine/framework/sampling/hf_sampler.h"
 
 #include <cstdint>
 #include <random>
@@ -68,8 +69,8 @@ class MossVoiceGenDelayDecoder {
     MossVoiceGenConfig config_;
     MossVoiceGenSamplingOptions sampling_;
     MossVoiceGenLengthBounds bounds_;
-    uint32_t seed_ = 0;
     std::mt19937 rng_;
+    engine::sampling::HfSamplerScratch sampler_scratch_;
     uint64_t sample_call_index_ = 0;
 
     int64_t step_index_ = 0;

include/engine/community_models/moss_voicegen/session.h

@@ -6,9 +6,9 @@
 #include "engine/community_models/moss_voicegen/heads.h"
 #include "engine/community_models/moss_voicegen/tokenizer_text.h"
 #include "engine/framework/core/execution_context.h"
+#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h"
+#include "engine/framework/modules/multi_codebook_embedding.h"
 #include "engine/framework/runtime/session_base.h"
-#include "engine/models/moss/shared/audio_tokenizer_decoder.h"
-#include "engine/models/moss/shared/token_rows.h"
 
 #include <cstddef>
 #include <memory>
@@ -63,10 +63,10 @@ class MossVoiceGenSession final
 
     // The execution context comes from RuntimeSessionBase; the runtimes below borrow it.
     std::unique_ptr<MossVoiceGenTextProcessor> text_processor_;
-    std::unique_ptr<moss::AudioCodebookEmbeddings> codebooks_;
+    std::unique_ptr<engine::modules::MultiCodebookEmbedding> codebooks_;
     std::unique_ptr<MossVoiceGenBackboneRuntime> backbone_;
     std::unique_ptr<MossVoiceGenHeadsRuntime> heads_;
-    std::unique_ptr<moss::MossAudioTokenizerDecoder> codec_;
+    std::unique_ptr<engine::codecs::MossAudioTokenizerCodecRuntime> codec_;
 };
 
 }  // namespace engine::models::moss_voicegen

include/engine/community_models/moss_voicegen/tokenizer_text.h

@@ -1,7 +1,7 @@
 #pragma once
 
 #include "engine/community_models/moss_voicegen/assets.h"
-#include "engine/models/moss/shared/token_rows.h"
+#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h"
 
 #include <memory>
 #include <optional>
@@ -23,7 +23,7 @@ class MossVoiceGenTextProcessor {
 
     // `instruction` describes the speaker to design. `language` must be the full language
     // name the model was trained on ("English", not "en"); an empty value renders "None".
-    moss::TokenRows build_generation_prefix(
+    engine::codecs::MossTokenRows build_generation_prefix(
         const std::string & text,
         const std::optional<std::string> & instruction,
         const std::optional<std::string> & language) const;

include/engine/community_models/vietneu_tts/session.h

@@ -10,7 +10,7 @@
 #include "engine/community_models/vietneu_tts/tokenizer_speech_encoder.h"
 #include "engine/community_models/vietneu_tts/tokenizer_text.h"
 
-#include "engine/models/moss/shared/audio_tokenizer_decoder.h"
+#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h"
 
 #include <cstddef>
 #include <cstdint>
@@ -80,7 +80,7 @@ class VietneuTTSSession final
     std::shared_ptr<const VietneuTalkerWeightsRuntime> talker_weights_;
     std::shared_ptr<VietneuTalkerStepRuntime> talker_step_;
     core::ExecutionContext voice_prompt_context_;
-    std::unique_ptr<engine::models::moss::MossAudioTokenizerDecoder> moss_speech_decoder_;
+    std::unique_ptr<engine::codecs::MossAudioTokenizerCodecRuntime> moss_speech_decoder_;
     std::unique_ptr<Qwen3SpeechTokenizerEncoderRuntime> speech_encoder_;
     std::unique_ptr<VietneuSpeakerEncoderRuntime> speaker_encoder_;
     runtime::CacheSlots<VoicePromptCacheKey, VoicePromptCacheEntry, VoicePromptCacheKeyEqual> voice_prompt_cache_;

include/engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h

@@ -0,0 +1,114 @@
+#pragma once
+
+#include "engine/framework/assets/tensor_source.h"
+#include "engine/framework/core/execution_context.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+namespace engine::codecs {
+
+struct MossAudioTokenizerTransformerStage {
+    int64_t input_dimension = 0;
+    int64_t output_dimension = 0;
+    int64_t model_dimension = 0;
+    int64_t num_heads = 0;
+    int64_t num_layers = 0;
+    int64_t feedforward_dimension = 0;
+    int64_t context_window = 0;
+    int64_t patch_size = 0;
+};
+
+struct MossAudioTokenizerQuantizerConfig {
+    int64_t codebook_size = 1024;
+    int64_t codebook_dim = 8;
+    int64_t rvq_dim = 512;
+    int64_t code_dim = 768;
+    int64_t num_quantizers = 12;
+};
+
+struct MossAudioTokenizerConfig {
+    int64_t sampling_rate = 48000;
+    int64_t samples_per_frame = 3840;
+    int64_t channels = 2;
+    MossAudioTokenizerQuantizerConfig quantizer;
+    std::vector<MossAudioTokenizerTransformerStage> encoder_stages;
+    std::vector<MossAudioTokenizerTransformerStage> decoder_stages;
+    int64_t encoder_final_patch = 1;
+    int64_t decoder_initial_patch = 1;
+    int64_t encoder_module_start = 1;
+    int64_t encoder_module_stride = 2;
+    int64_t decoder_module_start = 0;
+    int64_t decoder_module_stride = 2;
+};
+
+MossAudioTokenizerConfig moss_audio_tokenizer_v1_config();
+MossAudioTokenizerConfig moss_audio_tokenizer_v2_config();
+MossAudioTokenizerConfig moss_audio_tokenizer_nano_config();
+
+struct MossAudioTokenizerAudio {
+    int64_t sampling_rate = 0;
+    std::vector<std::vector<float>> channels;
+};
+
+struct MossAudioTokenizerCodes {
+    int64_t frames = 0;
+    std::vector<std::vector<int32_t>> codebooks;
+};
+
+struct MossTokenRows {
+    std::vector<int32_t> text_tokens;
+    std::vector<int32_t> audio_codes;
+};
+
+class MossTokenRowBuilder {
+public:
+    MossTokenRowBuilder(int64_t num_codebooks, int32_t audio_pad_token_id);
+
+    void push_text_token(int32_t token_id);
+    void push_text_tokens(const std::vector<int32_t> & token_ids);
+    void push_audio_row(int32_t text_slot_token_id, const int32_t * codes, int64_t num_codebooks);
+    void push_audio_row(int32_t text_slot_token_id, const std::vector<std::vector<int32_t>> & codes, int64_t frame);
+    MossTokenRows finish();
+
+private:
+    int64_t num_codebooks_ = 0;
+    int32_t audio_pad_token_id_ = 0;
+    MossTokenRows rows_;
+};
+
+struct MossAudioTokenizerCodecRuntimeOptions {
+    size_t weight_context_bytes = 256ull * 1024ull * 1024ull;
+    size_t encoder_graph_arena_bytes = 2048ull * 1024ull * 1024ull;
+    size_t decoder_graph_arena_bytes = 1536ull * 1024ull * 1024ull;
+    bool separate_encoder_context = false;
+};
+
+class MossAudioTokenizerCodecRuntime {
+public:
+    MossAudioTokenizerCodecRuntime(
+        std::shared_ptr<const assets::TensorSource> source,
+        core::ExecutionContext & execution_context,
+        int64_t num_quantizers,
+        MossAudioTokenizerCodecRuntimeOptions options,
+        MossAudioTokenizerConfig config = moss_audio_tokenizer_v2_config());
+    ~MossAudioTokenizerCodecRuntime();
+
+    MossAudioTokenizerCodecRuntime(const MossAudioTokenizerCodecRuntime &) = delete;
+    MossAudioTokenizerCodecRuntime & operator=(const MossAudioTokenizerCodecRuntime &) = delete;
+
+    int64_t sampling_rate() const noexcept;
+    void prepare_encoder();
+    void prepare_decoder();
+    MossAudioTokenizerCodes encode(const MossAudioTokenizerAudio & audio);
+    MossAudioTokenizerAudio decode(const MossAudioTokenizerCodes & codes);
+    void release_runtime_graphs();
+
+private:
+    struct Impl;
+    std::unique_ptr<Impl> impl_;
+};
+
+}  // namespace engine::codecs

include/engine/framework/modules/multi_codebook_embedding.h

@@ -0,0 +1,40 @@
+#pragma once
+
+#include "engine/framework/assets/tensor_source.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <string>
+#include <vector>
+
+namespace engine::modules {
+
+struct MultiCodebookEmbeddingSpec {
+    int64_t hidden_size = 0;
+    int64_t num_codebooks = 0;
+    int64_t vocab_size = 0;
+    int64_t pad_token_id = 0;
+    std::vector<int64_t> codebook_sizes;
+    std::string tensor_prefix = "audio_embeddings";
+};
+
+class MultiCodebookEmbedding {
+public:
+    MultiCodebookEmbedding(const assets::TensorSource & source, MultiCodebookEmbeddingSpec spec);
+
+    int64_t hidden_size() const noexcept { return hidden_size_; }
+    int64_t num_codebooks() const noexcept { return num_codebooks_; }
+    int32_t pad_token_id() const noexcept { return pad_token_id_; }
+    int64_t codebook_size(int64_t codebook) const;
+    const float * embedding(int64_t codebook, int32_t code) const;
+    void add_bias(const int32_t * codes, float * bias) const;
+    std::vector<float> bias_for(const int32_t * codes) const;
+
+private:
+    int64_t hidden_size_ = 0;
+    int64_t num_codebooks_ = 0;
+    int32_t pad_token_id_ = 0;
+    std::vector<std::vector<float>> embeddings_;
+};
+
+}  // namespace engine::modules

include/engine/models/moss/moss_tts_local/generator.h

@@ -3,7 +3,7 @@
 #include "engine/models/moss/moss_tts_local/assets.h"
 #include "engine/models/moss/moss_tts_local/backbone.h"
 #include "engine/models/moss/moss_tts_local/depth_transformer.h"
-#include "engine/models/moss/shared/token_rows.h"
+#include "engine/framework/modules/multi_codebook_embedding.h"
 #include "engine/framework/sampling/torch_random.h"
 
 #include <cstddef>
@@ -65,7 +65,7 @@ class MossGenerator {
     const MossDepthTransformer & depth_;
     int64_t hidden_size_ = 0;
     int64_t num_codebooks_ = 0;
-    std::unique_ptr<moss::AudioCodebookEmbeddings> audio_codebooks_;
+    std::unique_ptr<engine::modules::MultiCodebookEmbedding> audio_codebooks_;
     std::vector<float> local_text_head_;                // [2 * hidden]
     engine::sampling::TorchCudaSamplingPolicy sampling_policy_;
     struct ProjectionRuntime;

include/engine/models/moss/moss_tts_local/session.h

@@ -2,8 +2,7 @@
 
 #include "engine/framework/runtime/cache_slots.h"
 #include "engine/framework/runtime/session_base.h"
-#include "engine/models/moss/shared/audio_tokenizer_decoder.h"
-#include "engine/models/moss/shared/audio_tokenizer_encoder.h"
+#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h"
 #include "engine/models/moss/moss_tts_local/assets.h"
 #include "engine/models/moss/moss_tts_local/backbone.h"
 #include "engine/models/moss/moss_tts_local/depth_transformer.h"
@@ -38,8 +37,6 @@ class MossTTSLocalSession final
     runtime::TaskResult run(const runtime::TaskRequest & request) override;
 
 private:
-    moss::MossAudioTokenizerEncoder & encoder();
-
     struct ReferenceAudioCacheKey {
         uint64_t hash = 0;
         int sample_rate = 0;
@@ -62,11 +59,8 @@ class MossTTSLocalSession final
     std::unique_ptr<MossBackboneRuntime> backbone_;
     std::unique_ptr<MossDepthTransformer> depth_;
     std::unique_ptr<MossTextProcessor> processor_;
-    std::unique_ptr<moss::MossAudioTokenizerDecoder> codec_;
+    std::unique_ptr<engine::codecs::MossAudioTokenizerCodecRuntime> codec_;
     std::unique_ptr<MossGenerator> generator_;
-    // Lazily built the first time a speaker reference is provided (voice cloning).
-    std::unique_ptr<core::ExecutionContext> reference_encoder_execution_context_;
-    std::unique_ptr<moss::MossAudioTokenizerEncoder> encoder_;
     runtime::CacheSlots<ReferenceAudioCacheKey, ReferenceVoiceCacheEntry, ReferenceAudioCacheKeyEqual>
         reference_voice_cache_;
 };

include/engine/models/moss/moss_tts_local/tokenizer_text.h

@@ -1,6 +1,6 @@
 #pragma once
 
-#include "engine/models/moss/shared/token_rows.h"
+#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h"
 #include "engine/models/moss/moss_tts_local/assets.h"
 
 #include <cstdint>
@@ -14,7 +14,7 @@ namespace engine::models::moss_tts_local {
 // Decoder input for a generation request: the text channel (input_ids[..., 0]) plus the
 // n_vq audio channels flattened row-major as [seq, n_vq] (input_ids[..., 1:]). Every audio
 // slot of the prompt carries audio_pad_token_id, matching MossTTSLocalProcessor._build_text_rows.
-using MossGenerationPrefix = moss::TokenRows;
+using MossGenerationPrefix = engine::codecs::MossTokenRows;
 
 // Reproduces the direct-generation branch of MossTTSLocalProcessor: it renders the
 // <user_inst> template, byte-level BPE encodes each piece with the Qwen tokenizer, and
@@ -35,7 +35,7 @@ class MossTextProcessor {
         const std::optional<std::string> & language = std::nullopt) const;
 
     // Builds a voice-clone prompt. reference_codes is [num_codebooks][frames] as produced
-    // by MossAudioTokenizerEncoder for the reference speaker.
+    // by the MOSS audio tokenizer codec runtime for the reference speaker.
     MossGenerationPrefix build_clone_prefix(
         const std::string & text,
         const std::vector<std::vector<int32_t>> & reference_codes,

include/engine/models/moss/moss_tts_nano/local_frame_decoder.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "engine/framework/core/execution_context.h"
+#include "engine/framework/sampling/hf_sampler.h"
 #include "engine/framework/sampling/torch_random.h"
 #include "engine/models/moss/moss_tts_nano/assets.h"
 #include "engine/models/moss/moss_tts_nano/types.h"
@@ -42,6 +43,7 @@ class MossTTSNanoLocalFrameDecoderRuntime {
     std::shared_ptr<const MossTTSNanoAssets> assets_;
     core::ExecutionContext & execution_context_;
     engine::sampling::TorchCudaSamplingPolicy sampling_policy_;
+    engine::sampling::HfSamplerScratch sampler_scratch_;
     std::shared_ptr<Weights> weights_;
     size_t graph_arena_bytes_ = 0;
     std::unique_ptr<TextGraph> text_graph_;

include/engine/models/moss/moss_tts_nano/session.h

@@ -2,8 +2,7 @@
 
 #include "engine/framework/core/execution_context.h"
 #include "engine/framework/runtime/session_base.h"
-#include "engine/models/moss/shared/audio_tokenizer_decoder.h"
-#include "engine/models/moss/shared/audio_tokenizer_encoder.h"
+#include "engine/framework/codecs/moss_audio_tokenizer_codec_runtime.h"
 #include "engine/models/moss/moss_tts_nano/assets.h"
 #include "engine/models/moss/moss_tts_nano/generator.h"
 #include "engine/models/moss/moss_tts_nano/global_transformer.h"
@@ -33,7 +32,6 @@ class MossTTSNanoSession final
     runtime::TaskResult run(const runtime::TaskRequest & request) override;
 
 private:
-    moss::MossAudioTokenizerEncoder & encoder();
     MossTTSNanoAudioCodes encode_reference_audio(const runtime::AudioBuffer & audio, int64_t active_codebooks);
     runtime::AudioBuffer decode_generated_audio(const MossTTSNanoAudioCodes & codes, int64_t active_codebooks);
     MossTTSNanoRequest make_request(const runtime::TaskRequest & request) const;
@@ -56,9 +54,7 @@ class MossTTSNanoSession final
     MossTTSNanoGlobalTransformerRuntime global_transformer_;
     MossTTSNanoLocalFrameDecoderRuntime local_frame_decoder_;
     MossTTSNanoGenerator generator_;
-    moss::MossAudioTokenizerDecoder decoder_;
-    std::unique_ptr<core::ExecutionContext> reference_encoder_execution_context_;
-    std::unique_ptr<moss::MossAudioTokenizerEncoder> encoder_;
+    engine::codecs::MossAudioTokenizerCodecRuntime codec_;
     std::optional<runtime::AudioBuffer> prepared_prompt_audio_;
     std::optional<MossTTSNanoAudioCodes> prepared_reference_codes_;
 };