GGML源碼逐行調試（下）

前言

學習 UP 主比飛鳥貴重的多_HKL 的 GGML源碼逐行調試視頻，記錄下個人學習筆記，僅供自己參考😄

refer1：【大模型部署】GGML源碼逐行調試

refer2：llama.cpp源碼解讀–ggml框架學習

refer3：https://github.com/ggml-org/ggml

refer4：https://chatgpt.com/

1. 簡述

我們接著上篇文章來講，在上篇文章中我們梳理了 ggml 推理 gpt-2 時模型加載部分的總體流程，這篇文章我們就來看下剩余的流程具體都是怎么做的

2. 預分配計算圖內存

這個小節我們來看 ggml 框架是如何分配計算圖內存的，代碼如下：

ggml_gallocr_t allocr = NULL;
// allocate the compute buffer
{// create a graph allocator with the backend's default buffer typeallocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend));// create the worst case graph for memory usage estimationint n_tokens = std::min(model.hparams.n_ctx, params.n_batch);int n_past = model.hparams.n_ctx - n_tokens;struct ggml_cgraph * gf = gpt2_graph(model, n_past, n_tokens);// pre-allocate the compute buffer for the worst case (optional)ggml_gallocr_reserve(allocr, gf);size_t mem_size =  ggml_gallocr_get_buffer_size(allocr, 0);fprintf(stderr, "%s: compute buffer size: %.2f MB\n", __func__, mem_size/1024.0/1024.0);
}

整體流程如下：(from ChatGPT)

1. 創建圖內存分配器

allocr = ggml_gallocr_new(ggml_backend_get_default_buffer_type(model.backend));

使用后端默認的緩沖區類型，通過調用 ggml_gallocr_new 創建一個圖內存分配器（allocr）

2. 構建最壞情況的計算圖

int n_tokens = std::min(model.hparams.n_ctx, params.n_batch);
int n_past = model.hparams.n_ctx - n_tokens;
struct ggml_cgraph * gf = gpt2_graph(model, n_past, n_tokens);

為了估計最壞情況下所需的計算內存，代碼先根據模型上下文長度以及 batch 參數計算出當前 token 數和 past 數，然后調用 gpt2_graph 來構造一個最壞情況的計算圖

3. 預留計算圖內存

ggml_gallocr_reserve(allocr, gf);
size_t mem_size =  ggml_gallocr_get_buffer_size(allocr, 0);
fprintf(stderr, "%s: compute buffer size: %.2f MB\n", __func__, mem_size/1024.0/1024.0);

利用圖內存分配器和構造的最壞情況計算圖，調用 ggml_gallocr_reserve 對計算圖內所有節點所需的內存進行預分配

整個流程從后端緩沖區類型出發，創建一個專門用于計算圖內存管理的分配器，接著通過最壞情況的計算圖來精確估計計算過程中將使用的內存，然后預留這部分內存，最后查詢并輸出整個分配結果

下面我們就來對每個流程進行具體的分析

2.1 創建圖內存分配器

這部分代碼的作用是為后續計算圖內存預分配創建一個 “圖內存分配器”（graph allocator），也就是 ggml_gallocr 對象。該對象負責管理后端緩沖區（backend buffer）相關的信息以及對應的動態張量分配器（dynamic tensor allocator），為后續的計算圖節點內存分配提供統一的接口

1. 圖內存分配器創建接口函數 ggml_gallocr_new

ggml_gallocr_t ggml_gallocr_new(ggml_backend_buffer_type_t buft) {return ggml_gallocr_new_n(&buft, 1);
}

這是一個方便的封裝函數，用于創建只有一個緩沖區類型的圖內存分配器，它接收一個后端緩沖區類型 buft

2. 圖內存分配器創建實現函數 ggml_gallocr_new_n

