llama.cpp 源码完全拆解（二十五）：从 llama_batch 到 ubatch——一份 token 清单如何被切成能被 KV 同时装下的微批

先看全景：一次 llama_decode 里切分处在什么位置

llama_decode 的入口其实有两条。新接口 llama_process(ctx, LLAMA_PROCESS_TYPE_DECODE, &batch_ext) 直接吃 llama_batch_ext；老接口 llama_decode(ctx, batch) 吃老的 llama_batch，由 llama_context::decode(const llama_batch &) 先把它经 llama_batch_compat 翻译成 llama_batch_ext，再走同一条路：

1
int llama_context::decode(const llama_batch & batch_inp) {
2
    llama_batch_compat compat(this, batch_inp);
3
    return decode(*compat.batch_ext);
4
}

也就是说，切分前的「原始批次」统一是 llama_batch_ext，切分后的「计算单元」统一是 llama_ubatch，中间夹着一个 llama_batch_allocr 做记账。整个流程放在 llama_context::decode(const llama_batch_ext &) 里，可以压缩成下面几步：

1. allocr 初始化：balloc->init(batch_inp, vocab, output_all)。这一步做的是校验与规范化——把 seq_ids 集合展开、把位置数组按 M-RoPE 的要求展开成段优先布局、统计输出行数、检查位置是否单调。它不做任何切分。
2. 记忆模块切分：memory->init_batch(*balloc, cparams.n_ubatch, embd_all)。切分策略的选择权交给了记忆模块，因为只有记忆模块知道自己的容量与结构：统一 KV（单流）只需要「沿扁平顺序截断」，非统一 KV（每序列一个流）则需要「每个序列在每个 ubatch 里出一样多的 token」。
3. KV 预演：prepare(ubatches) 对每个 ubatch 跑一遍 find_slot，但只找位置不写单元。任何一个 ubatch 找不到位置，整批失败。
4. 逐 ubatch 执行：mctx->next() 循环里，每个 ubatch 经过 apply() 真正占用 KV 单元、复用或重建图、写输入、计算、取回 logits。
5. 输出行序还原：因为重排会让 ubatch 内 token 的顺序与调用方给的顺序不同，最后用 balloc->get_out_ids() 建一张 output_ids 反查表，并记录 output_swaps 供 llama_get_logits_ith 按需重排。

上图把五步画成了一条从左到右的流水。注意第三步的位置：切分是「先全部切完，再统一找位置」，而不是「切一个、放一个、算一个」。这个顺序上的选择后面会专门讲，它是 llama.cpp 在「一次 decode 失败时不留半截状态」这条不变式上的实现方式。图中右侧的分叉表示 split_equal 会把 token 重排成「先序列 0 的所有 token，再序列 1 的所有 token」，而执行完之后要通过 out_ids 把行序还原回调用方的顺序。

两份数据结构：输入侧的 llama_batch_ext 与输出侧的 llama_ubatch

理解切分，首先要理解它一边吃进去的是什么、另一边吐出来的是什么。这两个结构的设计意图截然不同。

1
struct llama_batch_ext {
2
    const size_t n_tokens_max;   // 上限，超过后 add_token 返回 -1
3
    const size_t n_embd_inp;     // 输入嵌入宽度（token 嵌入）
4
    const size_t n_embd_inp_enc; // 编码器输入宽度（如 EAGLE-3 / DFlash 提取出的特征）
5
    const size_t n_embd_state;   // MTP 等上下文携带的隐状态宽度，0 表示不带
6
    const llama_seq_id n_seq_max;
7
    llama_memory_i * mem;        // 用于位置推断
8
    const llama_token n_vocab;
9
    const size_t n_pos_per_embd; // M-RoPE 时为 3 或 4，文本模型为 1
10

11
    size_t n_embd = 0;           // 本批实际的嵌入行宽，由第一次 set_token_embd() 确定
12

13
    struct token {
14
        llama_token  id = LLAMA_TOKEN_NULL;
15
        bool         has_embd = false;
16
        size_t       embd_off = 0;
17
        bool         has_state = false;
18
        size_t       state_off = 0;
19
        bool         output = false;
20
        int32_t      decision_order = 0;
21
        std::unordered_set<llama_seq_id> seq_ids;  // 一个 token 可以同时属于多条序列
22
        std::array<llama_pos, GGML_MROPE_SECTIONS> pos = {0, 0, 0, 0};
23
    };
24
    std::vector<token> tokens;
25
    std::vector<float> embd;
26
    std::vector<float> state;
27
};

