编辑
2026-07-20
AI Infra
00

目录

推理变量源码阅读:llama.cpp 与 TensorRT-LLM
范围与版本
代表性场景
五个变量概览
一个具体的逻辑示例
端到端调用链
llama.cpp server 路径
TensorRT-LLM PyExecutor 路径
1. 初始隐藏状态
处理方式与对比
llama.cpp
TensorRT-LLM
差异
2. 位置 ID 与 RoPE
处理方式与对比
llama.cpp
TensorRT-LLM
差异
3. 注意力掩码
处理方式与对比
llama.cpp
TensorRT-LLM
差异
4. KV 缓存
处理方式与对比
llama.cpp
TensorRT-LLM
差异
5. 解码与采样
处理方式与对比
llama.cpp
TensorRT-LLM
差异
Prefill、Decode 与 Continuous Batching
llama.cpp 的 continuous batching
TensorRT-LLM 的 continuous batching
简明对比
附录:Qwen3.5 混合架构(GDN + 全注意力)
本附录的范围
A. TensorRT-LLM 模型加载后的运行流程(重点关注前后处理)
B. 示例输入:TensorRT-LLM 每一步的 shape 与 dtype
C. 该混合模型上的 llama.cpp 与 TensorRT-LLM 对比
C-1. 框架层
C-2. 输入模型格式
C-3. 计算图(差异最明显的部分)
附录:Serving 层优化与五个变量
1. Continuous(in-flight)batching
2. Speculative decoding
3. Paged memory management
为什么一块 GPU 可以同时服务多个用户
第一层:Continuous Batching
第二层:KV Cache 与显存管理
第三层:Speculative Decoding
三层优化分别解决什么问题
对应到五个核心变量
后续补充时的阅读规则
相关本地文档
待继续研究的问题

推理变量源码阅读:llama.cpp 与 TensorRT-LLM

状态:源码阅读整理文档

本文沿着两个已检出的运行时实现,追踪推理过程中的五个核心问题:

  1. 初始隐藏状态
  2. 位置 ID 与 RoPE
  3. 注意力掩码
  4. KV 缓存
  5. 解码与采样

本文的目标并不是强行用同一种抽象来描述两个系统,而是说明每个问题在运行时的哪个位置进入、使用什么对象表示,以及它们在 prefill、decode 与 continuous batching 之间如何变化。

范围与版本

项目选定内容
llama.cpp 版本c528416388b6363d2781d8b82a59d0fba4968740
TensorRT-LLM 版本5344a0ad729e3411c14a1ddac65c39f6d0050824
模型家族标准 Llama 因果解码器
llama.cpp 路径server -> llama_decode -> 标准 llama_kv_cache
TensorRT-LLM 路径PyExecutor -> PyTorch 模型 -> TRTLLM attention backend -> 标准 KV cache manager
Decode 场景每个活跃序列当前只有一个 token,beam width 为 1
不包含的内容speculative decoding、循环模型、M-RoPE、滑动窗口变体、beam search、encoder-decoder attention

TensorRT-LLM 的后端会动态选择受支持的 FMHA 实现。因此,本文只追踪到后端分发边界,不声称最终一定调用某一个特定 CUDA kernel。物理布局也可能随着 tensor parallelism、quantization 和后端配置而变化。

撰写本文时,llama.cpp 工作树中存在与本文无关的本地修改。下文的行为结论均基于检出版本中的源码路径,不依赖这些本地修改。

代表性场景

阅读各节时统一使用下面的场景:

  • 一个 Llama 模型收到一个请求。
  • Prefill 阶段处理该请求的 prompt。
  • Decode 阶段处理一个刚刚选出的新 token,此前 token 的 K/V 状态仍保留在缓存中。
  • 批处理章节再将同一个场景扩展到多个活跃请求,其中包括 prefill 与 decode 混合执行的情况。

[B, T, H] 等逻辑张量名称仅用于解释。运行时通常使用 packed 物理布局。标记为 推断(Inference) 的内容,表示该结论是根据邻近的多个源码操作综合得出,而不是来自某一条单独声明。

五个变量概览

下表是全文的核心对比。设一个请求已经缓存了 C 个 token,本轮有 T 个当前 token。完整 prefill 中,C = 0T = prompt_len;普通 decode 中,C = previous_sequence_lenT = 1

变量或状态单个请求的逻辑值llama.cpp 中的载体与处理方式TensorRT-LLM 中的载体与处理方式主要差异
初始隐藏状态X[0:T, 0:H],只为本轮当前 token 计算inp_tokensinp_embd -> build_inp_embd -> inpL;GGML 计算图通常使用 [H, T] 约定扁平化 input_ids -> embed_tokens -> hidden_states;packed 约定在逻辑上为 [T, H]embedding 结果相同,但物理维度习惯看起来互为转置,而且执行系统不同
位置 ID 与 RoPEP = [C, C+1, ..., C+T-1];RoPE 改变 Q 和 K,不改变 Vbatch.pos -> ubatch positions -> I32 inp_pos -> 显式 ggml_rope_ext(Q/K)packed position_ids;未融合时由 Python 侧 rotary_emb 处理,否则融合进 TRTLLM 后端显式计算图操作与可选后端融合之间的差异
注意力掩码逻辑保留条件:属于同一请求、是有效 cache cell,并且因果注意力满足 key_pos <= query_pos由 host 填充的 KQ_mask 计算图输入,保留写 0,屏蔽写 -INFcausal mask 枚举,加上 sequence lengths、request types、KV lengths 和 block offsets显式掩码数值与由 metadata/kernel 隐式推导掩码之间的差异
KV 缓存forward 前:保存 C 个旧 token 的 K/V;forward 后:保存 C+T 个 token 的 K/V根据 sequence ID 和 stream 选择 cache cells;计算图将当前 K/V 复制到 cell 索引,再读取 cache view根据 request ID 与 block-offset 行选择分页内存池;后端读取和写入选中的 pagescell/stream 所有权与 paged block table 所有权之间的差异
Decode/生成循环当前 token IDs -> hidden states -> logits -> sampled token -> 下一轮当前 token IDsserver slot 与 llama_batch -> llama_decode -> common_sampler_sample -> slot 更新scheduled request -> ModelEngine.forward -> async sampler -> request 更新递归逻辑相同,但编排状态不同;这一项是状态机,不是单个张量

一个具体的逻辑示例

下面的示例只用于说明逻辑,与具体后端无关。它展示了上述具体载体必须表达的数值关系。

步骤当前 token IDs本轮计算的隐藏状态位置逻辑上可见的 keyforward 后的缓存循环输出
prefill[a, b, c]三行/三个向量[0, 1, 2]query 0 看到 key 0;query 1 看到 key 0..1;query 2 看到 key 0..2保存位置 0..2 的 K/V从要求输出 logits 的 prompt 位置采样 d,通常是最后一个位置
decode 1[d]一行/一个向量[3]当前 query 看到缓存 key 0..2 和当前 key 3保存位置 0..3 的 K/V采样 e
decode 2[e]一行/一个向量[4]当前 query 看到缓存 key 0..3 和当前 key 4保存位置 0..4 的 K/V采样下一个 token