ggml_gallocr_t ggml_gallocr_new_n(ggml_backend_buffer_type_t * bufts, int n_bufs) {ggml_gallocr_t galloc = (ggml_gallocr_t)calloc(1, sizeof(struct ggml_gallocr));GGML_ASSERT(galloc != NULL);galloc->bufts = calloc(n_bufs, sizeof(ggml_backend_buffer_type_t));GGML_ASSERT(galloc->bufts != NULL);galloc->buffers = calloc(n_bufs, sizeof(ggml_backend_buffer_t));GGML_ASSERT(galloc->buffers != NULL);galloc->buf_tallocs = calloc(n_bufs, sizeof(struct ggml_dyn_tallocr *));GGML_ASSERT(galloc->buf_tallocs != NULL);for (int i = 0; i < n_bufs; i++) {galloc->bufts[i] = bufts[i];galloc->buffers[i] = NULL;// 檢查是否已經使用過該緩沖區類型，若是，則重用已有的動態張量分配器for (int j = 0; j < i; j++) {if (bufts[i] == bufts[j]) {galloc->buf_tallocs[i] = galloc->buf_tallocs[j];break;}}// 如果前面沒有匹配到，則新建一個動態張量分配器if (galloc->buf_tallocs[i] == NULL) {size_t alignment = ggml_backend_buft_get_alignment(bufts[i]);galloc->buf_tallocs[i] = ggml_dyn_tallocr_new(alignment);}}galloc->n_buffers = n_bufs;return galloc;
}

創建流程如下：

內存分配
- 通過 calloc 為 ggml_gallocr 結構體分配內存，并確保內存全部置 0
- 分別為存儲后端緩沖區類型的數組 bufts、后端緩沖區數組 buffers 和動態張量分配器數組 buf_tallocs 分配空間，空間大小都為 n_bufs
初始化每個緩沖區類型及其分配器
- 循環遍歷每個緩沖區
  - 將傳入的 bufts[i] 直接復制到 galloc->bufts[i] 中
  - 初始化對應的 buffers[i] 為 NULL
- 重用動態分配器
  - 內部循環檢查在前面是否已經遇到相同的緩沖區類型。如果相同，就直接將之前的動態張量分配器指針復用到當前 buf_tallocs[i] 中，這避免了對同一緩沖區類型重復創建分配器
- 新建動態張量分配器
  - 如果當前緩沖區類型第一次出現，則調用 ggml_backend_buft_get_alignment 獲取該緩沖區的對齊要求
  - 然后調用 ggml_dyn_tallocr_new 新建一個動態張量分配器，并存儲到 galloc->buf_tallocs[i] 中
設置緩沖區數量
- 最后將 n_bufs 保存在 galloc->n_buffers 字段中，并返回新建的圖內存分配器 galloc

這種設計使得圖內存分配器能夠根據實際后端的內存對齊和最大分配要求，對計算圖中可能產生的中間張量內存進行精確預估和統一管理

2.2 構建最壞情況的計算圖

這部分代碼（gpt2_graph(...)）的目的是在預分配計算圖內存之前，構造一個 “最壞情況” 下的計算圖，也就是在最極端內存需求下（例如最多 token）的圖結構，從而準確評估所需內存量。整個過程大體分為以下幾個步驟：(from ChatGPT)

1. 預留計算圖元數據緩沖區

static size_t buf_size = ggml_tensor_overhead()*GPT2_MAX_NODES + ggml_graph_overhead_custom(GPT2_MAX_NODES, false);
static std::vector<uint8_t> buf(buf_size);

為構建計算圖所需的 ggml 上下文與圖對象分配一個臨時內存緩沖區，該緩沖區只需要存儲 ggml_tensor 結構體和 ggml_cgraph 結構體，不包括實際張量數據

2. 初始化 ggml 上下文

struct ggml_init_params params = {/*.mem_size   =*/ buf_size,/*.mem_buffer =*/ buf.data(),/*.no_alloc   =*/ true,
};
struct ggml_context * ctx = ggml_init(params);

使用預先構造的內存緩沖區與初始化參數，調用 ggml_init 創建上下文。其中參數 no_alloc 為 true，表示后續張量數據將在圖內存分配器中進行統一分配，不在這里分配實際內存

3. 構造自定義計算圖對象

struct ggml_cgraph  * gf = ggml_new_graph_custom(ctx, GPT2_MAX_NODES, false);

使用函數 ggml_new_graph_custom 創建自定義計算圖節對象 gf，該對象用于記錄所有后續構造的運算節點

4. 構造輸入張量

struct ggml_tensor * embd = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, N);
ggml_set_name(embd, "embd");
ggml_set_input(embd);struct ggml_tensor * position = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, N);
ggml_set_name(position, "position");
ggml_set_input(position);

構造兩個一維張量：

embd：存放 token 索引，作為模型輸入
position：存放位置索引

并調用 ggml_set_input 標記為輸入張量

5. 構造嵌入查詢與位置嵌入加和

struct ggml_tensor * inpL =ggml_add(ctx,ggml_get_rows(ctx, model.wte, embd),ggml_get_rows(ctx, model.wpe, position));

對模型中預先加載的詞向量（wte）和位置編碼（wpe），通過 ggml_get_rows 分別取出對應于輸入 token 和位置的行，之后用 ggml_add 相加產生第一層的輸入張量 inpL

6. 進入 Transformer 層循環