1
struct llama_ubatch {
2
    uint32_t b_equal_seqs; // 布尔值，用 32 位存是为了对齐（否则 ASan 会报错）
3
    uint32_t n_tokens;     // 总 token 数 = n_seq_tokens * n_seqs
4
    uint32_t n_seq_tokens; // 每个序列集的 token 数（所有序列集相同）
5
    uint32_t n_seqs;       // 序列集个数
6
    uint32_t n_seqs_unq;   // 不同序列 id 的个数
7
    uint32_t n_pos;        // 每个 token 的位置段数
8

9
    llama_token    * token;      // [n_tokens]
10
    float          * embd;       // [n_embd, n_tokens]
11
    float          * embd_state; // [n_embd_state, n_tokens]
12
    llama_pos      * pos;        // [n_tokens*n_pos]，段优先
13
    int32_t        * n_seq_id;   // [n_tokens]
14
    llama_seq_id  ** seq_id;     // [n_tokens]，每行指向 seq_id_data 的一段
15
    llama_seq_id   * seq_id_unq; // [n_seqs_unq]
16
    int32_t        * seq_idx;    // [LLAMA_MAX_SEQ]，序列 id → 在 ubatch 内的下标
17
    int8_t         * output;     // [n_tokens]
18
    int8_t         * type;       // [n_tokens]，混合批次才有
19

20
    std::shared_ptr<data_t> data; // 所有指针指向这块内存，data_t 里有 batch_idxs
21
};

1
typedef std::bitset<LLAMA_MAX_SEQ> seq_set_t;

allocr 的 init：切分之前的校验与规范化

split_* 函数看似只是「挑 token」，但很多边界条件在 init() 里就已经被筛掉了。init() 做的事可以归纳为六件，按代码顺序：

1. 内容类型判定。 统计每行的「只带 token」「只带嵌入」「两者都带」的分布。规则是：如果存在「两者都带」的行，则必须全部行都是「两者都带」，不能混用；「只带 token」与「只带嵌入」可以混合，但要求模型 / 上下文支持 allow_mixed（llama_context 构造时传入 llm_arch_supports_mixed_batch(arch) && ctx_type == LLAMA_CONTEXT_TYPE_DEFAULT）。不支持时直接返回 false，decode 返回 -1。

2. 隐状态一致性。 has_state 取自第 0 行，然后逐行比对。MTP 等场景中隐状态要么全带要么全不带，不存在部分行携带的情况。

3. 扁平化 token / 嵌入。 只带 token 的行写进 token_vec；混合批次中嵌入行的 token id 填 0 占位（这一行的实际语义由 type[] 标记）。嵌入行按 embd_off 拷贝到 embd_vec 的对应位置。全部带嵌入时直接 embd_vec = batch_inp.embd 复用，避免大数组拷贝。

4. 位置展开成段优先。 这是读者最容易看错的一步：

1
pos.resize((size_t) n_tok*n_pos_per_embd);
2
for (int32_t i = 0; i < n_tok; ++i) {
3
    const auto & tok = batch_inp.tokens[i];
4
    const bool expand = tok.id != LLAMA_TOKEN_NULL;
5
    for (uint32_t j = 0; j < n_pos_per_embd; ++j) {
6
        llama_pos p = tok.pos[j];
7
        if (expand) {
8
            // 把 [p] 扩展成 [p, p, p, 0]，对应 M-RoPE 的文本位置
9
            p = j < 3 ? tok.pos[0] : 0;
10
        }
11
        pos[(size_t) j*n_tok + i] = p;
12
    }
13
}

对于 token 行（expand == true），M-RoPE 的三个文本位置段全部写成同一个值——文本没有图像上的二维位置，用同一个位置填满所有段就等价于标准 RoPE 的一维位置。嵌入行（expand == false）则按调用方给的三维位置原样保留，这正是图像 patch 所在行的语义：它们的二维位置要参与编码，从而让图像内部的注意力具有二维局部性。

5. 序列归属展开。 seq_id_data 把所有 token 的序列 id 顺序拼成一块数组，seq_id[i] 指向各自片段，n_seq_id[i] 记录长度。这一步必须做，因为原始输入里 seq_ids 是 unordered_set，没有稳定顺序，切分与后续的 seq_idx 构建都依赖「每行序列 id 连续存放」这个布局。

同时校验每个 seq_id 落在 [0, n_seq_max) 内。

6. 输出标志与位置单调性。 output 数组直接从输入的 output（老接口里是 logits）拷贝；若 output_all（嵌入计算或 embeddings 模式）为真而并非所有行都标了输出，会强制全部置 1 并打一条警告——池化嵌入必须看全部 token。随后逐序列检查 pos 非递减（seq_pos[seq_id] 记录该序列已见的最大位置），递减直接报错，因为 KV 缓存的单元是按位置排序的，乱序写入会让 find_slot 与 seq_rm 的区间假设失效。

init() 结束时调用 split_reset()，清空 used[]（一个与 token 数等长的布尔数组），并把 n_used 归零。此后每次切出一块 ubatch，used[] 里对应位置就置 1，n_used 累加；所有 token 都被用过后切分结束。

最后 init() 组装出一份内部的 llama_batch（仍然是扁平指针），字段全部指向自己的 vector：

1
batch.n_tokens = n_tok;
2
batch.token    = token_vec.data();   // 只在需要时有值
3
batch.pos      = pos.data();         // 段优先
4
batch.n_seq_id = n_seq_id.data();
5
batch.seq_id   = seq_id.data();
6
batch.logits   = output.data();