两个运行时都不需要把这些逻辑视图真正物化为稠密张量。下面各节将说明究竟由哪些对象承载这些信息。

端到端调用链

llama.cpp server 路径

text
server_context::update_slots -> server_batch::render -> server_context::decode -> llama_decode -> llama_context::decode -> llama_context::process_ubatch -> 构建计算图并填充输入 -> graph_compute -> server 侧采样 -> 将采样 token 加入下一轮 server batch

代码依据:

  • tools/server/server-context.cpp:2711 更新 slots,在 tools/server/server-context.cpp:2751 渲染 batch,并在 tools/server/server-context.cpp:2761 切分 batch 供 decode 使用。
  • server wrapper 在 tools/server/server-context.cpp:3528tools/server/server-context.cpp:3562 调用 llama_decode
  • llama_decodesrc/llama-context.cpp:4058 将调用委托给 context。
  • llama_context::decodesrc/llama-context.cpp:1680src/llama-context.cpp:1766src/llama-context.cpp:1822 准备 memory 并遍历各个 ubatch。
  • process_ubatchsrc/llama-context.cpp:1304src/llama-context.cpp:1337src/llama-context.cpp:1359src/llama-context.cpp:1364 构建计算图、填充计算图输入并执行计算。
  • 采样从 tools/server/server-context.cpp:3673 开始;在 tools/server/server-context.cpp:3724 选出 token,并通过 tools/server/server-context.cpp:448tools/server/server-context.cpp:2991 将其送入下一轮 batch。

TensorRT-LLM PyExecutor 路径

text
PyExecutor._executor_loop -> 调度 context 与 generation 请求 -> resource managers 准备请求资源 -> PyExecutor._forward_step -> ModelEngine.forward -> 准备 packed 输入与 attention metadata -> LlamaModel.forward -> attention backend 分发 -> 异步采样 -> 更新 request 状态

代码依据:

  • executor loop 在 TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:3933TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:3956TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:4004TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:4068TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:4106 依次完成调度、资源准备、forward、采样和 request 更新。
  • TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:5218 将 context 与 generation 请求分开调度。
  • model engine 在 TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:5340TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:5481TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:5547 准备输入并调用模型。
  • 采样与 request 更新位于 TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:6333TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:6379

1. 初始隐藏状态

处理方式与对比

