大模型底层架构与推理机制源码深度剖析

一、引言

大语言模型（LLM）之所以能够完成复杂的对话、推理与代码生成，依赖的不只是"海量参数"，更是一套精心设计的底层架构与高效的推理调度机制。本文从源码视角出发，剖析 Transformer 主干、注意力计算与推理引擎（以 llama.cpp 为例）的关键实现，帮助读者理解模型"运行起来"背后的工程细节。

二、Transformer 主干：从张量到语义

现代大模型的基础是仅解码器（Decoder-only）的 Transformer。核心组件包括嵌入层（Embedding）、多头自注意力（Multi-Head Attention）、前馈网络（FFN）以及层归一化（LayerNorm）与残差连接。

以 PyTorch 风格表达主干前向过程：

python 复制代码

import torch, torch.nn as nn, torch.nn.functional as F

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-5):
        super().__init__()
        self.eps, self.weight = eps, nn.Parameter(torch.ones(dim))

    def forward(self, x):
        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight

class TransformerBlock(nn.Module):
    def __init__(self, dim, n_heads, hidden_dim):
        super().__init__()
        self.attn_norm = RMSNorm(dim)
        self.attn = nn.MultiheadAttention(dim, n_heads, batch_first=True)
        self.ffn_norm = RMSNorm(dim)
        self.ffn = nn.Sequential(
            nn.Linear(dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, dim))

    def forward(self, x):
        x = x + self.attn(self.attn_norm(x), self.attn_norm(x), self.attn_norm(x))[0]
        return x + self.ffn(self.ffn_norm(x))

每层完成后经过残差相加，保证深层网络的梯度稳定流动。

三、注意力机制：因果掩码与 KV 缓存

自注意力计算 Q、K、V 三个投影，并用缩放点积求权重。由于是自回归生成，需用因果掩码屏蔽未来 token。推理阶段为避免重复计算历史 token 的 K、V，引入 KV 缓存：

python 复制代码

class AttentionWithCache(nn.Module):
    def __init__(self, dim, n_heads):
        super().__init__()
        self.n_heads, self.head_dim = n_heads, dim // n_heads
        self.wq, self.wk, self.wv = (nn.Linear(dim, dim) for _ in range(3))

    def forward(self, x, k_cache=None, v_cache=None):
        q = self.wq(x).view(-1, 1, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.wk(x).view(-1, 1, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.wv(x).view(-1, 1, self.n_heads, self.head_dim).transpose(1, 2)
        k_cache = torch.cat([k_cache, k], dim=1) if k_cache is not None else k
        v_cache = torch.cat([v_cache, v], dim=1) if v_cache is not None else v
        scores = (q @ k_cache.transpose(-2, -1)) / (self.head_dim ** 0.5)
        scores = scores.masked_fill(self._mask(scores), float("-inf"))
        out = F.softmax(scores, dim=-1) @ v_cache
        return out.transpose(1, 2).contiguous().view(-1, self.n_heads * self.head_dim), k_cache, v_cache

KV 缓存让每步只计算当前 token 的注意力，时间复杂度由 O(n²) 降为 O(n)，是高效推理的关键。

四、推理引擎：以 llama.cpp 为例

llama.cpp 用 C++ 重写推理，核心是 GGUF 量化格式与调度器。其 llama_decode 流程分为三阶段：tokenize → embedding → decode（逐层前向并更新 KV cache）。关键优化包括：

量化：权重以 4-bit（Q4_0 等）存储，显著降低显存占用。
批处理：Prefill 阶段一次处理整段 prompt，生成阶段逐步追加。
多线程：按层并行，利用 CPU/GPU 混合加速。

调度器伪代码：

cpp 复制代码

void llama_scheduler::process() {
  for (auto& seq : active_sequences) {
    llama_batch_add(seq.tokens);        // 追加新 token
    llama_decode(seq.id);               // 逐层前向 + 更新 KV cache
    seq.token = sample_top_p(seq.logits); // 采样输出
  }
}

五、采样策略：温度与 Top-p

生成并非简单取概率最大的 token，而是引入随机性：

python 复制代码

def sample(probs, temperature=0.8, top_p=0.9):
    probs = probs ** (1.0 / temperature)
    probs /= probs.sum()
    sorted_p, _ = torch.sort(probs, descending=True)
    cum_p = torch.cumsum(sorted_p, dim=-1)
    cutoff = sorted_p[cum_p - sorted_p > top_p].min()
    probs[probs < cutoff] = 0.0
    return torch.multinomial(probs / probs.sum(), 1).item()

温度控制概率分布"尖锐"程度，Top-p 则裁剪长尾，兼顾多样性（temperature）与稳定性（top_p）。

六、总结

从 Transformer 张量运算，到因果掩码、KV 缓存，再到量化与采样，每个环节都直接影响推理的速度与质量。掌握这些底层机制，不仅有助于读懂源码，更能为模型优化、部署调优与二次开发提供坚实的工程基础。

海量精选技术文档和实战案例持续更新，敬请关注【风骏时光少年】