注意这里的 pos 已经是段优先布局，而老 llama_batch 通过 llama_batch_compat 翻译进来时也是按 M-RoPE 布局读的——两种输入在 init() 出口处统一成了同一种布局。

三套切分策略：由记忆模块按自己的结构挑选

切分的调用点在 llama_kv_cache::init_batch 里，一行三元表达式决定用哪个策略：

1
llama_memory_context_ptr llama_kv_cache::init_batch(
2
            llama_batch_allocr & balloc,
3
            uint32_t n_ubatch,
4
            bool embd_all) {
5
    do {
6
        balloc.split_reset();
7

8
        std::vector<llama_ubatch> ubatches;
9
        while (true) {
10
            auto ubatch = n_stream == 1               ? balloc.split_simple(n_ubatch)
11
                        : embd_all                    ? balloc.split_seq(n_ubatch)
12
                        :                               balloc.split_equal(n_ubatch, true, 0);
13

14
            if (ubatch.n_tokens == 0) {
15
                break;
16
            }
17

18
            ubatches.push_back(std::move(ubatch));
19
        }
20

21
        if (balloc.get_n_used() < balloc.get_n_tokens()) {
22
            // 有 token 没被任何 ubatch 收走 = 切分失败
23
            break;
24
        }
25

26
        auto sinfos = prepare(ubatches);
27
        if (sinfos.empty()) {
28
            break;
29
        }
30

31
        return std::make_unique<llama_kv_cache_context>(
32
                this, std::move(sinfos), std::move(ubatches));
33
    } while (false);
34

35
    return std::make_unique<llama_kv_cache_context>(LLAMA_MEMORY_STATUS_FAILED_PREPARE);
36
}

三个条件分别对应三种 KV 缓存结构：

n_stream 在构造函数里定死：n_stream(unified ? 1 : n_seq_max)。kv_unified 是从 llama_context_params 透传下来的开关（-kvu 命令行选项），打开时全部序列共享一个连续 KV 空间，因此不存在「每个序列各有一个流、各放等量 token」的约束，split_simple 是自然选择。

循环的终止条件也值得注意：只要某个策略返回空 ubatch（n_tokens == 0），循环就结束。如果此时 n_used < n_tokens（还有 token 没被收走），说明剩下的 token 无法构成任何一个合法 ubatch——这不是错误返回值，而是「切分失败」，init_batch 会走到最后的 LLAMA_MEMORY_STATUS_FAILED_PREPARE。什么时候会这样？典型场景是 n_ubatch 太小，装不下某条序列的最小可分单位，或者某个序列集的 token 数与 n_ubatch 不成比例。这时 decode 会尝试一次 memory_update(true)（允许做代价高的缓存整理，比如平移、丢弃空闲 slot 的缓存），再重试；仍不行才返回 1（「找不到 KV 槽位」）——用户侧看到的「batch too large / context exceeded」报错大多来自这条路径。

1
llama_ubatch llama_batch_allocr::split_simple(uint32_t n_ubatch) {
2
    // 找到第一个还没被用掉的 token
3
    uint32_t cur_idx = 0;
4
    while (cur_idx < used.size() && used[cur_idx]) {
5
        ++cur_idx;
6
    }
7
    if (cur_idx >= used.size()) {
8
        return {};   // 全部用完
9
    }
10

11
    std::vector<int32_t> idxs;
12
    while (true) {
13
        idxs.push_back(cur_idx);
14
        used[cur_idx] = true;
15
        ++n_used;
16
        ++cur_idx;
17

18
        if (cur_idx >= used.size()) {
19
            break;
20
        }
21
        if (idxs.size() >= n_ubatch) {
22
            break;
23
        }
24
    }
25

26
    return ubatch_add(idxs, idxs.size(), false);
27
}

split_equal：常规路径的「公平切分」

非统一 KV 是 llama.cpp 的常规配置（默认 kv_unified = false），每条序列有自己的连续 KV 流。split_equal 的规则是：每个 ubatch 里，每个参与序列集的 token 数相同，且不同序列集互不相交。

它的实现分三步，全部围绕 used、seq_set、seq_set_map 展开：

第一步：挑选参与的序列集。 从第一个未使用的 token 开始扫，每当遇到一个新的序列集，检查它与已选中的序列集是否「有交集」（位集按位与是否非空）。有交集就跳过，无交集才收进 cur_seq_set。当 cur_seq_set.size() > n_ubatch 时停止——因为一个序列集在 ubatch 里至少要占 1 个 token，超过 n_ubatch 个序列集就必然装不下。

sequential = true 时还额外要求「序列 id 递增」：

1
bool add = true;
2
for (uint32_t s = 0; s < cur_seq_set.size(); ++s) {
3
    if (!(cur_seq_set[s] & seq_set[i]).none()) { add = false; break; }
4
}
5
if (sequential) {
6
    add = add && (cur_seq_set.empty() || batch.seq_id[i][0] == last_seq_id + 1);
7
}

