diff --git a/include/engine/framework/modules/activation_modules.h b/include/engine/framework/modules/activation_modules.h
index 7edf78a73..315c6897d 100644
--- a/include/engine/framework/modules/activation_modules.h
+++ b/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 {
diff --git a/include/engine/framework/modules/attention/cross_attention.h b/include/engine/framework/modules/attention/cross_attention.h
index 26113aeea..cbf575015 100644
--- a/include/engine/framework/modules/attention/cross_attention.h
+++ b/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,
diff --git a/include/engine/framework/modules/attention/feed_forward.h b/include/engine/framework/modules/attention/feed_forward.h
index 6a357102f..6372292a8 100644
--- a/include/engine/framework/modules/attention/feed_forward.h
+++ b/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 {
diff --git a/include/engine/framework/modules/attention/transformer_blocks.h b/include/engine/framework/modules/attention/transformer_blocks.h
index cd5d8c2bf..9c54923ff 100644
--- a/include/engine/framework/modules/attention/transformer_blocks.h
+++ b/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:
diff --git a/include/engine/framework/modules/conformer_modules.h b/include/engine/framework/modules/conformer_modules.h
index 142e670a8..543bc8af0 100644
--- a/include/engine/framework/modules/conformer_modules.h
+++ b/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 {
diff --git a/include/engine/framework/modules/norm_modules.h b/include/engine/framework/modules/norm_modules.h
index b0fa666c3..3c75111d8 100644
--- a/include/engine/framework/modules/norm_modules.h
+++ b/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);
diff --git a/include/engine/framework/modules/weight_binding.h b/include/engine/framework/modules/weight_binding.h
index e73c81b5a..05463b1a5 100644
--- a/include/engine/framework/modules/weight_binding.h
+++ b/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,
diff --git a/src/framework/modules/activation_modules.cpp b/src/framework/modules/activation_modules.cpp
index 1ef3f2d6c..8002d1e1f 100644
--- a/src/framework/modules/activation_modules.cpp
+++ b/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);
 
diff --git a/src/framework/modules/attention/attention_internal.h b/src/framework/modules/attention/attention_internal.h
index 9cacd6564..2d10fb8f2 100644
--- a/src/framework/modules/attention/attention_internal.h
+++ b/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));
 }
 
diff --git a/src/framework/modules/attention/cross_attention.cpp b/src/framework/modules/attention/cross_attention.cpp
index 138886561..89a052a6d 100644
--- a/src/framework/modules/attention/cross_attention.cpp
+++ b/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,7 +47,7 @@ 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 ||
@@ -54,8 +55,8 @@ void validate_cross_cache(
         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,
diff --git a/src/framework/modules/attention/transformer_blocks.cpp b/src/framework/modules/attention/transformer_blocks.cpp
index 6176af3b4..9a12d1379 100644
--- a/src/framework/modules/attention/transformer_blocks.cpp
+++ b/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;
 }
diff --git a/src/framework/modules/conformer_modules.cpp b/src/framework/modules/conformer_modules.cpp
index 65a14be84..7c748678a 100644
--- a/src/framework/modules/conformer_modules.cpp
+++ b/src/framework/modules/conformer_modules.cpp
@@ -2,6 +2,7 @@
 
 #include "tensor_layout_utils.h"
 #include "engine/framework/modules/activation_modules.h"
+#include "engine/framework/modules/primitive_modules.h"
 #include "engine/framework/modules/structural_modules.h"
 
 #include <stdexcept>
@@ -177,6 +178,53 @@ ConvSubsamplingOutputs ConvSubsamplingModule::build(
     };
 }
 
