Function edgetam_perceiver_layer_forward

sam3.cpp:5437–5539

Helper: build one perceiver layer on latents, given features x. latents: [D, N_lat, batch] (D=64, N_lat = n latent tokens) x: [D, N_x, batch] (D=64, N_x = n feature tokens) pos: [D, N_x, batch] or nullptr (added to K and V after projection) Returns updated latents [D, N_lat, batch].

// pos:     [D, N_x,   batch]   or nullptr (added to K and V after projection)
// Returns updated latents [D, N_lat, batch].
static struct ggml_tensor* edgetam_perceiver_layer_forward(
        struct ggml_context* ctx,
        struct ggml_tensor*  latents,
        struct ggml_tensor*  x,
        struct ggml_tensor*  pos,
        const edgetam_perceiver_layer& layer) {
    const int64_t D     = latents->ne[0];  // 64
    const int64_t N_lat = latents->ne[1];  // 256 (1D) or 1 (2D per window)
    const int64_t batch = latents->ne[2];  // 1 (1D) or 256 (2D)
    const int64_t N_x   = x->ne[1];       // H*W (1D) or 16 (2D per window)
    const float scale   = 1.0f / sqrtf((float)D);  // 0.125 for D=64

    // ── Cross-attention: latents attend to features ─────────────────────
    {
        auto* lat_n = sam3_layer_norm(ctx, latents, layer.ca_norm_latents_w,
                                      layer.ca_norm_latents_b);
        auto* x_n   = sam3_layer_norm(ctx, x, layer.ca_norm_x_w, layer.ca_norm_x_b);

        // Q from latents: [D, N_lat, batch] → [D, N_lat, batch]
        auto* q = ggml_mul_mat(ctx, layer.ca_q_w, lat_n);  // [D, N_lat, batch]

        // KV from features: [2*D, N_x, batch]
        auto* kv = ggml_mul_mat(ctx, layer.ca_kv_w, x_n);  // [2*D, N_x, batch]

        // Split into K and V via views — kv is contiguous [2*D, N_x, batch]
        auto* k_raw = ggml_view_3d(ctx, kv, D, N_x, batch,
                                   kv->nb[1], kv->nb[2], 0);               // [D, N_x, batch]
        auto* v_raw = ggml_view_3d(ctx, kv, D, N_x, batch,
                                   kv->nb[1], kv->nb[2], D * sizeof(float)); // [D, N_x, batch]

        // Add positional encoding to K and V
        struct ggml_tensor* k;
        struct ggml_tensor* v;
        if (pos) {
            k = ggml_add(ctx, k_raw, pos);
            v = ggml_add(ctx, v_raw, pos);
        } else {
            k = ggml_cont(ctx, k_raw);
            v = ggml_cont(ctx, v_raw);
        }

        // Reshape for flash_attn_ext (1 head):
        // Q: [D, N_lat, 1, batch]  K: [D, N_x, 1, batch]  V: [D, N_x, 1, batch]
        q = ggml_reshape_4d(ctx, q, D, N_lat, 1, batch);
        k = ggml_reshape_4d(ctx, k, D, N_x,   1, batch);
        v = ggml_reshape_4d(ctx, v, D, N_x,   1, batch);

        auto* attn = ggml_flash_attn_ext(ctx, q, k, v, nullptr, scale, 0.0f, 0.0f);
        // Output: [D, 1, N_lat, batch] (permuted) → reshape to [D, N_lat, batch]
        attn = ggml_reshape_3d(ctx, attn, D, N_lat, batch);

        // Output projection
        auto* ca_out = ggml_mul_mat(ctx, layer.ca_out_w, attn);  // [D, N_lat, batch]

        // Residual
        latents = ggml_add(ctx, latents, ca_out);
    }