「递增」有什么好处？它保证同一序列集内的 token 在原始批次里是连续的——这在注意力掩码、位置检查与流式回传里都更方便，尤其当多模态图像 token 与文本 token 混排时，递增条件会让同类 token 尽量归并到一起。代价是「遇到一个乱序 token 就跳过它」，如果整个批次乱序严重，参与的序列集会很少，每个 ubatch 装的 token 就少。这通常不是问题，因为 llama.cpp 的调用方（common/batch、服务端的 server-context）在组装 batch 时本身就按序列 id 排序。

第二步：等量扩展。 对每个选中的序列集，维护一个指针 cur_idx[s] 指向该序列集下一个未用 token。每轮从每个序列集各取一个 token 追加到 idxs_per_seq[s]，然后用 (idxs_per_seq[0].size() + 1) * n_seqs > n_ubatch 判断下一轮是否超限。注意 +1：这是「如果再加一轮，总 token 数会是多少」的预判，在超限前停下，保证每个序列集的 token 数严格一致。

这里有一个隐含的细节：cur_idx[s] 初始化时会跳过已经用过的 token（while (used[seq_set_map[cur_seq_set[s]][cur_idx[s]]]) ++cur_idx[s];）。所以同一个序列集在第二个、第三个 ubatch 里继续取 token 时，会从上次停下的地方接着走。

第三步（本例中传入 n_keep_tail = 0，因此不生效）：尾部保留。 当调用方希望「每个序列的最后 $k$ 个 token 不被拆到两个 ubatch」时（某个需要尾部窗口的编码、或者 DFlash 一类需要保留尾部的模型），n_keep_tail 会起作用：装不下的序列集会整体退回（把它们的 token 重新标记为未使用），只保留能在本 ubatch 里跑完或还剩 $\ge k$ 个 token 的序列集。如果所有序列集都违规（剩余 token 比 $k$ 少又没跑完），就把第一个序列集单独成全——要么在本 ubatch 里跑完，要么截短到只剩 $k$ 个尾部 token 留给下一个 ubatch。

最后是拼接。 按序列集的顺序把所有 idxs_per_seq[s] 首尾相接：

1
std::vector<int32_t> idxs;
2
for (uint32_t s = 0; s < n_seqs; ++s) {
3
    idxs.insert(idxs.end(), idxs_per_seq[s].begin(), idxs_per_seq[s].end());
4
}
5
return ubatch_add(idxs, n_seqs, true);   // equal_seqs = true

所以一个由 3 条序列、每条 4 个 token 组成的 ubatch，内存里是 [s0t0, s0t1, s0t2, s0t3, s1t0, ..., s2t3]，即「先序列 0 的全部，再序列 1 的全部」。这不一定是原始 batch 的顺序——如果原始 batch 是交错排的（t0、t1、t2 各来自不同序列），batch_idxs 就会记录真实来源，后面 out_ids 会修正输出行的映射。

上图展示了一个典型的 trace：原始 batch 里有序列 {0}、{1}、{2} 三种 token，前三个 token 分别是它们的第一个 token。split_equal 把它们全部选中，之后逐轮各取一个，最终形成 [s0t0,s0t1,s0t2, s1t0,s1t1,s1t2, s2t0,s2t1,s2t2]。图的下半部分是第二个 ubatch（继续取剩下的 token）。由于 n_ubatch 恰好能装下 9 个 token，第一轮就装完了全部。

如果 n_ubatch = 6，第一轮只会取 6 个（每个序列 2 个），第二轮再取剩下的 3 个。这时候第二轮中每个序列集只剩 1 个 token，cur_idx 会指向它们各自的下一个 token，没问题。

如果某个序列集只剩 1 个 token，但另一个序列集还有 5 个，第二轮仍然只能给每个序列集取 1 个，剩下的 4 个要第三轮。这就是 split_equal 的「公平但有浪费」：短序列会拉着所有序列一起变短，因为它要求每轮等量。这个浪费在解码（每序列每轮就 1 个 token）时不存在；在预填充（不同 prompt 长度不一）时确实存在，但预填充通常长 prompt 占大头，短 prompt 的浪费有限；而且 llama.cpp 的服务端在组装 batch 时会对 prompt 长度相近的 slot 做分组（can_batch_with 与 get_available_slot），部分缓解了这个问题。server-context 那一层如何分组，见 llama.cpp 源码完全拆解（八）：服务化与调度。

1
const bool output_all = cparams.embeddings;
2
const bool embd_all   = output_all || !cparams.causal_attn;

1
auto cur_seq_set = seq_set[cur_idx];
2
while (true) {
3
    idxs.push_back(cur_idx);
4
    used[cur_idx] = true;
5
    ++n_used;
6
    if (idxs.size() >= n_ubatch) break;
7

8
    do {
9
        ++cur_idx;
10
    } while (cur_idx < get_n_tokens() &&
11
             (used[cur_idx] || ((cur_seq_set & seq_set[cur_idx]) != seq_set[cur_idx])));
12

13
    if (cur_idx == get_n_tokens()) break;
14
    cur_seq_set = seq_set[cur_idx];
15
}
16
return ubatch_add(idxs, 1, true);   // n_seqs = 1

