Function sam3_twoway_block_forward

sam3.cpp:10442–10508

TwoWayAttentionBlock forward

// TwoWayAttentionBlock forward
static void sam3_twoway_block_forward(
    struct ggml_context* ctx,
    struct ggml_tensor*& queries,  // [D, N_q, B] — modified in place
    struct ggml_tensor*& keys,     // [D, N_kv, B] — modified in place
    struct ggml_tensor* query_pe,  // [D, N_q, B]
    struct ggml_tensor* key_pe,    // [D, N_kv, B]
    const sam3_twoway_block& blk,
    int n_heads,
    bool skip_first_layer_pe) {
    // 1. Self-attention on queries
    if (skip_first_layer_pe) {
        // Python: queries = self.self_attn(q=queries, k=queries, v=queries)
        // No residual connection when skipping first layer PE
        queries = sam3_sam_attention(ctx, queries, queries, queries, blk.self_attn, n_heads);
    } else {
        auto* q = ggml_add(ctx, queries, query_pe);
        auto* attn_out = sam3_sam_attention(ctx, q, q, queries, blk.self_attn, n_heads);
        queries = ggml_add(ctx, queries, attn_out);
    }
    queries = sam3_layer_norm(ctx, queries, blk.norm1_w, blk.norm1_b);
    if (skip_first_layer_pe) {
        ggml_set_name(queries, "dbg_twoway_skip_sa_norm");
        ggml_set_output(queries);
    }

    // 2. Cross-attention: tokens attending to image
    {
        auto* q = ggml_add(ctx, queries, query_pe);
        auto* k = ggml_add(ctx, keys, key_pe);
        auto* attn_out = sam3_sam_attention(ctx, q, k, keys, blk.ca_tok2img, n_heads);
        queries = ggml_add(ctx, queries, attn_out);
        queries = sam3_layer_norm(ctx, queries, blk.norm2_w, blk.norm2_b);
    }
    if (skip_first_layer_pe) {
        ggml_set_name(queries, "dbg_twoway_skip_ca_tok2img");
        ggml_set_output(queries);
    }

    // 3. MLP on queries (ReLU activation)
    {
        auto* mlp = ggml_mul_mat(ctx, blk.mlp_fc1_w, queries);
        mlp = ggml_add(ctx, mlp, blk.mlp_fc1_b);
        mlp = ggml_relu(ctx, mlp);
        mlp = ggml_mul_mat(ctx, blk.mlp_fc2_w, mlp);
        mlp = ggml_add(ctx, mlp, blk.mlp_fc2_b);
        queries = ggml_add(ctx, queries, mlp);
        queries = sam3_layer_norm(ctx, queries, blk.norm3_w, blk.norm3_b);
    }
    if (skip_first_layer_pe) {
        ggml_set_name(queries, "dbg_twoway_skip_mlp");
        ggml_set_output(queries);
    }

    // 4. Cross-attention: image attending to tokens
    {
        auto* q = ggml_add(ctx, queries, query_pe);
        auto* k = ggml_add(ctx, keys, key_pe);
        // Note: q and k are swapped — image (k) attends to tokens (q)