3、Transformer与生成式模型基础：注意力机制、类模型结构、KV Cache原理、生成式推理的典型流程

在边缘设备上部署生成式AI模型，首先需要理解其底层架构的核心逻辑。Transformer架构自2017年提出以来，已成为大语言模型（LLM）的事实标准。无论是系列、LLaMA还是ChatGLM，其本质都是基于Transformer解码器（Decoder-Only）的变体。本章将深入拆解注意力机制的计算本质、类模型的层级结构、KV Cache的缓存原理，以及一次完整的生成式推理流程。这些知识是后续学习TensorRT-LLM优化（如Paged KV Cache、In-Flight Batching、量化）的必备前提。

3.1 注意力机制：从自注意力到因果掩码

3.1.1 核心原理：查询-键-值的数学本质

注意力机制的核心思想是让序列中的每个都能“关注”到其他，并根据相关性加权聚合信息。在自注意力（Self-Attention）中，输入序列 \( X \in \mathbb{R}^{L \times d} \)（L为序列长度，d为隐藏维度）通过三个可学习的权重矩阵 \( W_Q, W_K, W_V \) 分别投影为查询（Query）、键（Key）和值（Value）矩阵：

Q = X @ W_Q   # 形状 [L, d_k]
K = X @ W_K   # 形状 [L, d_k]
V = X @ W_V   # 形状 [L, d_v]

