LLM 大模型推理机制全解析：从 Transformer 到推理优化的完整技术栈

flowchart TD
    subgraph "单层 Decoder Block"
        A[输入 Token Embedding] --> B[位置编码]
        B --> C[Masked Self-Attention]
        C --> D[残差连接 + LayerNorm]
        D --> E[Feed-Forward Network]
        E --> F[残差连接 + LayerNorm]
        F --> G[下一层]
    end

词嵌入维度: |V| × d_model
V = 词表大小（约 32k ~ 128k）
d_model = 隐藏层维度（4096 或更高）

# 简化代码
logits = output @ embedding_matrix.T  # [batch, seq_len, d_model] @ [d_model, vocab]
probs = softmax(logits[:, -1, :])     # 取最后一个位置的分布
next_token = sample(probs)

output = LayerNorm(x + Sublayer(x))

def layer_norm(x, gamma, beta, eps=1e-5):
    mean = x.mean(dim=-1, keepdim=True)
    var  = x.var(dim=-1, keepdim=True)
    return gamma * (x - mean) / sqrt(var + eps) + beta

def attention(Q, K, V, mask=None):
    # Q, K, V: [batch, num_heads, seq_len, d_k]
    scores = Q @ K.transpose(-2, -1)  # [batch, h, seq, seq]
    scores = scores / sqrt(d_k)       # 缩放

    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)

    attn_weights = softmax(scores, dim=-1)
    output = attn_weights @ V         # [batch, h, seq, d_k]
    return output

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        self.num_heads = num_heads
        self.d_k = d_model // num_heads

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, x, mask=None):
        batch_size, seq_len, _ = x.shape

        Q = self.W_q(x).reshape(batch, seq, h, d_k).transpose(1,2)
        K = self.W_k(x).reshape(batch, seq, h, d_k).transpose(1,2)
        V = self.W_v(x).reshape(batch, seq, h, d_k).transpose(1,2)

        out = scaled_dot_attention(Q, K, V, mask)
        out = out.transpose(1,2).reshape(batch, seq, -1)
        return self.W_o(out)

flowchart LR
    subgraph "多头机制"
        X[输入向量] --> P[线性投影]
        P --> H1[头1: 捕捉局部语义]
        P --> H2[头2: 捕捉长距离依赖]
        P --> H3[头3: 捕捉语法信息]
        P --> Hn[头n: 捕捉其他模式]
        H1 --> Concat[拼接]
        H2 --> Concat
        H3 --> Concat
        Hn --> Concat
        Concat --> O[输出投影]
    end

def create_causal_mask(seq_len):
    # 上三角矩阵，位置 i 只能看到位置 0..i
    mask = torch.tril(torch.ones(seq_len, seq_len))
    return mask  # [[1,0,0], [1,1,0], [1,1,1]]

# ❌ 无 KV Cache：每次重新计算整段历史
def generate_without_cache(model, prompt, max_len):
    for i in range(max_len):
        # 每次都从 0 到 i 计算所有位置的 K/V
        logits = model(prompt[:, :i])
        next_token = sample(logits[:, -1, :])
        prompt = torch.cat([prompt, next_token], dim=-1)

# ✅ 有 KV Cache：只计算新 token 的 K/V
def generate_with_cache(model, prompt, max_len):
    # 第一次：计算所有 prompt token 的 K/V
    logits, kv_cache = model.forward_with_cache(prompt)

    for i in range(max_len):
        next_token = sample(logits[:, -1, :])
        # 只传入新 token，利用缓存中的历史 K/V
        logits, kv_cache = model.forward_with_cache(
            next_token, kv_cache)

KV Cache 内存 = 2 × num_layers × num_heads × d_k × seq_len × precision
              = 2 × L × h × d_k × n × bytes_per_elem

示例：LLaMA-7B（L=32, h=32, d_k=128, FP16=2bytes）
seq=2048: 2×32×32×128×2048×2 ≈ 1.07 GB
seq=8192: 2×32×32×128×8192×2 ≈ 4.29 GB
seq=32768: ~17.2 GB

flowchart TD
    subgraph "MHA: 标准多头"
        MHA_Q[Q: 32 heads] --> MHA_K[K: 32 heads]
        MHA_Q --> MHA_V[V: 32 heads]
    end
    subgraph "GQA: 分组查询"
        GQA_Q[Q: 32 heads] --> GQA_K[K: 8 heads]
        GQA_Q --> GQA_V[V: 8 heads]
    end
    subgraph "MQA: 单键值"
        MQA_Q[Q: 32 heads] --> MQA_K[K: 1 head]
        MQA_Q --> MQA_V[V: 1 head]
    end

def temperature_scale(logits, temperature):
    # temperature > 1: 分布更均匀（更随机）
    # temperature = 1: 原始分布
    # temperature < 1: 分布更尖锐（更确定）
    # temperature → 0: 退化为argmax
    return logits / temperature

# 示例
logits = torch.tensor([1.0, 2.0, 3.0, 0.5])

print(softmax(logits / 0.5))   # 更尖锐: [0.04, 0.29, 0.65, 0.02]
print(softmax(logits / 1.0))   # 原始:   [0.11, 0.24, 0.52, 0.13]
print(softmax(logits / 2.0))   # 更均匀: [0.19, 0.26, 0.35, 0.20]

def top_k_sampling(logits, k=50):
    # 只保留 top-k 的 token
    values, indices = torch.topk(logits, k)

    # 将其他 token 的概率设为 -inf
    filtered = torch.full_like(logits, float('-inf'))
    filtered[indices] = logits[indices]

    probs = softmax(filtered, dim=-1)
    return torch.multinomial(probs, 1)

def top_p_sampling(logits, p=0.9):
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    sorted_probs = softmax(sorted_logits, dim=-1)

    # 计算累积概率
    cumsum = torch.cumsum(sorted_probs, dim=-1)

    # 去掉累积超过 p 的 token
    sorted_probs[cumsum > p] = 0.0

    # 重归一化
    return torch.multinomial(sorted_probs / sorted_probs.sum(), 1)

flowchart LR
    subgraph "标准 Attention"
        S1[Q,K] --> S2[Scores O(n²)] --> S3[Softmax O(n²)] --> S4[Output O(n²)]
    end
    subgraph "Flash Attention"
        F1[分块 Q/K] --> F2[片上 SRAM 计算] --> F3[增量更新 Output]
        F2 -->|在线 softmax 重计算| F2
    end

# 伪代码：量化过程
def quantize_to_int8(weights):
    # 计算缩放因子
    scale = weights.abs().max() / 127.0
    # 量化
    q_weights = (weights / scale).round().to(torch.int8)
    return q_weights, scale

# 反量化
def dequantize(q_weights, scale):
    return q_weights * scale

sequenceDiagram
    participant D as 草稿模型 (小)
    participant T as 目标模型 (大)

    D->>D: 快速生成 K 个候选 token
    D->>T: 提交候选序列
    T->>T: 并行验证所有候选
    T->>T: 发现第 3 个 token 概率低
    T->>D: 接受前 2 个，拒绝后续
    Note over T: 一次推理 = 接受多个 token