char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Pedro Cuencamodel: Muse Glimmer Support (#26841)62bf73d25

char/aliceai-t5-moe
4.1 KiB88 linesraw
1#include "models.h"
2
3// MuseGlimmer vision encoder: 50-layer ViT with 2D RoPE, sparse block-diagonal
4// window attention (every 4th + last layer global), pixel-shuffle downsample, then
5// adapter MLP + LLM's vision_projection.
6//
7// Several quantities are precomputed on host and fed as named graph inputs (filled in
8// clip.cpp set_input, PROJECTOR_TYPE_MUSE_GLIMMER branch):
9//   muse_glimmer_pos_w/_h [n_tok] i32         : 1-indexed RoPE positions (sparse-permuted order)
10//   muse_glimmer_sp_perm  [n_tok] i32         : window grouping permutation (applied after ln_pre)
11//   muse_glimmer_inv_perm [n_tok] i32         : inverse of sp_perm (applied after blocks)
12//   muse_glimmer_ds_perm  [n_tok] i32         : pixel-shuffle gather (original order)
13//   muse_glimmer_sp_mask  [n_tok, n_tok] f32  : block-diagonal window mask (sparse layers)
14ggml_cgraph * clip_graph_muse_glimmer::build() {
15    const int ds = hparams.n_merge;              // downsample factor (2)
16    const int sf = hparams.muse_glimmer_sparse_factor;   // 4
17    const int n_tok     = n_patches;
18    const int n_out     = (n_patches_x / ds) * (n_patches_y / ds);
19    const float rope_base = hparams.rope_theta;  // 10000
20
21    auto inp_i32 = [&](const char * name, int64_t n) {
22        ggml_tensor * t = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n);
23        ggml_set_name(t, name);
24        ggml_set_input(t);
25        return t;
26    };
27
28    ggml_tensor * pos_w    = inp_i32("muse_glimmer_pos_w",    n_tok);
29    ggml_tensor * pos_h    = inp_i32("muse_glimmer_pos_h",    n_tok);
30    ggml_tensor * sp_perm  = inp_i32("muse_glimmer_sp_perm",  n_tok);
31    ggml_tensor * inv_perm = inp_i32("muse_glimmer_inv_perm", n_tok);
32    ggml_tensor * ds_perm  = inp_i32("muse_glimmer_ds_perm",  n_tok);
33
34    ggml_tensor * sp_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_tok, n_tok);
35    ggml_set_name(sp_mask, "muse_glimmer_sp_mask");
36    ggml_set_input(sp_mask);
37
38    // patchify via build_inp (conv2d over raw pixels) + bilinear-resized learned pos-emb
39    ggml_tensor * x = build_inp();                                                     // [n_embd, n_tok, 1]
40    x = ggml_add(ctx0, x, resize_position_embeddings(GGML_SCALE_MODE_BILINEAR));
41    cb(x, "after_posemb", -1);
42
43    // group patches into pgrid x pgrid windows (sparse attention order)
44    x = ggml_get_rows(ctx0, x, sp_perm);
45    cb(x, "after_sp_perm", -1);
46
47    // per-layer mask: sparse layers get sp_mask, global layers (every sf-th and last) get none
48    std::vector<ggml_tensor *> attn_mask_layers(n_layer);
49    for (int il = 0; il < n_layer; ++il) {
50        const bool is_global = (il == n_layer - 1) || ((il + 1) % sf == 0);
51        attn_mask_layers[il] = is_global ? nullptr : sp_mask;
52    }
53
54    // 2D RoPE: first half of head_dim uses width pos, second half uses height pos
55    auto add_pos = [&](ggml_tensor * cur, const clip_layer &) {
56        return build_rope_2d(ctx0, cur, pos_w, pos_h, rope_base, false);
57    };
58
59    build_vit_opts opts;
60    opts.attn_mask_layers = std::move(attn_mask_layers);
61
62    // pre_ln, per-layer transformer, post_ln (all inside build_vit); reference uses exact (erf) GELU
63    x = build_vit(x, n_tok, NORM_TYPE_NORMAL, FFN_GELU_ERF, nullptr, add_pos, opts);
64
65    // un-permute back to original grid order
66    x = ggml_get_rows(ctx0, x, inv_perm);
67    cb(x, "after_inv_perm", -1);
68
69    // pixel-shuffle downsample: gather f*f spatial neighbors then concat channel-outer.
70    // out[c*(ds*ds)+s, o] = x[ds_perm gathered][o*(ds*ds)+s, c]
71    x = ggml_get_rows(ctx0, x, ds_perm);                 // [n_embd, n_tok], grouped
72    x = ggml_reshape_3d(ctx0, x, n_embd, ds * ds, n_out);// [c, s, o]
73    x = ggml_permute(ctx0, x, 1, 0, 2, 3);               // [s, c, o]
74    x = ggml_cont(ctx0, x);
75    x = ggml_reshape_2d(ctx0, x, n_embd * ds * ds, n_out); // [6144, n_out]
76    cb(x, "encoder_out", -1);
77
78    // adapter (6144->4096->4096, exact GELU each) + LLM vision_projection (4096->6656)
79    x = build_mm(model.mm_0_w, x);
80    x = ggml_gelu_erf(ctx0, x);
81    x = build_mm(model.mm_1_w, x);
82    x = ggml_gelu_erf(ctx0, x);
83    x = build_mm(model.mm_2_w, x);                       // [6656, n_out]
84    cb(x, "projected", -1);
85
86    ggml_build_forward_expand(gf, x);
87    return gf;
88}