1
GGML_ASSERT((cparams.causal_attn || cparams.n_ubatch >= n_tokens_all) &&
2
            "non-causal attention requires n_ubatch >= n_tokens");

ubatch_add：重排之后的所有指针都要「修一遍」

三个 split_* 最后都调用 ubatch_add(idxs, n_seqs, equal_seqs)。传入的 idxs 是「按 ubatch 内顺序排列的原始 token 下标」，ubatch_add 负责把它变成一份自洽的 llama_ubatch：为每个 token 重新分配 token id、嵌入、位置、序列、输出，并把所有指针都指到自己的内存。

它必须自己做内存，不能复用 balloc 里的 token_vec / pos / seq_id_data，原因是：balloc 只在一次 llama_decode 的生命周期里存在，且其数组是按原始批次索引的；ubatch 需要在后续的 apply()、set_inputs()、graph_compute() 中独立存活，还可能被 mctx 持有直到下一轮。所以 ubatch_add 为每个 ubatch 分配一个 std::shared_ptr<data_t>，把所有数组放进去：

1
auto udata = std::make_shared<llama_ubatch::data_t>();
2
udata->token.resize(n_tokens);
3
udata->embd.resize(n_embd_all);
4
udata->embd_state.resize(n_state_all);
5
udata->pos.resize(n_pos_all);
6
udata->n_seq_id.resize(n_tokens);
7
udata->seq_id.resize(n_tokens);
8
udata->seq_id_unq.resize(0);
9
udata->seq_idx.resize(LLAMA_MAX_SEQ, -1);
10
udata->output.resize(n_tokens);
11
udata->type.resize(mixed ? n_tokens : 0);
12
udata->batch_idxs = idxs;   // 记录原始下标

随后一个循环把每个下标 idxs[i] 的数据拷到新位置 i：

1
for (size_t i = 0; i < idxs.size(); ++i) {
2
    if (use_token) udata->token[i] = batch.token[idxs[i]];
3
    ...
4
    for (size_t j = 0; j < n_pos_per_embd; ++j) {
5
        udata->pos[j*n_tokens + i] = batch.pos[j*batch.n_tokens + idxs[i]];
6
    }
7
    udata->n_seq_id[i] = batch.n_seq_id[idxs[i]];
8
    udata->output[i]   = batch.logits[idxs[i]];
9
    for (int s = 0; s < udata->n_seq_id[i]; ++s) {
10
        udata->seq_id_data.push_back(batch.seq_id[idxs[i]][s]);
11
    }
12
    ...
13
}

位置数组的索引要单独看。 原 batch.pos 是段优先布局，batch.pos[j*batch.n_tokens + idxs[i]] 是新布局同段位置的正确来源；新数组的 pos[j*n_tokens + i] 把 $j$ 段搬到新位置。两个 n_tokens 不一样：前者是原批次的 batch.n_tokens，后者是本 ubatch 的 n_tokens。这一步是整段代码里最容易写错的地方之一，写反了就会让 M-RoPE 把图像位置读错，表现为多模态理解质量下降而不是报错。

第二级指针的修复单独占用一段循环：

1
llama_seq_id * seq_id_ptr = udata->seq_id_data.data();
2
for (size_t i = 0; i < idxs.size(); ++i) {
3
    udata->seq_id[i] = seq_id_ptr;
4
    seq_id_ptr += udata->n_seq_id[i];
5
}

因为 seq_id_data 在循环里不断 push_back，std::vector 可能扩容并移动内存，所以不能在 push 的同时保存 &seq_id_data.back()。正确做法是先全部 push 完，再统一计算每行的起始地址。这种「两遍法」在 C++ 里处理「指针指向正在增长的容器」时是标准做法，llama.cpp 里许多类似的地方也采用同一思路。

seq_id_unq 与 seq_idx： 遍历 seq_set_unq 位集，把置位序列 id 依次追加进 seq_id_unq，并在 seq_idx 里记录它的下标。seq_idx[seq_id] 是后面 decode 里做池化嵌入提取时的索引：

1
const llama_seq_id seq_id  = ubatch.seq_id_unq[s];
2
const int32_t      seq_idx = ubatch.seq_idx[seq_id];
3
embd_seq_out[seq_id].resize(n_embd_out);
4
ggml_backend_tensor_get_async(backend_embd, t_embd,
5
    embd_seq_out[seq_id].data(), (n_embd_out*seq_idx)*sizeof(float),
6
    n_embd_out*sizeof(float));

即：池化嵌入在 GPU 上按「本 ubatch 的第 seq_idx 个序列」排列，seq_idx 把「原始序列 id」映射到「本 ubatch 内的序列槽」。这个映射之所以必要，是因为一个 ubatch 可能只包含原始批次的一小部分序列（split_equal 第二轮只含少数序列），序列 id 与槽位不再一一对应。

prepare：只找位置不写单元，以及为什么必须整体成功

init_batch 切出全部 ubatch 之后，调用 prepare(ubatches)：