for (int il = 0; il < n_layer; ++il) {struct ggml_tensor * cur;// norm: 歸一化并加權偏置{cur = ggml_norm(ctx, inpL, hparams.eps);cur = ggml_add(ctx,ggml_mul(ctx,cur,model.layers[il].ln_1_g),model.layers[il].ln_1_b);}// attn: 對輸入進行全連接變換并加偏置{cur = ggml_mul_mat(ctx, model.layers[il].c_attn_attn_w, cur);cur = ggml_add(ctx, cur, model.layers[il].c_attn_attn_b);}// self-attention: 從 cur 中按塊拆分出 Q、K、V，再進行計算與 KV 緩存存儲{struct ggml_tensor * Qcur = ggml_view_2d(ctx, cur, n_embd, N, cur->nb[1], 0*sizeof(float)*n_embd);struct ggml_tensor * Kcur = ggml_view_2d(ctx, cur, n_embd, N, cur->nb[1], 1*sizeof(float)*n_embd);struct ggml_tensor * Vcur = ggml_view_2d(ctx, cur, n_embd, N, cur->nb[1], 2*sizeof(float)*n_embd);if (N >= 1) {struct ggml_tensor * k = ggml_view_1d(ctx, model.memory_k, N*n_embd, (ggml_element_size(model.memory_k)*n_embd)*(il*n_ctx + n_past));struct ggml_tensor * v = ggml_view_1d(ctx, model.memory_v, N*n_embd, (ggml_element_size(model.memory_v)*n_embd)*(il*n_ctx + n_past));ggml_build_forward_expand(gf, ggml_cpy(ctx, Kcur, k));ggml_build_forward_expand(gf, ggml_cpy(ctx, Vcur, v));}// 對 Q、K 進行 reshape、permutestruct ggml_tensor * Q = ggml_permute(ctx,ggml_cont_3d(ctx, Qcur, n_embd/n_head, n_head, N),0, 2, 1, 3);struct ggml_tensor * K = ggml_permute(ctx,ggml_reshape_3d(ctx,ggml_view_1d(ctx, model.memory_k, (n_past + N)*n_embd, il*n_ctx*ggml_element_size(model.memory_k)*n_embd),n_embd/n_head, n_head, n_past + N),0, 2, 1, 3);// 注意力計算: K*Q, 縮放、mask、softmaxstruct ggml_tensor * KQ = ggml_mul_mat(ctx, K, Q);struct ggml_tensor * KQ_scaled = ggml_scale(ctx, KQ, 1.0f/sqrtf(float(n_embd)/n_head));struct ggml_tensor * KQ_masked = ggml_diag_mask_inf(ctx, KQ_scaled, n_past);struct ggml_tensor * KQ_soft_max = ggml_soft_max(ctx, KQ_masked);// V_trans: 對 V 進行相應的變換struct ggml_tensor * V_trans = ggml_cont_3d(ctx,ggml_permute(ctx,ggml_reshape_3d(ctx,ggml_view_1d(ctx, model.memory_v, (n_past + N)*n_embd, il*n_ctx*ggml_element_size(model.memory_v)*n_embd),n_embd/n_head, n_head, n_past + N),1, 2, 0, 3),n_past + N, n_embd/n_head, n_head);// 最后計算 KQV，即 V_trans * attention權重struct ggml_tensor * KQV = ggml_mul_mat(ctx, V_trans, KQ_soft_max);struct ggml_tensor * KQV_merged = ggml_permute(ctx, KQV, 0, 2, 1, 3);cur = ggml_cont_2d(ctx, KQV_merged, n_embd, N);}// attention projection：線性變換（矩陣乘法+加偏置）{cur = ggml_mul_mat(ctx, model.layers[il].c_attn_proj_w, cur);cur = ggml_add(ctx, cur, model.layers[il].c_attn_proj_b);}// 添加輸入的殘差連接：將計算結果與上層輸入 inpL 相加cur = ggml_add(ctx, cur, inpL);// Feed Forward 部分（前饋網絡）：先歸一化、全連接、GELU激活，再做投影和加殘差struct ggml_tensor * inpFF = cur;{{cur = ggml_norm(ctx, inpFF, hparams.eps);cur = ggml_add(ctx,ggml_mul(ctx, cur, model.layers[il].ln_2_g),model.layers[il].ln_2_b);}cur = ggml_mul_mat(ctx, model.layers[il].c_mlp_fc_w, cur);cur = ggml_add(ctx, cur, model.layers[il].c_mlp_fc_b);cur = ggml_gelu(ctx, cur);cur = ggml_mul_mat(ctx, model.layers[il].c_mlp_proj_w, cur);cur = ggml_add(ctx, cur, model.layers[il].c_mlp_proj_b);}// 再次殘差連接：將前饋網絡的輸出與輸入 inpFF 相加inpL = ggml_add(ctx, cur, inpFF);
}

對于每個 Transformer 層（總層數為 n_layer），執行如下主要操作：

LayerNorm 與殘差連接（前置）
多頭自注意力
- 先將上一層輸出（inpL）經過歸一化（ggml_norm）
- 經過全連接變換生成 attention 輸入
- 使用 ggml_view_2d 從 cur 中分離出 Q、K、V 向量
- 利用 ggml_view_1d 將當前層的 K 與 V 存入模型內部的 kv cache 中（這里利用 n_ctx、n_past 參數定位到正確的偏移）
- 對 Q、K 進行 reshape 和 permute 操作，準備進行注意力點積計算
- 計算注意力得分（ $QK^T$ ）、進行縮放（ $\frac{1}{\sqrt{D}}$ ）、mask 掩碼、softmax 得到注意力權重
- 對 V 進行變換，最終通過矩陣乘法得到注意力層輸出
- 對輸出進行排列（permute）和變形（contiguous 轉為 2D 張量）
Attention Projection
- 將多頭自注意力輸出進行線性變換
殘差連接與后續 FFN
- 將輸出結果與輸入相加（殘差連接）
- 接著進行前饋網絡操作：LayerNorm、全連接、GELU 激活，再做一次全連接投影，并同樣進行殘差連接
更新 inpL 為當前層輸出，為下一層做輸入

7. 最終歸一化與輸出映射

{inpL = ggml_norm(ctx, inpL, hparams.eps);inpL = ggml_add(ctx,ggml_mul(ctx, inpL, model.ln_f_g),model.ln_f_b);
}
inpL = ggml_mul_mat(ctx, model.lm_head, inpL);
ggml_set_name(inpL, "logits");
ggml_set_output(inpL);

在所有 Transformer 層之后，先對最終的 inpL 進行歸一化，并使用最后的歸一化參數 ln_f_g 和 ln_f_b 對其進行變換。接著將歸一化后的張量與模型的 lm_head 做矩陣乘法，生成 logits 輸出，并將該張量設置為 graph 的輸出。

8. 構建前向傳播圖并清理上下文

ggml_build_forward_expand(gf, inpL);
ggml_free(ctx);
return gf;

最后通過 ggml_build_forward_expand(...) 將計算圖的最后一個節點（logits）添加至圖中完成前向傳播路徑構建，然后調用 ggml_free(ctx) 釋放上下文中的臨時內存，返回構造好的計算圖 gf

通過構建這樣一個最壞情況的計算圖（采用最壞的 tokens 數量、最大 n_ctx），上層代碼就能通過預分配接口準確估算計算過程中內存需求，為后續推理的高效執行做好內存預留工作

Note：在這里整個 graph 計算圖是手動構建的，這有點類似于 tensorrtx 這個 repo，tensorrtx 是通過 TensorRT 的 C++ API 手動來構建模型

2.3 預留計算圖內存

