Add opt-in encoder-decoder framework modules

include/engine/framework/modules/activation_modules.h

@@ -114,11 +114,20 @@ class SoftmaxModule {
     static const core::ModuleSchema & static_schema() noexcept;
 };
 
+struct GLUConfig {
+    bool contiguous_gate = false;
+};
+
 class GLUModule {
 public:
+    GLUModule() = default;
+    explicit GLUModule(GLUConfig config);
     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:
+    GLUConfig config_;
 };
 
 struct Snake1dConfig {

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

@@ -43,6 +43,16 @@ class CrossAttentionModule {
         const core::TensorValue * attention_prior = nullptr,
         core::TensorValue * last_attention = nullptr) const;
 
+    // Opt-in flash path. KV is [B,H,K,D]; mask is contiguous F16 [B|1,H|1,Q,K],
+    // with additive scores (0 for allowed, -infinity for excluded positions).
+    // Every query must have at least one allowed key. No attention prior/output.
+    core::TensorValue build_cached_flash(
+        core::ModuleBuildContext & ctx,
+        const core::TensorValue & query,
+        const CrossAttentionKeyValue & key_value,
+        const AttentionWeights & weights,
+        const core::TensorValue & attention_mask) const;
+
     CrossAttentionKeyValue build_key_value(
         core::ModuleBuildContext & ctx,
         const core::TensorValue & memory,

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

@@ -9,12 +9,18 @@
 
 namespace engine::modules {
 
+enum class FeedForwardActivation {
+    Gelu,
+    Relu,
+};
+
 struct FeedForwardConfig {
     int64_t hidden_size = 0;
     int64_t intermediate_size = 0;
     bool use_bias = true;
     GeluApproximation gelu_approximation = GeluApproximation::ExactErf;
     ggml_prec projection_precision = GGML_PREC_DEFAULT;
+    FeedForwardActivation activation = FeedForwardActivation::Gelu;
 };
 
 struct FeedForwardWeights {

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

@@ -225,6 +225,10 @@ struct TransformerDecoderBlockConfig {
     int64_t intermediate_size = 0;
     float eps = 1e-5f;
     bool use_bias = true;
+    FeedForwardActivation activation = FeedForwardActivation::Gelu;
+    bool use_packed_qkv = false;
+    bool use_packed_kv = false;
+    bool use_flash_cross_attention = false;
 };
 
 struct TransformerDecoderBlockWeights {
@@ -249,6 +253,21 @@ class TransformerDecoderBlockModule {
         const core::TensorValue & memory,
         const TransformerDecoderBlockWeights & weights) const;
 
+    // Requires packed QKV/KV opt-ins. Caches and masks are caller-owned, using
+    // SelfAttentionModule::build_cached_tail and CrossAttentionModule::build_cached layouts.
+    // When use_flash_cross_attention is enabled, memory_mask instead follows
+    // CrossAttentionModule::build_cached_flash's additive F16 mask contract.
+    core::TensorValue build_cached_tail(
+        core::ModuleBuildContext & ctx,
+        const core::TensorValue & input,
+        const TransformerDecoderBlockWeights & weights,
+        const core::TensorValue & self_key_cache,
+        const core::TensorValue & self_value_cache,
+        const core::TensorValue & cache_slot,
+        const core::TensorValue & causal_mask,
+        const CrossAttentionKeyValue & memory_key_value,
+        const core::TensorValue & memory_mask) const;
+
     static const core::ModuleSchema & static_schema() noexcept;
 
 private:

include/engine/framework/modules/conformer_modules.h

@@ -14,6 +14,7 @@ struct ConformerConvModuleConfig {
     bool use_bias = true;
     float eps = 1e-5f;
     int64_t cache_drop_size = 0;
+    bool contiguous_glu_gate = false;
 };
 
 struct ConvSubsamplingConfig {
@@ -37,6 +38,41 @@ struct ConvSubsamplingOutputs {
     core::TensorValue lengths;
 };
 
+struct DepthwiseConvSubsamplingConfig {
+    int64_t input_features = 0;
+    int64_t output_features = 0;
+    int64_t conv_channels = 0;
+    int kernel_size = 3;
+    int stride = 2;
+    int padding = 1;
+    bool use_bias = true;
+};
+
+struct DepthwiseConvSubsamplingStageWeights {
+    Conv2dWeights depthwise;
+    Conv2dWeights pointwise;
+};
+
+struct DepthwiseConvSubsamplingWeights {
+    Conv2dWeights input_conv;
+    std::vector<DepthwiseConvSubsamplingStageWeights> stages;
+    LinearWeights projection;
+};
+
+class DepthwiseConvSubsamplingModule {
+public:
+    explicit DepthwiseConvSubsamplingModule(DepthwiseConvSubsamplingConfig config);
+    // Input is [batch, time, features]; optional masks cover each downsampling stage.
+    core::TensorValue build(
+        core::ModuleBuildContext & ctx,
+        const core::TensorValue & input,
+        const DepthwiseConvSubsamplingWeights & weights,
+        const std::vector<core::TensorValue> & stage_keep_masks = {}) const;
+
+private:
+    DepthwiseConvSubsamplingConfig config_;
+};
+
 class ConvSubsamplingModule {
 public:
     explicit ConvSubsamplingModule(ConvSubsamplingConfig config);
@@ -54,10 +90,7 @@ struct ConformerConvModuleWeights {
     NormWeights norm;
     LinearWeights pointwise_in;
     DepthwiseConv1dWeights depthwise;
-    struct {
-        core::TensorValue scale;
-        core::TensorValue bias;
-    } depthwise_norm;
+    ChannelAffineWeights depthwise_norm;
     LinearWeights pointwise_out;
 };
 
@@ -103,6 +136,7 @@ struct ConformerBlockConfig {
     int64_t left_context = -1;
     int64_t right_context = -1;
     int64_t cache_drop_size = 0;
+    bool contiguous_glu_gate = false;
 };
 
 struct ConformerBlockWeights {

include/engine/framework/modules/norm_modules.h

@@ -19,6 +19,11 @@ struct NormWeights {
     std::optional<core::TensorValue> bias;
 };
 
+struct ChannelAffineWeights {
+    core::TensorValue scale;
+    core::TensorValue bias;
+};
+
 class LayerNormModule {
 public:
     explicit LayerNormModule(NormConfig config);

include/engine/framework/modules/weight_binding.h

@@ -20,6 +20,28 @@
 
 namespace engine::modules::binding {
 
+template <typename Store>
+ChannelAffineWeights batch_norm_eval_from_source(
+    Store & store,
+    const assets::TensorSource & source,
+    const std::string & prefix,
+    int64_t channels,
+    float eps,
+    assets::TensorStorageType storage_type = assets::TensorStorageType::F32) {
+    const auto gamma = source.require_f32(prefix + ".weight", {channels});
+    const auto beta = source.require_f32(prefix + ".bias", {channels});
+    const auto mean = source.require_f32(prefix + ".running_mean", {channels});
+    const auto variance = source.require_f32(prefix + ".running_var", {channels});
+    std::vector<float> scale(static_cast<size_t>(channels)), bias(static_cast<size_t>(channels));
+    for (size_t i = 0; i < scale.size(); ++i) {
+        scale[i] = gamma[i] / std::sqrt(variance[i] + eps);
+        bias[i] = beta[i] - mean[i] * scale[i];
+    }
+    const auto shape = core::TensorShape::from_dims({channels});
+    return {store.make_from_f32(shape, storage_type, std::move(scale)),
+        store.make_from_f32(shape, storage_type, std::move(bias))};
+}
+
 inline LinearConfig linear_config(
     int64_t in_features,
     int64_t out_features,

src/framework/modules/activation_modules.cpp

@@ -513,6 +513,8 @@ const core::ModuleSchema & SoftmaxModule::static_schema() noexcept {
     return kSoftmaxSchema;
 }
 
+GLUModule::GLUModule(GLUConfig config) : config_(config) {}
+
 const core::ModuleSchema & GLUModule::schema() const noexcept {
     return static_schema();
 }
@@ -540,6 +542,9 @@ core::TensorValue GLUModule::build(core::ModuleBuildContext & ctx, const core::T
         ggml_view_2d(ctx.ggml, flat.tensor, hidden, flat.shape.dims[0], flat.tensor->nb[1], hidden * sizeof(float)),
         core::TensorShape::from_dims({flat.shape.dims[0], hidden}),
         GGML_TYPE_F32);
+    if (config_.contiguous_gate) {
+        rhs = core::wrap_tensor(ggml_cont(ctx.ggml, rhs.tensor), rhs.shape, GGML_TYPE_F32);
+    }
     rhs = core::wrap_tensor(ggml_sigmoid(ctx.ggml, rhs.tensor), rhs.shape, GGML_TYPE_F32);
     auto output = core::wrap_tensor(ggml_mul(ctx.ggml, lhs.tensor, rhs.tensor), lhs.shape, GGML_TYPE_F32);
 

src/framework/modules/attention/attention_internal.h

@@ -426,7 +426,11 @@ inline core::TensorValue build_feed_forward_impl(
     const LinearModule fc2({config.intermediate_size, config.hidden_size, config.use_bias, config.projection_precision});
 
     auto hidden = fc1.build(ctx, input, make_linear_weights(weights.fc1_weight, weights.fc1_bias));
-    hidden = gelu.build(ctx, hidden);
+    switch (config.activation) {
+        case FeedForwardActivation::Gelu: hidden = gelu.build(ctx, hidden); break;
+        case FeedForwardActivation::Relu: hidden = ReluModule().build(ctx, hidden); break;
+        default: throw std::runtime_error("Unsupported feed-forward activation");
+    }
     return fc2.build(ctx, hidden, make_linear_weights(weights.fc2_weight, weights.fc2_bias));
 }
 

src/framework/modules/attention/cross_attention.cpp

@@ -1,4 +1,5 @@
 #include "attention_internal.h"
+#include "engine/framework/modules/attention/scaled_dot_product_attention.h"
 
 namespace engine::modules {
 
@@ -46,16 +47,16 @@ void validate_cross_memory(const core::TensorValue & memory, const AttentionConf
 void validate_cross_cache(
     const CrossAttentionKeyValue & key_value,
     const core::TensorValue & query,
-    const core::TensorValue & memory_mask,
+    int64_t memory_frames,
     const AttentionConfig & config) {
     const int64_t head_dim = cross_head_dim(config);
     if (key_value.key.shape.rank != 4 || key_value.value.shape.rank != 4 ||
         key_value.key.shape.dims[0] != query.shape.dims[0] ||
         key_value.value.shape.dims[0] != query.shape.dims[0] ||
         key_value.key.shape.dims[1] != config.num_heads ||
         key_value.value.shape.dims[1] != config.num_heads ||
-        key_value.key.shape.dims[2] != memory_mask.shape.dims[1] ||
-        key_value.value.shape.dims[2] != memory_mask.shape.dims[1] ||
+        key_value.key.shape.dims[2] != memory_frames ||
+        key_value.value.shape.dims[2] != memory_frames ||
         key_value.key.shape.dims[3] != head_dim ||
         key_value.value.shape.dims[3] != head_dim) {
         throw std::runtime_error("CrossAttentionModule cached KV shape is invalid");
@@ -200,7 +201,7 @@ core::TensorValue CrossAttentionModule::build_cached(
         throw std::runtime_error("CrossAttentionModule cached path requires packed KV");
     }
     validate_cross_query(query, config_);
-    validate_cross_cache(key_value, query, memory_mask, config_);
+    validate_cross_cache(key_value, query, memory_mask.shape.dims[1], config_);
     auto query_heads = build_cross_query(ctx, query, config_, weights);
     auto probs = build_cross_probabilities(
         ctx,
@@ -213,6 +214,38 @@ core::TensorValue CrossAttentionModule::build_cached(
     return build_cross_output(ctx, query, probs, key_value.value, config_, weights);
 }
 
+core::TensorValue CrossAttentionModule::build_cached_flash(
+    core::ModuleBuildContext & ctx,
+    const core::TensorValue & query,
+    const CrossAttentionKeyValue & key_value,
+    const AttentionWeights & weights,
+    const core::TensorValue & attention_mask) const {
+    if (!config_.use_packed_kv) {
+        throw std::runtime_error("CrossAttentionModule cached flash path requires packed KV");
+    }
+    validate_cross_query(query, config_);
+    core::validate_rank_between(attention_mask, 4, 4, "cross_attention.flash_mask");
+    const auto & shape = attention_mask.shape;
+    if (attention_mask.type != GGML_TYPE_F16 || !ggml_is_contiguous(attention_mask.tensor) ||
+        (shape.dims[0] != 1 && shape.dims[0] != query.shape.dims[0]) ||
+        (shape.dims[1] != 1 && shape.dims[1] != config_.num_heads) ||
+        shape.dims[2] != query.shape.dims[1]) {
+        throw std::runtime_error("CrossAttentionModule flash mask must be contiguous F16 [B|1,H|1,Q,K]");
+    }
+    validate_cross_cache(key_value, query, shape.dims[3], config_);
+    const auto query_heads = build_cross_query(ctx, query, config_, weights);
+    const auto precision = config_.attention_precision == GGML_PREC_DEFAULT
+        ? GGML_PREC_F32 : config_.attention_precision;
+    auto context = ScaledDotProductAttentionModule({cross_head_dim(config_),
+        ScaledDotProductAttentionLowering::Flash, precision})
+        .build(ctx, query_heads, key_value.key, key_value.value, attention_mask);
+    context = core::reshape_tensor(ctx, context,
+        core::TensorShape::from_dims({query.shape.dims[0], query.shape.dims[1], cross_attention_size(config_)}));
+    return LinearModule({cross_attention_size(config_), config_.hidden_size,
+        config_.use_bias, config_.projection_precision})
+        .build(ctx, context, make_linear_weights(weights.out_weight, weights.out_bias));
+}
+
 CrossAttentionKeyValue CrossAttentionModule::build_key_value(
     core::ModuleBuildContext & ctx,
     const core::TensorValue & memory,

src/framework/modules/attention/transformer_blocks.cpp

@@ -333,12 +333,17 @@ core::TensorValue TransformerDecoderBlockModule::build(
     validate_sequence_input(memory, config_.hidden_size, "memory");
 
     const LayerNormModule norm1(make_norm_config(config_.hidden_size, config_.eps));
-    const SelfAttentionModule self_attention({config_.hidden_size, config_.num_heads, config_.use_bias});
+    AttentionConfig self_config{config_.hidden_size, config_.num_heads, config_.use_bias};
+    self_config.use_packed_qkv = config_.use_packed_qkv;
+    const SelfAttentionModule self_attention(self_config);
     const LayerNormModule norm2(make_norm_config(config_.hidden_size, config_.eps));
-    const CrossAttentionModule cross_attention({config_.hidden_size, config_.num_heads, config_.use_bias});
+    AttentionConfig cross_config{config_.hidden_size, config_.num_heads, config_.use_bias};
+    cross_config.use_packed_kv = config_.use_packed_kv;
+    const CrossAttentionModule cross_attention(cross_config);
     const LayerNormModule norm3(make_norm_config(config_.hidden_size, config_.eps));
     const FeedForwardModule feed_forward(
-        {config_.hidden_size, config_.intermediate_size, config_.use_bias, GeluApproximation::ExactErf});
+        {config_.hidden_size, config_.intermediate_size, config_.use_bias, GeluApproximation::ExactErf,
+            GGML_PREC_DEFAULT, config_.activation});
     const ResidualAddModule add;
 
     auto cur = norm1.build(ctx, input, weights.norm1);
@@ -354,6 +359,39 @@ core::TensorValue TransformerDecoderBlockModule::build(
     return add.build(ctx, cur, ff_out);
 }
 
+core::TensorValue TransformerDecoderBlockModule::build_cached_tail(
+    core::ModuleBuildContext & ctx,
+    const core::TensorValue & input,
+    const TransformerDecoderBlockWeights & weights,
+    const core::TensorValue & self_key_cache,
+    const core::TensorValue & self_value_cache,
+    const core::TensorValue & cache_slot,
+    const core::TensorValue & causal_mask,
+    const CrossAttentionKeyValue & memory_key_value,
+    const core::TensorValue & memory_mask) const {
+    validate_sequence_input(input, config_.hidden_size, "input");
+    AttentionConfig self_config{config_.hidden_size, config_.num_heads, config_.use_bias};
+    self_config.use_packed_qkv = config_.use_packed_qkv;
+    self_config.causal = true;
+    AttentionConfig cross_config{config_.hidden_size, config_.num_heads, config_.use_bias};
+    cross_config.use_packed_kv = config_.use_packed_kv;
+    const LayerNormModule norm(make_norm_config(config_.hidden_size, config_.eps));
+    const ResidualAddModule add;
+    auto x = norm.build(ctx, input, weights.norm1);
+    x = SelfAttentionModule(self_config).build_cached_tail(ctx, x, weights.self_attention,
+        self_key_cache, self_value_cache, cache_slot, causal_mask).output;
+    auto cur = add.build(ctx, input, x);
+    x = norm.build(ctx, cur, weights.norm2);
+    x = config_.use_flash_cross_attention
+        ? CrossAttentionModule(cross_config).build_cached_flash(ctx, x, memory_key_value, weights.cross_attention, memory_mask)
+        : CrossAttentionModule(cross_config).build_cached(ctx, x, memory_key_value, weights.cross_attention, memory_mask);
+    cur = add.build(ctx, cur, x);
+    x = norm.build(ctx, cur, weights.norm3);
+    x = FeedForwardModule({config_.hidden_size, config_.intermediate_size, config_.use_bias,
+        GeluApproximation::ExactErf, GGML_PREC_DEFAULT, config_.activation}).build(ctx, x, weights.feed_forward);
+    return add.build(ctx, cur, x);
+}
+
 const core::ModuleSchema & TransformerDecoderBlockModule::static_schema() noexcept {
     return kTransformerDecoderBlockSchema;
 }