Function build_wavlm_self_attention

src/framework/modules/wavlm_encoder.cpp:407–472

namespace {

core::TensorValue build_wavlm_self_attention(
    core::ModuleBuildContext & ctx,
    const core::TensorValue & hidden_btc,
    const core::TensorValue & position_bias,
    const core::TensorValue & attention_mask,
    const WavlmEncoderWeights & weights,
    int64_t layer_index) {
    const auto & config = weights.config;
    const int64_t batch = hidden_btc.shape.dims[0];
    const int64_t tokens = hidden_btc.shape.dims[1];
    const int64_t head_dim = config.hidden_size / config.num_attention_heads;
    const std::string prefix = "encoder.layers." + std::to_string(layer_index) + ".attention";

    auto q = LinearModule({config.hidden_size, config.hidden_size, true, GGML_PREC_F32})
                 .build(ctx, hidden_btc, linear_weights(weights, prefix + ".q_proj"));
    auto k = LinearModule({config.hidden_size, config.hidden_size, true, GGML_PREC_F32})
                 .build(ctx, hidden_btc, linear_weights(weights, prefix + ".k_proj"));
    auto v = LinearModule({config.hidden_size, config.hidden_size, true, GGML_PREC_F32})
                 .build(ctx, hidden_btc, linear_weights(weights, prefix + ".v_proj"));

    q = core::reshape_tensor(ctx, contiguous(ctx, q), core::TensorShape::from_dims({batch, tokens, config.num_attention_heads, head_dim}));
    k = core::reshape_tensor(ctx, contiguous(ctx, k), core::TensorShape::from_dims({batch, tokens, config.num_attention_heads, head_dim}));
    v = core::reshape_tensor(ctx, contiguous(ctx, v), core::TensorShape::from_dims({batch, tokens, config.num_attention_heads, head_dim}));
    q = TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, q);
    k = TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, k);
    v = TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, v);

    auto query_layer = core::reshape_tensor(
        ctx,
        contiguous(ctx, hidden_btc),
        core::TensorShape::from_dims({batch, tokens, config.num_attention_heads, head_dim}));
    query_layer = TransposeModule({{0, 2, 1, 3}, 4}).build(ctx, query_layer);
    auto gates = LinearModule({head_dim, 8, true, GGML_PREC_F32})
                     .build(ctx, query_layer, linear_weights(weights, prefix + ".gru_rel_pos_linear"));
    auto gate_a = SliceModule({3, 0, 4}).build(ctx, gates);
    auto gate_b = SliceModule({3, 4, 4}).build(ctx, gates);
    gate_a = SigmoidModule().build(ctx, ReduceSumModule({3}).build(ctx, gate_a));
    gate_b = SigmoidModule().build(ctx, ReduceSumModule({3}).build(ctx, gate_b));
    gate_a = core::reshape_tensor(ctx, contiguous(ctx, gate_a), core::TensorShape::from_dims({batch, config.num_attention_heads, tokens, 1}));
    gate_b = core::reshape_tensor(ctx, contiguous(ctx, gate_b), core::TensorShape::from_dims({batch, config.num_attention_heads, tokens, 1}));
    auto rel_const = core::reshape_tensor(
        ctx,
        require_tensor(weights, prefix + ".gru_rel_pos_const"),
        core::TensorShape::from_dims({1, config.num_attention_heads, 1, 1}));
    rel_const = RepeatModule({gate_b.shape}).build(ctx, rel_const);
    auto gated = MulModule().build(ctx, gate_b, rel_const);
    gated = add_scalar(ctx, gated, -1.0F);
    gated = MulModule().build(ctx, gate_a, gated);
    gated = add_scalar(ctx, gated, 2.0F);
    gated = RepeatModule({position_bias.shape}).build(ctx, gated);
    auto rel_bias = MulModule().build(ctx, gated, position_bias);

    const auto k_t = TransposeModule({{0, 1, 3, 2}, 4}).build(ctx, k);
    auto scores = MatMulModule().build(ctx, q, k_t);
    scores = mul_scalar(ctx, scores, static_cast<float>(1.0 / std::sqrt(static_cast<double>(head_dim))));
    scores = add_same(ctx, scores, rel_bias);
    auto attn = core::wrap_tensor(
        ggml_soft_max_ext(ctx.ggml, contiguous(ctx, scores).tensor, attention_mask.tensor, 1.0F, 0.0F),