這部分代碼的主要目標是為構建好的計算圖（graph）預先計算并分配好所有中間節點和葉子節點所需內存，利用圖內存分配器（gallocr）對計算圖中各個張量的內存需求進行 “預留”，代碼通過以下幾個步驟來實現這一目標：

1. 初始化 Hash 表用于內存分配記錄

size_t min_hash_size = graph->n_nodes + graph->n_leafs;
min_hash_size += min_hash_size / 4;  // 增加 25% 余量以避免哈希沖突

計算需要的 hash 表大小，即總節點數（計算節點和葉子節點之和）再加上 25% 的余量，旨在降低哈希沖突的概率

2. 初始化并重置 Hash 表

if (galloc->hash_set.size < min_hash_size) {ggml_hash_set_free(&galloc->hash_set);galloc->hash_set = ggml_hash_set_new(min_hash_size);GGML_ASSERT(galloc->hash_set.keys != NULL);free(galloc->hash_values);galloc->hash_values = malloc(sizeof(struct hash_node) * galloc->hash_set.size);GGML_ASSERT(galloc->hash_values != NULL);
}

檢查當前圖內存分配器中的 hash set 是否足夠大。如果不足，則先釋放舊的，再創建一個新哈希集合，保證其大小能容納所有節點信息

3. 重置所有動態張量分配器

for (int i = 0; i < galloc->n_buffers; i++) {ggml_dyn_tallocr_reset(galloc->buf_tallocs[i]);
}

對每個后端緩沖區對應的動態張量分配器進行重置，將內部的空閑塊初始化為初始狀態

4. 分配圖中各個節點和葉子的內存

ggml_gallocr_alloc_graph_impl(galloc, graph, node_buffer_ids, leaf_buffer_ids);

這個函數內部會遍歷計算圖中所有的節點（nodes）和葉子（leafs），為每個 tensor 計算出對應的內存分配（buffer_id, offset 等），并將這些信息記錄到前面分配的 Hash 表中

5. 建立 node_allocs 數組

if (galloc->n_nodes < graph->n_nodes) {free(galloc->node_allocs);galloc->node_allocs = calloc(graph->n_nodes, sizeof(struct node_alloc));GGML_ASSERT(galloc->node_allocs != NULL);
}
galloc->n_nodes = graph->n_nodes;
for (int i = 0; i < graph->n_nodes; i++) {struct ggml_tensor * node = graph->nodes[i];struct node_alloc * node_alloc = &galloc->node_allocs[i];if (node->view_src || node->data) {node_alloc->dst.buffer_id = -1;node_alloc->dst.offset = SIZE_MAX;node_alloc->dst.size_max = 0;} else {struct hash_node * hn = ggml_gallocr_hash_get(galloc, node);node_alloc->dst.buffer_id = hn->buffer_id;node_alloc->dst.offset    = hn->offset;node_alloc->dst.size_max  = ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], node);}for (int j = 0; j < GGML_MAX_SRC; j++) {struct ggml_tensor * src = node->src[j];if (!src || src->view_src || src->data) {node_alloc->src[j].buffer_id = -1;node_alloc->src[j].offset = SIZE_MAX;node_alloc->src[j].size_max = 0;} else {struct hash_node * hn = ggml_gallocr_hash_get(galloc, src);node_alloc->src[j].buffer_id = hn->buffer_id;node_alloc->src[j].offset   = hn->offset;node_alloc->src[j].size_max = ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], src);}}
}

為計算圖中的每個 “計算節點”（非葉子張量）建立一個 node_alloc 結構，記錄該節點的輸出（dst）在內存緩沖區中的分配情況，以及其依賴的各個輸入（src）的內存分配信息。具體操作如下：

若 node 已經有數據或屬于視圖，則標記其輸出（dst）為無需重新分配
否則利用 ggml_gallocr_hash_get 從 Hash 表中查找該 tensor 的分配記錄，將 buffer_id、offset、size_max 寫入對應的 node_alloc
對每個節點的每個輸入源做類似處理

6. 建立 leaf_allocs 數組

if (galloc->n_leafs < graph->n_leafs) {free(galloc->leaf_allocs);galloc->leaf_allocs = calloc(graph->n_leafs, sizeof(galloc->leaf_allocs[0]));GGML_ASSERT(galloc->leaf_allocs != NULL);
}
galloc->n_leafs = graph->n_leafs;
for (int i = 0; i < graph->n_leafs; i++) {struct ggml_tensor * leaf = graph->leafs[i];struct hash_node * hn = ggml_gallocr_hash_get(galloc, leaf);if (leaf->view_src || leaf->data) {galloc->leaf_allocs[i].leaf.buffer_id = -1;galloc->leaf_allocs[i].leaf.offset = SIZE_MAX;galloc->leaf_allocs[i].leaf.size_max = 0;} else {galloc->leaf_allocs[i].leaf.buffer_id = hn->buffer_id;galloc->leaf_allocs[i].leaf.offset = hn->offset;galloc->leaf_allocs[i].leaf.size_max = ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], leaf);}
}

為計算圖中的每個 “葉子節點”（通常是輸入、常量或權重）建立一個 leaf_alloc 結構，記錄其在內存中的分配信息。具體操作如下：

同樣利用 ggml_gallocr_hash_get 獲取 leaf 在哈希表中的分配記錄，并寫入 leaf_alloc
如果 leaf 已經是一個 view 或者已經分配數據，則標記為無需分配

7. 根據當前分配器信息重新分配后端緩沖區