1
llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(
2
        const std::vector<llama_ubatch> & ubatches) {
3
    llama_kv_cache::slot_info_vec_t res;
4

5
    struct state_t {
6
        slot_info sinfo;
7
        std::vector<uint32_t> v_heads_old;
8
        std::vector<llama_kv_cells> v_cells;
9
    };
10
    std::vector<state_t> states;
11
    bool success = true;
12

13
    for (const auto & ubatch : ubatches) {
14
        // 只找合适的位置，暂不修改单元
15
        const auto sinfo_new = find_slot(ubatch, false);
16
        if (sinfo_new.empty()) {
17
            success = false;
18
            break;
19
        }
20

21
        res.push_back(sinfo_new);
22

23
        // 把单元的旧状态存入恢复栈
24
        state_t state = { sinfo_new, v_heads, {} };
25
        ...
26
    }
27

28
    ...
29
}

要点有三个。

第一，find_slot(ubatch, false) 的第二参数 cont = false。 这个参数控制「找不到连续区间时，是否把碎片拼起来」。prepare 不传 cont=true，因为这只是预演，没必要为了预演去做碎片拼接（拼接会改变 slot_info 的 strm 分布，使恢复变复杂）。真正 apply 时的 find_slot 也不传 cont=true 用于常规路径，cont 只在某些特定的重试路径里使用。

第二，每个 ubatch 找到位置后，旧单元的 pos、seq 归属、序列流头指针都要备份。 因为 find_slot 虽然只「找」不「写」，但 find_slot 的 cont=true 路径和某些「把已有 token 平移」的操作会改状态；prepare 选择最简单的方式：前后完全隔离——先备份，然后全部 apply_ubatch 一遍，最后恢复。这等价于「尝试性执行」，即使 apply_ubatch 改动了很多东西，结束时也一定会复原。

第三，任何一个 ubatch 失败即整批失败。 不存在「前两个 ubatch 放了、第三个放不下，就只跑前两个」的降级。为什么？因为 init_batch 的调用语义是「这一批 token 要么全处理、要么全不处理」：调用方会认为 llama_decode 返回非零时整个批次未生效。如果允许部分成功，调用方在重试时就会重复写前两个 ubatch，KV 里出现重复位置；如果不允许，调用方只能整批减少批次大小（服务端做的是 n_batch /= 2 重试），代价是丢掉一部分吞吐，但状态是确定的。

具体地，apply_ubatch 在写单元时还会做两件事：

• 对被覆盖的非空单元（统一 KV 或环形复用情况下，一个位置可能已被别的 token 用过），记录 seq_pos_max_rm[seq_id] = max(...)；
• 写完所有单元后，对每个被覆盖的序列，把 [pos_min, pos_max] 内小于被覆盖位置的所有单元删除（purge），以维持「一个序列在缓存里的位置区间是连续的」这条不变式。

这条不变式的来源是 llama.cpp PR #13746 讨论中提出的「位置区间内不能有洞」约束——注意力掩码与 RoPE 的假设都依赖它。prepare 的恢复栈之所以必须保存 llama_kv_cells（而不只是 v_heads），就是因为 apply_ubatch 的 purge 可能删除大量单元，恢复必须把单元结构（位置、序列占用）完整还原。

执行侧：输出行序的还原与 batch_idxs 的最终用途

切分、预演、执行都完成后，decode 要做一件事：让调用方按原始批次顺序拿结果。这个过程分三层：

第一层，out_ids。 ubatch_add 在拷贝 output 标志的同时，把「本 ubatch 中哪些行是输出」记录到 out_ids（所有 ubatch 共享同一个 out_ids，按 ubatch 的出现顺序追加）。decode 在末尾用它检查输出数量：

1
auto & out_ids = balloc->get_out_ids();
2
GGML_ASSERT(out_ids.size() == (size_t) n_outputs);

第二层，output_ids。 建立「原始行号 → logits 缓冲行号」的反查表：

1
for (int64_t i = 0; i < n_outputs; ++i) {
2
    output_ids[out_ids[i]] = i;
3
}

logits 缓冲区是按 ubatch 内顺序填的：第一个 ubatch 中第一个输出行在缓冲区第 0 行，第二个 ubatch 的第一个输出行紧随其后。而调用方给的 llama_get_logits_ith(ctx, i) 里 i 指向原始批次。output_ids 正是这两套编号的桥。

第三层，output_swaps。 当 out_ids 本身不是单调递增时（比如统一 KV 下 split_simple 恰好按顺序、或者非统一 KV 下 split_equal 的重排让输出顺序变了），代码用选择排序把 out_ids 排成 [0,1,2,...] 的样子，同时记录每次交换：

1
std::swap(out_ids[i], out_ids[j_min]);
2
output_swaps.push_back({ i, j_min });

这些交换不在采样时执行，而是在调用方读取 logits 时懒执行。这样做的原因是：采样通常只关心少数几行的 logits，懒交换可以把实际交换次数降到最少。文件注释说得很直白——这是「mostly relevant for recurrent models atm」。

batch_idxs 的最终用途在嵌入提取分支：

