Function sam3_text_block_forward

sam3.cpp:3957–4033

Single text encoder block forward pass. Input x: [E, L] where E=text_width=1024, L=seq_len (typically 32). causal_mask: [L, L] F16 additive mask for ggml_flash_attn_ext. Returns: [E, L]

// causal_mask: [L, L] F16 additive mask for ggml_flash_attn_ext.
// Returns: [E, L]
static struct ggml_tensor* sam3_text_block_forward(struct ggml_context* ctx,
                                                   struct ggml_tensor* x,
                                                   const sam3_text_block& blk,
                                                   const sam3_hparams& hp,
                                                   struct ggml_tensor* causal_mask,
                                                   int block_idx) {
    const int E = hp.text_width;   // 1024
    const int NH = hp.text_heads;  // 16
    const int HD = E / NH;         // 64
    const int64_t L = x->ne[1];    // sequence length

    auto* shortcut = x;
    x = sam3_layer_norm(ctx, x, blk.ln1_w, blk.ln1_b);
    sam3_name_tensorf(x, "text_block_%02d_after_ln1", block_idx);

    auto* qkv = ggml_mul_mat(ctx, blk.attn_in_proj_w, x);
    qkv = ggml_add(ctx, qkv, blk.attn_in_proj_b);
    sam3_name_tensorf(qkv, "text_block_%02d_qkv", block_idx);

    // [3*E, L] → [E, 3, L] → permute → [E, L, 3]
    qkv = ggml_reshape_3d(ctx, qkv, E, 3, L);
    qkv = ggml_cont(ctx, ggml_permute(ctx, qkv, 0, 2, 1, 3));
    // qkv: [E, L, 3]

    auto* Q = ggml_view_2d(ctx, qkv, E, L, qkv->nb[1], 0);
    auto* K = ggml_view_2d(ctx, qkv, E, L, qkv->nb[1], 1 * qkv->nb[2]);
    auto* V = ggml_view_2d(ctx, qkv, E, L, qkv->nb[1], 2 * qkv->nb[2]);

    // [E, L] → [HD, NH, L] → permute → [HD, L, NH, 1]
    Q = ggml_reshape_3d(ctx, Q, HD, NH, L);
    Q = ggml_cont(ctx, ggml_permute(ctx, Q, 0, 2, 1, 3));
    Q = ggml_reshape_4d(ctx, Q, HD, L, NH, 1);

    K = ggml_reshape_3d(ctx, K, HD, NH, L);
    K = ggml_cont(ctx, ggml_permute(ctx, K, 0, 2, 1, 3));
    K = ggml_reshape_4d(ctx, K, HD, L, NH, 1);

    V = ggml_reshape_3d(ctx, V, HD, NH, L);
    V = ggml_permute(ctx, V, 0, 2, 1, 3);  // non-contiguous; flash_attn uses strides

    float scale = 1.0f / sqrtf((float)HD);
    auto* attn_out = ggml_flash_attn_ext(ctx, Q, K, V, causal_mask, scale, 0.0f, 0.0f);
    x = ggml_reshape_2d(ctx, attn_out, E, L);

    x = ggml_mul_mat(ctx, blk.attn_out_proj_w, x);
    x = ggml_add(ctx, x, blk.attn_out_proj_b);

    if (blk.ls1) {
        x = ggml_mul(ctx, x, blk.ls1);
    }
    sam3_name_tensorf(x, "text_block_%02d_attn_out", block_idx);

    x = ggml_add(ctx, shortcut, x);
    sam3_name_tensorf(x, "text_block_%02d_after_attn_residual", block_idx);

    shortcut = x;
    x = sam3_layer_norm(ctx, x, blk.ln2_w, blk.ln2_b);
    sam3_name_tensorf(x, "text_block_%02d_after_ln2", block_idx);