Add pre-0.7 framework building blocks

CMakeLists.txt

@@ -327,6 +327,7 @@ add_library(engine_core OBJECT
     src/framework/modules/speech_encoders/whisper_embedding.cpp
     src/framework/modules/speech_encoders/whisper_frontend.cpp
     src/framework/modules/speech_encoders/campplus_encoder.cpp
+    src/framework/modules/codecs/nemo_nano_codec.cpp
     src/framework/modules/pitch_extractors/rmvpe_pitch_extractor.cpp
     src/framework/modules/vocoders/bigvgan_vocoder.cpp
     src/framework/modules/vocoders/hift_vocoder.cpp

include/engine/framework/codecs/mimi_codec_runtime.h

@@ -203,6 +203,8 @@ class MimiEncoderRuntime {
     ~MimiEncoderRuntime();
 
     std::vector<int32_t> encode(const runtime::AudioBuffer & audio);
+    void reset_streaming();
+    std::vector<int32_t> encode_streaming(const runtime::AudioBuffer & audio, bool flush);
 
 private:
     struct Impl;

include/engine/framework/modules/activation_modules.h

@@ -21,6 +21,23 @@ class ReluModule {
     static const core::ModuleSchema & static_schema() noexcept;
 };
 
+struct LeakyReluConfig {
+    float negative_slope = 0.01F;
+};
+
+class LeakyReluModule {
+public:
+    explicit LeakyReluModule(LeakyReluConfig config = {});
+
+    const LeakyReluConfig & config() const noexcept;
+    const core::ModuleSchema & schema() const noexcept;
+    core::TensorValue build(core::ModuleBuildContext & ctx, const core::TensorValue & input) const;
+    static const core::ModuleSchema & static_schema() noexcept;
+
+private:
+    LeakyReluConfig config_;
+};
+
 class SigmoidModule {
 public:
     const core::ModuleSchema & schema() const noexcept;

include/engine/framework/modules/attention/cross_attention.h

@@ -3,8 +3,15 @@
 #include "engine/framework/core/module.h"
 #include "engine/framework/modules/attention/types.h"
 
+#include <optional>
+
 namespace engine::modules {
 
+struct CrossAttentionKeyValue {
+    core::TensorValue key;
+    core::TensorValue value;
+};
+
 class CrossAttentionModule {
 public:
     explicit CrossAttentionModule(AttentionConfig config);
@@ -18,6 +25,29 @@ class CrossAttentionModule {
         const core::TensorValue & memory,
         const AttentionWeights & weights) const;
 
+    core::TensorValue build(
+        core::ModuleBuildContext & ctx,
+        const core::TensorValue & query,
+        const core::TensorValue & memory,
+        const AttentionWeights & weights,
+        const core::TensorValue & memory_mask,
+        const core::TensorValue * attention_prior = nullptr,
+        core::TensorValue * last_attention = nullptr) const;
+
+    core::TensorValue build_cached(
+        core::ModuleBuildContext & ctx,
+        const core::TensorValue & query,
+        const CrossAttentionKeyValue & key_value,
+        const AttentionWeights & weights,
+        const core::TensorValue & memory_mask,
+        const core::TensorValue * attention_prior = nullptr,
+        core::TensorValue * last_attention = nullptr) const;
+
+    CrossAttentionKeyValue build_key_value(
+        core::ModuleBuildContext & ctx,
+        const core::TensorValue & memory,
+        const AttentionWeights & weights) const;
+
     static const core::ModuleSchema & static_schema() noexcept;
 
 private:

include/engine/framework/modules/attention/feed_forward.h

@@ -2,6 +2,7 @@
 
 #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 <optional>
@@ -45,6 +46,20 @@ struct GatedFeedForwardWeights {
     LinearWeights down_proj;
 };
 
+struct ConvFeedForwardConfig {
+    int64_t hidden_size = 0;
+    int64_t intermediate_size = 0;
+    int64_t kernel_size = 0;
+    bool causal = false;
+    bool use_bias = false;
+    GeluApproximation gelu_approximation = GeluApproximation::Tanh;
+};
+
+struct ConvFeedForwardWeights {
+    Conv1dWeights proj;
+    Conv1dWeights out;
+};
+
 class FeedForwardModule {
 public:
     explicit FeedForwardModule(FeedForwardConfig config);
@@ -99,4 +114,22 @@ class GatedFeedForwardModule {
     GatedFeedForwardConfig config_;
 };
 
+class ConvFeedForwardModule {
+public:
+    explicit ConvFeedForwardModule(ConvFeedForwardConfig config);
+
+    const ConvFeedForwardConfig & config() const noexcept;
+    const core::ModuleSchema & schema() const noexcept;
+
+    core::TensorValue build(
+        core::ModuleBuildContext & ctx,
+        const core::TensorValue & input,
+        const ConvFeedForwardWeights & weights) const;
+
+    static const core::ModuleSchema & static_schema() noexcept;
+
+private:
+    ConvFeedForwardConfig config_;
+};
+
 }  // namespace engine::modules

include/engine/framework/modules/attention/self_attention.h

@@ -2,6 +2,7 @@
 
 #include "engine/framework/core/module.h"
 #include "engine/framework/modules/attention/types.h"
+#include "engine/framework/modules/optimizations/fast_kv_modules.h"
 
 #include <optional>
 
@@ -19,6 +20,22 @@ class SelfAttentionModule {
         const core::TensorValue & input,
         const AttentionWeights & weights) const;
 
+    core::TensorValue build(
+        core::ModuleBuildContext & ctx,
+        const core::TensorValue & input,
+        const AttentionWeights & weights,
+        const std::optional<core::TensorValue> & attention_mask) const;
+
+    StreamingAttentionOutputs build_cached_tail(
+        core::ModuleBuildContext & ctx,
+        const core::TensorValue & input,
+        const AttentionWeights & weights,
+        const core::TensorValue & cache_key,
+        const core::TensorValue & cache_value,
+        const core::TensorValue & cache_slot,
+        const core::TensorValue & attention_mask,
+        FastKVSetRowsMode set_rows_mode = FastKVSetRowsMode::BackendViewOptimized) const;
+
     static const core::ModuleSchema & static_schema() noexcept;
 
 private:

include/engine/framework/modules/attention/types.h

@@ -18,6 +18,12 @@ struct AttentionConfig {
     ggml_prec projection_precision = GGML_PREC_DEFAULT;
     ggml_prec attention_precision = GGML_PREC_DEFAULT;
     AttentionPrefixCacheLayout prefix_cache_layout = AttentionPrefixCacheLayout::SequenceHeads;
+    bool use_packed_qkv = false;
+    bool causal = false;
+    bool use_packed_kv = false;
+    int64_t key_value_size = 0;
+    int64_t attention_size = 0;
+    int64_t head_dim = 0;
 };
 
 struct RelativeAttentionConfig {

include/engine/framework/modules/codecs/nemo_nano_codec.h

@@ -0,0 +1,53 @@
+#pragma once
+
+#include "engine/framework/assets/tensor_source.h"
+#include "engine/framework/core/backend.h"
+#include "engine/framework/core/execution_context.h"
+#include "engine/framework/runtime/session.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+namespace engine::modules {
+
+struct NemoNanoCodecConfig {
+    int64_t sample_rate = 22050;
+    int64_t input_dim = 32;
+    int64_t base_channels = 864;
+    int64_t audio_codebooks = 0;
+    std::vector<int64_t> upsample_rates;
+    std::vector<int64_t> resblock_kernel_sizes;
+    std::vector<int64_t> resblock_dilation_sizes;
+    std::vector<int32_t> fsq_num_levels;
+    std::vector<int32_t> fsq_dim_base_index;
+};
+
+struct NemoNanoCodecRuntimeOptions {
+    size_t graph_arena_bytes = 1024ull * 1024ull * 1024ull;
+    size_t weight_context_bytes = 2048ull * 1024ull * 1024ull;
+    assets::TensorStorageType weight_storage_type = assets::TensorStorageType::Native;
+};
+
+class NemoNanoCodecRuntime {
+public:
+    NemoNanoCodecRuntime(
+        std::shared_ptr<const assets::TensorSource> source,
+        core::ExecutionContext & execution,
+        NemoNanoCodecConfig config,
+        NemoNanoCodecRuntimeOptions options);
+    ~NemoNanoCodecRuntime();
+
+    NemoNanoCodecRuntime(const NemoNanoCodecRuntime &) = delete;
+    NemoNanoCodecRuntime & operator=(const NemoNanoCodecRuntime &) = delete;
+
+    runtime::AudioBuffer decode_codes(const std::vector<int32_t> & codes);
+    void release_runtime_graph();
+
+private:
+    struct Impl;
+    std::unique_ptr<Impl> impl_;
+};
+
+}  // namespace engine::modules

include/engine/framework/modules/streaming_conv_modules.h

@@ -60,6 +60,12 @@ enum class StreamingPadMode {
     Replicate,
 };
 
+enum class StreamingConv1dPaddingMode {
+    StreamingSame,
+    StrictCausal,
+    Explicit,
+};
+
 struct StreamingConv1dConfig {
     int64_t in_channels = 0;
     int64_t out_channels = 0;
@@ -68,6 +74,9 @@ struct StreamingConv1dConfig {
     int dilation = 1;
     bool use_bias = true;
     StreamingPadMode pad_mode = StreamingPadMode::Constant;
+    StreamingConv1dPaddingMode padding_mode = StreamingConv1dPaddingMode::StreamingSame;
+    int64_t explicit_left = 0;
+    int64_t explicit_right = 0;
 };
 
 using StreamingConv1dWeights = Conv1dWeights;
@@ -84,6 +93,12 @@ class StreamingConv1dModule {
     StreamingConv1dConfig config_;
 };
 
+using CausalConv1dPadMode = StreamingPadMode;
+using CausalConv1dPaddingMode = StreamingConv1dPaddingMode;
+using CausalConv1dConfig = StreamingConv1dConfig;
+using CausalConv1dWeights = StreamingConv1dWeights;
+using CausalConv1dModule = StreamingConv1dModule;
+
 struct DepthwiseConvTranspose1dConfig {
     int64_t channels = 0;
     int64_t kernel_size = 0;

include/engine/framework/modules/transformers/qwen_causal_decode_runtime.h

@@ -26,6 +26,7 @@ struct QwenCausalDecodeRuntimeConfig {
     QwenCausalDecodeOutputMode output_mode = QwenCausalDecodeOutputMode::Logits;
     bool return_hidden = false;
     std::optional<ggml_type> readback_round_type;
+    std::vector<int32_t> logits_readback_token_ids;
 };
 
 struct QwenCausalDecodeRuntimeWeights {
@@ -41,6 +42,12 @@ struct QwenCausalPrefillResult {
     runtime::TransformerKVState state;
 };
 
+struct QwenCausalBatchedPrefillResult {
+    std::vector<float> logits;
+    std::vector<float> hidden;
+    runtime::TransformerBatchedKVState state;
+};
+
 struct QwenCausalDecodeStepResult {
     std::vector<float> logits;
     std::vector<float> hidden;
@@ -60,11 +67,31 @@ class QwenCausalDecodeRuntime {
     QwenCausalPrefillResult prefill_tokens(const std::vector<int32_t> & token_ids);
     QwenCausalPrefillResult prefill_embeddings(const std::vector<float> & embeddings, int64_t steps);
 
+    QwenCausalBatchedPrefillResult prefill_tokens_batched(
+        const std::vector<int32_t> & token_ids,
+        int64_t batch_size,
+        int64_t steps);
+    QwenCausalBatchedPrefillResult prefill_embeddings_batched(
+        const std::vector<float> & embeddings,
+        int64_t batch_size,
+        int64_t steps);
+
     void start_decode_tokens(const runtime::TransformerKVState & state, int64_t required_cache_steps);
     void start_decode_embeddings(const runtime::TransformerKVState & state, int64_t required_cache_steps);
     QwenCausalDecodeStepResult decode_token(int32_t token);
     QwenCausalDecodeStepResult decode_embedding(const std::vector<float> & embedding);
 
+    void start_decode_tokens_batched(
+        const runtime::TransformerBatchedKVState & state,
+        int64_t required_cache_steps);
+    void start_decode_embeddings_batched(
+        const runtime::TransformerBatchedKVState & state,
+        int64_t required_cache_steps);
+    QwenCausalDecodeStepResult decode_tokens_batched(const std::vector<int32_t> & tokens);
+    QwenCausalDecodeStepResult decode_embeddings_batched(
+        const std::vector<float> & embeddings,
+        int64_t batch_size);
+
     int64_t decode_cache_steps() const noexcept;
     int64_t decode_current_end() const noexcept;
     int64_t decode_valid_steps() const noexcept;