1
bool extract_all_idxs = extract_layer_inputs(res, n_tokens_prev, ubatch.n_tokens);
2
...
3
if (extract_all_idxs) {
4
    const auto & batch_idxs = ubatch.data->batch_idxs;
5
    embd_batch_idxs.insert(embd_batch_idxs.end(), batch_idxs.begin(), batch_idxs.end());
6
}

在需要逐层输入嵌入的模型里（MTP 钩子批次、EAGLE-3 / DFlash 提取目标特征），每个 ubatch 的嵌入行要按实际来源映射回原始批次，供投机解码的验证阶段配对。

三种策略的实际表现：什么时候会浪费

用一个具体算例把三种策略的差异摆出来。假设有一批次，4 条序列，每条 5 个 token，n_ubatch = 8，n_stream > 1：

• split_equal：第一轮挑选 4 个序列集（全部不相交），每轮各取 1 个，一共能取 2 轮（8 个 token，(2+1)*4=12 > 8 所以在第 2 轮后停），形成 [8 token × 2 ubatch]。剩下的 12 个 token 构成第二个批次单元：第一轮取 4 个序列集，每序列 1 个，8 个 token 用完；第二个 ubatch 同样 8 个 token；最后剩 4 个 token，第三轮取 4 个序列集，每个 1 个——结果 3 个 ubatch。
• split_simple（假想统一 KV）：直接 8、8、4，形状不均匀但总 ubatch 数相同。
• split_seq：整条序列不能被拆分，每条 5 个 token 都能整除 n_ubatch 之外不强制要求，实际上会得到 4 个 ubatch（$n\_ubatch - 5 = 3$ 个空位无法装下第二个序列）。

从这个例子能看出，split_equal 的「公平」在等量序列上表现良好；但若 4 条序列长度是 5、5、5、1，split_equal 的第二轮就只给每条序列 1 个（因为第 4 条只剩 1 个），第一轮每序列 1 个、第二轮每序列 1 个，剩给第三轮的是 4 条序列各 4、4、4、0 个，第三轮无法凑到 4 个序列集（第 4 条没了），所以只能让前 3 条进第三轮、每条 2 个（(2+1)*3=9 > 8），第四轮再处理最后 2、2、2 个 → 仍然 5 个 ubatch。split_simple 只需要 $\lceil 16/8 \rceil = 2$ 个 ubatch。

这就是 split_equal 在「长短混排」时的浪费：短序列会形成新的「轮次」，每轮的 token 数被 n_ubatch / 序列集数 压低。但这只是 ubatch 数量增多，并没有丢数据，位置与序列正确性由每个 ubatch 自带的 pos、seq_id 保证。

实际 llama.cpp 服务端在预填充时通常把 prompt 长度接近的 slot 分到同一轮（见 server-context 的批组装逻辑），因此长短混排不如「一次塞满所有 slot」常见；单流解码时每条序列每个 ubatch 也就 1 个 token，不会出现浪费。

与站内其他框架的对照

llama.cpp 这套「一份扁平 batch 切成若干 ubatch」的思路与 vLLM、SGLang 的调度结构有本质不同，值得对照着看。

vLLM 的 V1 调度器直接把请求按「每轮出 1 个 token」组织成 ScheduledSequenceGroup，每个请求在本轮要追加 1 个 token（预填充是若干块），InputBatch 负责把不同请求的「本轮 token」拼成一份 token_ids 与 positions 的扁平 buffer。它没有显式的 n_ubatch 概念——调度器本身就是切分器：准入阶段决定 max_num_seqs、块预算，准入后的请求就是本轮的 ubatch，形状固定为「$N$ 条序列 × 1 个 token」或「若干预填充请求的组合」。对照阅读可参考 vLLM 源码完全拆解（二十五）：请求准入的五层关卡 和 KV Cache 的页粒度：128 token 一块的账，以及五个推理框架为什么选了五个不同答案。

SGLang 的调度器走的是「PrefillAdder 预算 KV、切分长提示、组装 Prefill Batch」的路子，ForwardBatch 承载本轮要跑的所有 token；对于超出 KV 预算的 prompt，它会切分成多个 extend 请求，每个请求下一次处理一段。这与 llama.cpp 的 split_equal 在目的上相似——都是「让长 prompt 不独占一轮」——但 SGLang 的切分在调度器层面完成，llama.cpp 的切分在上下文层完成：调用方只管递扁平 batch，切分是上下文对调用方透明地做的。可参考 SGLang 源码完全拆解（四）：PrefillAdder 如何预算 KV、切分长提示并组装 Prefill Batch。

更贴切的对照是 vLLM 的 Dual Batch Overlap 与 SGLang 的 Two-batch Overlap——它们「一次前向切成两半、一半算、一半等通信」的切法，在概念上与 ubatch 切分相似（都把一轮逻辑执行切成多个可独立推进的段），但目的不同：后两者是为了重叠通信与计算，而 ubatch 切分是为了适应固定形状的计算图。可参考 vLLM 源码完全拆解（二十七）：Dual Batch Overlap 与 SGLang 源码完全拆解（二十二）：Two-batch Overlap。