for (int i = 0; i < galloc->n_buffers; i++) {// 如果同一個動態分配器被重復使用，則復用已有的緩沖區for (int j = 0; j < i; j++) {if (galloc->buf_tallocs[j] == galloc->buf_tallocs[i]) {galloc->buffers[i] = galloc->buffers[j];break;}}size_t cur_size = galloc->buffers[i] ? ggml_backend_buffer_get_size(galloc->buffers[i]) : 0;size_t new_size = ggml_dyn_tallocr_max_size(galloc->buf_tallocs[i]);// 即使當前該緩沖區中沒有 tensor 被分配，也需要分配以便初始化視圖 tensorif (new_size > cur_size || galloc->buffers[i] == NULL) {
#ifndef NDEBUGGGML_LOG_DEBUG("%s: reallocating %s buffer from size %.02f MiB to %.02f MiB\n", __func__, ggml_backend_buft_name(galloc->bufts[i]), cur_size / 1024.0 / 1024.0, new_size / 1024.0 / 1024.0);
#endifggml_backend_buffer_free(galloc->buffers[i]);galloc->buffers[i] = ggml_backend_buft_alloc_buffer(galloc->bufts[i], new_size);if (galloc->buffers[i] == NULL) {GGML_LOG_ERROR("%s: failed to allocate %s buffer of size %zu\n", __func__, ggml_backend_buft_name(galloc->bufts[i]), new_size);return false;}ggml_backend_buffer_set_usage(galloc->buffers[i], GGML_BACKEND_BUFFER_USAGE_COMPUTE);}
}

根據預先通過動態分配器（ggml_dyn_tallocr）的累計信息，計算出每個緩沖區所需的內存新大小。具體操作如下：

對每個后端緩沖區（galloc->buffers[i]），先檢查是否存在同一動態分配器已復用的情況，如果有則復用相同的緩沖區
獲取當前緩沖區大小（cur_size），并調用 ggml_dyn_tallocr_max_size 得到該緩沖區內所有 tensor 分配所需的最大內存空間（new_size）
如果 new_size 大于當前分配大小或者當前沒有分配，則釋放舊緩沖區，并通過 ggml_backend_buft_alloc_buffer 分配一個新的緩沖區，大小為 new_size
分配成功后，調用 ggml_backend_buffer_set_usage 標記該緩沖區的使用場景為 COMPUTE 計算

8. 返回成功標志

函數最后返回 true 表示預留計算圖內存成功完成

這一系列流程確保了計算圖在執行前所有內存都已被 “預留”，即各個節點和葉子的內存位置與大小都已經確定，后續在實際執行圖的運算（前向傳播）時，只需要按照預留信息訪問連續的后端緩沖區，即可獲得高效且正確的內存訪問

至此，我們完成了預分配計算圖內存的全部代碼分析，下面我們來看分詞

3. 分詞

這個小節我們來看 ggml 框架在推理 gpt-2 時是如何做分詞的：

在這里插入圖片描述

gpt_tokenize 函數負責將用戶輸入的 prompt 文本轉化為一系列 token（即 vocab_id）以便后續模型推理時作為輸入，下面我們來具體分析：

1. 文本拆分為單詞和特殊 token

std::vector<std::string> words;// first split the text into words
{std::string str = text;// Generate the subpattern from the special_tokens vector if it's not emptyif (!vocab.special_tokens.empty()) {const std::regex escape(R"([\[\\\^\$\.\|\?\*\+\(\)\{\}])");std::string special_tokens_subpattern;for (const auto & token : vocab.special_tokens) {if (!special_tokens_subpattern.empty()) {special_tokens_subpattern += "|";}special_tokens_subpattern += std::regex_replace(token, escape, R"(\$&)");}std::regex re(special_tokens_subpattern);std::smatch m;// Split the text by special tokens.while (std::regex_search(str, m, re)) {// Split the substrings in-between special tokens into words.gpt_split_words(m.prefix(), words);// Add matched special tokens as words.for (auto x : m) {words.push_back(x);}str = m.suffix();}// Remaining text without special tokens will be handled below.}gpt_split_words(str, words);
}

這部分代碼首先將輸入的 text 拷貝給局部變量 str，然后檢查 vocab.special_tokens 是否為空。如果有特殊 token，則構造一個正則表達式 special_tokens_subpattern 來匹配這些 token，并使用 std::regex_search 將特殊 token 從文本中提取出來

在匹配過程中，通過 m.prefix() 調用 gpt_split_words 拆分特殊 token 前面的內容，然后直接把匹配到的特殊 token 加入 words 向量

最后，對剩余的字符串 str（即沒有特殊 token 的部分），再次調用 gpt_split_words 進一步拆分，確保所有單詞、標點、數字等都被拆分出來

gpt_split_words 函數實現如下：

void gpt_split_words(std::string str, std::vector<std::string>& words) {const std::string pattern = R"('s|'t|'re|'ve|'m|'ll|'d| ?[[:alpha:]]+| ?[[:digit:]]+| ?[^\s[:alpha:][:digit:]]+|\s+(?!\S)|\s+)";const std::regex re(pattern);std::smatch m;while (std::regex_search(str, m, re)) {for (auto x : m) {words.push_back(x);}str = m.suffix();}
}

首先預定義正則表達式模式，包含以下幾類匹配規則：

匹配特定的縮寫形式，如 's、't、're 等
匹配可選空格后的一串字母（英文單詞）
匹配可選空格后的一串數字
匹配可選空格后非空格、非字母、非數字的符號
匹配連續空白字符（但不跟隨非空白字符）以及其他空白組合

該正則表達式設計的目的是將文本細致地拆分成組成部分，使得后續的 token greedy 分割可以更準確地匹配詞匯表中的 token

