Function decoder_layer_with_static_cache

src/models/miotts/causal_lm.cpp:285–336

}

DecoderLayerOutputs decoder_layer_with_static_cache(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & input,
    const core::TensorValue & positions,
    const TextLayerWeights & weights,
    const MioTTSConfig & config,
    const core::TensorValue & cache_key,
    const core::TensorValue & cache_value,
    const core::TensorValue & cache_slot,
    const core::TensorValue & attention_mask) {
    const int64_t dim = head_dim(config);
    const modules::LinearModule q_proj({config.hidden_size, config.num_attention_heads * dim, false});
    const modules::LinearModule k_proj({config.hidden_size, config.num_key_value_heads * dim, false});
    const modules::LinearModule v_proj({config.hidden_size, config.num_key_value_heads * dim, false});
    const modules::LinearModule o_proj({config.num_attention_heads * dim, config.hidden_size, false});
    const modules::RMSNormModule hidden_norm({config.hidden_size, config.rms_norm_eps, true, false});

    auto x_norm = hidden_norm.build(ctx, input, {weights.input_norm, std::nullopt});
    auto q = q_proj.build(ctx, x_norm, {weights.q_proj, std::nullopt});
    auto k = k_proj.build(ctx, x_norm, {weights.k_proj, std::nullopt});
    auto v = v_proj.build(ctx, x_norm, {weights.v_proj, std::nullopt});
    q = modules::RMSNormModule({dim, config.rms_norm_eps, true, false})
            .build(ctx, reshape_heads(ctx, q, config.num_attention_heads, dim), {weights.q_norm, std::nullopt});
    k = modules::RMSNormModule({dim, config.rms_norm_eps, true, false})
            .build(ctx, reshape_heads(ctx, k, config.num_key_value_heads, dim), {weights.k_norm, std::nullopt});
    v = reshape_heads(ctx, v, config.num_key_value_heads, dim);
    q = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, q, positions);
    k = modules::RoPEModule({dim, GGML_ROPE_TYPE_NEOX, config.rope_theta}).build(ctx, k, positions);

    const modules::FastKVSetRowsModule set_rows;
    auto updated_cache_key = set_rows.build(ctx, cache_key, k, cache_slot);
    auto updated_cache_value = set_rows.build(ctx, cache_value, v, cache_slot);

    auto q_heads = modules::TransposeModule({{0, 2, 1, 3}, q.shape.rank}).build(ctx, q);
    auto k_heads = modules::TransposeModule({{0, 2, 1, 3}, updated_cache_key.shape.rank}).build(ctx, updated_cache_key);
    auto v_heads = modules::TransposeModule({{0, 2, 1, 3}, updated_cache_value.shape.rank}).build(ctx, updated_cache_value);
    auto context = flash_attention_from_grouped_heads(ctx, q_heads, k_heads, v_heads, dim, attention_mask);
    context = core::ensure_backend_addressable_layout(ctx, context);
    context = core::reshape_tensor(ctx, context, core::TensorShape::from_dims({1, 1, config.num_attention_heads * dim}));
    auto x = modules::AddModule{}.build(ctx, input, o_proj.build(ctx, context, {weights.o_proj, std::nullopt}));

    auto ff_in = hidden_norm.build(ctx, x, {weights.post_norm, std::nullopt});
    auto gate = modules::LinearModule({config.hidden_size, config.intermediate_size, false})
                    .build(ctx, ff_in, {weights.gate_proj, std::nullopt});
    gate = modules::SiluModule{}.build(ctx, gate);
    auto up = modules::LinearModule({config.hidden_size, config.intermediate_size, false})
                  .build(ctx, ff_in, {weights.up_proj, std::nullopt});
    auto gated = modules::MulModule{}.build(ctx, gate, up);
    auto ff = modules::LinearModule({config.intermediate_size, config.hidden_size, false})
                  .build(ctx, gated, {weights.down_proj, std::nullopt});
    return {modules::AddModule{}.build(ctx, x, ff), k, v};
}

MioTTSCausalLMWeights load_weights(
    const MioTTSAssets & assets,