+DepthwiseConvSubsamplingModule::DepthwiseConvSubsamplingModule(DepthwiseConvSubsamplingConfig config)
+    : config_(config) {
+    if (config.input_features <= 0 || config.output_features <= 0 || config.conv_channels <= 0 ||
+        config.kernel_size <= 0 || config.stride <= 0 || config.padding < 0) {
+        throw std::runtime_error("DepthwiseConvSubsampling requires positive dimensions and nonnegative padding");
+    }
+}
+
+core::TensorValue DepthwiseConvSubsamplingModule::build(
+    core::ModuleBuildContext & ctx,
+    const core::TensorValue & input,
+    const DepthwiseConvSubsamplingWeights & weights,
+    const std::vector<core::TensorValue> & stage_keep_masks) const {
+    core::validate_rank_between(input, 3, 3, "subsampling.input");
+    core::validate_last_dim(input, config_.input_features, "subsampling.input");
+    if (!stage_keep_masks.empty() && stage_keep_masks.size() != weights.stages.size() + 1) {
+        throw std::runtime_error("DepthwiseConvSubsampling requires one keep mask per downsampling stage");
+    }
+    const auto channels = config_.conv_channels;
+    const int k = config_.kernel_size, s = config_.stride, p = config_.padding;
+    auto x = core::reshape_tensor(ctx, input,
+        core::TensorShape::from_dims({input.shape.dims[0], 1, input.shape.dims[1], input.shape.dims[2]}));
+    x = Conv2dModule({1, channels, k, k, s, s, p, p, 1, 1, config_.use_bias}).build(ctx, x, weights.input_conv);
+    x = ReluModule().build(ctx, x);
+    if (!stage_keep_masks.empty()) {
+        x = TimeMask4dModule().build(ctx, x, stage_keep_masks[0]);
+    }
+    for (size_t i = 0; i < weights.stages.size(); ++i) {
+        x = DepthwiseConv2dModule({channels, k, k, s, s, p, p, 1, 1, config_.use_bias})
+            .build(ctx, x, weights.stages[i].depthwise);
+        if (!stage_keep_masks.empty()) {
+            x = TimeMask4dModule().build(ctx, x, stage_keep_masks[i + 1]);
+        }
+        x = Conv2dModule({channels, channels, 1, 1, 1, 1, 0, 0, 1, 1, config_.use_bias})
+            .build(ctx, x, weights.stages[i].pointwise);
+        x = ReluModule().build(ctx, x);
+        if (!stage_keep_masks.empty()) {
+            x = TimeMask4dModule().build(ctx, x, stage_keep_masks[i + 1]);
+        }
+    }
+    x = tensor_layout::swap_channel_time_axes_4d(ctx, x);
+    x = core::wrap_tensor(ggml_cont(ctx.ggml, x.tensor), x.shape, x.type);
+    const int64_t flat_features = x.shape.dims[2] * x.shape.dims[3];
+    x = core::reshape_tensor(ctx, x, core::TensorShape::from_dims({x.shape.dims[0], x.shape.dims[1], flat_features}));
+    return LinearModule({flat_features, config_.output_features, config_.use_bias}).build(ctx, x, weights.projection);
+}
+
 ConformerConvModule::ConformerConvModule(ConformerConvModuleConfig config) : config_(config) {}
 
 core::TensorValue ConformerConvModule::build(
@@ -186,7 +234,7 @@ core::TensorValue ConformerConvModule::build(
     const std::optional<core::TensorValue> & keep_mask) const {
     auto x = LayerNormModule({config_.hidden_size, config_.eps, true, true}).build(ctx, input, weights.norm);
     x = LinearModule({config_.hidden_size, config_.hidden_size * 2, config_.use_bias}).build(ctx, x, weights.pointwise_in);
-    x = GLUModule().build(ctx, x);
+    x = GLUModule({config_.contiguous_glu_gate}).build(ctx, x);
     if (keep_mask.has_value()) {
         x = MaskingModule().build(ctx, x, *keep_mask);
     }
@@ -209,7 +257,7 @@ StreamingConformerConvOutputs StreamingConformerConvModule::build(
     const std::optional<core::TensorValue> & keep_mask) const {
     auto x = LayerNormModule({config_.hidden_size, config_.eps, true, true}).build(ctx, input, weights.norm);
     x = LinearModule({config_.hidden_size, config_.hidden_size * 2, config_.use_bias}).build(ctx, x, weights.pointwise_in);
-    x = GLUModule().build(ctx, x);
+    x = GLUModule({config_.contiguous_glu_gate}).build(ctx, x);
     if (keep_mask.has_value()) {
         x = MaskingModule().build(ctx, x, *keep_mask);
     }
@@ -272,7 +320,7 @@ core::TensorValue ConformerBlockModule::build(
     y = SelfAttentionModule({config_.hidden_size, config_.num_heads, config_.use_bias}).build(ctx, y, weights.self_attention);
     x = core::wrap_tensor(ggml_add(ctx.ggml, x.tensor, y.tensor), x.shape, GGML_TYPE_F32);
 
-    y = ConformerConvModule({config_.hidden_size, config_.kernel_size, config_.use_bias, config_.eps, 0}).build(ctx, x, weights.conv);
+    y = ConformerConvModule({config_.hidden_size, config_.kernel_size, config_.use_bias, config_.eps, 0, config_.contiguous_glu_gate}).build(ctx, x, weights.conv);
     x = core::wrap_tensor(ggml_add(ctx.ggml, x.tensor, y.tensor), x.shape, GGML_TYPE_F32);
 
     y = LayerNormModule({config_.hidden_size, config_.eps, true, true}).build(ctx, x, weights.norm2);
@@ -315,7 +363,7 @@ core::TensorValue RelativeConformerBlockModule::build(
     }).build(ctx, y, pos_emb, weights.self_attention, attention_mask, query_keep_mask, projected_pos_emb);
     x = core::wrap_tensor(ggml_add(ctx.ggml, x.tensor, y.tensor), x.shape, GGML_TYPE_F32);
 
-    y = ConformerConvModule({config_.hidden_size, config_.kernel_size, config_.use_bias, config_.eps, 0}).build(ctx, x, weights.conv, keep_mask);
+    y = ConformerConvModule({config_.hidden_size, config_.kernel_size, config_.use_bias, config_.eps, 0, config_.contiguous_glu_gate}).build(ctx, x, weights.conv, keep_mask);
     x = core::wrap_tensor(ggml_add(ctx.ggml, x.tensor, y.tensor), x.shape, GGML_TYPE_F32);
 
     y = LayerNormModule({config_.hidden_size, config_.eps, true, true}).build(ctx, x, weights.norm2);
@@ -350,7 +398,7 @@ StreamingConformerBlockOutputs StreamingConformerBlockModule::build(
     y = SelfAttentionModule({config_.hidden_size, config_.num_heads, config_.use_bias}).build(ctx, y, weights.self_attention);
     x = core::wrap_tensor(ggml_add(ctx.ggml, x.tensor, y.tensor), x.shape, GGML_TYPE_F32);
 
-    auto conv = StreamingConformerConvModule({config_.hidden_size, config_.kernel_size, config_.use_bias, config_.eps, config_.cache_drop_size}).build(
+    auto conv = StreamingConformerConvModule({config_.hidden_size, config_.kernel_size, config_.use_bias, config_.eps, config_.cache_drop_size, config_.contiguous_glu_gate}).build(
         ctx,
         x,
         weights.conv,
@@ -403,7 +451,7 @@ StreamingConformerBlockOutputs StreamingRelativeConformerBlockModule::build(
         attention_mask);
     x = core::wrap_tensor(ggml_add(ctx.ggml, x.tensor, attn.output.tensor), x.shape, GGML_TYPE_F32);
 
-    auto conv = StreamingConformerConvModule({config_.hidden_size, config_.kernel_size, config_.use_bias, config_.eps, config_.cache_drop_size}).build(
+    auto conv = StreamingConformerConvModule({config_.hidden_size, config_.kernel_size, config_.use_bias, config_.eps, config_.cache_drop_size, config_.contiguous_glu_gate}).build(
         ctx,
         x,
         weights.conv,
diff --git a/tests/unittests/test_encoder_modules.cpp b/tests/unittests/test_encoder_modules.cpp
index 4abae5687..00e3040a6 100644
--- a/tests/unittests/test_encoder_modules.cpp
+++ b/tests/unittests/test_encoder_modules.cpp
@@ -1,9 +1,13 @@
 #include "engine/framework/core/backend.h"
+#include "engine/framework/core/backend_weight_store.h"
 #include "engine/framework/modules/attention_modules.h"
 #include "engine/framework/modules/conformer_modules.h"
+#include "engine/framework/modules/primitive_modules.h"
+#include "engine/framework/modules/weight_binding.h"
 #include "engine/framework/runtime/graph_optimizer.h"
 
 #include <cmath>
+#include <cstring>
 #include <iostream>
 #include <memory>
 #include <optional>
@@ -11,6 +15,7 @@
 #include <stdexcept>
 #include <string>
 #include <vector>
+#include <unordered_map>
 
 namespace {
 
@@ -33,7 +38,7 @@ void require_allclose(
     }
     for (size_t i = 0; i < actual.size(); ++i) {
         const float diff = std::fabs(actual[i] - expected[i]);
-        if (diff > atol) {
+        if (!std::isfinite(diff) || diff > atol) {
             std::ostringstream oss;
             oss << label << " mismatch at " << i << ": expected " << expected[i] << ", got " << actual[i]
                 << ", diff=" << diff;
@@ -788,10 +793,310 @@ void test_relative_attention_specialized_flash_matches_reference_on_realistic_sh
     }
 }
 
+void test_glu_contiguous_gate_opt_in() {
+    using namespace engine;
+    modules::GLUModule legacy = {};
+    legacy = []() -> modules::GLUModule { return {}; }();
+    const std::vector<float> data{1, -2, 3, 4, 0, 1, -1, 2, -3, 2, 1, -1, -2, 0, 2, 1};
+    std::vector<float> expected;
+    for (size_t row = 0; row < 2; ++row) {
+        for (size_t col = 0; col < 4; ++col) {
+            expected.push_back(data[row * 8 + col] / (1.0f + std::exp(-data[row * 8 + col + 4])));
+        }
+    }
+    require(!modules::GLUConfig{}.contiguous_gate, "GLU layout change must be opt-in");
+    require(!modules::ConformerBlockConfig{}.contiguous_glu_gate, "Conformer layout change must be opt-in");
+    for (const auto type : {core::BackendType::Cpu, core::BackendType::Cuda}) {
+        if (type == core::BackendType::Cuda && !backend_is_available(type)) {
+            continue;
+        }
+        for (const bool opt_in : {false, true}) {
+            if (type == core::BackendType::Cuda && !opt_in) {
+                continue;  // Legacy strided sigmoid is not supported on CUDA.
+            }
+            ModuleRunner runner;
+            set_runner_backend(runner, type);
+            runner.ctx.backend_type = type;
+            auto input = runner.make_f32(core::TensorShape::from_dims({1, 2, 8}));
+            const auto module = opt_in ? modules::GLUModule(modules::GLUConfig{true}) : legacy;
+            auto output = module.build(runner.ctx, input);
+            auto graph = ggml_new_graph_custom(runner.ggml, kTestGraphNodes, false);
+            ggml_build_forward_expand(graph, output.tensor);
+            bool found_sigmoid = false;
+            for (int i = 0; i < ggml_graph_n_nodes(graph); ++i) {
+                auto node = ggml_graph_node(graph, i);
+                if (node->op == GGML_OP_UNARY && ggml_get_unary_op(node) == GGML_UNARY_OP_SIGMOID) {
+                    require(ggml_is_contiguous(node->src[0]) == opt_in, "GLU gate layout opt-in");
+                    found_sigmoid = true;
+                }
+            }
+            require(found_sigmoid, "GLU sigmoid node");
+            runner.allocate_tensors();
+            core::write_tensor_f32(input, data);
+            require_allclose(runner.run_f32(output), expected, 1e-5f, "GLU output");
+        }
+    }
+}
+
+void test_depthwise_subsampling_stage_masks() {
+    using namespace engine;
+    ModuleRunner runner;
+    auto input = runner.make_f32(core::TensorShape::from_dims({1, 8, 4}));
+    auto first_mask = runner.make_i32(core::TensorShape::from_dims({1, 4}));
+    auto second_mask = runner.make_i32(core::TensorShape::from_dims({1, 2}));
+    auto kernel = runner.make_f32(core::TensorShape::from_dims({1, 1, 1, 1}));
+    auto bias = runner.make_f32(core::TensorShape::from_dims({1}));
+    auto depth_bias = runner.make_f32(core::TensorShape::from_dims({1}));
+    auto point_kernel = runner.make_f32(core::TensorShape::from_dims({1, 1, 1, 1}));
+    auto point_bias = runner.make_f32(core::TensorShape::from_dims({1}));
+    auto proj = runner.make_f32(core::TensorShape::from_dims({1, 1}));
+    auto proj_bias = runner.make_f32(core::TensorShape::from_dims({1}));
+    modules::DepthwiseConvSubsamplingWeights weights{
+        {kernel, bias}, {{{kernel, depth_bias}, {point_kernel, point_bias}}}, {proj, proj_bias}};
+    const modules::DepthwiseConvSubsamplingModule module({4, 1, 1, 1, 2, 0, true});
+    auto masked = module.build(runner.ctx, input, weights, {first_mask, second_mask});
+    auto unmasked = module.build(runner.ctx, input, weights);
+    require(masked.shape.dims[1] == 2, "two-stage subsampling output length");
+    bool rejected = false;
+    try {
+        module.build(runner.ctx, input, weights, {first_mask});
+    } catch (const std::runtime_error &) {
+        rejected = true;
+    }
+    require(rejected, "subsampling rejects missing stage masks");
+    runner.allocate_tensors();
+    std::vector<float> data(32);
+    for (size_t i = 0; i < data.size(); ++i) data[i] = static_cast<float>(i + 1);
+    core::write_tensor_f32(input, data);
+    core::write_tensor_f32(kernel, {1});
+    core::write_tensor_f32(bias, {1});
+    core::write_tensor_f32(depth_bias, {2});
+    core::write_tensor_f32(point_kernel, {3});
+    core::write_tensor_f32(point_bias, {5});
+    core::write_tensor_f32(proj, {2});
+    core::write_tensor_f32(proj_bias, {7});
+    core::write_tensor_i32(first_mask, {1, 1, 1, 0});
+    core::write_tensor_i32(second_mask, {1, 0});
+    require_allclose(runner.run_f32(masked), {41, 7}, 1e-6f, "subsampling tail mask");
+    require_allclose(runner.run_f32(unmasked), {41, 137}, 1e-6f, "subsampling unmasked");
+    core::write_tensor_i32(first_mask, {1, 1, 0, 0});
+    core::write_tensor_i32(second_mask, {1, 1});
+    require_allclose(runner.run_f32(masked), {41, 29}, 1e-6f, "subsampling earlier stage mask");
+}
+
+class AffineTestSource final : public engine::assets::TensorSource {
+public:
+    std::unordered_map<std::string, std::vector<float>> values;
+    const std::filesystem::path & source_path() const noexcept override { return path_; }
+    bool has_tensor(std::string_view name) const noexcept override { return values.count(std::string(name)) != 0; }
+    engine::assets::TensorMetadata require_metadata(std::string_view name) const override {
+        return {std::string(name), "f32", {static_cast<int64_t>(values.at(std::string(name)).size())}};
+    }
+    std::vector<engine::assets::TensorMetadata> tensors() const override {
+        std::vector<engine::assets::TensorMetadata> result;
+        for (const auto & entry : values) result.push_back(require_metadata(entry.first));
+        return result;
+    }
+    engine::assets::RawTensorData require_tensor_data(std::string_view name) const override {
+        engine::assets::RawTensorData result;
+        result.metadata = require_metadata(name);
+        const auto & data = values.at(std::string(name));
+        result.bytes.resize(data.size() * sizeof(float));
+        std::memcpy(result.bytes.data(), data.data(), result.bytes.size());
+        return result;
+    }
+    std::vector<float> require_f32(std::string_view name,
+        const std::optional<std::vector<int64_t>> & shape) const override {
+        require(!shape || *shape == require_metadata(name).shape, "affine source shape");
+        return values.at(std::string(name));
+    }
+    std::optional<std::vector<float>> optional_f32(std::string_view name,
+        const std::optional<std::vector<int64_t>> & shape) const override {
+        return has_tensor(name) ? std::optional<std::vector<float>>(require_f32(name, shape)) : std::nullopt;
+    }
+    int64_t require_i64_scalar(std::string_view) const override { throw std::runtime_error("not an integer tensor"); }
+private:
+    std::filesystem::path path_{"affine-test"};
+};
+
+void test_batch_norm_eval_binding() {
+    using namespace engine;
+    ModuleRunner runner;
+    core::BackendWeightStore store(runner.backend, core::BackendType::Cpu, "affine-test", 1024 * 1024);
+    AffineTestSource source;
+    source.values = {{"bn.weight", {2, -3}}, {"bn.bias", {1, 2}},
+        {"bn.running_mean", {4, -2}}, {"bn.running_var", {3, 8}}};
+    auto weights = modules::binding::batch_norm_eval_from_source(store, source, "bn", 2, 1.0f);
+    store.upload();
+    std::vector<float> scale, bias;
+    core::read_tensor_f32_into(weights.scale.tensor, scale);
+    core::read_tensor_f32_into(weights.bias.tensor, bias);
+    require_allclose(scale, {1, -1}, 1e-6f, "batch norm eval scale");
+    require_allclose(bias, {-3, 0}, 1e-6f, "batch norm eval bias");
+}
+
+void test_feed_forward_activation_opt_in() {
+    using namespace engine;
+    ModuleRunner runner;
+    auto input = runner.make_f32(core::TensorShape::from_dims({1, 1, 4}));
+    auto identity = runner.make_f32(core::TensorShape::from_dims({4, 4}));
+    modules::FeedForwardWeights weights{identity, std::nullopt, identity, std::nullopt};
+    modules::FeedForwardConfig config{4, 4, false};
+    require(config.activation == modules::FeedForwardActivation::Gelu, "default activation remains GELU");
+    auto legacy = modules::FeedForwardModule(config).build(runner.ctx, input, weights);
+    config.activation = modules::FeedForwardActivation::Relu;
+    auto relu = modules::FeedForwardModule(config).build(runner.ctx, input, weights);
+    runner.allocate_tensors();
+    const std::vector<float> data{-2, -1, 1, 2};
+    core::write_tensor_f32(input, data);
+    core::write_tensor_f32(identity, {1,0,0,0, 0,1,0,0, 0,0,1,0, 0,0,0,1});
+    std::vector<float> expected;
+    for (float v : data) expected.push_back(0.5f * v * (1.0f + std::erf(v / std::sqrt(2.0f))));
+    require_allclose(runner.run_f32(legacy), expected, 1e-6f, "legacy GELU");
+    require_allclose(runner.run_f32(relu), {0,0,1,2}, 1e-6f, "opt-in ReLU");
+}
+
+void test_cached_decoder_block_matches_composition(bool flash_cross) {
+    using namespace engine;
+    using core::TensorShape;
+    constexpr int64_t hidden = 32;
+    ModuleRunner runner;
+    std::vector<std::pair<core::TensorValue, std::vector<float>>> initializers;
+    auto param = [&](const TensorShape & shape) {
+        auto tensor = runner.make_f32(shape);
+        initializers.emplace_back(tensor, make_patterned_f32(shape.num_elements(), 0.37f, 0.07f));
+        return tensor;
+    };
+    auto norm = [&]() -> modules::NormWeights {
+        auto weight = param(TensorShape::from_dims({hidden}));
+        initializers.back().second.assign(hidden, 1.0f);
+        return {weight, param(TensorShape::from_dims({hidden}))};
+    };
+    modules::TransformerDecoderBlockWeights w;
+    w.norm1 = norm(); w.norm2 = norm(); w.norm3 = norm();
+    w.self_attention.qkv_weight = param(TensorShape::from_dims({3 * hidden, hidden}));
+    w.self_attention.qkv_bias = param(TensorShape::from_dims({3 * hidden}));
+    w.self_attention.out_weight = param(TensorShape::from_dims({hidden, hidden}));
+    w.self_attention.out_bias = param(TensorShape::from_dims({hidden}));
+    w.cross_attention.q_weight = param(TensorShape::from_dims({hidden, hidden}));
+    w.cross_attention.q_bias = param(TensorShape::from_dims({hidden}));
+    w.cross_attention.out_weight = param(TensorShape::from_dims({hidden, hidden}));
+    w.cross_attention.out_bias = param(TensorShape::from_dims({hidden}));
+    w.feed_forward = {param(TensorShape::from_dims({64, hidden})), param(TensorShape::from_dims({64})),
+        param(TensorShape::from_dims({hidden, 64})), param(TensorShape::from_dims({hidden}))};
+    auto input = runner.make_f32(TensorShape::from_dims({1, 1, hidden}));
+    auto slot = runner.make_i32(TensorShape::from_dims({1}));
+    auto mask = core::make_tensor(runner.ctx, GGML_TYPE_F16, TensorShape::from_dims({1, 8}));
+    auto memory_mask = runner.make_i32(TensorShape::from_dims({1, 3}));
+    auto flash_mask = core::make_tensor(runner.ctx, GGML_TYPE_F16, TensorShape::from_dims({1, 1, 1, 3}));
+    modules::CrossAttentionKeyValue memory{param(TensorShape::from_dims({1, 1, 3, hidden})),
+        param(TensorShape::from_dims({1, 1, 3, hidden}))};
+    std::vector<core::TensorValue> caches;
+    for (int i = 0; i < 4; ++i) caches.push_back(core::make_tensor(runner.ctx, GGML_TYPE_F16,
+        TensorShape::from_dims({1, 8, 1, hidden})));
+    modules::TransformerDecoderBlockConfig config{hidden, 1, 64};
+    require(!config.use_flash_cross_attention, "cached cross flash remains opt-in");
+    config.use_flash_cross_attention = flash_cross;
+    config.activation = modules::FeedForwardActivation::Relu;
+    config.use_packed_qkv = true;
+    config.use_packed_kv = true;
+    auto actual = modules::TransformerDecoderBlockModule(config).build_cached_tail(runner.ctx, input, w,
+        caches[0], caches[1], slot, mask, memory, flash_cross ? flash_mask : memory_mask);
+    modules::AttentionConfig attention{hidden, 1, true};
+    attention.causal = true;
+    attention.use_packed_qkv = true;
+    auto y = modules::LayerNormModule({hidden}).build(runner.ctx, input, w.norm1);
+    y = modules::SelfAttentionModule(attention).build_cached_tail(runner.ctx, y, w.self_attention,
+        caches[2], caches[3], slot, mask).output;
+    auto expected = modules::AddModule().build(runner.ctx, input, y);
+    y = modules::LayerNormModule({hidden}).build(runner.ctx, expected, w.norm2);
+    modules::AttentionConfig cross_attention{hidden, 1, true};
+    cross_attention.use_packed_kv = true;
+    y = modules::CrossAttentionModule(cross_attention).build_cached(runner.ctx, y, memory, w.cross_attention, memory_mask);
+    expected = modules::AddModule().build(runner.ctx, expected, y);
+    y = modules::LayerNormModule({hidden}).build(runner.ctx, expected, w.norm3);
+    y = modules::LinearModule({hidden, 64, true}).build(runner.ctx, y, {w.feed_forward.fc1_weight, w.feed_forward.fc1_bias});
+    y = modules::ReluModule().build(runner.ctx, y);
+    y = modules::LinearModule({64, hidden, true}).build(runner.ctx, y, {w.feed_forward.fc2_weight, w.feed_forward.fc2_bias});
+    expected = modules::AddModule().build(runner.ctx, expected, y);
+    runner.allocate_tensors();
+    for (const auto & item : initializers) core::write_tensor_f32(item.first, item.second);
+    for (const auto & cache : caches) core::write_tensor_f16(cache, std::vector<float>(8 * hidden, 0));
+    core::write_tensor_i32(memory_mask, {1, 1, 0});
+    core::write_tensor_f16(flash_mask, {0, 0, -std::numeric_limits<float>::infinity()});
+    for (int32_t step = 0; step < 3; ++step) {
+        core::write_tensor_f32(input, make_patterned_f32(hidden, 0.4f + step, 0.3f));
+        core::write_tensor_i32(slot, &step, 1);
+        std::vector<float> causal(8, -std::numeric_limits<float>::infinity());
+        std::fill_n(causal.begin(), step + 1, 0.0f);
+        core::write_tensor_f16(mask, causal);
+        require_allclose(runner.run_f32(actual), runner.run_f32(expected), 1e-5f, "cached decoder composition");
+    }
+}
+
+void test_cached_cross_flash_masks() {
+    using namespace engine;
+    using core::TensorShape;
+    ModuleRunner runner;
+    modules::AttentionConfig config{64, 2, false};
+    config.use_packed_kv = true;
+    const modules::CrossAttentionModule module(config);
+    auto query = runner.make_f32(TensorShape::from_dims({2, 2, 64}));
+    auto weight = runner.make_f32(TensorShape::from_dims({64, 64}));
+    modules::AttentionWeights weights;
+    weights.q_weight = weight;
+    weights.out_weight = weight;
+    modules::CrossAttentionKeyValue kv{
+        runner.make_f32(TensorShape::from_dims({2, 2, 3, 32})),
+        runner.make_f32(TensorShape::from_dims({2, 2, 3, 32}))};
+    auto keep = runner.make_i32(TensorShape::from_dims({2, 3}));
+    auto mask = core::make_tensor(runner.ctx, GGML_TYPE_F16, TensorShape::from_dims({2, 1, 2, 3}));
+    auto shared_mask = core::make_tensor(runner.ctx, GGML_TYPE_F16, TensorShape::from_dims({1, 1, 2, 3}));
+    auto per_head_mask = core::make_tensor(runner.ctx, GGML_TYPE_F16, TensorShape::from_dims({2, 2, 2, 3}));
+    auto expected = module.build_cached(runner.ctx, query, kv, weights, keep);
+    auto actual = module.build_cached_flash(runner.ctx, query, kv, weights, mask);
+    auto shared = module.build_cached_flash(runner.ctx, query, kv, weights, shared_mask);
+    auto per_head = module.build_cached_flash(runner.ctx, query, kv, weights, per_head_mask);
+    for (const auto & invalid : {keep,
+        runner.make_f32(TensorShape::from_dims({2, 1, 2, 3})),
+        core::make_tensor(runner.ctx, GGML_TYPE_F16, TensorShape::from_dims({2, 1, 1, 3})),
+        core::make_tensor(runner.ctx, GGML_TYPE_F16, TensorShape::from_dims({2, 1, 2, 4}))}) {
+        bool rejected = false;
+        try { module.build_cached_flash(runner.ctx, query, kv, weights, invalid); }
+        catch (const std::runtime_error &) { rejected = true; }
+        require(rejected, "cached flash rejects invalid mask contracts");
+    }
+    runner.allocate_tensors();
+    std::vector<float> identity(64 * 64, 0);
+    for (size_t i = 0; i < 64; ++i) identity[i * 64 + i] = 1;
+    core::write_tensor_f32(weight, identity);
+    core::write_tensor_f32(query, make_patterned_f32(256, 0.2f, 0.3f));
+    core::write_tensor_f32(kv.key, make_patterned_f32(384, 0.4f, 0.2f));
+    core::write_tensor_f32(kv.value, make_patterned_f32(384, 0.7f, 0.1f));
+    const float excluded = -std::numeric_limits<float>::infinity();
+    core::write_tensor_i32(keep, {1, 1, 0, 1, 0, 1});
+    core::write_tensor_f16(mask, {0, 0, excluded, 0, 0, excluded, 0, excluded, 0, 0, excluded, 0});
+    require_allclose(runner.run_f32(actual), runner.run_f32(expected), 1e-5f, "batched cached flash padding");
+    core::write_tensor_i32(keep, {1, 1, 0, 1, 1, 0});
+    core::write_tensor_f16(shared_mask, {0, 0, excluded, 0, 0, excluded});
+    std::vector<float> per_head_values(24);
+    for (size_t i = 0; i < per_head_values.size(); ++i) per_head_values[i] = i % 3 == 2 ? excluded : 0;
+    core::write_tensor_f16(per_head_mask, per_head_values);
+    require_allclose(runner.run_f32(shared), runner.run_f32(expected), 1e-5f, "shared cached flash mask");
+    require_allclose(runner.run_f32(per_head), runner.run_f32(expected), 1e-5f, "per-head cached flash mask");
+}
+
 }  // namespace
 
 int main() {
     try {
+        test_glu_contiguous_gate_opt_in();
+        test_depthwise_subsampling_stage_masks();
+        test_batch_norm_eval_binding();
+        test_feed_forward_activation_opt_in();
+        test_cached_decoder_block_matches_composition(false);
+        test_cached_decoder_block_matches_composition(true);
+        test_cached_cross_flash_masks();
         test_graph_optimizer_elides_metadata_nodes_without_changing_output();
         test_graph_optimizer_two_sided_broadcast_binary_matches_repeat();
         test_graph_optimizer_unary_broadcast_scale_matches_repeat();