问题llama.cppTensorRT-LLM
生产者build_inp_embd(model.tok_embd) 选择 token lookup 或调用方传入的 embeddings(src/llama-graph.cpp:2130src/models/llama.cpp:108LlamaModel.forward 调用 embed_tokens(input_ids)TensorRT-LLM/tensorrt_llm/_torch/models/modeling_llama.py:1096TensorRT-LLM/tensorrt_llm/_torch/models/modeling_llama.py:1112
Prefill 数值当前 prompt 中每个 token 对应一个 hidden vector;逻辑上 T = prompt_chunk_len扁平化后的每个当前 prompt token 对应一个 hidden vector;sequence_lengths 保留请求边界(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3389
Decode 数值普通单 token decode 中,每个活跃序列对应一个 hidden vector每个活跃 generation 序列对应一个 hidden vector;generation 的 sequence_lengths 初始化为 1(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3693
具体载体I32 inp_tokens 或 F32 inp_embd;选中的结果成为 inpL,物理形状为 [n_embd, n_tokens]src/llama-graph.cpp:2138src/llama-graph.cpp:2143src/llama-graph.cpp:2190扁平 GPU input_ids,随后变为 inputs_embeds/hidden_states;**推断:**逻辑形状为 [num_scheduled_tokens, hidden_size],这与扁平输入打包方式及后端对 token-major Q 形状的检查一致(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3838TensorRT-LLM/tensorrt_llm/_torch/attention_backend/trtllm.py:1640TensorRT-LLM/tensorrt_llm/_torch/attention_backend/trtllm.py:1648
第一个消费者第一层归一化,以及 layer loop 中的 Q/K/V 投影(src/models/llama.cpp:126第一个 LlamaDecoderLayer,它将 hidden_states 传入 self-attention(TensorRT-LLM/tensorrt_llm/_torch/models/modeling_llama.py:1118
迭代之间发生的变化新采样 token 取代 prompt chunk 成为当前输入;已缓存 token 不会重新计算隐藏状态相同;只对本轮调度的当前 token 做 embedding,旧历史由 KV metadata/cache 表示
核心差异GGML 计算图张量,序列身份位于 hidden-state 维度之外packed PyTorch/CUDA 张量,请求边界位于 hidden-state 维度之外

llama.cpp

build_inp_embd 会创建 token-ID 输入或显式 embedding 输入。对于 token ID,ggml_get_rows 从模型 token embedding table 中取出对应行(src/llama-graph.cpp:2130src/llama-graph.cpp:2138src/llama-graph.cpp:2155)。调用方直接提供 embedding 的备用路径位于 src/llama-graph.cpp:2180

Llama 计算图从该结果开始,将其送入每一层 transformer,最后执行 final normalization 与 output projection(src/models/llama.cpp:98src/models/llama.cpp:108src/models/llama.cpp:126src/models/llama.cpp:231)。

物理约定:GGML 将 embedding 输入表示为 [n_embd, n_tokens]。这里不会保留稠密 batch 维度;序列身份通过独立的 batch 与 memory metadata 传递。

TensorRT-LLM

模型接收 packed input IDs。LlamaModel.forward 调用 embed_tokens(input_ids),将结果赋值给 hidden_states,然后依次送入 decoder layers(TensorRT-LLM/tensorrt_llm/_torch/models/modeling_llama.py:1096TensorRT-LLM/tensorrt_llm/_torch/models/modeling_llama.py:1112TensorRT-LLM/tensorrt_llm/_torch/models/modeling_llama.py:1118)。

输入准备过程会把所有已调度请求的 token 扁平化到 input_ids 中,并复制到 GPU buffer(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3288TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3838)。

**推断:**普通 packed hidden-state 的逻辑形状为 [num_scheduled_tokens, hidden_size]。请求边界通过 lengths 和 metadata 重建,而不是依赖稠密 [batch, sequence, hidden] 张量。

差异

两个系统都从 embedding lookup 开始,并让 hidden states 流经概念上相同的 transformer stack。表面差异主要来自表示方式:llama.cpp 围绕 [H, N] 构建 GGML 计算图,而 TensorRT-LLM 围绕 [N, H] 准备 packed PyTorch/CUDA 输入,并配合更丰富的 attention metadata。

2. 位置 ID 与 RoPE

处理方式与对比

问题llama.cppTensorRT-LLM
位置生产者若调用方未提供位置,llama_batch_allocrmemory->seq_pos_max(seq_id) + 1 开始,并针对每个序列递增(src/llama-batch.cpp:90context 位置使用 range(begin_compute, begin_compute + len(prompt_tokens));generation 使用 past_seen_token_numTensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3368TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3725
Prefill 数值通常为 [C, ..., C+T-1];对于分块或复用 context,C 可以不为 0每个 context 请求附加相同的逻辑位置范围,随后再打包
Decode 数值每个序列一个位置,通常为 [C]每个 beam/request 一个 position_id,通常等于此前已经处理的 token 数
具体载体batch.pos -> ubatch.pos -> I32 计算图输入 inp_possrc/llama-batch.cpp:780src/llama-graph.cpp:2219Python position_ids 列表 -> packed CUDA position tensor(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:4092
消费者/效果ggml_rope_ext 旋转 QcurKcur;V 不参与 RoPE(src/models/llama.cpp:146src/models/llama.cpp:153rope_fusion 为 false 时,apply_rope 使用 rotary_emb 旋转 Q/K;否则由后端消费 positions(TensorRT-LLM/tensorrt_llm/_torch/modules/attention.py:1011TensorRT-LLM/tensorrt_llm/_torch/modules/attention.py:1062
核心差异positions 与 RoPE 都是 GGML 计算图中可见的输入/操作positions 始终显式存在,但其对 Q/K 的作用可以移动到融合后端中

llama.cpp

如果调用方没有提供 positions,batch allocator 会让每个序列从 memory->seq_pos_max(seq_id) + 1 开始,并在该序列内部逐 token 递增(src/llama-batch.cpp:90)。位置随后在 src/llama-batch.cpp:748src/llama-batch.cpp:780 被复制进各个 ubatch。

计算图将 positions 暴露为显式 I32 输入(src/llama-graph.cpp:2219)。Llama 计算图构建该输入,并把同一组 positions 提供给 Q 和 K 的 RoPE 操作(src/models/llama.cpp:110src/models/llama.cpp:146src/models/llama.cpp:153)。

因此,位置状态从每个序列的 memory state 推导出来,被物化为计算图输入,再由显式的图级 RoPE 操作消费。

TensorRT-LLM

对于 context token,输入准备阶段会创建与所选 prompt chunk 对应的位置范围(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3368)。对于普通 generation,它会追加最新 token,并根据已处理 token 数附加一个位置(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3725)。packed position tensor 在 TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:4092 被复制到 CUDA 并返回。

Attention 接收 position_ids。当 RoPE fusion 被关闭时,Python 侧 attention 会在后端分发前应用 RoPE;当 fusion 被启用时,该操作归后端负责(TensorRT-LLM/tensorrt_llm/_torch/modules/attention.py:1011TensorRT-LLM/tensorrt_llm/_torch/modules/attention.py:1062)。TRTLLM 后端在 TensorRT-LLM/tensorrt_llm/_torch/attention_backend/trtllm.py:1763 声明支持融合 RoPE。

差异

即使 token 已经被打包,两个系统仍然都按逻辑序列计算位置。llama.cpp 把 RoPE 明确表示为计算图操作;TensorRT-LLM 可以在 Python 模块中应用 RoPE,也可以将其融合进所选 attention backend。

3. 注意力掩码

处理方式与对比

问题llama.cppTensorRT-LLM
逻辑规则只保留属于当前 query 序列的 cache cell;因果注意力还必须满足 key_pos <= query_possrc/llama-kv-cache.cpp:1625src/llama-kv-cache.cpp:1641预定义 causal mask,加上每个请求有效的 query/KV 长度与该请求的 cache blocks(TensorRT-LLM/tensorrt_llm/_torch/modules/attention.py:938TensorRT-LLM/tensorrt_llm/_torch/attention_backend/trtllm.py:1674
Prefill 行为prompt 中每个 query 只能看到本序列中被允许的更早位置和当前位置context request lengths 与 cumulative/packed metadata 告诉后端每个 query 和 KV 的有效范围
Decode 行为唯一的当前 query 行可以保留所有允许的旧 cache cells 和当前 cellgeneration metadata 表示一个当前 query,以及 cached_len + current_len 个有效 KV 位置
显式数值host 代码向 KQ mask tensor 写入 0mask_keep)或 -INFINITYmask_drop)(src/llama-kv-cache.cpp:1546src/llama-kv-cache.cpp:1669src/llama-kv-cache.cpp:1674代表性 TRTLLM 路径不会构建稠密 mask 数值;attention_mask 是预定义枚举,约束由 metadata/kernel 规则表达
具体载体计算图输入,形状由 n_kv、当前 tokens 和 streams 决定(src/llama-graph.cpp:38);由 memory context 填充(src/llama-graph.cpp:467attention_maskseq_lenskv_lensrequest_typesrequest_idskv_cache_block_offsetsTensorRT-LLM/tensorrt_llm/_torch/attention_backend/interface.py:140TensorRT-LLM/tensorrt_llm/_torch/attention_backend/trtllm.py:547
消费者build_attn_mha(q, k, v, ..., kq_mask, ...)src/llama-graph.cpp:2656TRTLLM attention backend/FMHA 分发(TensorRT-LLM/tensorrt_llm/_torch/attention_backend/trtllm.py:1740
核心差异逻辑掩码被物化为计算图输入数值逻辑掩码由后端根据枚举与 metadata 重建

llama.cpp

计算图会创建显式 KQ mask 输入。其形状由 KV cells、tokens 与 streams 决定,并不是传统的稠密 [B, T, T] 张量(src/llama-graph.cpp:38)。memory context 通过 graph callback 填充该输入(src/llama-graph.cpp:467)。

对于标准 KV cache,mask 构建逻辑会拒绝不属于当前序列的 cell,并屏蔽因果意义上的未来位置(src/llama-kv-cache.cpp:1516src/llama-kv-cache.cpp:1625src/llama-kv-cache.cpp:1641)。可选的 sliding-window 限制在 src/llama-kv-cache.cpp:1659 附近应用。

Llama 计算图将该 mask 传入 multi-head attention(src/llama-graph.cpp:2575src/llama-graph.cpp:2656)。Flash-attention 与 non-flash 路径可能使用不同的 mask 数据类型,但 mask 始终是一个显式计算图输入。

TensorRT-LLM

代表性的 TRTLLM 路径不会构建稠密 attention-mask tensor,而是传递预定义 mask type 与 sequence metadata。模块默认使用 causal mask(TensorRT-LLM/tensorrt_llm/_torch/modules/attention.py:938),后端再把该枚举映射为面向 kernel 的 mask type(TensorRT-LLM/tensorrt_llm/_torch/attention_backend/trtllm.py:1369)。

metadata 保存 sequence lengths、context/generation 数量、累计长度信息、request IDs 与 block offsets(TensorRT-LLM/tensorrt_llm/_torch/attention_backend/interface.py:140TensorRT-LLM/tensorrt_llm/_torch/attention_backend/interface.py:205TensorRT-LLM/tensorrt_llm/_torch/attention_backend/trtllm.py:154)。后端在分发阶段使用每个序列的 KV lengths 与 prompt lengths(TensorRT-LLM/tensorrt_llm/_torch/attention_backend/trtllm.py:1674)。

差异

llama.cpp 将允许和禁止的注意力关系物化为由 host 填充的计算图张量。TensorRT-LLM 则用 causal mask 枚举、packed sequence metadata 与 paged-cache metadata 表达相同逻辑约束,并由 kernel 解释这些约束。

4. KV 缓存

处理方式与对比

问题llama.cppTensorRT-LLM
分配者memory->init_batch 与标准 llama_kv_cache 为 ubatch 查找 cells(src/llama-context.cpp:1771src/llama-kv-cache.cpp:698PyExecutor 的 KVCacheManager 分配 paged block pools 与 offsets(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:266TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:643
Prefill 动作分配 T 个位置,写入每层当前 K/V,再让后续 prompt queries 可以读取它们收集并加入 context sequences,为它们分配 blocks,并把对应 block rows 暴露给 attention(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:793TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:797
Decode 动作为新位置找到 cell,为每个活跃序列/层写入一组 K/V,并读取旧缓存与当前缓存为每个 generation request 增加当前 token 容量,刷新 block offsets,并读取旧 pages 与新 page(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:823TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:828
写入位置mctx_cur->cpy_kmctx_cur->cpy_vk_cur/v_cur 复制到选定索引(src/llama-graph.cpp:2646后端设置 update_kv_cache,并携带 manager 管理的 cache pool pointers 分发;具体 FMHA kernel 动态选择,本文不固定到某一个实现(TensorRT-LLM/tensorrt_llm/_torch/attention_backend/trtllm.py:1516TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:644
读取位置mctx_cur->get_kget_v 向 attention 返回 cache views(src/llama-graph.cpp:2653TRTLLM attention 在后端分发前接收有效 KV lengths 与 request 选中的 block offsets(TensorRT-LLM/tensorrt_llm/_torch/attention_backend/trtllm.py:572TensorRT-LLM/tensorrt_llm/_torch/attention_backend/trtllm.py:589
请求到存储的映射seq_id -> seq_to_stream -> cache cells;cells 同时记录 sequence ownership(src/llama-kv-cache.cpp:982request_id -> batch block-offset row -> paged pool blocksTensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:2226
增长/释放每个 ubatch 查找并应用 cells;后续生命周期操作可以移动、复制或释放序列拥有的 cellsmanager 增加 tokens/blocks、更新或回退 sequences,并释放 request resources(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:1019TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:1066
核心差异以 cell/stream 为中心的 cache,mask 会直接检查 sequence ownership以 page/block table 为中心的 cache,attention metadata 选择属于请求的 block rows

llama.cpp

计算图执行前,context 会要求 memory object 初始化 batch(src/llama-context.cpp:1766src/llama-context.cpp:1771)。标准 KV cache 将工作拆成 ubatches,查找可用 cells 并应用分配结果(src/llama-kv-cache.cpp:698src/llama-kv-cache.cpp:747src/llama-kv-cache.cpp:894)。

计算图将当前 K/V 写入选中的 cache 位置,然后读取 cached K/V views 供 attention 使用(src/llama-graph.cpp:2646src/llama-graph.cpp:2653)。序列隔离由 sequence-to-stream 映射,以及 attention mask 中的 sequence-ownership 检查共同保证(src/llama-kv-cache.cpp:982src/llama-kv-cache.cpp:1625)。

TensorRT-LLM

标准 PyExecutor KV manager 根据 layer、head、head-dimension 与 tokens-per-block 配置 paged block pools(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:266)。它在 TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:643 分配 cache pools 与 block-offset storage。

资源准备期间,context requests 被加入管理器,generation requests 则扩展 sequence,为新 token 增加容量(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:793TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:823)。manager 在 TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:2226 把每个请求的 block offsets 复制到 batch rows;attention metadata 再在 TensorRT-LLM/tensorrt_llm/_torch/attention_backend/trtllm.py:589 根据 request ID 选择对应行。

差异

两个系统都把 packed compute order 与持久化 sequence history 分开。llama.cpp 在构建计算图时暴露逻辑 cache cells 与 sequence-to-stream ownership;TensorRT-LLM 向 attention backend 暴露 paged block pools 与 per-request block-offset tables。不能在假设 K/V 维度与轴顺序相同的前提下直接比较两者物理布局。

5. 解码与采样

第五项并不是一个模型张量,而是产生前四项下一轮取值的状态转换过程。

处理方式与对比

问题llama.cppTensorRT-LLM
循环输入状态server slot 向 server_batch 提供 current token、position、sequence ID 与 output flag(tools/server/server-context.cpp:104tools/server/server-context.cpp:136tools/server/server-context.cpp:448scheduled request 提供 latest token、position、request ID、cached length 与 output gather index(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3696TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3717TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3725
解码器计算llama_decode -> llama_context::decode -> ubatch graph -> transformer layers(src/llama-context.cpp:4058src/llama-context.cpp:1680_forward_step -> ModelEngine.forward -> LlamaModel.forwardTensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:6207TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:5340
模型输出选中的 output hidden states -> final norm -> LM head/logits(src/models/llama.cpp:231model engine 在 model forward 后 gather logits(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:5598
采样server 调用 common_sampler_sample 并接受 token(tools/server/server-context.cpp:3724tools/server/server-context.cpp:3729_sample_async 调用 sampler(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:6333
停止/状态更新server 处理 token 与停止条件,然后更新 slot(tools/server/server-context.cpp:3755_update_requests 应用 sampler 结果并更新 request state(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:6379
下一轮输入构建下一次 update_slots/batch build 时,将 sampled token 放到下一位置(tools/server/server-context.cpp:2991下一轮 scheduled generation input preparation 从已更新 request 中收集 latest token
核心差异围绕 llama_decode 的 slot-centric、表面同步式 server 编排request-centric scheduler,将 model forward 与 asynchronous sampling 分阶段执行

llama.cpp

普通 decode step 中,下一枚 sampled token 会在下一位置加入 server batch(tools/server/server-context.cpp:448)。llama_decode 可以把组合后的 server batch 拆成 ubatches、更新 memory、计算 logits,并把结果返回 server。server 调用 common_sampler_sample,接受采样结果、检查停止条件,随后再把它加入下一轮 batch(tools/server/server-context.cpp:3724tools/server/server-context.cpp:3729tools/server/server-context.cpp:3755)。

TensorRT-LLM

普通 generation input preparation 为每个活跃 request 提供 latest token,并记录其 past-seen-token count(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3717)。forward step 之后,_sample_async 调用 sampler,_update_requests 再把结果写入 request state(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:6333TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:6379)。

差异

两者递归关系相同:

text
selected token -> next position -> 携带缓存历史执行 model forward -> logits -> sampling -> selected token

llama.cpp 通过 server slots、llama_batch 与 GGML graph execution 暴露该递归过程。TensorRT-LLM 则通过 scheduler requests、packed model inputs、asynchronous sampling 与 request-state updates 表达该过程。

Prefill、Decode 与 Continuous Batching

关注点PrefillDecode
当前 tokensprompt chunk,通常每个请求包含多个 token通常每个活跃序列一个 token
positions每个 prompt chunk 对应一段连续范围缓存历史之后的一个位置
attention在当前 prompt 内执行 causal attention,同时可读取已有 cache当前 query 读取被允许的 cached history
KV 动作分配并写入多个位置分配并写入新位置
输出用途通常只有 prompt 最后位置需要用于采样的 logits每个活跃序列都需要 next-token logits

llama.cpp 的 continuous batching

server 先为处于 generation 状态的 slots 加入 sampled tokens;启用 continuous batching 时,还可以加入待处理 prompt tokens(tools/server/server-context.cpp:2991)。server_batch 为每个 token 保存 slot ID,并在 render 时把它作为该 token 的 sequence ID(tools/server/server-context.cpp:104tools/server/server-context.cpp:136)。来自多个 slots 的 prompt tokens 在 tools/server/server-context.cpp:3411 被加入。

底层 cache 将 sequence IDs 映射到 streams/cells,mask 构建过程会检查某个 cache cell 是否属于当前 token 的 sequence。因此,把多个请求的计算合并到一起,并不代表它们会共享请求历史。

TensorRT-LLM 的 continuous batching

scheduler 在同一个 scheduled step 中返回 context 与 generation request groups(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:5218)。输入准备阶段把它们的当前 tokens 扁平化,同时保留 request IDs、sequence lengths、cached-token counts 与 paged-cache block offsets(TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3288TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:4141TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:4168)。

request ID 会在 attention 前选择该请求的 block-offset rows。与 llama.cpp 一样,请求打包只合并计算工作,不会合并 KV 历史。

简明对比

关注点llama.cppTensorRT-LLM PyExecutor + TRTLLM
hidden statesGGML 计算图,通常为 [H, N]packed PyTorch/CUDA tensor,逻辑上为 [N, H]
positions显式计算图输入packed CUDA tensor 与 attention metadata
RoPE可见的计算图操作Python 操作或后端融合
causal masking显式 KQ mask tensormask 枚举,加 lengths 与 cache metadata
请求隔离sequence IDs、streams/cells 与 ownership maskrequest IDs 与 paged block-offset rows
KV cachecache cells 暴露给计算图构建过程paged pools 暴露给 attention backend
batchingserver slots -> token-level sequence IDs -> ubatchesscheduled requests -> flat tokens -> metadata
sampling loopserver 侧 sampler 与 slot 更新async sampler 与 request 更新

附录:Qwen3.5 混合架构(GDN + 全注意力)

上文主体限定为标准 Llama 因果解码器,并排除了循环模型与 M-RoPE。本附录把相同的五变量阅读方法扩展到此前特意未纳入范围的 Qwen3.5 混合架构。本附录内容自洽,不会改变上文任何结论。

本附录的范围

项目选定内容
模型Qwen/Qwen3.5-0.8B,纯文本
形状信息hidden 1024,8 个 query heads,2 个 KV heads,head_dim 256,24 层 = 18 个 GDN(linear_attention)+ 6 个 full_attention,vocab 248320,partial_rotary_factor = 0.25 -> rotary_dim 64
GDN 信息每个 GDN 层:causal-conv 状态宽度为 conv_kernel(4) - 1 = 3;循环状态为 [num_v_heads=16, head_k_dim=128, head_v_dim=128]
阶段prefill 后的普通 decode;beam width 1;TP=1;不使用 speculative decoding
llama.cpp 载体src/models/qwen35moe.cpp 计算图构建 + hybrid/recurrent memory
TensorRT-LLM 载体Qwen3NextModel(Qwen3.5 复用 Qwen3Next 混合层)

下文所有 tensor shape 都属于 推断(Inference):它们由配置和源码推导,并非从实际运行中采集,符合阅读规则 7。

A. TensorRT-LLM 模型加载后的运行流程(重点关注前后处理)

flowchart TD
    subgraph PRE["前处理(CPU,编排逻辑,不执行模型数学)"]
        A1["scheduler.schedule_request"] --> A2["resource_manager.prepare_resources\nKVCache.add_token / Mamba state slot"]
        A2 --> A3["_prepare_tp_inputs\ninput_ids=[y0](只保留最后一个 token)\nposition_id = past_seen_token_num = 8\nseq_lens / kv_lens / block_offsets / request_ids"]
    end
    A3 --> B0
    subgraph FWD["模型前向(GPU,执行模型数学)"]
        B0["embedding lookup:input_ids[1] -> hidden[1,1024] BF16"] --> B1{"层类型(共 24 层)"}
        B1 -->|"GDN x18(忽略 position)"| G1["causal_conv1d_update\n+ fused sigmoid-gating delta-rule\n(读写 conv 状态与循环状态)"]
        B1 -->|"full-attn x6"| F1["qkv_proj -> 拆分 Q/gate/K/V\nfused_qk_norm_rope(QK-norm + 旋转前 64 维)"]
        F1 --> F2["FMHA:读取历史 KV(0..7) + 当前 K/V(8)\ncausal 由 metadata 隐式表达(无稠密 mask)\n把当前 K/V 写回 page slot"]
        F2 --> F3["out * sigmoid(gate) -> o_proj"]
        G1 --> M1["residual + MoE FFN"]
        F3 --> M1
        M1 --> B9{"还有下一层?"}
        B9 -->|是| B1
        B9 -->|否| B10["final hidden[1,1024] BF16"]
    end
    B10 --> C0
    subgraph POST["后处理(GPU -> CPU,编排逻辑)"]
        C0["LogitsProcessor:每个请求取最后一个 token\nlm_head -> logits[1,248320] FP32"] --> C1["TorchSampler.sample_async(greedy=argmax)-> y1"]
        C1 --> C2["_update_requests:将 y1 写回 request"]
        C2 --> C3["下一轮:input=[y1],position=9,kv_lens=10"]
    end
    C3 -.下一轮迭代.-> A1

该引擎中五个变量的处理规则,以及前后状态如下:

变量处理前(生产者)计算规则处理后(消费者)
input hidden statesembedding lookup weight[id]只保留最后一个 token;宽度保持 1024,sequence axis 对应 packed tokens依次经过 24 层更新
position / RoPEposition = past_seen_token_num不是加法操作;RoPE 只旋转每个 head 前 head_dim*0.25 = 64 维,角度为 pos / base^(2i/64);仅由 6 个 full-attn 层消费,GDN 忽略 positionQ/K rotation
attention maskcausal 枚举 + kv_lens / block offsets该路径不会物化稠密矩阵;causal 与 request isolation 都存在于 metadata 中FMHA kernel
KV cacheqkv_proj 生成当前 K/Vfull-attn:分页读取历史并写入当前 K/V;GDN:原地推进固定大小的 conv 状态与循环状态下一 decode step
decode整个 forwardfull-attn = fused FMHA;GDN = causal_conv1d_update + delta-rule recurrencesampler 输出 y1

代码依据:_prepare_tp_inputs 中 position 位于 TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3725,mRoPE 三轴展开位于 TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:3769;embedding 位于 TensorRT-LLM/tensorrt_llm/_torch/models/modeling_qwen3_next.py:942;融合 QK-norm+RoPE 位于 TensorRT-LLM/tensorrt_llm/_torch/modules/qk_norm_attention.py:237;full-attn forward 位于 TensorRT-LLM/tensorrt_llm/_torch/modules/attention.py:973;FMHA 分发位于 TensorRT-LLM/tensorrt_llm/_torch/attention_backend/trtllm.py:1740;GDN decode 位于 TensorRT-LLM/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py:594

B. 示例输入:TensorRT-LLM 每一步的 shape 与 dtype

假设 prompt 被 tokenize 为 8 个 token(位置 0..7);prefill 采样得到 y0,随后进入 decode。

Prefill(T = 8C = 0):

变量shapedtype说明
input_ids[8]int32完整 prompt
hidden(embedding 后)[8, 1024]BF16sequence axis 长度为 8
position_ids[3, 1, 8]int32mRoPE 三个轴;纯文本场景中三个轴都为 0..7
attention mask不存在稠密张量;kv_lens=[8]-通过 metadata 表达三角形 causal 约束
KV cache(每个 full-attn 层)K/V 逻辑形状分别为 [8, 2, 256]BF16物理分页形状为 [pages, 2, 2, 32, 256]
GDN state(每个 GDN 层)conv [conv_dim, 3] / recurrent [16,128,128]BF16扫描 8 个 token 后的最终状态
logits[1, 248320]FP32只取最后一个 token

Decode 第 1 步(T = 1C = 8,position = 8):

变量输入 shape/dtype计算输出 shape/dtype
input_ids[1] int32 = [y0]只保留最后一个 token-
hidden-embedding lookup[1,1024] BF16
position_ids[3,1,1] int32,数值为 8fused_qk_norm_ropeQ/K 前 64 维原地旋转(BF16)
Q / gate / K / Vqkv_proj [1,1024] -> [1,5120] BF16拆分Q [1,2048]->[1,8,256];gate [1,2048];K/V [1,512]->[1,2,256]
attention maskkv_lens=9(读取 0..8)隐式 causal不产生 mask tensor
KV cache读取历史 [8,2,256] + 当前 K/VFMHA 读取并写入写入 slot 8;逻辑形状变为 [9,2,256]
GDN staterecurrent [16,128,128]执行 1 步 delta-rule形状不变,数值推进
attn out[1,2048] BF16* sigmoid(gate) -> o_proj[1,1024] BF16
logitsfinal hidden [1,1024]lm_head + cast[1,248320] FP32
sampled-argmaxy1 int

Decode 第 2 步(position = 9)的结构完全相同,只有三个标量发生变化:input_ids=[y1]position=9kv_lens=10。所有非缓存 tensor 都保持第 1 步的 shape 与 dtype。只有 KV cache 的逻辑长度增长为 [10,2,256];GDN recurrent state 甚至不会增长,因为其大小固定。正是这种形状稳定性,使 decode forward 可以被捕获为可重复 replay 的 CUDA graph。

C. 该混合模型上的 llama.cpp 与 TensorRT-LLM 对比

C-1. 框架层

维度llama.cppTensorRT-LLM
语言/后端C++ + ggml;CPU/Metal/CUDA/VulkanPython+C++(nanobind)+ PyTorch/CUDA;仅支持 NVIDIA
执行模型每个 ubatch 重建静态计算图(build_arch_graph),由 ggml_backend_sched 调度PyTorch eager + fused custom ops(torch.ops.trtllm.*)+ 可选的整段 forward CUDA graph replay
混合层标记GGUF KV 中的 is_recr_impl / full_attention_intervalsrc/models/qwen35moe.cpp:24根据 HF config 调用 get_qwen3_hybrid_layer_types()
KV/状态存储llama_kv_cache + memory_recurrent + memory_hybridKVCacheManager(paged)+ MambaHybridCacheManager
采样CPU 上的 common/samplingGPU 上的 TorchSampler

C-2. 输入模型格式

llama.cppTensorRT-LLM(PyTorch 后端)
载体单个 GGUF 文件;需要离线执行 convert_hf_to_gguf.pyHF 目录(config.json + safetensors + tokenizer);无需离线转换
内容weights + hparams(例如 ssm_d_convssm_n_group)+ tokenizer,统一打包weights 与 config 分离,在运行时读取
Tensor 命名blk.{i}.attn_norm.weightblk.{i}.ssm_conv1d.weightsrc/models/qwen35moe.cpp:72PyTorch state_dict 中的 model.layers.{i}.self_attn.*
加载/量化mmap;直接读取 GGUF block quant(Q4_K/Q5_K/Q8_0)safetensors + weight_mapper;运行时通过 ModelOpt 使用 FP8/NVFP4/AWQ

C-3. 计算图(差异最明显的部分)

同一个 GDN linear-attention 层,在两个框架中呈现为非常不同的图结构。

TensorRT-LLM(gdn_mixer.py:594)主要由两个粗粒度融合算子组成:

text
causal_conv1d_update(...) # 一个融合 kernel:更新 conv window + 计算 q/k/v fused_sigmoid_gating_delta_rule_update(...) # 一个融合 kernel:gating + delta recurrence + 状态读写

llama.cpp(src/models/qwen35moe.cpp:362build_layer_attn_linear)则由大约 10 个显式 ggml 节点组成:

text
ggml_ssm_conv -> ggml_silu -> view q/k/v -> ggml_l2_norm x2 -> build_recurrent_attn(delta net)-> build_norm_gated(rms_norm + silu + mul) -> ggml_reshape -> ssm_out mul_mat

full-attention 层同样体现出这种差异。TensorRT-LLM:qkv_proj -> fused_qk_norm_rope -> FMHA -> gate -> o_proj,大约 5 个操作,其中 norm+RoPE 已融合。llama.cpp(src/models/qwen35moe.cpp:281):wq/wk/wv three mul_mat -> build_norm(Q) -> build_norm(K) -> ggml_rope_multi x2 -> build_attn -> sigmoid -> mul -> wo,大约 12 个节点,norm 与 RoPE 分别是独立节点。

最重要的三个计算图层面差异如下:

关注点llama.cppTensorRT-LLM
attention mask物化 attn_inp_kq_mask tensor,形状 [n_kv, n_tokens, 1, n_stream],dtype 为 F16/F32,由 host 填充 0/-INF(src/llama-graph.cpp:38src/llama-graph.cpp:450不构建稠密 mask;causal 由 metadata + kernel 隐式表达
position / RoPEI32 一维 tensor,长度为 n_tokens * n_pos_per_embdsrc/llama-graph.cpp:2224);ggml_rope_multi 是独立节点,QK-norm 在它之前单独计算int32 [3,1,N]fused_qk_norm_rope 把 norm + RoPE 融合成一个操作
fusion / schedulingggml_backend_sched 调度细粒度节点;每个 ubatch 重建计算图粗粒度融合算子 + 运行时 FMHA 分发 + 整段 forward CUDA graph replay

一句话总结:llama.cpp 是“显式静态图 + 细粒度 ggml 节点 + 物化 mask + 离线 GGUF”;TensorRT-LLM 是“eager + 粗粒度融合 kernel + metadata 隐式 mask + 直接读取 HF,并在运行时完成量化/CUDA graph”。两者数学含义等价,只是抽象边界画在了不同位置。

附录:Serving 层优化与五个变量

现代 serving framework(vLLM、TensorRT-LLM、TGI)会在 decode loop 外再包一层工程系统。真正让一块 GPU 同时服务几十个并发用户的,主要是 continuous batchingspeculative decodingpaged memory management。这些优化都不会改变模型数学,但会改变五个变量的组织方式生命周期。本附录分别从三种优化的角度重新观察五个变量。下文 shape 均属于由配置推导的 推断(Inference),代码引用沿用文档开头的版本。

1. Continuous(in-flight)batching

系统不再等待一整个 batch 全部结束后才开始下一批,而是在每次迭代中重新组成运行 batch:新请求加入,完成的请求退出,所有活跃请求的当前 token 被压入一个扁平 buffer。

变量影响
input hidden states不再是稠密 [B, T, H],而是 [各请求当前 token 数之和, H] 的扁平 packed tensor。用户 A 与用户 B 的 token 可以连续存放;边界由 seq_lens / num_contexts 恢复。一个 mixed batch 可以同时包含多 token 的 context request 与单 token 的 generation request。
position / RoPEposition 属于每个 token,而不是属于整个 batch;每个 token 都携带其所属请求的 past_seen_token_num,因此交错执行不会造成冲突。
attention mask这是请求隔离的关键:A 的 query 不能读取 B 的 keys。TRT-LLM 通过每请求 kv_lens + block table,llama.cpp 通过 cache cell 的 sequence-ownership,并结合 varlen/cu_seqlens attention 实现,而不是构建稠密 mask。
KV cache每个请求拥有自己的 blocks;packing 合并的是计算,不是历史(copy_batch_block_offsets 选择每个请求自己的 offset rows)。
decode循环不再把一个请求完整执行完才处理下一个请求;每轮让所有活跃请求前进一步,同时允许新的 context requests 加入。scheduler 每一步都重新决定 batch 成员。

代码依据:executor loop 在 TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/py_executor.py:3933 每轮执行调度、forward、采样与更新;保留边界的 packed inputs 位于 TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/model_engine.py:4141;每请求 block-offset rows 位于 TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:2226。llama.cpp 通过 server_batch 保留每个 token 的 sequence ID,并在 src/llama-kv-cache.cpp:1625 使用 ownership mask 完成隔离。

2. Speculative decoding

draft 模型先提出 K 个 token,target 模型在一次 forward 中验证 K+1 个候选,并保留最长的正确前缀。这会把 decode 再次变成一个“小型 prefill”,因此在普通 decode 中退化为简单形式的变量,在这里重新变得复杂。

变量影响
input hidden states一个 generation request 不再输入 1 个 token,而是输入 draft_len + 1 个候选 token;hidden 变为 [num_req * (draft_len+1), H]。此时 decode 的形态重新接近 prefill。
position / RoPEdraft tokens 使用连续位置 past + 0..K;tree-based drafting 根据 tree offset 分配位置。这正是 _adjust_position_ids_for_spec_decspec_decoding_position_offsets 存在的原因。position 不再是单个标量。
attention maskdecode 阶段真正复杂的 mask 再次出现:链式 drafting 使用覆盖 K 个 token 的小型 causal mask;tree drafting 则使用自定义 tree mask,每个候选只能关注其祖先。这是“decode mask 会退化”的例外。
KV cacheK 个 token 的 KV 会被暂时写入;若部分候选被拒绝,则需要把 cache 回退(rewind) 到被接受的长度。循环状态(GDN/Mamba)更复杂:它不能简单原地覆盖,因此需要保存每一步的 intermediate states,最后选择被接受位置对应的状态。
decode一次 target forward 验证 K 个 draft,并在一次迭代中接受多个 token。Qwen3.5 把 draft head 放在模型内部(内置 MTP),因此不一定需要单独的 draft model。

代码依据:spec-dec position adjustment 位于 TensorRT-LLM/tensorrt_llm/_torch/modules/attention.py:1008;GDN verify 路径在 TensorRT-LLM/tensorrt_llm/_torch/modules/mamba/gdn_mixer.py:459 写入 intermediate_conv_states / intermediate_ssm_states,并延迟到 update_mamba_states() 再选择最终状态;内置 MTP head Qwen3NextMTP 与单模型 MTP 接线位于 TensorRT-LLM/tensorrt_llm/_torch/models/modeling_qwen3_next.py:704TensorRT-LLM/tensorrt_llm/_torch/models/modeling_qwen3_next.py:966。llama.cpp 在 src/models/qwen35moe.cpp:552graph_mtp 中构建 Qwen3.5 MTP draft graph。

3. Paged memory management

并发能力最终由内存决定,而 KV cache 是其中的核心。

变量影响
KV cachePagedAttention:KV 存放在固定大小的 blocks 中(TRT-LLM 的 tokens_per_block=32,vLLM 默认 16);逻辑上连续的序列通过 block table 映射到物理上离散的 pages。优势包括:无需按 max_seq_len 预分配、减少碎片、通过 add_token 按需增长、请求完成后释放,以及 prefix/block reuse。Qwen3.5 混合模型会关闭 block reuse,因为 GDN state 没有等价的 block-sharing 语义。
GDN / recurrent state与 KV cache 对比:attention KV 会分页并持续增长;recurrent state 则是每请求固定大小的 slab,通过 state_indices 管理,不分页,大小恒定。
attention mask / isolationblock table 本身就是隔离机制,即使用 metadata 而不是 mask matrix。
input hidden / position它们是每一步的临时数据,不由 paged memory 直接管理;但正是分页所支持的 packing,使许多请求能够共享同一个计算 buffer。
decodescheduler 根据空闲 block 数决定是否接纳请求;内存不足时执行 preemption/eviction,例如 recompute 或 offload。这是“一块 GPU 支持几十个并发用户”的内存侧基础。

代码依据:tokens_per_block 默认值位于 TensorRT-LLM/tensorrt_llm/llmapi/llm_args.py:3435;paged pools 与 block offsets 位于 TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:643add_token 增长位于 TensorRT-LLM/tensorrt_llm/_torch/pyexecutor/resource_manager.py:823;hybrid block reuse 在 TensorRT-LLM/tensorrt_llm/_torch/models/modeling_qwen3_next.py:989 被禁用;固定大小 Mamba state slots 由 CppMambaHybridCacheManagerget_state_indices 管理。llama.cpp 则把对应职责拆分到 llama_kv_cachememory_recurrentmemory_hybrid 中。

为什么一块 GPU 可以同时服务多个用户

现代 serving framework 的核心价值,不是把单次 forward() 写得多么特殊,而是把模型 forward() 放进一个由 请求调度、KV 管理、批处理重排与结果回填 组成的长期运行循环中。

普通单请求推理主循环通常是:

text
用户 A 输入 prompt → prefill → decode token 1 → decode token 2 → decode token 3 → ...

现代 serving framework 的主循环则更接近:

text
请求队列中同时存在 A / B / C / D ... → scheduler 在本轮选择可以执行的请求与 token → 将 context chunk 和 decode token 组织为 packed batch → 为每个请求分配或查询 KV cache block → 执行一次模型 forward → 完成采样、候选验证或状态回填 → 已完成请求退出,新请求进入 → 开始下一轮调度

因此,llama.cpp 的简单示例或 PyTorch 手写生成循环更像“单请求推理流程”,而 vLLM、TensorRT-LLM 和 TGI 更像一个面向多请求的“推理操作系统”。它们不仅执行模型计算,还管理请求生命周期、token 级调度、KV cache、显存容量、流式输出和性能监控。

第一层:Continuous Batching

Continuous batching 也称为 in-flight batchingiteration-level batching

传统静态 batching 要求同一批请求一起开始,并且通常要等待该批中最慢的请求结束。这样会产生两个问题:

  1. 已经完成的短请求仍然占据 batch slot;
  2. 新请求必须等待当前整批结束后才能进入 GPU。

Continuous batching 则在每一轮推理后重新调度:

text
第 N 轮:A、B、C、D 一起执行 第 N+1 轮:B 已完成并退出,E 立即加入 第 N+2 轮:A、C、D、E 继续执行

调度粒度从“整条请求”下降到“本轮需要处理的 token”。同一轮中可以同时出现:

  • 某些请求的 context/prefill chunk;
  • 某些请求的单 token decode;
  • 某些请求的 speculative verification token。

TensorRT-LLM 将 in-flight batching 描述为 continuous batching / iteration-level batching,并支持在同一调度体系内处理 context phase 与 generation phase 请求。vLLM 和 TGI 同样把 continuous batching 作为提高 GPU 利用率与并发吞吐的核心机制。(TensorRT-LLM) (vLLM) (TGI)

第二层:KV Cache 与显存管理

Decode 阶段中,每个请求每轮通常只增加一个新 token,但 attention 仍需读取该请求此前所有 token 对应的 K/V。

若为每个请求预留一块足以容纳最大序列长度的连续显存,会产生大量浪费:

text
请求实际长度:128 tokens 预留最大容量:8192 tokens 未使用空间:8064 tokens 对应的 KV 显存

Paged KV cache / PagedAttention 将 KV cache 切分为固定大小的 block/page,并通过 block table 完成映射:

text
逻辑 token 位置 ↓ block table 物理 KV block/page

因此,一个请求的逻辑 KV 序列在物理显存中不必连续:

text
请求 A:block 3 → block 9 → block 15 请求 B:block 1 → block 7 请求 C:block 4 → block 5 → block 12

Cache manager 负责:

  • 分配新 block;
  • 跟踪逻辑位置与物理 block 的对应关系;
  • 在请求结束时回收 block;
  • 在容量不足时执行 eviction、preemption 或 offloading;
  • 在条件满足时进行 prefix/KV cache reuse。

TensorRT-LLM 的 paged KV cache 将 KV 分解为 block,并由 cache manager 负责分配、跟踪和回收;其 KV cache 系统还提供跨请求复用、offloading 与优先级驱逐等能力。(TensorRT-LLM IFB) (TensorRT-LLM KV Cache)

第三层:Speculative Decoding

普通自回归 decode 通常遵循:

text
一次 target model forward → 确认 1 个 token

由于每轮 decode 的 query token 很少,这一阶段往往受到显存带宽、kernel launch 和串行依赖限制。Speculative decoding 使用轻量 drafting 机制预先提出多个候选 token,再由 target model 一次验证多个候选:

text
Draft model / Medusa / N-gram / EAGLE / MTP → 草拟 token 1、2、3、4 → Target model 一次 forward 验证 → 接受连续正确前缀 → 从第一个不接受的位置继续生成

其核心目标不是改变 target model 的输出分布,而是减少必须串行执行的 target-model forward 次数。

TensorRT-LLM 将 speculative decoding 描述为由轻量 drafting 机制提出候选、target model 在一次 forward 中验证多个 token。vLLM 强调该技术主要用于降低中低 QPS、memory-bound 场景下的 inter-token latency;TGI 则将 speculative decoding、Medusa 与 assisted generation 归为“先提出候选,再由大模型检查”的同类方法。(TensorRT-LLM Speculative Decoding) (vLLM Speculative Decoding) (TGI Speculation)

三层优化分别解决什么问题

优化主要缓解的瓶颈对五个变量影响最明显的部分
continuous batchingGPU 利用率与请求排队延迟input_ids/hidden states 的 packed 组织、请求级 position、metadata mask、scheduler-driven decode
paged KV cache显存碎片、并发容量与请求生命周期管理KV cache 从连续 tensor 变成 page/block pool,并通过 block table 间接寻址
speculative decoding自回归串行 forward 次数与 inter-token latency一轮输入多个候选 token、tree/verification mask、KV 接受与回滚

对应到五个核心变量

变量在现代 serving framework 中的变化
input_ids / input hidden states不再只是单用户连续序列,而是多个请求的 context chunk、当前 decode token 或候选 token 组成的 packed batch。Embedding 后通常形成 [sum_tokens, hidden_size] 一类的紧凑表示。
position_ids / position embedding每个 token 必须携带所属请求中的逻辑位置。即使多个请求在 packed batch 中交错,RoPE 位置也不能按照物理 batch 下标重新编号。
attention mask除 causal 约束外,还必须保证请求隔离;请求 A 的 query 不能读取请求 B 的 KV。该约束往往不再表现为完整稠密 mask,而由 seq_lens、request mapping、block table 和 kernel metadata 隐式编码。
KV cache从模型内部的普通连续 tensor,演变为由 cache manager 管理的 block/page pool。其生命周期跨越多轮 forward,并与请求加入、暂停、迁移、完成和淘汰绑定。
decode不再是单请求 while 循环,而是 scheduler 驱动的多请求 token-level loop;每轮都会重新选择请求、组织 batch、执行 forward、采样并更新状态。

这五个变量的数学身份并没有改变,但它们的组织方式、寻址方式和生命周期发生了变化:hidden states 被打包,positions 归属于逻辑 token,mask 通过 metadata 保证请求隔离,KV 由 paged cache 管理,decode 则由 scheduler 编排。

一句话概括:

一块 GPU 能同时服务几十个并发用户,不是因为单个请求的计算速度提高了几十倍,而是 serving framework 将许多用户零散的“一 token decode 小任务”持续合并为 GPU 更擅长执行的批量计算,同时利用分页 KV cache 控制显存碎片、容量与请求生命周期。

后续补充时的阅读规则

扩展本文时应遵循以下规则:

  1. 明确模型、后端、阶段、batching mode 与 cache implementation。
  2. 从公开执行入口开始,一直追踪到 backend boundary。
  3. 分别记录 logical shape 与 physical packed shape。
  4. 区分显式 tensor 与通过 metadata 编码的约束。
  5. 必须证明 request isolation,而不能默认 batching 天然保持隔离。
  6. 后端动态分发时,应把具体 backend behavior 标记为尚未确定。
  7. 多个源码位置综合得出的结论标记为 推断(Inference)

相关本地文档

待继续研究的问题

  • 在固定 GPU architecture、dtype、quantization 与 runtime configuration 后,继续追踪一个确定的 TensorRT-LLM FMHA 实现。
  • 对比 llama.cpp unified/non-unified KV 模式与 TensorRT-LLM cache reuse/offloading 模式。
  • 把 speculative decoding 作为独立场景加入主文;它会同时改变 current-token packing 与 cache acceptance/rewind 行为。
  • 使用具体模型采集运行时 tensor shapes,验证上文的逻辑 shape 推断。

本文作者:WarF

本文链接:

版权声明:本博客所有文章除特别声明外,均采用 BY-NC-SA 许可协议。转载请注明出处!