注意力分数的计算为 \( \text{Attention}(Q,K,V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V \)。除以 \( \sqrt{d_k} \) 是为了防止点积结果过大导致Softmax梯度消失。多头注意力（Multi-Head Attention）则将Q、K、V切分为多个头（Head），每个头独立计算注意力，最后拼接并线性变换，使模型能够从不同子空间捕捉信息。

3.1.2 因果掩码：类模型的单向约束

类模型采用因果（Causal）注意力，即每个只能关注它之前的，不能看到未来信息。实现方式是在Softmax之前，将注意力矩阵的上三角部分（包括对角线）加上一个极大的负值（如 \(-\infty\) 或 \(-1e9\)），使得Softmax后这些位置的权重趋近于0。这一设计保证了自回归生成时，第 \( t \) 个的预测只依赖于前 \( t-1 \) 个，符合概率链式法则 \( P(x_1...x_T) = \prod_{t=1}^T P(x_t | x_{

3.1.3 案例拆解：一个简单序列的注意力计算

假设输入序列为 ["我", "爱", "AI"]，嵌入维度 d=4，单头注意力。经过投影后得到Q、K、V矩阵（数值仅为示意）：

计算“爱”位置的注意力分数（未缩放）：与“我”的点积为 0.4×0.1+0.3×0.2+0.1×0.4+0.2×0.3=0.20；与“爱”自身的点积为 0.4×0.3+0.3×0.1+0.1×0.2+0.2×0.4=0.31；与“AI”的点积为 0.4×0.2+0.3×0.3+0.1×0.1+0.2×0.4=0.32。由于因果掩码，“爱”不能看到“AI”，因此“AI”的分数被置为 \(-\infty\)。Softmax后，“我”和“爱”的权重约为0.47和0.53，最终输出为两者V的加权和。这个例子说明注意力机制本质上是基于内容寻址的软检索。

3.1.4 常见误区与复盘

• 误区1：认为多头注意力中每个头关注不同位置范围。实际上，头与头之间没有显式的位置分工，这种分工是训练中自动涌现的。
• 误区2：忽略缩放因子 \( \sqrt{d_k} \)。当维度较大时，点积结果方差变大，Softmax会进入饱和区，梯度极小，训练不稳定。
• 误区3：在边缘设备上使用过大的上下文长度。注意力计算复杂度为 \( O(L^2) \)，L=4096时单层计算量是L=1024的16倍，且KV Cache显存也线性增长。

3.2 类模型结构：解码器堆叠与参数分布

3.2.1 核心原理：Decoder-Only架构

（Generative Pre-trained Transformer）采用Decoder-Only架构，即只保留Transformer的解码器部分，并移除编码器-解码器交叉注意力层。一个标准的层包含：Masked Multi-Head Self-Attention、Feed-Forward Network（FFN）、两个LayerNorm和两个残差连接。FFN通常包含两个线性层，中间使用GELU激活函数，隐藏维度一般为模型维度的4倍（如d=768时，FFN中间维度为3072）。

3.2.2 模型结构参数与计算量估算

以 1.5B（15亿参数）为例，其结构参数如下：

单次前向传播的计算量（FLOPs）约为 \( 2 \times \text{参数量} \times \text{数} \)。对于1.5B模型，处理1个需要约3 GFLOPs（乘加操作计为2 FLOPs）。在边缘设备（如Jetson Orin 64GB，FP16算力约128 TFLOPS）上，理论峰值可支持约40K /s，但实际受内存带宽限制，通常只能达到2-5K /s。

3.2.3 实战场景：边缘设备上的模型选型

在边缘设备部署时，模型结构直接影响推理性能。以Jetson Orin NX 16GB为例，对比不同规模模型：

• 124M：FP16权重约248MB，KV Cache（上下文1024）约16MB，适合实时交互场景，生成速度可达50+ /s。
• LLaMA-7B：FP16权重约14GB，超出内存，需INT4量化（约3.5GB）才能运行，生成速度约10-15 /s。
• Phi-3-mini 3.8B：INT4量化后约2.2GB，配合Paged KV Cache可在16GB设备上流畅运行，适合边缘RAG场景。

3.2.4 常见误区与复盘

• 误区1：认为层数越多模型越强。实际上，层数增加会线性增加推理延迟，且边缘设备的内存带宽有限，深层模型可能因访存瓶颈导致速度下降。
• 误区2：忽略嵌入层（Embedding）的显存占用。词表为50K、维度为1600时，嵌入矩阵大小为 50K×1600×2字节 ≈ 160MB，在INT4量化时这部分通常保留FP16精度以保证输出质量。
• 误区3：混淆训练和推理的显存需求。训练需要保存梯度、优化器状态和激活值，显存需求是推理的3-5倍；边缘设备只做推理，应使用推理专用优化。

3.3 KV Cache原理：用显存换算力的核心机制

3.3.1 核心原理：为什么需要KV Cache

在自回归生成中，每生成一个新，都需要重新计算所有历史的注意力。如果不做缓存，第 \( t \) 步需要重新计算前 \( t-1 \) 个的K和V矩阵，计算量随序列长度平方增长。KV Cache的核心思想是：在每一步生成时，只计算当前的K和V，并将其追加到缓存中；注意力计算时，直接使用缓存的K、V与当前的Q进行运算。这样，单步生成的计算量从 \( O(t^2) \) 降为 \( O(t) \)。

3.3.2 缓存结构与内存占用分析

KV Cache的存储结构为 [2, num_layers, batch_size, num_heads, seq_len, head_dim]。以LLaMA-7B（32层，32头，head_dim=128）为例，FP16精度下，每个每层需要存储 2×32×128×2字节 = 16KB。上下文长度为2048时，单个序列的KV Cache占用 16KB×2048 = 32MB。若并发8个请求，则需256MB显存。TensorRT-LLM引入了Paged KV Cache，将缓存划分为固定大小的块（如每块16个），按需分配，避免预分配整个上下文长度的连续显存。

# KV Cache 更新伪代码（单层单头）
def update_kv_cache(cache, new_k, new_v, step):
    # cache 形状: [2, max_len, head_dim]
    cache[0, step, :] = new_k  # 写入K
    cache[1, step, :] = new_v  # 写入V
    return cache

# 注意力计算时
K = cache[0, :step+1, :]  # 读取历史K
V = cache[1, :step+1, :]  # 读取历史V
scores = Q @ K.T / sqrt(d_k)
scores = mask_causal(scores, step)
weights = softmax(scores)
output = weights @ V

3.3.3 案例拆解：KV Cache对延迟的影响

在Jetson Orin NX上运行 124M（12层，12头，head_dim=64），对比有无KV Cache的生成延迟（上下文长度512，生成100个）：

可以看到，KV Cache将后续的生成延迟降低了约71%。这是因为无缓存时，每步需要重新计算所有历史的K、V，计算量随序列长度线性增长；而有缓存时，每步只需计算当前的K、V，计算量恒定。

3.3.4 常见误区与复盘

• 误区1：认为KV Cache只影响显存不影响计算。实际上，Decode阶段是访存密集型，性能瓶颈在于从HBM读取KV Cache的带宽，而非计算单元利用率。
• 误区2：忽略KV Cache的精度。FP16 KV Cache比INT8精度更高，但显存翻倍。TensorRT-LLM支持INT8 KV Cache，在精度损失可接受的情况下可减少50%显存。
• 误区3：在Prefill阶段也使用KV Cache。Prefill阶段需要并行处理所有输入，此时应直接计算完整注意力，而不是逐写入缓存。

3.4 生成式推理的典型流程：Prefill与Decode

3.4.1 核心原理：两阶段推理

生成式推理分为两个阶段：Prefill（预填充）和Decode（解码）。Prefill阶段处理完整的输入提示（Prompt），并行计算所有输入的注意力，生成初始的KV Cache。Decode阶段逐生成输出，每步只处理一个，并更新KV Cache。Prefill是计算密集型（Compute-Bound），Decode是访存密集型（Memory-Bound）。

3.4.2 操作流程：一次完整的生成过程

1. 输入预处理：将用户输入文本通过Tokenizer转换为 ID序列，并添加特殊（如 <bos>、<eos>）。
2. Prefill阶段：将整个序列一次性输入模型，经过所有解码器层，计算并缓存每层的K、V矩阵。此阶段输出最后一个的隐藏状态。
3. LM Head与采样：将隐藏状态通过LM Head（词表大小的线性层）映射为词汇表上的概率分布，使用采样策略（如Top-K、Top-P、温度）选择下一个。
4. Decode阶段循环：将新生成的作为输入，重复执行：前向传播（仅计算当前的Q、K、V）→ 更新KV Cache → 计算注意力 → 输出概率 → 采样。直到生成 <eos> 或达到最大长度。
5. 后处理：将生成的 ID序列通过Tokenizer解码为文本。

3.4.3 实战场景：边缘设备上的推理参数调优

在Jetson Orin上使用TensorRT-LLM部署时，关键参数配置如下：

# 构建TensorRT引擎时的关键参数
trtllm-build --model_dir ./ \
  --max_batch_size 4 \
  --max_input_len 512 \
  --max_seq_len 1024 \
  --max_num_tokens 2048 \
  --kv_cache_type paged \
  --kv_cache_block_size 16 \
  --dtype float16

# 运行时参数
# --max_num_tokens 控制每次推理的最大数（Prefill+Decode）
# --kv_cache_free_gpu_mem_fraction 0.3  # 限制KV Cache最多使用30%显存

实际测试中，当并发4个请求、上下文长度512时，设置 max_num_tokens=2048 可以充分利用GPU的并行能力，吞吐量比逐请求推理提升约2.3倍。但需要注意，增大 max_num_tokens 会增加Prefill阶段的延迟，因为需要等待更多一起处理。

3.4.4 常见误区与复盘

• 误区1：认为Prefill阶段越快越好。实际上，Prefill阶段的计算量是固定的，但可以通过连续批处理（Continuous Batching）将多个请求的Prefill合并执行，提高GPU利用率。
• 误区2：忽略采样参数对性能的影响。贪婪解码（Greedy）比Top-K采样更快，但输出多样性差。在边缘场景中，如果对实时性要求高，可先用贪婪解码，再在需要时切换采样策略。
• 误区3：不监控KV Cache的命中率。在长对话场景中，如果历史被频繁访问，KV Cache的局部性较好；如果用户频繁切换话题，缓存利用率下降，应考虑使用滑动窗口或摘要压缩。