还有一种对照值得注意：vLLM 的 max_num_batched_tokens 与 llama.cpp 的 n_batch 名字很像，但语义不同。前者衡量「一个调度轮次能装多少 token」，后者衡量「一次 llama_decode 调用能装多少 token」。在 llama.cpp 里，一次 llama_decode 调用的输入由服务端 update_slots 组装，可能包含多个 slot 的解码 token 与一段 prompt；n_batch 是这条链路上单次调用能承受的总量。而 n_ubatch 是 n_batch 之下、GPU 图执行单元的上限。换句话说，llama.cpp 的切分是两次：服务端在组装 batch 时受 n_batch 约束，上下文在切成 ubatch 时受 n_ubatch 约束。

切分对参数选择的影响

最后给一组实用结论：读者在调参时看到的 -b、-ub 两个参数分别影响这条链路的不同环节。

-b（n_batch）决定一次 llama_decode 的逻辑容量，直接影响：

• 服务端能一次塞进多少 prompt token；
• 图复用率——n_batch 越大，每次 decode 包含的形状越多样，图复用越难，但单次调用能吞的 token 更多；
• 错误回滚粒度——失败时回滚的是整个 llama_decode 涉及的序列位置。

-ub（n_ubatch）决定每次图计算实际能装下的 token 数，直接影响：

• 计算缓冲区大小（按 n_ubatch 预留，注意力中间量是大头）；
• 长 prompt 被切成的段数（split_equal 或 split_simple 切成几轮）；
• split_seq 模式下必须满足的硬约束（非因果注意力要求 n_ubatch >= n_tokens）。

两个参数的经验取值：默认情况下 n_batch = n_ubatch = 512 或 2048（随版本与设备而异），预填充时增大 n_batch 可以提升吞吐，但 n_ubatch 别跟着盲目变大——它直接决定显存峰值；解码时 n_batch 只需要覆盖「最大并发数 × 1」，n_ubatch 也只需 ≥ 该值即可。对显存较小的设备，缩 n_ubatch 换取更大 n_batch（更多并发、更少 KV 碎片）往往是正确的取舍。

小结

把这一篇的要点压成几条：

第一，切分不在调用方，在上下文与记忆模块之间。 调用方递进来的是扁平 llama_batch（或新的 llama_batch_ext），llama_context::decode 通过 llama_batch_allocr 做校验与规范化，然后交给 llama_memory_i::init_batch 做切分。切分策略的选择权在记忆模块，因为只有它知道 n_stream 与容量结构。

第二，三种策略各有适用面。 split_simple 适合统一 KV 的扁平顺切；split_equal 适合非统一 KV 的「序列集不相交 + 等量」；split_seq 适合「整条序列必须同批」的非因果注意力与嵌入计算。split_equal 的尾部保留机制（n_keep_tail）为需要尾部窗口的模型提供额外保证。

第三，ubatch_add 负责所有指针的修复。 位置段优先布局在两份不同 n_tokens 之间搬运，seq_id 的两级指针采用「先 push 完再算起点」的两遍法，batch_idxs 为后续行序还原与嵌入提取留存原始下标。

第四，prepare 的「预演 + 恢复」是正确性保证。 全部 ubatch 必须同时装得下，任何一个失败即整批失败，不留半截状态。恢复栈保存单元内容，使 purge 也能被完整还原。

第五，输出侧的还原是懒执行的。 output_ids 提供反查表，output_swaps 记录选择排序过程中的交换，在读取 logits 时才真正执行。

下一篇进入 llama.cpp 源码完全拆解（二十六）：词嵌入与输出头——llama_context::get_embeddings 如何把图输出里的 logits 与嵌入按原始顺序交还给调用方。这一篇讲了切分之后「哪个 ubatch 里放了什么」，下一篇要讲切分之后的结果如何被 output_ids、output_swaps 与 embd 缓冲区还原回调用方视角的编号——也就是 llama_get_logits_ith、llama_get_embeddings_ith、池化嵌入与逐层输入这几条路径的完整实现。

参考资料

1. llama.cpp 官方仓库（src/ 与 ggml/ 源码）
2. src/llama-batch.cpp：llama_batch_allocr 的 init、三套 split_*、ubatch_add、llama_batch_compat
3. src/llama-batch.h：llama_ubatch、llama_batch_ext 与 llama_batch_allocr 接口
4. src/llama-context.cpp：decode 主循环、process_ubatch、graph_params 与输出行序还原
5. src/llama-kv-cache.cpp：init_batch 的策略选择、prepare 的预演与 find_slot
6. src/llama-memory.h：llama_memory_i 与 llama_memory_context_i 接口
7. src/llama-cparams.h：n_batch、n_ubatch、kv_unified 等运行时参数
8. tests/test-batch.cpp：llama.cpp 自带的批次切分单元测试
9. llama.cpp Issue #13746 讨论：KV 缓存位置区间的连续不变式
10. Efficient Memory Management for Large Language Model Serving with PagedAttention
11. vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention
12. M-RoPE: Multimodal Rotary Position Embedding and Modality-Agnostic Embeddings（Qwen2-VL）