接著拆分文本，利用 std::regex_search 對輸入字符串進行迭代匹配，將每次匹配到的部分（smatch 中的各個子匹配）加入到 word 向量中，直至字符串完全匹配結束

2. 采用貪心匹配策略為每個單詞找到最長 token

// find the longest token that forms each word in words:
std::vector<gpt_vocab::id> tokens;
for (const auto & word : words) {for (int i = 0; i < (int) word.size(); ){for (int j = word.size() - 1; j >= i; j--){auto cand = word.substr(i, j-i+1);auto it = vocab.token_to_id.find(cand);if (it != vocab.token_to_id.end()){ // word.substr(i, j-i+1) in vocabtokens.push_back(it->second);i = j + 1;break;}else if (j == i){ // word.substr(i, 1) has no matchingfprintf(stderr, "%s: unknown token '%s'\n", __func__, word.substr(i, 1).data());i++;}}}
}

遍歷 words 向量中每一個字符串 word，對于每個 word，以索引 i 開始，嘗試取從 i 到 j 的子串，其中 j 從單詞末尾開始向 i 遍歷，以保證能匹配最長可能的子串

若找到在 vocab.token_to_id 中存在的子串（即 token），則將對應的 token id 加入返回的 tokens 向量，并將 i 移動到匹配結束后的位置，從而實現貪心匹配。

如果連一個字符都匹配不到時（j == i），則輸出錯誤并將 i 加 1

3. 返回 token 序列

return tokens;

經過貪心匹配后，得到完整的 token 序列，存儲在 tokens 向量中，函數最后返回這一向量作為最終的 token 化結果

該函數完成了將原始輸入文本轉換成一系列 token id，為后續模型推理提供輸入表示

至此，我們完成了分詞的全部代碼分析，下面我們來看最后一部分即模型的推理與生成

4. 模型推理與生成

這個小節我們來看 ggml 框架在推理 gpt-2 時是如何進行前向傳播推理與文本生成的：

在這里插入圖片描述

這里我們逐 token 來進行文本生成，先通過調用模型推理函數 gpt2_eval 來計算 logits，接著通過采樣函數 gpt_sample_top_k_top_p 得到下一個 token，如此反復循環直至達到預測 token 數或者遇到結束符退出，下面我們來具體分析這兩個函數

4.1 模型推理

gpt2_eval 函數的主要任務是：

根據當前上下文構造計算圖
為計算圖分配內存，并設置好輸入張量
調用后端執行圖計算，得到輸出張量（logits）
從 logits 中提取最后一個 token 的預測結果（embd_w），返回預測成功的狀態

函數簽名及參數說明：

輸入參數：
- model：模型結構體，包含模型上下文大小、層數、詞匯表大小等信息
- allocr：圖內存分配器（ggml_gallocr）對象，用于為計算圖申請計算內存
- n_threads：設置后端計算使用的線程數
- n_past：上下文中已有的 token 數
- embd_inp：當前輸入的 token 序列（以詞匯 id 表示），它們將用于填充計算圖中輸入的張量
輸出參數：
- embd_w：輸出的 logits 數組
返回值
- 返回 true 表示推理計算完成

下面我們來具體分析整個流程：(from ChatGPT)

1. 計算輸入序列長度與參數準備

const int N = embd_inp.size();const auto & hparams = model.hparams;const int n_vocab = hparams.n_vocab;

將輸入 token 序列的長度保存為 N，引用模型中的超參數，獲取詞匯表大小 n_vocab

2. 構造計算圖

struct ggml_cgraph * gf = gpt2_graph(model, n_past, embd_inp.size());

調用 gpt2_graph 函數，傳入當前歷史 token 數 n_past 和本次輸入 token 數來構造當前上下文下的計算圖 gf

3. 分配計算圖中的所有張量的內存

// allocate the graph tensors
ggml_gallocr_alloc_graph(allocr, gf);

利用圖內存分配器 allocr 為計算圖 gf 內部所有尚未分配內存的張量進行內存分配

4. 設置輸入張量

struct ggml_tensor * embd = ggml_graph_get_tensor(gf, "embd");
ggml_backend_tensor_set(embd, embd_inp.data(), 0, N*ggml_element_size(embd));

設置 token 嵌入輸入，通過 ggml_graph_get_tensor(gf, "embd") 獲取計算圖中名為 “embd” 的輸入張量，利用 ggml_backend_tensor_set 將實際的輸入 token 序列寫入該張量的內存中

struct ggml_tensor * position = ggml_graph_get_tensor(gf, "position");
for (int i = 0; i < N; ++i) {int32_t v = n_past + i;ggml_backend_tensor_set(position, &v, i*sizeof(int32_t), sizeof(v));
}

設置位置編碼輸入，獲取計算圖中名為 “position” 的位置張量，循環為每個輸入 token 設置位置索引

5. 后端配置與線程設置

if (ggml_backend_is_cpu(model.backend)) {ggml_backend_cpu_set_n_threads(model.backend, n_threads);
}

判斷當前使用的后端是否為 CPU 后端，如果是，則調用 ggml_backend_cpu_set_n_threads 設置使用的線程數為 n_threads

6. 執行計算圖

// run the computation
ggml_backend_graph_compute(model.backend, gf);

調用后端接口 ggml_backend_graph_compute 執行構造好的計算圖 gf，模型推理的核心步驟就在這個函數。它會根據傳入的 model.backend 參數調用不同后端的算子實現去完成整個 graph 的計算，對于當前示例的 CUDA 后端而言，最終調用的函數是 ggml_cuda_compute_forward，如下圖所示：

在這里插入圖片描述

在 ggml_cuda_compute_forward 函數中我們會調用不同 op 算子的 CUDA 實現來完成計算，這是我們要學習的重點內容，看看每個算子的 CUDA 實現具體是怎么做的，不過在 ggml 框架中我們先跳過這部分吧，因為要分析的算子比較多，內容也比較長，以后有機會的話我們在 llama.cpp 中再來分析吧😄

7. 獲取輸出 logits

struct ggml_tensor * logits = ggml_graph_get_tensor(gf, "logits");// return result just for the last token
embd_w.resize(n_vocab);
ggml_backend_tensor_get(logits, embd_w.data(), (n_vocab*(N-1))*sizeof(float), sizeof(float)*n_vocab);//embd_w.resize(n_vocab*N);
//ggml_backend_tensor_get(logits, embd_w.data(), 0, sizeof(float)*n_vocab*N);// return result just for the last token
embd_w.resize(n_vocab);
ggml_backend_tensor_get(logits, embd_w.data(), (n_vocab*(N-1))*sizeof(float), sizeof(float)*n_vocab);

通過 ggml_graph_get_tensor(gf, "logits") 獲取計算圖中名為 “logits” 的輸出張量，該張量表示對每個 token 預測的 logits 分布，之后獲取最后一個 token 的輸出：

分配 embd_w 向量大小為 n_vocab，即每個 token 對應一個 n_vocab 維度的 logits
通過 ggml_backend_tensor_get 從 logits 張量中讀取最后一個 token 的 logits 數據

在自回歸生成中，我們使用整個序列輸出中最后一個 token 的 logits 分布來預測下一個 token，因此這里我們只取最后一個 token 的 logits 即可

8. 返回成功狀態

return true;

這樣 gpt2_eval 函數便完成了從輸入 token 到生成預測 logits 的整個推理過程

4.2 采樣

gpt_sample_top_k_top_p 采樣函數結合了 Top-K 與 Top-P（Nucleus Sampling）策略，對模型計算出來的 logits 進行溫度縮放、選擇候選 token、計算概率、歸一化、截斷概率質心集合，最后根據離散分布隨機采樣一個 token。

下面我們來具體分析整個流程：(from ChatGPT)

1. 初始化候選列表

int n_logits = vocab.id_to_token.size();std::vector<std::pair<double, gpt_vocab::id>> logits_id;
logits_id.reserve(n_logits);{const double scale = 1.0 / temp;for (int i = 0; i < n_logits; ++i) {logits_id.push_back(std::make_pair(logits[i] * scale, i));}
}

獲取候選數量：根據詞匯表中 token 的數量 n_logits 來初始化一個 vector，用于存儲每個 token 縮放后對應的 logit 值以及它的 id
溫度縮放：使用給定的溫度值 temp 對 logit 進行縮放（示例中給定的溫度值為 0.9），公式為 logits[i] * (1 / temp)
- 當 temp < 1 時，縮放后 logits 差異會更大，生成更 “確定” 的分布
- 當 temp > 1 時，則分布更平坦，生成較為隨機的選擇
結果存儲：每個 token 的 logit 和對應 id 存入到 logits_id 向量中，供后續排序和采樣使用

2. Top-K 篩選

std::partial_sort(logits_id.begin(),logits_id.begin() + top_k,logits_id.end(),[](const std::pair<double, gpt_vocab::id> & a, const std::pair<double, gpt_vocab::id> & b) {return a.first > b.first;
});logits_id.resize(top_k);

部分排序：使用 std::partial_sort 對 logits_id 向量進行排序，僅排序出前 top_k 個最高的候選 token（示例中 top-k 為 40）
- 排序基于 a.first > b.first，即按照縮放后的 logit 值從大到小排序
截斷候選集：調用 resize(top_k) 截斷 vector，使得后續只處理這 top_k 個候選 token

3. 計算 softmax 概率

double maxl = -INFINITY;
for (const auto & kv : logits_id) {maxl = std::max(maxl, kv.first);
}std::vector<double> probs;
probs.reserve(logits_id.size());double sum = 0.0;
for (const auto & kv : logits_id) {double p = exp(kv.first - maxl);probs.push_back(p);sum += p;
}// normalize the probs
for (auto & p : probs) {p /= sum;
}

數值穩定性處理：遍歷 Top-K 候選 token 計算最大值 maxl。計算 softmax 時先減去最大值，防止指數函數計算中的數值溢出，也就是我們前面文章中提到的 safe softmax
計算 softmax 概率：遍歷 Top-K 個候選 token 分別計算它們的 softmax 概率，公式為 p = exp(scaled_logit - maxl) / sum

4. Top-P 截斷（Nucleus Sampling）

if (top_p < 1.0f) {double cumsum = 0.0f;for (int i = 0; i < top_k; i++) {cumsum += probs[i];if (cumsum >= top_p) {top_k = i + 1;probs.resize(top_k);logits_id.resize(top_k);break;}}cumsum = 1.0 / cumsum;for (int i = 0; i < (int) probs.size(); i++) {probs[i] *= cumsum;}
}

Nucleus Sampling：如果 top_p 小于 1（示例中 top-p 為 0.9），表示需要保留累積概率質量達到 top_p 部分，丟棄剩余低概率 token
累加與截斷：逐個累加 Top-K 個 token 的概率 probs，直到累計和 cumsum 達到或超過設定的閾值 top_p。此時，將 top_k 更新為當前索引加 1，即只保留這部分候選 token，并將 probs 和 logits_id 調整為新大小
重新歸一化：將截斷后累積總概率的倒數乘以每個概率，使得新的候選集的概率和重新歸一化為 1

5. 離散采樣

std::discrete_distribution<> dist(probs.begin(), probs.end());
int idx = dist(rng);return logits_id[idx].second;

構造離散分布：利用標準庫的 std::discrete_distribution<> 構造一個離散概率分布，分布參數為前面經過 top_p 截斷歸一化后的概率 probs 數組
隨機采樣：使用提供的隨機數生成器 rng 從該離散分布中采樣一個索引 idx，該索引指向候選 token 集中的某個 token
返回采樣結果：根據采樣到的索引從 logits_id 數組中獲取對應的 token id 并返回，作為最終的采樣結果

通過以上步驟，函數最終返回了一個 token id，即下一個生成的 token

下面我們以一個具體的例子來說明采樣流程，假設：

詞匯表大小：10
溫度：0.9
top-p：0.9
top-k：5

原始 logits 數組如下：

Token id	原始 Logit 值
0	0.5
1	2.0
2	1.5
3	0.0
4	1.0
5	-0.5
6	3.0
7	0.2
8	2.5
9	1.8

下面具體說明各個步驟：

1. 溫度縮放

溫度設為 0.9，則縮放因子為：

$\frac{1}{0.9} \approx 1.1111$

對每個原始 logit 乘以該因子，得到縮放后的 logit：

Token id	原始 Logit 值	縮放后 Logit
0	0.5	0.5556
1	2.0	2.2222
2	1.5	1.6667
3	0.0	0.0
4	1.0	1.1111
5	-0.5	-0.5556
6	3.0	3.3333
7	0.2	0.2222
8	2.5	2.7778
9	1.8	2.0000

2. Top-K 篩選

接下來使用 partial_sort 僅保留 logit 最高的前 5 個候選 token，對上面縮放后的 logit 值排序，從高到低，我們得到前 5 個候選為：

1. Token 6：3.3333
2. Token 8：2.7778
3. Token 1：2.2222
4. Token 9：2.0000
5. Token 2：1.6667

其他 token 的得分較低，不被考慮

3. softmax 概率計算

為了得到概率分布，首先確定候選集合中最大的值： $ma x l = 3.3333$ （來自 Token 6）

然后對每個候選 token 計算：

$p_i = \exp(scaled\_logit_i - maxl)$

計算結果（近似值）：

Token 6： $\exp(3.3333 - 3.3333) = \exp(0) = 1.0$
Token 8： $\exp(2.7778 - 3.3333) = \exp(-0.5555) \approx 0.5738$
Token 1： $\exp(2.2222 - 3.3333) = \exp(-1.1111) \approx 0.3293$
Token 9： $\exp(2.0000 - 3.3333) = \exp(-1.3333) \approx 0.2636$
Token 2： $\exp(1.6667 - 3.3333) = \exp(-1.6666) \approx 0.1889$

總和 $\approx 2.3556$

歸一化后，每個 token 的概率為：

Token 6： $\approx 0.4244$
Token 8： $\approx 0.2435$
Token 1： $\approx 0.1398$
Token 9： $\approx 0.1119$
Token 2： $\approx 0.0802$

4. Top-P 截斷

設置的 top-p 等于 0.9，表示我們只希望候選 token 的累積概率至少達到 90%

從上述排序后的候選開始累積概率

加入 Token 6：累積概率 = 0.4244
加入 Token 8：累積概率 = 0.4244 + 0.2435 = 0.6679
加入 Token 1：累積概率 = 0.6679 + 0.1398 = 0.8077
加入 Token 9：累積概率 = 0.8077 + 0.1119 = 0.9196

當加入 Token 9 后累積概率超過 0.9，因此，將候選集合截斷為前 4 個 token：Token 6，Token 8，Token 1，Token 9。Token 2 被舍棄

接下來，對剩下的 4 個 token 的概率重新歸一化。它們原累積概率為 0.9196，歸一化因子為 $\approx 1.0870$ ：

Token 6： $0.4244 \times 1.0870 \approx 0.4617$
Token 8： $0.2435 \times 1.0870 \approx 0.2650$
Token 1： $0.1398 \times 1.0870 \approx 0.1519$
Token 9： $0.1119 \times 1.0870 \approx 0.1219$

最終總體候選概率約為：

Token 6：46.17%
Token 8：26.50%
Token 1：15.19%
Token 9：12.19%

5. 離散分布采樣

利用上面的 4 個候選 token 的歸一化概率構造一個離散分布。函數調用標準庫 std::discrete_distribution<>，并使用傳入的隨機數生成器 rng 根據該分布采樣出一個索引

假設采樣結果得到的索引為 0，那么對應的候選 token 就是 Token 6，函數最后返回此 token id 即 6

OK，這就是模型推理與文本采樣生成的整體過程了

至此，我們算簡單過了一遍 ggml 框架推理 gpt-2 的整體流程，主要包括模型加載、預分配計算圖內存、分詞以及模型推理與生成

結語

這里我們完整的過了一遍 ggml 推理 gpt-2 的整體流程，之后再來閱讀 llama.cpp 時會相對輕松些

其實博主對 ggml 整個框架的設計并不了解，只是表面上隨著代碼完整的 debug 了一遍，那如果想要深入的學習了解 ggml 這個框架，可能需要從框架設計方式出發去理解它，為什么要這么設計，有什么優點

博主能力有限，所有也就帶著大家簡單過了一遍，大部分都是 ChatGPT 分析的結果，真正理解的少之又少，其中的一些內容分析可能有誤，博主也沒有完全理解，更多的細節可能需要大家自己深入研究分析了

大家感興趣的可以看看 UP 的視頻講解，還是非常不錯的🤗