harleyszhang Blog

Pytorch显存管理机制与显存占用分析方法

2025-05-20

VGG
ResNet
Inceptionv3
Resnetv2
ResNeXt
Darknet53
DenseNet
CSPNet
VoVNet
一些结论
参考资料

VGG

VGG网络结构参数表如下图所示。

ResNet

ResNet 模型比 VGG 网络具有更少的滤波器数量和更低的复杂性。比如 Resnet34 的 FLOPs 为 3.6G，仅为 VGG-19 19.6G 的 18%。

注意，论文中算的 FLOPs，把乘加当作 1 次计算。

ResNet 和 VGG 的网络结构连接对比图，如下图所示。

不同层数的 Resnet 网络参数表如下图所示。

看了后续的 ResNeXt、ResNetv2、Densenet、CSPNet、VOVNet 等论文，越发觉得 ResNet 真的算是 Backone 领域划时代的工作了，因为它让深层神经网络可以训练，基本解决了深层神经网络训练过程中的梯度消失问题，并给出了系统性的解决方案（两种残差结构），即系统性的让网络变得更“深”了。而让网络变得更“宽”的工作，至今也没有一个公认的最佳方案（Inception、ResNeXt 等后续没有广泛应用），难道是因为网络变得“宽”不如“深”更重要，亦或是我们还没有找到一个更有效的方案。

Inceptionv3

Inception v3 是一种图像识别模型，经证实可以对 ImageNet 数据集实现 78.1% 以上的准确率。该模型是数年来多位研究人员提出的诸多想法积淀的成果。它以 Szegedy 等人发表的《Rethinking the Inception Architecture for Computer Vision》原创性论文为理论依据。

模型本身由对称和非对称构建块组成，包括卷积、平均池化、最大池化、串联、丢弃、全连接层。批量归一化也在模型中广泛应用，同时用于激活输入。损失是通过 Softmax 计算的。

以下是该模型的简要图示：

常见的一种 Inception Modules 结构如下：

论文地址 https://arxiv.org/pdf/1512.00567.pdf。

Resnetv2

作者总结出恒等映射形式的快捷连接和预激活对于信号在网络中的顺畅传播至关重要的结论。

ResNeXt

ResNeXt 的卷积block 和 Resnet 对比图如下所示。

ResNeXt 和 Resnet 的模型结构参数对比图如下图所示。

Darknet53

Darknet53 模型结构连接图，如下图所示。

DenseNet

作者 Gao Huang 于 2018 年发表的论文 Densely Connected Convolutional Networks。

在密集块（DenseBlock）结构中，每一层都会将前面所有层 concate 后作为输入。DenseBlock（类似于残差块的密集块结构）结构的 3 画法图如下所示：

可以看出 DenseNet 论文更侧重的是 DenseBlock 内各个卷积层之间的密集连接（dense connection）关系，另外两个则是强调每层的输入是前面所有层 feature map 的叠加，反映了 feature map 数量的变化。

CSPNet

CSPDenseNet 的一个阶段是由局部密集块和局部过渡层组成（a partial dense block and a partial transition layer）。

CSP 方法可以减少模型计算量和提高运行速度的同时，还不降低模型的精度，是一种更高效的网络设计方法，同时还能和 Resnet、Densenet、Darknet 等 backbone 结合在一起。

VoVNet

One-Shot Aggregation（只聚集一次）是指 OSA 模块的 concat 操作只进行一次，即只有最后一层($1\times 1$ 卷积)的输入是前面所有层 feature map 的 concat（叠加）。OSA 模块的结构图如图 1(b) 所示。

在 OSA module 中，每一层产生两种连接，一种是通过 conv 和下一层连接，产生 receptive field 更大的 feature map，另一种是和最后的输出层相连，以聚合足够好的特征。通过使用 OSA module，5 层 43 channels 的 DenseNet-40 的 MAC 可以被减少 30%（3.7M -> 2.5M）。

基于 OSA 模块构建的各种 VoVNet 结构参数表如下。

作者认为 DenseNet 用更少的参数与 Flops 而性能却比 ResNet 更好，主要是因为concat 比 add 能保留更多的信息。但是，实际上 DenseNet 却比 ResNet要慢且消耗更多资源。

GPU 的计算效率：

GPU 特性是擅长 parallel computation，tensor 越大，GPU 使用效率越高。
把大的卷积操作拆分成碎片的小操作将不利于 GPU 计算。
设计 layer 数量少的网络是更好的选择。
1x1 卷积可以减少计算量，但不利于 GPU 计算。

在 CenterMask 论文提出了 VoVNetv2，其卷积模块结构图如下：

一些结论

当卷积层的输入输出通道数相等时，内存访问代价（MAC）最小。
影响 CNN 功耗的主要因素在于内存访问代价 MAC，而不是计算量 FLOPs。
GPU 擅长并行计算，Tensor 越大，GPU 使用效率越高，把大的卷积操作拆分成碎片的小操作不利于 GPU 计算。
1x1 卷积可以减少计算量，但不利于 GPU 计算。

参考资料

VGG/ResNet/Inception/ResNeXt/CSPNet 论文
深度学习论文: An Energy and GPU-Computation Efficient Backbone Network for Object Detection及其PyTorch

DeepseekMoE 结构详解和代码实现

2025-02-12

1. 基础 MOE 结构介绍
2. DeepseekMOE 结构介绍
- 2.1 Gate 网络与 DeepseekMOE 计算流程
3. DeepseekMOE 结构代码实现
参考资料

1. 基础 MOE 结构介绍

Mixtral 8x7B (announcement, model card) 是高质量的混合专家模型 (Mixed Expert Models，简称 MoEs) 的 Transformer 模型，或者说是一种稀疏的 mixture-of-experts 模型，采用纯解码器结构，并使用 MOE 结构，替换原始的 FFN 结构。在每一层，对每个 token，存在一个 router network 会挑选两组 “experts”(即参数量更小的 FFN）来分别处理该 token，并通过加法方式融合两组 “experts” 的输出。

基础版的（稀疏）MOE 结构图如下图所示:

MOE 通常由两部分组成：

门控或 Router 网络：模块负责根据输入 token 的特征动态选择激活哪些专家，路由器是由带学习的参数组成的网络。
“experts” 网络（小型 FFN）：每层 MOE 都包含若干个（稀疏）专家网络，其通常是小型的 FFN，在实际推理中只有部分专家(通常 8 个)会被激活参与计算。

2. DeepseekMOE 结构介绍

和基础 MOE 结构的区别是：

更精细地划分专家网络，提升每个专家的专业性，提高知识表达的准确度。
引入部分共享专家，减少不同专家间的知识冗余，提升计算效率；所有 tokens 都会经过的共享专家，每个 token 会用计算的 Router 权重，来选择 topK 个专家，然后和共享的专家的输出一起加权求和。

DeepseekMOE 其实是有两类专家的：

共享专家（Shared Expert）：1 个共享专家，用于捕捉通用、全局的特征信息。
路由专家（Routed Experts）：每个 MoE 层都包含 256 个路由专家，负责精细化处理输入 tokens 的专业特征。

2.1 Gate 网络与 DeepseekMOE 计算流程

当一个 token 的向量传入 MoE 层时，首先会经过一个专门的 Gate 网络，该网络负责计算 token 与各个路由专家之间的匹配得分。具体流程如下：

计算 tokens 和专家的匹配得分
- Gate 网络通过线性变换计算每个 token 与所有路由专家的兼容性得分。得分可以反映 token 和各专家“契合”的程度。
选择 Top-K 专家
- 基于得分，Gate 网络为每个 token 选择 Top-K 个最合适的路由专家。在 DeepSeek‐V3 中，每个 token 通常选择 8 个路由专家（在一些实现中还可能对跨节点路由做限制，如最多路由到 4 个不同节点），从而只激活极少数专家进行计算。
专家处理与加权聚合
- 被选中的专家各自对 token 进行独立处理（专家实际上是小型 FFN 模块，类似于 Transformer 中的 FFN 模块），并产生各自的输出。然后，这些专家的输出会根据 Gate 网络给出的得分权重进行加权聚合，最后再和共享专家的输出进行融合，形成当前 MoE 层的最终输出表示。

DeepseekV2 模型的 MOE 参数如下：

{
  // 部分参数省略
  "hidden_act": "silu",
  "hidden_size": 5120,
  "initializer_range": 0.02,
  "intermediate_size": 12288,
  "model_type": "deepseek_v2",
  "moe_intermediate_size": 1536,
  "moe_layer_freq": 1,
  "n_group": 8,
  "n_routed_experts": 160,
  "n_shared_experts": 2,
  "norm_topk_prob": false,
  "num_experts_per_tok": 6,
  "num_hidden_layers": 60,
  "num_key_value_heads": 128,
  "topk_group": 3,
  "topk_method": "group_limited_greedy",
}

混合专家（MoE）参数说明：

3. DeepseekMOE 结构代码实现

这里只考虑推理模式下的 DeepseekMOE 结构实现，且分步实现。

3.1 DeepseekV2MLP 实现

专家其实就是参数量更少的 FFN/MLP 结构，和 llama 中结构一样，只是参数量和计算量更少了，DeepseekV2MLP 代码如下所示。

class DeepseekV2MLP(nn.Module):
    def __init__(self, config, hidden_size=None, intermediate_size=None):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
        self.intermediate_size = (
            config.intermediate_size if intermediate_size is None else intermediate_size
        )

        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
        self.act_fn = ACT2FN[config.hidden_act] # silu 激活函数

    def forward(self, x):
        mlp_out = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
        return mlp_out

3.2 门控/路由网络实现

门控网络的作用是，根据输入 tokens 动态的选择 Top-K 个专家，并为每个 Token 分配权重。关键流程如下：

门控分数计算：通过线性层 + Softmax 生成专家选择概率分布。
Top-K 专家选择：支持两种模式（贪婪选择 vs 分组限制贪婪选择），贪婪模式直接使用 torch.topk 函数选取分数张量中的前 k 个分数。
权重归一化：对 Top-K 权重进行归一化或缩放。

代码实现如下所示:

import torch
import torch.nn as nn
import torch.nn.functional as F
import math 
from dataclasses import dataclass

class MoEGate(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.top_k = config.num_experts_per_tok
        self.n_routed_experts = config.n_routed_experts
        self.routed_scaling_factor = config.routed_scaling_factor
        self.scoring_func = config.scoring_func
        self.topk_method = config.topk_method
        self.n_group = config.n_group
        self.topk_group = config.topk_group
        self.norm_topk_prob = config.norm_topk_prob
        
        # 静态化推理配置（假设配置固定）
        self.inference_norm = self.norm_topk_prob and (self.top_k > 1)
        self.use_group_limited = (self.topk_method == "group_limited_greedy")

        # 门控权重
        self.gating_dim = config.hidden_size
        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim)))
        self.reset_parameters()

    def reset_parameters(self):
        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))

    @torch.inference_mode()  # 禁用梯度与训练逻辑
    def forward(self, hidden_states):
        bsz, seq_len, h = hidden_states.shape
        hidden_states = hidden_states.reshape(-1, h)
        
        # 门控分数计算（保持原始数据类型）
        logits = F.linear(hidden_states, self.weight)  # [n_tokens, n_experts]
        scores = logits.softmax(dim=-1)  # 自动推断 dtype

        # Top-K 选择（静态分支）
        if self.use_group_limited:
            # 分组限制逻辑优化
            group_scores = scores.view(bsz * seq_len, self.n_group, -1).max(dim=-1).values
            group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
            group_mask = torch.zeros_like(group_scores).scatter_(1, group_idx, 1)
            score_mask = group_mask.unsqueeze(-1).expand(-1, -1, self.n_routed_experts // self.n_group).reshape(bsz * seq_len, -1)
            scores = scores.masked_fill(~score_mask.bool(), 0.0)
        
        topk_weight, topk_idx = torch.topk(scores, k=self.top_k, dim=-1, sorted=False)

        # 权重归一化（静态分支）
        if self.inference_norm:
            topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20)
        else:
            topk_weight = topk_weight * self.routed_scaling_factor

        return topk_idx, topk_weight, None  # aux_loss 始终为 None

@dataclass
class DeepseekV2Config:
    # 1, Position Config
    max_position_embeddings: int = 163840
    vocab_size: int = 102400

    # 2, MLA Config
    # down_linear config
    q_lora_rank: int = 1536
    kv_lora_rank: int = 512

    # head_dim、heads and hidden_size config
    v_head_dim: int = 128
    qk_nope_head_dim: int = 128
    qk_rope_head_dim: int = 64
    hidden_size: int = 5120
    num_attention_heads: int = 128
    num_key_value_heads: int = 128
    
    attention_bias: bool = False

    attention_dropout: float = 0.1
    # rope config
    rope_theta: float = 10000

    # 3, MOE Config
    n_group: int = 8
    n_routed_experts: int = 160
    num_experts_per_tok: int = 6
    topk_group: int = 3
    routed_scaling_factor: float = 1.0
    scoring_func: str="softmax"
    topk_method: str="greedy"
    norm_topk_prob: bool = True

# 初始化配置
config = DeepseekV2Config()

# 模拟输入，CPU 电脑可直接跑，去除了 cuda 设备限制代码
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
hidden_states = torch.randn(32, 64, 5120, device=device)

# 创建模块
moe_gate = MoEGate(config)  # 半精度推理

# gate 网络推理
topk_idx, topk_weight, _ = moe_gate(hidden_states)

print("topk_idx shape ", topk_idx.shape) # 32 * 64 = 2048 个 tokens
print("topk_weight shape", topk_weight.shape)

"""
# 输出如下，表示每个 token 会激活 6 个专家参与计算
topk_idx shape  torch.Size([2048, 6]) 
topk_weight shape torch.Size([2048, 6])
"""

3.3 DeepseekMOE 实现

门控计算
- 调用门控网络（self.gate），对输入 hidden_states 计算得到 top‑k 专家索引（topk_idx）、对应权重（topk_weight）以及辅助损失（aux_loss，推理时不参与梯度计算）。
数据重排
- 将输入 hidden_states 展平为二维张量（形状 [B * T, d]），并将 topk_idx 也展平。
- 在推理模式下，通常不需要像训练时那样对每个 token 进行 repeat_interleave，因为每个 token 只会由对应专家处理一次。
专家计算
- 根据展平后的 topk_idx，依次对每个专家负责的 token 子集进行计算。
- 由于这里可能存在多个 token 被分配给不同专家，实际实现中需要将每个专家的输出按顺序记录下来。
输出重构与加权融合
- 将所有专家计算的输出进行合并。通过将输出重新整理（排序）回原始 token 顺序，并按照 topk_weight 对各个专家输出进行加权求和，从而获得最终输出。
- 整个过程保证最终输出形状与原始输入保持一致，即 [B, T, d]。

代码实现如下所示：

# 为了单元测试，模拟不使用分布式（ep_size默认为1）
class DeepseekV2MoE(nn.Module):
    """
    A mixed expert module containing shared experts.
    """
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.num_experts_per_tok = config.num_experts_per_tok

        self.experts = nn.ModuleList(
            [
                DeepseekV2MLP(
                    config, intermediate_size=config.moe_intermediate_size
                )
                for i in range(config.n_routed_experts)
            ]
        )
        self.gate = MoEGate(config)
        if config.n_shared_experts is not None:
            intermediate_size = config.moe_intermediate_size * config.n_shared_experts
            self.shared_experts = DummyMLP(config=config, intermediate_size=intermediate_size)

    # 此处为简化实现，仅做推理示例，不涉及分布式通信
    @torch.no_grad()
    def moe_infer(self, x, topk_ids, topk_weight):
        # x: [batch * seq_len, hidden_size]
        # 对每个 token 依然采用与训练类似的方式进行专家计算
        outputs = []
        flat_topk_ids = topk_ids.view(-1)
        for i, expert in enumerate(self.experts):
            mask = (flat_topk_ids == i)
            if mask.sum() == 0:
                continue
            outputs.append(expert(x[mask]))
        # 简单拼接，不做复杂排序和 all-to-all 操作
        outs = torch.cat(outputs, dim=0)
        new_x = torch.empty_like(outs)
        # 这里直接返回加权求和的结果（实际实现更复杂）
        final_out = (outs.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim=1)
        return final_out

参考资料

MLA 结构代码实现及优化

2025-02-10

1. MLA 实现拆解
2. 标准 MLA 模块的代码实现
3 MLA 模块的代码优化-Projection Absorption
- 3.1 CC (CacheCompressed）
- 3.2 A_CC（AbsorbCacheCompressed）
参考资料

1. MLA 实现拆解

DeepDeekv2 的模型配置如下所示:

1.1 Q 向量计算

大部分参考 DeepSeek-V2高性能推理优化笔记：MLA优化，部分细节做了修改和优化， MLA 结构图以及这章节的公式更多的是给出 MLA 过程和细节，实际的代码实现没有一一对应。

1，在 DeepSeek-V2 中，Q 向量也采用了低秩压缩的方式。首先，将输入向量投影到一个 1536（对应模型配置文件中的 q_lora_rank 参数）维的低维空间，得到 Latent $c_t^Q$。

\[c_t^Q = W^{DQ} h_t \in \mathbb{R}^{B \times L \times 1536}\]

2，然后，再将其投影到 $\mathbb{R}^{H \times 128}$ 的多头向量空间上（其中 $H=128$ 是 heads 数，对应配置文件中的 qk_nope_head_dim 参数），得到了 Q 向量的第一部分: $q_t^C$。

\[q_t^C = W^{UQ} c_t^Q \in \mathbb{R}^{B \times L \times H \times 128}\]

3，再将其投影到 $\mathbb{R}^{H \times 64}$（对应模型配置文件中的 qk_rope_head_dim 参数）上，并使用 RoPE 嵌入位置信息，得到 Q 向量的第二部分: $q_t^R$。

\[q_t^R = \mathrm{RoPE}(W^{QR} h_t) \in \mathbb{R}^{B \times L \times H \times 64}\]

4，最后，将这两部分进行 concat 拼接得到最终的 $Q$ 向量：$q_t$。

\[q_t = [q_t^C, q_t^R] \in \mathbb{R}^{B \times L \times H \times 192}\]

其中：

$B$: batch_size 批量大小；
$L$: seq_len 序列长度；
$H$: heads 注意力头数；
$\mathbb{R}$ 的最后一维是 head_dim。

1.2 KV 向量计算

1，计算 $KV$ 向量时，首先，将输入向量投影到一个 $512$（对应模型配置文件中的 kv_lora_rank 参数）维的低维空间，得到 Latent $c_t^{KV}$。

\[c_t^{KV} = W^{DKV} h_t \in \mathbb{R}^{B \times L \times 512}\]

2，然后，和 $Q$ 向量的计算过程类似，再将其投影到 $\mathbb{R}^{H \times 128}$ 的多头向量空间上（其中 $H=128$ 是 heads 数，$128$ 对应模型配置文件中的 qk_rope_head_dim 参数，得到了 $K$ 向量的第一部分 $k_t^C$。

\[k_t^C = W^{UK}c_t^{K} \in \mathbb{R}^{B\times L\times H\times 128}\]

3，将输入向量投影到 $64$（对应模型配置文件中的 qk_rope_head_dim 参数）维向量空间，并应用 RoPE 嵌入位置信息得到 $K$ 向量的第二部分： $k_t^R$。

\[k_t^R = \mathrm{RoPE}(W^{KR} h_t) \in \mathbb{R}^{B \times L \times 1 \times 64}\]

4，最后，和 $Q$ 不同的是，完整的 $K$ 是将 $k_t^R$ 广播到每个 head 后与 $k_t^C$ concate 拼接得到：

\[k_t = \begin{bmatrix} k_{t,1}^C & k_t^R \\ k_{t,2}^C & k_t^R \\ \vdots & \vdots \\ \end{bmatrix} \in \mathbb{R}^{B \times L \times H \times 192}\]

上述广播后拼接的方式意味着，每个 head 的 RoPE 部分是完全相同的。

$V$ 向量因为不需要执行 ROPE 操作，所以它的的计算较为简单，直接将 $c_t^{KV}$ 解压缩（升维）到 $\mathbb{R}^{H \times 128}$ 即可：

\[\mathbf{v}_t = W^{UV} c_t^{KV} \in \mathbb{R}^{B \times L \times H \times 128}\]

注意: $k_t^R$ 和 $c_t^{KV}$ 是需要缓冲的向量。前面计算得到 $q_t$、$k_t$ 和 $\mathbf{v}_t$ 用来执行 self-attention 计算。

1.3 Self-Attention 计算

Self-Attention 的计算过程和传统的 MHA 一模一样。同样也是首先计算 attention score：

$p = \mathrm{softmax}\left(\frac{q_t^\top k_t + \mathrm{Mask}}{\sqrt{192}}\right) = \mathrm{softmax}\left(\frac{{q_t^C}^\top k_t^C + {q_t^R}^\top k_t^R + \mathrm{Mask}}{\sqrt{128 + 64}} \right) \mathrm{softmax}\left(\frac{{q_t^C}^\top k_t^C + {q_t^R}^\top k_t^R + \mathrm{Mask}} {\sqrt{128 + 64}} \right) \in \mathbb{R}^{B \times L \times H \times L}$

计算对 $V$的加权和，并将所有 heads 压平（即 heads * head_dim），得到 Attention 输出：

\[o = p \cdot \mathbf{v}_t \in \mathbb{R}^{B \times L \times H \times 128} \cong \mathbb{R}^{B \times L \times 16384}\]

其中，$16384 = 128 \times 128 = \text{num\;attention\;heads * v\;head\;dim}$。最后，经过另一个注意力输出矩阵的投影（5120 是 hidden_size），就能得到 MLA 的最终输出：

\[u = W^O o \in \mathbb{R}^{B \times L \times 5120}\]

2. 标准 MLA 模块的代码实现

transformers 库中的 modeling_deepseek.py 是没有经过推理加速优化的原始实现，我参考其实现给出了一个更为精简和更易看懂的版本，完整代码在这里。

# 从 LlamaAttention 修改而来，适配 DeepseekV2 模型的注意力模块，简单版本不带 kv cache
class DeepseekV2MLA(nn.Module):
    def __init__(self, config: DeepseekV2Config):
        super().__init__()
        # MHA 初始化相关
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.v_head_dim = config.v_head_dim

        self.o_proj = nn.Linear(
            self.v_head_dim * self.num_heads, 
            self.hidden_size,
            bias=config.attention_bias,
        )

        self.attention_dropout = config.attention_dropout
        self.training = False
        self.qk_nope_head_dim = config.qk_nope_head_dim
        self.qk_rope_head_dim = config.qk_rope_head_dim

        # MLA 相关 part1: 压缩
        self.q_lora_rank = config.q_lora_rank
        self.kv_lora_rank = config.kv_lora_rank

        self.q_down_proj = nn.Linear(self.hidden_size, self.q_lora_rank)
        self.q_down_rmsnorm = DeepseekV2RMSNorm(self.q_lora_rank)
        
        self.kv_down_proj = nn.Linear(
            self.hidden_size, 
            self.kv_lora_rank + config.qk_rope_head_dim
        )
        self.kv_down_rmsnorm = DeepseekV2RMSNorm(self.kv_lora_rank)
        
        # MLA 相关 part2: 解压缩
        self.q_head_dim = self.qk_nope_head_dim  + self.qk_rope_head_dim
        self.q_up_proj = nn.Linear(
            self.q_lora_rank, 
            self.num_heads * self.q_head_dim,
            bias=False,
        )
        # qk_nope_head_dim = q_head_dim - qk_rope_head_dim
        self.kv_up_proj = nn.Linear(
            self.kv_lora_rank, 
            self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
            bias=False,
        )
        
        # MLA 相关 part3: 切片 q k 张量，以及 rope 旋转位置编码
        self.rotary_emb = DeepseekV2RotaryEmbedding(
            config.qk_rope_head_dim,
            config.max_position_embeddings,
            config.rope_theta,
        ) 

    def forward(self, hidden_states, position_ids, casual_mask=None):
        batch_size, q_len, hidden_size = hidden_states.shape

        # 1，q 压缩和解压缩，以及 split to q_nope, q_rope
        q = self.q_up_proj(
            self.q_down_rmsnorm(self.q_down_proj(hidden_states))
        )

        q = q.view(batch_size, q_len, self.num_heads, self.q_head_dim).transpose(1,2)
        q_nope, q_rope = torch.split(
            q,
            [self.qk_nope_head_dim, self.qk_rope_head_dim],
            dim = -1,
        )

        # 2, kv 压缩和解压缩
        kv_down = self.kv_down_proj(hidden_states)
        
        # compressed_kv 压缩后的 kv 张量
        compressed_kv, k_rope = torch.split(
            kv_down,
            [self.kv_lora_rank, self.qk_rope_head_dim],
            dim = -1,
        )
        # num_heads = 1 后续广播其它 heads 上
        k_rope = k_rope.view(batch_size, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)

        # 对 compressed_kv 解压缩
        kv = (
            self.kv_up_proj(self.kv_down_rmsnorm(compressed_kv))
            .view(batch_size, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)
            .transpose(1, 2)
        )

        k_nope, value_states = torch.split(
            kv,
            [self.qk_nope_head_dim, self.v_head_dim],
            dim = -1,
        )

        # 3, 计算 cos 和 sin，并应用 rope 旋转位置编码
        kv_seq_len = value_states.shape[-2] # shape (b, nums_head, seq_len, v_head_dim)
        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
        
        q_rope, k_rope = apply_rotary_pos_emb(q_rope, k_rope, cos, sin, position_ids)

        # 4, 执行 self-attention 计算
        query_states = torch.concat([q_nope, q_rope], dim=-1)
        key_states = torch.concat(
            [k_nope, k_rope.expand(-1, self.num_heads, -1, -1)], 
            dim=-1
        )
        # qk^t
        scores = (
            torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.q_head_dim)
        )

        if casual_mask is not None:
            scores = scores.masked_fill(casual_mask == 0, float('-inf'))
        
        attn_weights = F.softmax(scores, dim=-1).to(query_states.dtype)
        attn_weights = F.dropout(
            attn_weights, p=self.attention_dropout, training=self.training
        ) # attn_weights shape: [bs, num_heads, seq_len, seq_len]
        
        attn_output = torch.matmul(attn_weights, value_states) # shape: [bs, num_heads, seq_len, head_dim]
        attn_output = attn_output.transpose(1, 2).contiguous().reshape(batch_size, q_len, self.num_heads * self.v_head_dim)

        # 5, MLA 输出映射
        output = self.o_proj(attn_output)

        return output, attn_weights

3 MLA 模块的代码优化-Projection Absorption

3.1 CC (CacheCompressed）

在 transformers 的最新开源版本中， MLA 算子改为缓存压缩后的 KV Cache，并将 RoPE 后的 k_pe 一并缓存入 KV Cache 中，与缓存完整的 KV Cache 相比，这将大大减少每个 token 的每层Cache 大小。

3.2 A_CC（AbsorbCacheCompressed）

上述 CacheCompressed 的实现代码其实并不能实质减少 KV Cache 过大的问题，因为在计算 MLA 的时候，仍然需要存储解压后的完整的 KV Cache（中间激活），这很可能引起 OOM 崩溃。

DeepSeek-V2 论文中提出，可以将 KV 的解压缩矩阵吸收到Q-projection 和 Out-projection 中，从而可以在不解压缩 KV Cache的情况下直接计算最终的 Attention 结果。

1，对于 K 的吸收（吸收进 self-attention 算子中，相当于算子合并），在 Attention Score 的计算公式中，K 向量的非 RoPE 部分的可以做如下展开：

\[{q_t^C}^\top k_t^C = (W^{UQ} c_t^Q)^{\top} W^{UK} c_t^{KV} = {c_t^Q}^{\top}{W^{UQ}}^{\top} W^{UK} c_t^{KV} = ({c_t^Q}^{\top}{W^{UQ}}^{\top} W^{UK}) c_t^{KV}\]

即通过矩阵乘法结合律，可以改为计算 $({c_t^Q}^{\top}{W^{UQ}}^{\top} W^{UK})$，避免了解压缩出完整的 $K$ 矩阵。另外，在原始版本的解压缩的过程中，由于每个 token 的 key 都需要与 $W^{UK}$ 相乘才能得到，因此计算量较大；矩阵吸收后，$W^{UK}$ 只需要对 $q_t^C$ 这一个向量相乘，也大大减少了浮点计算量。

总结：A_CC 相比于 CC，把原来属于单 kv 的计算量转移到 q 上了，而 q 的 seq_len=1，可减少计算量。

其中，$c_t^{KV}$ 是我们实际保存的 KV cache。

2，$V$ 的吸收，其实现更为复杂。为了更方便表述，采用 Einstein 求和约定描述该过程：

v_t = einsum('hdc,blc->blhd', W_UV, c_t_KV) # (1) 生成值向量 v_t
o   = einsum('bqhl,blhd->bqhd', a, v_t)     # (2) 加权求和得到 o
u   = einsum('hdD,bhqd->bhD', W_o, o)       # (3) 投影到最终输出 u

# 将上述三式合并，得到总的计算过程
u   = einsum('hdc,blc,bqhl,hdD->bhD', W_UV, c_t_KV, a, W_o)

# 利用结合律改变计算顺序
o_  = einsum('bhql,blc->bhqc', a, c_t_KV) # (4) 避免显式生成 v_t，减少存储 (b, l, h, d) 的开销。
o   = einsum('bhqc,hdc->bhqd', o_, W_UV)  # (5) 延迟投影操作，减少计算量。
u   = einsum('hdD,bhqd->bhD', W_o, o)     # (6)

其中生成值向量 v_t：输入：

W_UV：权重矩阵，形状为 (h, d, c)，其中 h 是注意力头数，d 是值向量维度，c 是输入特征维度。
c_t_KV：键值上下文向量，形状为 (b, l, c)，其中 b 是批次大小，l 是序列长度。

操作：

将每个位置的 c 维特征通过 W_UV 投影到 d 维，生成多头值向量 v_t，形状为 (b, l, h, d)。

改变计算顺序的优化: 通过结合律调整计算顺序，减少中间张量的内存占用：先计算加权上下文 o_:

操作：

将注意力权重 attn_weights 直接作用于原始上下文 c_t_KV，生成中间结果 o_，形状为 (b, h, q, c)。

意义：

避免显式生成 v_t，减少存储 (b, l, h, d) 的开销。

上述优化方法的实现和对比测试代码如下所示:

import torch
import time

# 配置参数
b, q, l, h, d, c, D = 32, 64, 128, 64, 64, 128, 256  # 将 h 调整为 64
n_warmup = 10   # 预热次数
n_trials = 100  # 正式测试次数

# 初始化张量（GPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

W_UV = torch.randn(h, d, c, device=device)
c_t_KV = torch.randn(b, l, c, device=device)
attn_weights = torch.randn(b, q, h, l, device=device)
W_o = torch.randn(h, d, D, device=device)  # h 维度为 64

# 预热 GPU
for _ in range(n_warmup):
    _ = torch.einsum('hdc,blc->blhd', W_UV, c_t_KV)
    _ = torch.einsum('bqhl,blhd->bqhd', attn_weights, _)
    _ = torch.einsum('hdD,bhqd->bhD', W_o, _)

# 原始分步实现
def original_method():
    v_t = torch.einsum('hdc,blc->blhd', W_UV, c_t_KV)
    o = torch.einsum('bqhl,blhd->bqhd', attn_weights, v_t)
    u = torch.einsum('hdD,bhqd->bhD', W_o, o.permute(0, 2, 1, 3))

# 优化后实现
def optimized_method():
    o_ = torch.einsum('bhql,blc->bhqc', attn_weights.permute(0, 2, 1, 3), c_t_KV)
    o = torch.einsum('bhqc,hdc->bhqd', o_, W_UV)
    u = torch.einsum('hdD,bhqd->bhD', W_o, o)

# 测量时间
def benchmark(func):
    times = []
    for _ in range(n_trials):
        start = time.time()
        func()
        end = time.time()
        times.append(end - start)
    return sum(times) / n_trials

# 执行测试
time_original = benchmark(original_method) * 1000  # 转换为毫秒
time_optimized = benchmark(optimized_method) * 1000

# 打印结果
print(f"原始方法平均时间: {time_original:.3f} ms")
print(f"优化方法平均时间: {time_optimized:.3f} ms")
print(f"速度提升: {time_original / time_optimized - 1:.1%}")

# 验证等价性
def validate_equivalence():
    v_t_orig = torch.einsum('hdc,blc->blhd', W_UV, c_t_KV)
    o_orig = torch.einsum('bqhl,blhd->bqhd', attn_weights, v_t_orig)
    u_orig = torch.einsum('hdD,bhqd->bhD', W_o, o_orig.permute(0, 2, 1, 3))
    
    v_t_opt = torch.einsum('hdc,blc->blhd', W_UV, c_t_KV)
    o_opt = torch.einsum('bqhl,blhd->bqhd', attn_weights, v_t_opt)
    u_opt = torch.einsum('hdD,bhqd->bhD', W_o, o_opt.permute(0, 2, 1, 3))
    
    # 检查是否等价
    assert torch.allclose(u_orig, u_opt, atol=1e-4), "两种方法结果不一致！"
    print("两种方法结果一致，验证通过。")

# 调用验证函数
validate_equivalence()

"""
原始方法平均时间: 28.649 ms
优化方法平均时间: 20.378 ms
速度提升: 40.6%
两种方法结果一致，验证通过。
"""

参考资料

DeepSeekV2 论文解读

2025-02-07

1. 介绍
2. 架构
- 2.1 多头潜变量注意力（MLA）：提升推理效率
- 2.2 DeepSeekMoE：以经济成本训练强大的模型
参考资料

1. 介绍

DeepSeek-V2 是一种高效的开源混合专家（MoE）语言模型，基于创新的 Transformer 架构，实现了经济的训练和高效的推理。DeepSeek-V2 具有 2360 亿个参数(236B)，每个 token 激活 21 亿个参数，支持 128K tokens 的上下文长度。

和 DeepSeekV1 模型结构沿用 llama 结构不同，DeepSeekV2 提出了多头潜在注意力（MLA）和 DeepSeekMoE，旨在优化 Transformer 框架中的注意力模块和前馈网络（FFNs）。

MLA: Multi-head Latent Attention 结构，通过低秩键值联合压缩，减少了推理时的 KV 缓存，从而提高了推理效率。
DeepSeekMoE: FFN（标准 MOE）的优化版。
- 细粒度专家划分(Routed Expert)：相比标准 MOE，DeepSeekMoE 在保持参数量不变的前提下，通过减小每个 Expert 的 FFN 维度，来增加 Expert 数量，进行更细粒度专家划分。
- 共享专家隔离(Shared Expert): 用于表示 Routed Expert 中的共用知识信息，减少 Routed Expert 的知识冗余问题。

DeepSeek-V2 架构图如下所示：

2. 架构

本节介绍 MLA 和 DeepSeekMoE 的详细设计。

2.1 多头潜变量注意力（MLA）：提升推理效率

传统的 Transformer 模型通常采用多头注意力（MHA），但在生成（generation）过程中，其庞大的 Key-Value（KV）缓存会成为限制推理效率的瓶颈。为减少 KV 缓存占用，研究者提出了多查询注意力（MQA）（Shazeer, 2019）和分组查询注意力（GQA）（Ainslie et al., 2023）。这两种方法虽然减少了 KV 缓存需求，但在性能上仍无法与 MHA 相媲美（关于 MHA、GQA 和 MQA 的消融实验见附录 D.1）。

DeepSeek-V2 引入了一种全新的注意力机制多头潜变量注意力（MLA）。MLA 结合了低秩键值联合压缩（low-rank key-value joint compression,），在推理时大幅降低 KV 缓存需求，同时在性能上超越 MHA。

MLA 本质上是通过低秩转换的思路减少 head 的维度，即换为一个压缩的 QKV，存储的KV 的维度显著减小，而不是 GQA 方法减少 kv heads 的数量。

2.1.1 Standard Multi-Head Attention

先回顾下标准的多头注意力（MHA）机制。设 $d$ 为嵌入维度，$n_h$ 为注意力头数，$d_h$ 为单个注意力头的维度，$h_t \in R_d$ 表示第 $t$ 个 token 进入注意力层的输入向量。

在 MHA 机制中，我们通过三个投影矩阵 $W_Q、W_K、W_V \in \mathbb{R}^{n_h d_h\times d}$ 分别计算得到查询向量、键向量和值向量（$q_t、k_t、v_t \in \mathbb{R}^{n_h d_h}$），QKV 向量的线性变换公式如下所示：

QKV 的线性变换的权重矩阵的第二个维度大小一定为嵌入维度 $d$。

\[\mathbf{q}_t = W^Q \mathbf{h}_t, \tag{1}\] \[\mathbf{k}_t = W^K \mathbf{h}_t, \tag{2}\] \[\mathbf{v}_t = W^V \mathbf{h}_t, \tag{3}\]

然后，$q_t$, $k_t$, $v_t$ 将被切分为 $n_h$ 个头（heads），用于多头注意力计算：

\[[\mathbf{q}_{t,1}; \mathbf{q}_{t,2}; \dots; \mathbf{q}_{t,n_h}] = \mathbf{q}_t \tag{4}\] \[[\mathbf{k}_{t,1}; \mathbf{k}_{t,2}; \dots; \mathbf{k}_{t,n_h}] = \mathbf{k}_t \tag{5}\] \[[\mathbf{v}_{t,1}; \mathbf{v}_{t,2}; \dots; \mathbf{v}_{t,n_h}] = \mathbf{v}_t \tag{6}\] \[\mathbf{o}_{t,i} = \sum_{j=1}^{t} \text{Softmax}_j \left( \frac{\mathbf{q}_{t,i}^T \mathbf{k}_{j,i}}{\sqrt{d_h}} \right) \mathbf{v}_{j,i} \tag{7}\] \[\mathbf{u}_t = W^O [\mathbf{o}_{t,1}; \mathbf{o}_{t,2}; \dots; \mathbf{o}_{t,n_h}] \tag{8}\]

其中，$q_{t,i}$, $k_{t,i}$, $v_{t,i} \in \mathbb{R}^{d_h}$ 分别表示第 $i$ 个注意力头的查询（query）、键（key）和值（value）；$W_O \in \mathbb{R}^{d \times d_h n_h}$ 表示输出投影矩阵。在推理过程中，key 和 value 需要被缓存，以加速计算，避免重复计算。

标准 MHA 每个 token 的 kv 缓冲大小 = $2n_hd_h l$，单位为字节 byte；如果使用了 GQA 优化技术，每个 token 的 kv 缓冲大小变为 $2n_{kv}d_h l = 2n_hd_h l/\text{groups}$ 个元素。下标 $t$ 表示第几个 token，下标 $[1, n_h]$ 表示注意力头数，$l$ 表示 decoder layers 数目。

在模型部署时，这种庞大的 KV 缓存成为了一个主要的瓶颈，限制了最大批量大小（batch size）和序列长度（sequence length）。

2.1.2 Low-Rank Key-Value Joint Compression

MLA 的核心是对键（keys）和值（values）进行低秩联合压缩（low-rank joint compression），以减少 KV 缓存（KV cache）的占用：

\[\mathbf{c}_t^{KV} = W^{DKV} \mathbf{h}_t, \tag{9}\] \[\mathbf{k}_t^{C} = W^{UK} \mathbf{c}_t^{KV}, \tag{10}\] \[\mathbf{v}_t^{C} = W^{UV} \mathbf{c}_t^{KV}, \tag{11}\]

KV 向量的生成是先投影到一个低维（5120 -> 512）的 compressed_kv 向量（$\mathbf{c}_t^{KV}$）再升维展开得到 $\mathbf{k}_t^{C}$ 和 $\mathbf{v}_t^{C}$。上述公式的各个变量定义：

$\mathbf{c}_t^{KV}$ 是 keys 和 values 的压缩后的潜在向量（latent vector）；
$d_c (\ll d_h n_h)$ 代表 KV 压缩维度（KV compression dimension）
$W^{DKV} \in \mathbb{R}^{d_c \times d}$ 是降维投影矩阵（down-projection matrix）；
$W^{UK}, W^{UV} \in \mathbb{R}^{d_h n_h \times d_c}$ 分别是 keys 和 values 的升维投影矩阵（up-projection matrices）。

另外，虽然不能减少 KV Cache 的占用，但是为了减少训练时的激活内存（activation memory），同样也对查询（queries）也进行了低秩压缩（low-rank compression）。同样也是先投影到一个低维（5120 -> 1536）的 compressed_kv 向量（$\mathbf{c}_t^{Q}$）再升维展开得到 $\mathbf{q}_t^{C}$:

\[\mathbf{c}_t^{Q} = W^{DQ} \mathbf{h}_t, \tag{12}\] \[\mathbf{q}_t^{C} = W^{UQ} \mathbf{c}_t^{Q}, \tag{13}\]

类比前面的公式可知:

$\mathbf{c}_t^{Q} \in \mathbb{R}^{d’_c}$ 是查询的压缩潜在向量（compressed latent vector for queries）；
$d’_c (\ll d_h n_h)$ 表示查询的压缩维度（query compression dimension）；
$W^{DQ} \in \mathbb{R}^{d’_c \times d}$ 是查询的降维投影矩阵；
$W^{UQ} \in \mathbb{R}^{d_h n_h \times d’_c}$ 是查询的升维投影矩阵（up-projection matrix）。

2.1.3 Decoupled Rotary Position Embedding

和 DeepSeek 67B（DeepSeek-AI, 2024）类似，作者也计划在 DeepSeek-V2 中使用旋转位置编码（RoPE, Rotary Position Embedding）（Su et al., 2024）。但是，RoPE 与低秩 KV 压缩（low-rank KV compression）并不兼容。

具体来说，RoPE 使键（Key）和查询（Query）都具备位置敏感性（position sensitivity）。如果我们在压缩后的键 $\mathbf{k}_t^{C}$ 上应用 ROPE，那么实际上我们得到的键表示会是这样的形式：

\[k_t^R = \text{ROPE}(W^{UK} \mathbf{c}_t^{KV})\]

很明显式（10）中的 $W^{UK}$ 和 RoPE 旋转矩阵在计算过程中“耦合”在一起—这意味着 $W^{UK}$ 输出的结果会始终被那个依赖于具体位置的旋转矩阵所“修正”或“调制”。

这样会导致在执行 atten weight（$QK^T$）的计算优化中，无法像原本设想的那样，把 $W^{UK}$ 吸收到 $W^Q$ 中，因为当前生成 token 相关的 RoPE 矩阵位于 $W^Q$ 和 $W^{UK}$ 之间，而矩阵乘法不满足交换律（commutative law）。这直接导致在推理过程中，我们必须重新计算所有 prefix token 的键（keys），这将显著降低推理效率。

为了解决这个问题，作者提出了一种解耦 RoPE（decoupled RoPE）的策略，通过额外引入多头查询（multi-head queries）$q_{t,i}^R \in \mathbb{R}^{d^R_h}$和采用一个共享键（shared key) $k_t^R \in \mathbb{R}^{d^R_h}$ 来承载 RoPE 信息。其中 $d^R_h$ 代表解耦查询和键的每头维度（per-head dimension of the decoupled queries and key）。

在使用解耦 RoPE 策略后，MLA 的计算过程变成如下所示：

\[\left[ \mathbf{q}_{t,1}^{R}; \mathbf{q}_{t,2}^{R}; \dots; \mathbf{q}_{t,n_h}^{R} \right] = \mathbf{q}_t^{R} = \text{RoPE}(W^{QR} \mathbf{c}_t^{Q}), \tag{14}\] \[\mathbf{k}_t^{R} = \text{RoPE}(W^{KR} \mathbf{h}_t), \tag{15}\] \[\mathbf{q}_{t,i} = \left[ \mathbf{q}_{t,i}^{C}; \mathbf{q}_{t,i}^{R} \right], \tag{16}\] \[\mathbf{k}_{t,i} = \left[ \mathbf{k}_{t,i}^{C}; \mathbf{k}_{t,i}^{R} \right], \tag{17}\] \[\mathbf{o}_{t,i} = \sum_{j=1}^{t} \text{Softmax}_j \left( \frac{\mathbf{q}_{t,i}^{T} \mathbf{k}_{j,i}}{\sqrt{d_h + d_h^{R}}} \right) \mathbf{v}_{j,i}^{C}, \tag{18}\] \[\mathbf{u}_t = W^{O} \left[ \mathbf{o}_{t,1}; \mathbf{o}_{t,2}; \dots; \mathbf{o}_{t,n_h} \right], \tag{19}\]

其中:

$W^{QR} \in \mathbb{R}^{d^R_h n_h \times d’_c}$ 表示生成解耦查询（decoupled queries）矩阵
$W^{KR} \in \mathbb{R}^{d^R_h \times d}$ 表示解耦键（decoupled key）的矩阵。
$\text{RoPE}(\cdot)$ 表示应用 RoPE 矩阵的操作;
$\cdot ; \cdot$ 表示拼接（concatenation）操作。

在推理过程中，解耦后的键（decoupled key）也需要缓存。因此，DeepSeek-V2 的 KV 缓存总大小为 $(d_c + d^R_h)l$ 个元素。

很明显和前面公式相比，多了 $\mathbf{q}_t^{R}$ 和 $\mathbf{k}_t^{R}$ 两个变量的计算过程，它们用于单独承载 ROPE 信息，并和前面计算得到的 $\mathbf{q}_t^{C}$ 和 $\mathbf{k}_t^{C}$ 做拼接后得到新的 $q、k$，再执行 atten weight 计算（$qk^t$）。

最后，总结下完成的 MLA 计算过程如下所示：

MLA 结构的可视化图如下所示：

2.1.4 kv cache 大小的比较

多头注意力（MHA）、分组查询注意力（GQA）、多查询注意力（MQA）和多头潜在注意力（MLA）的简化示意图对比如下图 3 所示。通过将键（Key）和值（Value）压缩到一个潜在向量中，MLA 在推理时大幅减少了对 KV 缓存的需求。

下表 1 中对比了不同注意力机制下，每个 token 需要的 KV 缓存大小。MLA 仅需要少量的 KV 缓存，其大小相当于仅有 $2.25$ 组（groups）的 GQA，但其性能却强于 MHA。

表 1｜不同注意力机制下，每个 token 需要的 KV 缓存对比。其中，

$n_h$ 表示注意力头的数量，
$d_h$ 表示每个注意力头的维度，
$l$ 表示模型层数，
$n_g$ 表示 GQA 的组数，
$d_c$ 和 $d^R_h$ 分别表示 MLA 中的 KV 压缩维度和解耦查询与键的 per-head 维度。

KV 缓存的数量以元素(elements)个数计算，不考虑存储精度（storage precision）。对于 DeepSeek-V2，

$d_c$ 设定为 $4d_h$。
$d^R_h$ 设定为 $\frac{d_h}{2}$。

因此，DeepSeek-V2 只需要相当于 GQA $2.25$ 组的 KV 缓存，但相比 MHA 仍能提供更强的性能。

2.1.5 总结

MLA 虽然增大了计算量，但 KV Cache 的减少也降低了显存和带宽的压力，且 llm 推理的 decode 阶段是受限于带宽瓶颈和显存瓶颈，因此 MLA 的引入理论上能明显提高 Generation 的速度。

2.2 DeepSeekMoE：以经济成本训练强大的模型

对于 FFN（前馈网络），我们采用 DeepSeekMoE 架构（Dai et al., 2024）。DeepSeekMoE 主要包含两个关键思想：

更精细地划分专家网络，提升每个专家的专业性，提高知识表达的准确度。
引入部分共享专家，减少不同专家间的知识冗余，从而提升计算效率。

相比传统的 MoE 架构（如 GShard，Lepikhin et al., 2021），DeepSeekMoE 在相同的专家参数量和激活参数量下，能显著提升模型性能。

设 $u_t$ 为第 $t$ 个 token 的 FFN 输入，其 FFN 输出 $h’_t$ 计算如下：

\[\mathbf{h}_t' = \mathbf{u}_t + \sum_{i=1}^{N_s} \text{FFN}_i^{(s)} (\mathbf{u}_t) + \sum_{i=1}^{N_r} g_{ij,t} \text{FFN}_i^{(r)} (\mathbf{u}_t), \tag{20}\] \[g_{ij,t} = \begin{cases} s_{ij,t}, & s_{ij,t} \in \text{Topk}(\{s_{ij,t}| 1 \leq j \leq N_r\}, K_r), \\ 0, & \text{otherwise}, \end{cases} \tag{21}\] \[s_{ij,t} = \text{Softmax}_i (\mathbf{u}_t^T e_i), \tag{22}\]

上述公式中：

$N_s$ 和 $N_r$ 分别表示共享专家和路由专家的数量；
$FFN(s) i(·)$ 和 $FFN(r) 𝑖(·)$ 分别表示第 $i$ 个共享专家和第 $i$ 个路由专家的计算过程；
$K_r$ 表示激活的路由专家数量；
$g_{i, t}$ 是第 $i$ 个专家的门控值，用来决定该专家是否激活；
$s_{i, t}$ 是 token 到专家的亲和度值，表示 token 和专家之间的相关性；
$e_i$ 是该层第 $i$ 个路由专家的质心，用于表示专家的聚合特征；
$\text{Topk}(·, K)$ 表示从第 $t$ 个 token 计算的所有路由专家的亲和度分数中，选择出 $K$ 个最高的值，并将这些分数组成一个集合。

2.2.2. 设备受限路由（Device-Limited Routing）

作者设计了一种设备受限的路由机制，用于限制 MoE 相关的通信成本。

当采用专家并行（expert parallelism）时，路由专家（routed experts）会分布在多个设备上。对于每个 token，它的 MoE 相关通信频率与其目标专家所涉及的设备数量成正比。由于 DeepSeekMoE 采用了精细的专家划分策略，激活的专家数量可能较多，如果直接应用专家并行，会导致更高的 MoE 相关通信成本。

在 DeepSeek-V2 中，除了直接选择得分最高的 $K$ 个路由专家（top-K selection）之外，我们还确保每个 token 的目标专家最多分布在 $M$ 台设备上。具体而言，对于每个 token，我们首先选择拥有最高亲和度分数的 $M$ 台设备，然后在这 $M$ 台设备上的专家中执行 top-K 选择。实践中，我们发现当 $M \geq 3$ 时，设备受限路由的效果可以大致匹配不受限的 top-K 路由。

2.2.3. 负载均衡的辅助损失（Auxiliary Loss for Load Balance）

作者在自动学习的路由策略中引入了负载均衡机制。

负载不均衡会导致路由塌陷（routing collapse）（Shazeer et al., 2017），即部分专家可能无法得到充分训练和利用。
在专家并行（expert parallelism）机制下，负载不均衡会降低计算效率。

在 DeepSeek-V2 训练过程中，我们设计了三种辅助损失（auxiliary losses），分别用于控制：

专家级负载均衡（ $L_{\text{ExpBal}}$ ），
设备级负载均衡（ $L_{\text{DevBal}}$ ），
通信均衡（ $L_{\text{CommBal}}$ ）。

这三种损失函数协同作用，确保 DeepSeek-V2 在计算资源受限的条件下，仍能高效训练高性能 MoE 模型。

2.2.4. Token-Dropping 策略

虽然平衡损失有助于实现负载平衡，但它无法完全保证负载的严格平衡。为了解决负载不平衡带来的计算浪费，作者在训练中引入了设备级的 Token-Dropping 策略。该策略首先计算每个设备的平均计算预算，将每个设备的容量因子设为 1.0。然后，借鉴 Riquelme 等人（2021）的思路，我们会丢弃每个设备上亲和度最低的 token，直到达成计算预算。此外，我们还确保约 10% 的训练序列中的 token 不会被丢弃。这样，在推理过程中，我们可以灵活地根据效率需求选择是否丢弃 token，同时保持训练和推理的一致性。

参考资料

DeepSeekV3 简单概述

2025-02-06

1. 介绍

DeepSeek-V3 是一款强大的 Mixture-of-Experts (MoE) 语言模型，总参数量为 671B，每个 token 激活的参数为 37B。为了实现高效推理和成本效益的训练，DeepSeek-V3 采用了在 DeepSeek-V2 中经过充分验证的 Multi-head Latent Attention (MLA) 和 DeepSeekMoE 架构。此外，DeepSeek-V3 首创了无辅助损失的负载平衡策略，并引入了 Multi-Token Prediction 技术来实现多 token 预测的训练目标，用以提高性能。我们在 14.8 万亿多样且高质量的 token 上进行了 DeepSeek-V3 的预训练，随后通过监督微调和强化学习阶段，充分发挥其能力。全面评估结果表明，DeepSeek-V3 在性能上超越了其他开源模型，且与领先的闭源模型的性能相当。尽管表现出色，DeepSeek-V3 仅需 2.788M H800 GPU 小时进行全程训练。此外，其训练过程非常稳定。在整个训练过程中，我们没有遇到任何不可恢复的损失峰值，也没有进行任何回滚。

2. 模型总结

架构：创新的负载平衡策略和训练目标

在 DeepSeek-V2 高效架构的基础上，我们首创了一种无辅助损失的负载平衡策略，最大限度减少了因鼓励负载平衡而带来的性能下降。我们研究了一种多 token 预测（MTP）目标，并证明它对模型性能有益。它还可以用于推测性解码，从而加速推理过程。

预训练：迈向极致的训练效率

我们设计了一个 FP8 混合精度训练框架，并首次验证了在超大规模模型上进行 FP8 训练的可行性和有效性。通过算法、框架和硬件的联合设计，我们克服了跨节点 MoE 训练中的通信瓶颈，几乎实现了计算与通信的完全重叠。这显著提高了我们的训练效率并降低了训练成本，使得我们能够在没有额外开销的情况下进一步扩大模型规模。以仅 2.664M H800 GPU 小时的经济成本，我们完成了 DeepSeek-V3 在 14.8T tokens 上的预训练，产生了目前最强的开源基础模型。预训练后的后续训练阶段仅需要 0.1M GPU 小时。

后训练：来自 DeepSeek-R1 的知识蒸馏

我们引入了一种创新的方法，将推理能力从长链推理（CoT）模型，特别是来自 DeepSeek R1 系列模型的能力，蒸馏到标准 LLM 中，尤其是 DeepSeek-V3。我们的管道巧妙地将 R1 的验证和反思模式融入 DeepSeek-V3，显著提高了其推理能力。同时，我们还保持对 DeepSeek-V3 输出风格和长度的控制。

参考资料

github-DeepSeek-V3

triton 内核编译流程解析

2025-01-25

一 Triton 概述
二 Triton 编译（JIT）入口
三内核编译函数 compile
- AST -> TritonIR
四内核 compile 流程解析
参考资料

一 Triton 概述

Triton 既是语言，可用于编写高性能计算 GPU kernel，使用 Triton 编写核函数必须按照其提供的 DSL 进行编写，否则运行失败。也是编译器（根据配置自动生成 kernel），采用 just-in-time 机制进行编译。如果你编写了一个用 triton.jit 修饰的核函数，那么在运行时其会被编译成底层 GPU 二进制（PTX / SASS），并马上执行。

Triton编译器的架构，包括前端、优化器以及后端等部件。

前端（Frontend）：用于将用户使用 Python 编写的 kernel 或者 Pytorch 2.0 中通过 Inductor 生成的 TritonKernel 转换为对应的 Triton IR。
优化器（Optimizer）：通过各类 pass 将 Triton IR 逐步转换并优化为 TritonGPU IR。
后端（Backend）：将 TritonGPU IR 逐步转换为 LLVM IR，对于 NVIDIA GPU 最终会被编译为 cubin 文件。

二 Triton 编译（JIT）入口

Triton kernel 入口：python 代码中使用 triton.jit 装饰器修饰的函数，被视为“Triton Kernel”，函数内部可使用 Triton 的 DSL（域专用语言）来编写 GEMM、reduce 算子等。

下面以官方给出的经典张量 Add 的 kernel 来简单剖析 triton 编译流程：

@triton.jit
def add_kernel(x_ptr,  # *Pointer* to first input vector.
               y_ptr,  # *Pointer* to second input vector.
               output_ptr,  # *Pointer* to output vector.
               n_elements,  # Size of the vector.
               BLOCK_SIZE: tl.constexpr,  # Number of elements each program should process.
               # NOTE: `constexpr` so it can be used as a shape value.
               ):
    pid = tl.program_id(axis=0)  # We use a 1D launch grid so axis is 0.
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    # Create a mask to guard memory operations against out-of-bounds accesses.
    mask = offsets < n_elements
    # Load x and y from DRAM, masking out any extra elements in case the input is not a multiple of the block size.
    x = tl.load(x_ptr + offsets, mask=mask)
    y = tl.load(y_ptr + offsets, mask=mask)
    output = x + y
    # Write x + y back to DRAM.
    tl.store(output_ptr + offsets, output, mask=mask) 

def add(x: torch.Tensor, y: torch.Tensor):
    output = torch.empty_like(x) # preallocate the output.
    assert x.device == DEVICE and y.device == DEVICE and output.device == DEVICE
    n_elements = output.numel()
    # In this case, we use a 1D grid where the size is the number of blocks:
    grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']), )
    # NOTE:
    #  - Each torch.tensor object is implicitly converted into a pointer to its first element.
    #  - `triton.jit`'ed functions can be indexed with a launch grid to obtain a callable GPU kernel.
    #  - Don't forget to pass meta-parameters as keywords arguments.
    add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024)
    
    return output 

2.1 jit() 函数

python/triton/runtime/jit.py: jit() 装饰器函数将原始函数 fn 封装为 JIT 函数: JITFunction，JITFunction 被调用的时候，是以 fn[grid](*args, **kwargs) 的形式被调用的。jit() 函数执行流程：

decorator 装饰器定义：检查环境变量 TRITON_INTERPRET：
- 如果设置为 “1”，返回 InterpretedFunction （用于调试）
- 否则返回 JITFunction （正常编译模式）
装饰器将原始函数封装为 JIT 函数并返回。

2.2 JITFunction 类

JITFunction 类继承自 KernelInterface，KernelInterface 类定义了内核函数的基本接口和启动机制，其代码实现如下所示，很明显，通过 __getitem__ 方法允许使用方括号语法 [grid] 来设置内核的执行网格，这里的网格参数定义了 kernel 的并行执行结构。

__getitem_ 是 Python 的魔法方法，通常用来定义 obj[key] 的行为。KernelInterface 类继承自 Generic：泛型类型的抽象基类。Generic 作为所有泛型类的基类，提供类型参数化的基础设施。

class KernelInterface(Generic[T]):
    run: T

    def __getitem__(self, grid) -> T:
        """
        A JIT function is launched with: fn[grid](*args, **kwargs).
        Hence JITFunction.__getitem__ returns a callable proxy that
        memorizes the grid.
        """
        return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
        # return cast(T, functools.partial(cast(Callable, self.run), grid=grid))

__getitem__ 会返回一个匿名函数（lambda），其调用了 self.run 函数，并记住了 grid 参数。self.run() 函数在子类 JITFunction 中实现，函数执行流程包括:

获取/设置编译、执行环境: 通过 driver.active.get_current_device() 与 driver.active.get_current_stream(device) 来获取当前设备和流（stream）。
调用预处理钩子（pre_run_hooks）: 在运行或编译前，可以执行一些用户自定义的操作（例如日志、数据检查、统计）。
绑定参数、生成内核缓存键: 与 binder 交互，将传入的 *args, **kwargs 绑定成特定的参数结构并生成 “specialization” 信息。然后使用 specialization + options 组合成一个字符串 str 类型的缓存键 key。
通过 kernel = kernel_cache.get(key, None) 查找或编译内核: 如果缓存中已有对应内核，则复用；否则就走编译流程并把新编译的内核写回缓存。
执行内核（可选）: 如果 warmup=False，则真正将内核提交到 GPU 执行，传入相应的网格大小（grid）及参数；否则仅完成编译预热（不执行），返回编译后的 kernel。
返回编译（或缓存）得到的内核对象。

这里重点分析下 triton 编译内核的过程代码:

# 当内核未缓存时执行编译流程
if kernel is None:
    # -------------------- 选项处理 --------------------
    # 从kwargs解析后端编译选项（如GPU架构参数等）
    options = backend.parse_options(kwargs)
    
    # -------------------- 签名构建 --------------------
    # 从类参数中获取参数名称列表
    sigkeys = [x.name for x in self.params]
    # 从特化参数中获取参数值列表（第一个元素为参数值）
    sigvals = [x[0] for x in specialization]
    # 构建参数名-值映射的签名字典
    signature = {k: v for (k, v) in zip(sigkeys, sigvals)}
    
    # -------------------- 参数校验 --------------------
    # 检查是否使用了已弃用的设备相关参数
    assert "device_type" not in kwargs, "device_type选项已弃用，将使用当前目标设备"
    assert "device" not in kwargs, "device选项已弃用，将使用当前设备"
    assert "stream" not in kwargs, "stream选项已弃用，将使用当前流"
    # 验证所有传入参数是否合法
    for k in kwargs:
        if k not in options.__dict__ and k not in sigkeys:
            raise KeyError(f"检测到未识别的关键字参数: {k}")

    # -------------------- 常量表达式处理 --------------------
    # 查找标记为constexpr的常量参数路径
    constexprs = find_paths_if(sigvals, lambda _, val: val == "constexpr")
    # 获取这些常量参数的实际值（用于编译时优化）
    constexprs = {path: get_iterable_path(list(bound_args.values()), path) for path in constexprs}

    # -------------------- 属性处理 --------------------
    # 从特化参数中获取属性值（第二个元素为属性）
    attrvals = [x[1] for x in specialization]
    # 查找字符串类型的属性路径（如GPU内存配置参数）
    attrs = find_paths_if(attrvals, lambda _, x: isinstance(x, str))
    # 解析属性配置（转换为后端需要的格式）
    attrs = {k: backend.parse_attr(get_iterable_path(attrvals, k)) for k in attrs}

    # -------------------- 编译前钩子 --------------------
    # 执行预编译钩子函数（如缓存检查、日志记录等）
    if self._call_hook(key, signature, device, constexprs, options, [attrs], warmup, before=True):
        return None  # 如果钩子返回True则中止编译

    # -------------------- 编译流程 --------------------
    # 生成抽象语法树（AST）形式的源代码
    src = self.ASTSource(self, signature, constexprs, attrs)
    # 调用后端编译器生成可执行内核
    kernel = self.compile(src, target=target, options=options.__dict__)
    # 将编译结果缓存
    kernel_cache[key] = kernel
    
    # -------------------- 编译后钩子 --------------------
    # 执行后编译钩子函数（如性能分析、资源注册等）
    self._call_hook(key, signature, device, constexprs, options, [attrs], warmup, before=False)

上述代码看着很复杂，但大部分都是预处理流程代码，真正的内核编译代码就三行：通过抽象语法树 AST 生成中间代码，调用编译器函数 self.compile，并缓存结果。

src = self.ASTSource(self, signature, constexprs, attrs)
# print(src.code) 可直接打印生成的 GPU 代码
kernel = self.compile(src, target=target, options=options.__dict__)
kernel_cache[key] = kernel

2.3 ASTSource 类

triton/compiler/compiler.py: ASTSource 是 Triton 编译器中的一个关键类，负责管理和处理源代码的抽象语法树(AST)。

JITFunction.run 函数调用 ASTSource 类的代码是：src = self.ASTSource(self, signature, constexprs, attrs)。

这段代码的作用是将内核的参数、常量表达式和属性配置转换为中间表示（IR）的抽象语法树（AST），这是实现跨平台代码生成和即时编译（JIT）的关键步骤，其步骤包括：语义解析、类型推导、常量替换、属性注入、多后端适配（CUDA/ROCm/Metal）。这里举个例子理解其作用：

假设用户编写如下 Triton 核函数：

@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M: tl.constexpr, N: tl.constexpr, K: tl.constexpr
):
    # ... 矩阵乘法逻辑

上述内核函数对应的输入参数结构：

signature（函数签名）：描述参数名称和类型。

{
    "a_ptr": "pointer",
    "b_ptr": "pointer",
    "c_ptr": "pointer",
    "M": "int32",
    "N": "int32",
    "K": "int32"
}

constexprs（编译时常量）: 标记需要在编译时确定的参数值。

{"M": 1024, "N": 512, "K": 2048}  # 矩阵维度

attrs（硬件属性）: 配置 GPU 核函数的执行参数。

{
    "num_warps": 1,     # 每个线程块的 warp 数量
    "num_stages": 1,    # 流水线阶段数（用于异步拷贝优化）
    "static_shared": 4096  # 预分配的静态共享内存大小
}

ASTSource.make_ir 方法会调用 ast_to_ttir 函数将 Python AST 转换为目标平台中间表示（如 MLIR 中间表示），如下所示：

// Triton 生成的 MLIR 中间表示（简化版）
module attributes {
  "triton_gpu.num-warps" = 8,
  "triton_gpu.threads-per-warp" = 32,
  "triton_gpu.shared" = 4096
} {
  func.func @matmul_kernel(
      %a_ptr: !tt.ptr<f32>,
      %b_ptr: !tt.ptr<f32>,
      %c_ptr: !tt.ptr<f32>
  ) {
    // 常量替换后的参数（M=1024, N=512, K=2048）
    %M = arith.constant 1024 : i32
    %N = arith.constant 512 : i32
    %K = arith.constant 2048 : i32

    // 计算逻辑（伪代码）
    %row = tt.get_program_id(0) : i32
    %col = tt.get_program_id(1) : i32
    %a = tt.load %a_ptr[%row, %col] : !tt.ptr<f32>
    %b = tt.load %b_ptr[%col, %row] : !tt.ptr<f32>
    %c = tt.dot %a, %b : f32
    tt.store %c_ptr[%row, %col] = %c : !tt.ptr<f32>
    return
  }
}

三内核编译函数 compile

triton/compiler/compiler.py: compile() 函数最终返回 CompiledKernel 对象，输入参数 src 为 ASTSource 对象，ASTSource 主要作用是将 JITFunction 从 ast tree 转换为 Triton IR。

compile() 函数其简化版代码如下所示:

def compile(src, target=None, options=None):
    if target is None:
        target = driver.active.get_current_target()
    assert isinstance(target, GPUTarget), "target must be of GPUTarget type"
    backend = make_backend(target)
    ir_source = not isinstance(src, ASTSource)
    #################省略部分代码#################
    # try:
    # 1. AST -> TritonIR
    module = src.make_ir(options, codegen_fns, module_map, context)
    # except Exception as e:
    #     filter_traceback(e)
    #     raise
    use_ir_loc = os.environ.get("USE_IR_LOC", None)
    # 2. 遍历执行 backend 中定义的 stages 的 pass
    for ext, compile_ir in list(stages.items())[first_stage:]:
        next_module = compile_ir(module, metadata)
        ir_filename = f"{file_name}.{ext}"
        if (fn_override_manager is not None and (full_name := fn_override_manager.get_file(ir_filename)) is not None):
            print(f"\nOverriding kernel with file {full_name}")
            next_module = parse(full_name, ext, context)
        metadata_group[ir_filename] = fn_cache_manager.put(next_module, ir_filename)
        if fn_dump_manager is not None:
            fn_dump_manager.put(next_module, ir_filename)
        # use an env variable to parse ir from file
        if use_ir_loc == ext:
            ir_full_name = fn_cache_manager.get_file(ir_filename)
            next_module.create_location_snapshot(ir_full_name)
            print(f"Creating new locations for {ir_full_name}")
        module = next_module
    # write-back metadata
    metadata_group[metadata_filename] = fn_cache_manager.put(json.dumps(metadata, default=vars), metadata_filename,
                                                             binary=False)
    fn_cache_manager.put_group(metadata_filename, metadata_group)
    # Compilation completed, disabling multithreading in context.
    # This is needed to safely finalize threads pool inside context: if current process forks before
    # python GC deletes context object, thread pool in child process will be invalid, which could
    # lead to child crash or hang.
    #
    # However disabling multithreading causes the code to hang if the ASAN pass is enabled
    # this is likely due to the llvm-symbolizer forking a process
    # TODO: Reconcile the difference here between the ASAN and non-ASAN path with enabling
    # multithreading in the MLIR context
    if not os.environ.get("TRITON_ENABLE_ASAN", "0") == "1":
        context.disable_multithreading()
    # return handle to compiled kernel
    return CompiledKernel(src, metadata_group, hash)

complie 函数的核心流程如下，最终的返回值就是编译完成后的内核，其内部通常包含 GPU 二进制（PTX/Cubin）或相应后端需要的可执行形式，也可能包括若干中间表征和完整的元数据。

AST -> TritonIR
遍历执行 backend 中定义的 stages 的 pass

AST -> TritonIR

make_ir 函数实现如下所示：

class ASTSource:
    def __init__(self, fn, signature, constexprs=None, attrs=None) -> None:
        ############## 省略代码 ##############

    def make_ir(self, options, codegen_fns, module_map, context):
        return ast_to_ttir(self.fn, self, context=context, options=options, codegen_fns=codegen_fns,
                            module_map=module_map)

make_ir 函数的第一个参数 self 实际是 JITFunction 对象，其作用是将 JITFunction 从 ast tree 转换为 Triton IR，通过调用 ast_to_ttir 实现。

ast_to_ttir 函数定义如下：

def ast_to_ttir(fn, src, context, options, codegen_fns, module_map):
    arg_types = list(map(str_to_ty, src.signature.values()))
    prototype = ASTFunction([], arg_types, src.constants, src.attrs)
    file_name, begin_line = get_jit_fn_file_line(fn)
    # query function representation
    from collections import namedtuple
    leaves = filter(lambda v: len(v) == 1, src.constants)
    constants = {fn.arg_names[i[0]]: src.constants[i] for i in leaves}
    signature = src.signature
    proxy = namedtuple("SpecializationProxy", ["constants", "signature"])(constants, signature)
    generator = CodeGenerator(context, prototype, gscope=fn.__globals__.copy(), function_name=fn.repr(proxy), jit_fn=fn,
                              is_kernel=True, file_name=file_name, begin_line=begin_line, options=options,
                              codegen_fns=codegen_fns, module_map=module_map)
    generator.visit(fn.parse())
    ret = generator.module
    # module takes ownership of the context
    ret.context = context
    return ret

函数实现最关键的部分是初始化 CodeGenerator 对象，和使用 JITFunction.parse() 方法会将 kernel 代码转换为 ast tree，返回ast tree 的 root 节点。

四内核 compile 流程解析

Backend 系统

Triton 在高层提供了一个 compile(src, target, options) 这样的统一编译函数(在 python/triton/compiler/compiler.py 文件中)，但具体如何把 ASTSource 转成可执行文件（PTX / Cubin / GCN / etc.）则由各个后端模块在 backends 目录下完成，如英伟达 GPU 的后端实现目录 :vthird_party/nvidia/backend。

python/triton/backends 目录下的 __init__.py、compiler.py（BaseBackend 类）、driver.py（DriverBase 类）代码实现了跨硬件平台支持，提供了一套子类，定义了一系列抽象方法 abstractmethod 统一接口，用于将 Triton 的中间表示（IR）或 AST 编译成不同 GPU 平台所需的可执行内核。

Triton 定义了一个抽象后端 BaseBackend 类，定义所有硬件后端必须实现的接口：

class Backend(ABC):
    @abstractmethod
    def add_stages(self, stages: dict, options):
        """定义编译流水线阶段（如 TTIR→LLVMIR→PTX）"""
        
    @abstractmethod
    def get_codegen_implementation(self, options):
        """返回平台特定的代码生成回调函数"""

    @abstractmethod
    def get_module_map(self):
        """获取依赖的预编译模块（如数学函数库）"""

    ################## 省略代码 ####################

Triton 内核的通用编译流程总结如下图所示:

图片来源 Triton概念与编程入门笔记（以Matmul为例）

不同的GPU 硬件 backend 定义了不同的编译 stages：

third_party/nvidia/backend/compiler.py 文件: class CUDABackend(BaseBackend):

def add_stages(self, stages, options):
    capability = self._parse_arch(options.arch)
    stages["ttir"] = lambda src, metadata: self.make_ttir(src, metadata, options)
    stages["ttgir"] = lambda src, metadata: self.make_ttgir(src, metadata, options, capability)
    stages["llir"] = lambda src, metadata: self.make_llir(src, metadata, options, capability)
    stages["ptx"] = lambda src, metadata: self.make_ptx(src, metadata, options, self.target.arch)
    stages["cubin"] = lambda src, metadata: self.make_cubin(src, metadata, options, self.target.arch)

很明显英伟达 cuda 后端的编译流程分为 5 个 stages，将 kernel 函数的中间表示（IR）转换为最终的二进制 cubin 代码。

third_party/amd/backend/compiler.py 文件: class HIPBackend(BaseBackend):

def add_stages(self, stages, options):
    stages["ttir"] = lambda src, metadata: self.make_ttir(src, metadata, options)
    stages["ttgir"] = lambda src, metadata: self.make_ttgir(src, metadata, options)
    stages["llir"] = lambda src, metadata: self.make_llir(src, metadata, options)
    stages["amdgcn"] = lambda src, metadata: self.make_amdgcn(src, metadata, options)
    stages["hsaco"] = lambda src, metadata: self.make_hsaco(src, metadata, options)

下面以 CUDABackend 后端类为例简单描述下 kernel 编译流程和对应函数作用。

4.1 make_ttir

make_ttir 函数会构建并将 py 文件转换 Triton IR 文件，Triton IR 是 Triton 编译器的高级中间表示，用于表示深度学习模型的计算图，是硬件无关的。其特点包括：

高级抽象：使开发者用接近于高级深度学习框架的方式来描述计算图。
操作表示：包含了一系列的操作，如矩阵乘法、卷机、激活函数等。
优化：编译器可以应用一些高级优化，如死代码消除（DCE）、共子表达式消除（CSE）等。
转换：可以被转换为更接近硬件特性的低阶表示：Triton GPU IR。

ttir 的全称是 Triton Tensor IR。

@staticmethod
def make_ttir(mod, metadata, options):
    pm = ir.pass_manager(mod.context)
    pm.enable_debug()
    passes.common.add_inliner(pm)
    passes.ttir.add_rewrite_tensor_pointer(pm)
    passes.common.add_canonicalizer(pm)
    passes.ttir.add_combine(pm)
    passes.ttir.add_reorder_broadcast(pm)
    passes.common.add_cse(pm)
    passes.common.add_licm(pm)
    passes.common.add_symbol_dce(pm)
    passes.ttir.add_loop_unroll(pm)
    pm.run(mod)
    return mod

输入：

mod: 一个 Module（通常是 MLIR 中定义的 Triton IR 模块）。
metadata: 编译过程中的一些上下文信息（可能包含目标架构、调试选项、或其他元数据）。
options: 进一步控制 Pass 管线行为的编译选项。

输出：

优化/转换过的同一 mod 对象（MLIR Module）。
返回给调用方后，可继续流向下一步编译或直接生成 GPU 二进制。

4.2 make_ttgir

make_ttgir 首先将 TritonIR 转换为 TritonGPUIR（专门针对 GPU 硬件优化的中间表示），然后对 TritonGPUIR 进行优化：

架构特定优化根据 GPU 计算能力(capability)启用不同优化:
通用优化：内存访问优化（合并内存访问、数据预取、优化线程局部性）、计算相关优化（矩阵乘法加速、点积运算优化、指令流水线优化）等。

函数返回依然是优化后的同一个 Module 对象。

@staticmethod
def make_ttgir(mod, metadata, opt, capability):
    # 1. 准备 cluster_info 结构，用于描述集群（cluster）维度（X, Y, Z）。
    cluster_info = nvidia.ClusterInfo()
    if opt.cluster_dims is not None:
        cluster_info.clusterDimX = opt.cluster_dims[0]
        cluster_info.clusterDimY = opt.cluster_dims[1]
        cluster_info.clusterDimZ = opt.cluster_dims[2]

    # 2. 如果环境变量里开启了 MLIR_ENABLE_REMARK，则设置源码诊断（Diagnostic）
    if os.environ.get("MLIR_ENABLE_REMARK", "0") == "1":
        srcMgr = llvm.source_mgr()
        _ = ir.source_mgr_diag(srcMgr, mod.context)
        mod.context.printOpOnDiagnostic(True)

    # 3. 构建一个 PassManager，用于把 TTIR 转成 TTGIR，并对其进行各种优化
    pm = ir.pass_manager(mod.context)
    pm.enable_debug()

    # 3.1 把 TTIR 转换为 TTGPU IR
    #     - 这里指定了 “cuda:{capability}”，以及 opt.num_warps, 32, opt.num_ctas 等参数。
    passes.ttir.add_convert_to_ttgpuir(pm, f"cuda:{capability}", opt.num_warps, 32, opt.num_ctas)

    # 4. 优化 TTGIR：coalesce（合并张量操作）、f32_dot_tc（在capability>=8.0时启用f32 TensorCore等）
    passes.ttgpuir.add_coalesce(pm)
    if capability // 10 >= 8:
        passes.ttgpuir.add_f32_dot_tc(pm)

    # 4.1 PlanCTA：将 CTA（thread block）按照 cluster_info 做布局规划
    #     注意这里由 nvidia.passes.ttnvgpuir 提供
    nvidia.passes.ttnvgpuir.add_plan_cta(pm, cluster_info)

    # 4.2 进一步优化：移除多余 layout 转换、优化线程局部性、加速矩阵乘法等
    passes.ttgpuir.add_remove_layout_conversions(pm)
    passes.ttgpuir.add_optimize_thread_locality(pm)
    passes.ttgpuir.add_accelerate_matmul(pm)
    passes.ttgpuir.add_remove_layout_conversions(pm)

    # 4.3 优化点乘操作（点积）；capability >= 80（即8.0以上）时可能有专门的硬件指令
    passes.ttgpuir.add_optimize_dot_operands(pm, capability >= 80)

    # 4.4 通用优化：CSE（消除公共子表达式）
    passes.common.add_cse(pm)

    # 4.5 如果是 8.0（Ampere）或更高：做一些更高级的优化（累加器初始化、select和if合并、loop调度与pipeline等）
    if capability // 10 >= 8:
        passes.ttgpuir.add_optimize_accumulator_init(pm)
        passes.ttgpuir.add_combine_tensor_select_and_if(pm)
        passes.ttgpuir.add_loop_scheduling(pm, opt.num_stages)
        passes.ttgpuir.add_pipeline(pm, opt.num_stages)

    # 4.6 继续做预取（prefetch）、coalesce async copy、指令重排等，提升内存访问性能
    passes.ttgpuir.add_prefetch(pm)
    passes.ttgpuir.add_optimize_dot_operands(pm, capability >= 80)
    passes.ttgpuir.add_coalesce_async_copy(pm)
    passes.ttgpuir.add_remove_layout_conversions(pm)
    passes.ttgpuir.add_reduce_data_duplication(pm)
    passes.ttgpuir.add_reorder_instructions(pm)

    # 4.7 再做一波公共子表达式消除 & 符号死代码删除
    passes.common.add_cse(pm)
    passes.common.add_symbol_dce(pm)

    # 5. 如果是 9.0（Hopper）或更高架构：插入 fence & TMA（Tensor Memory Accelerator）等相关 pass
    if capability // 10 >= 9:
        nvidia.passes.ttnvgpuir.add_fence_insertion(pm)
        nvidia.passes.ttnvgpuir.add_tma_lowering(pm)

    # 6. 最后做一次 canonicalize（规范化），并运行整个 Pass 管线
    passes.common.add_canonicalizer(pm)
    pm.run(mod)

    # 7. 将 cluster_dims 信息记录到 metadata，便于后续阶段查询
    metadata["cluster_dims"] = (cluster_info.clusterDimX,
                                cluster_info.clusterDimY,
                                cluster_info.clusterDimZ)

    # 8. 返回处理好的 MLIR Module
    return mod

4.3 make_llir

make_llir 是 Triton 编译流水线中的 LLVM IR 生成与优化阶段，负责将 Triton GPU IR（TTGPUIR）转换为优化后的 LLVM IR，并然后再把 MLIR 表示的 LLVM IR 进一步转化成实际的 LLVM 模块（通过 llvm.to_module），并最终返回一段 LLVM IR 字符串。它同时会根据硬件“compute capability”、编译选项（options）等信息做相应的优化与元数据记录。其核心任务包括：

低级优化：将平台无关的 TTGPUIR 转换为 NVIDIA GPU 专用的 LLVM IR
内存管理：分配共享内存（Shared Memory）和全局暂存内存（Global Scratch Memory）
架构适配：根据 GPU 计算能力（如 sm_80）和 PTX 版本生成目标指令
性能调优：设置寄存器限制、链接外部库、应用 O3 级优化

函数输出是：一段 LLVM IR 源码（字符串形式）。示例：

; 优化后的 LLVM IR
target triple = "nvptx64-nvidia-cuda"
target datalayout = "e-i64:64-i128:128-v16:16-v32:32-n16:32:64"

define void @kernel(ptr %arg0) #0 {
  %shared = alloca [128 x float], align 4
  %val = load float, ptr %arg0
  store float %val, ptr %shared
  ret void
}

attributes #0 = { "target-features"="+ptx80,+sm_80" "maxnreg"="128" }

make_llir 函数定义：

@staticmethod
def make_llir(src, metadata, options, capability):
    """
    将 Triton GPU IR (TTGIR) 转换为 LLVM IR (MLIR 表示)，
    再进一步转换为实际的 LLVM 模块并返回其字符串形式。
    同时根据编译选项与硬件 capability 做相应处理。
    """

    # 1. 从编译选项中获取 PTX 版本（如果是 NVIDIA 后端）
    ptx_version = get_ptx_version_from_options(options)

    # 2. 根据属性 "ttg.num-warp-groups-per-cta" 动态调整 num_warps
    #    (warp-specialization 机制)
    num_warp_groups = src.get_int_attr("ttg.num-warp-groups-per-cta")
    if num_warp_groups is not None:
        metadata["num_warps"] *= num_warp_groups

    mod = src  # mod 就是 TTGIR Module (MLIR)

    # 3. 构建 MLIR PassManager 来执行 TTGIR -> LLVM-IR (MLIR) 的转换与优化
    pm = ir.pass_manager(mod.context)
    pm.enable_debug()

    # 3.1 若环境变量启用 MLIR 诊断，将在编译时输出额外信息
    if os.environ.get("MLIR_ENABLE_REMARK", "0") == "1":
        srcMgr = llvm.source_mgr()
        _ = ir.source_mgr_diag(srcMgr, mod.context)
        mod.context.printOpOnDiagnostic(True)

    # 3.2 对 TTGIR 进行一些转换与优化 Pass
    passes.ttgpuir.add_combine_tensor_select_and_if(pm)
    passes.convert.add_scf_to_cf(pm)
    passes.convert.add_index_to_llvmir(pm)
    passes.ttgpuir.add_allocate_shared_memory(pm)
    passes.ttgpuir.add_allocate_global_scratch_memory(pm)
    # 3.3 将 TTGIR 转换成 MLIR 中的 LLVM Dialect IR (兼顾 GPU capability 和 ptx_version)
    nvidia.passes.ttgpuir.add_to_llvmir(pm, capability, ptx_version)
    # 3.4 如果有 nvidia GPU 定制化 pass，也在此时调用
    nvidia.passes.ttnvgpuir.add_nvgpu_to_llvm(pm)
    passes.convert.add_arith_to_llvmir(pm)

    # 3.5 做一些通用 IR 优化 & 清理
    passes.common.add_canonicalizer(pm)
    passes.common.add_cse(pm)
    passes.common.add_symbol_dce(pm)

    # 3.6 如果环境变量允许行号信息，则添加 debug scope
    if os.environ.get("TRITON_DISABLE_LINE_INFO", "0") == "0":
        passes.llvmir.add_di_scope(pm)

    # 3.7 运行所有 Pass，得到 LLVM-IR (仍是 MLIR 表示)
    pm.run(mod)

    # 4. 初始化 LLVM，准备将 MLIR -> LLVM 原生模块
    llvm.init_targets()
    context = llvm.context()

    # 4.1 不支持 ASAN 时的提示
    if os.environ.get("TRITON_ENABLE_ASAN", "0") == "1":
        raise RuntimeError(
            "Address Sanitizer Error: Address sanitizer is only supported on AMD backend"
        )

    # 4.2 把 MLIR 中的 LLVM Dialect Module 转为一个真实的 LLVM 模块 (llvmlite)
    llvm_mod = llvm.to_module(mod, context)

    # 4.3 根据 compute capability 设置目标 (sm_80 / sm_90a 等) 与特性
    proc = 'sm_90a' if capability == 90 else f'sm_{capability}'
    features = get_features(options)
    triple = 'nvptx64-nvidia-cuda'
    llvm.attach_datalayout(llvm_mod, triple, proc, features)

    # 4.4 在 LLVM 模块上设定 “FTZ (flush-to-zero)” 等配置
    nvidia.set_nvvm_reflect_ftz(llvm_mod)

    # 4.5 如果用户在 options 里指定了 maxnreg，则给内核函数附加相应的 NVVM 属性
    if options.maxnreg is not None:
        for k in llvm_mod.get_functions():
            if not k.is_declaration() and k.is_external_linkage():
                k.set_nvvm_maxnreg(options.maxnreg)

    # 4.6 若需要链接外部库，则在此操作
    if options.extern_libs:
        paths = [path for (name, path) in options.extern_libs]
        llvm.link_extern_libs(llvm_mod, paths)

    # 4.7 优化 LLVM 模块到 O3
    llvm.optimize_module(llvm_mod, llvm.OPTIMIZE_O3)

    # 5. 提取更多 metadata 信息 (共享内存大小、全局临时区大小/对齐等)
    metadata["shared"] = src.get_int_attr("ttg.shared")
    metadata["global_scratch_size"] = src.get_int_attr("ttg.global_scratch_memory_size")
    metadata["global_scratch_align"] = src.get_int_attr("ttg.global_scratch_memory_alignment")

    # 6. 输出最终的 LLVM IR 字符串
    ret = str(llvm_mod)

    # 7. 释放资源，然后返回 LLVM IR 字符串
    del llvm_mod
    del context
    return ret

4.4 make_ptx

make_ptx 用于将此前编译产生的 LLVM IR (在字符串形式) 翻译成PTX 汇编代码，并对生成的 PTX 进行一定的后处理（如插入正确的 .version、去除 debug 标记等）。它主要依赖 llvmlite 或 Triton 自带的 llvm.translate_to_asm 接口，将 LLVM IR 转换为可在 NVIDIA GPU 上使用的 PTX 代码。

@staticmethod
def make_ptx(src, metadata, opt, capability):
    """
    将 'src' (LLVM IR 字符串) 翻译成 PTX 汇编，并进行必要的后处理和信息提取。
    :param src:         字符串形式的 LLVM IR
    :param metadata:    字典类型的元数据，用于存放编译后的一些信息，如 kernel 名称等
    :param opt:         编译选项对象（包含 arch、fp_fusion 等）
    :param capability:  GPU 架构（如 80=sm_80, 90=sm_90a 等）
    :return:            最终生成的 PTX 字符串
    """

    # 1. 从选项 'opt' 中获取所需 PTX 版本 (例如 63 对应 PTX 6.3)
    ptx_version = get_ptx_version_from_options(opt)

    # 2. 设定编译目标三元组 (triple) 和目标 GPU (proc)，以及特性 (features)
    #    - triple 固定为 "nvptx64-nvidia-cuda"
    #    - 若 capability==90 则用 "sm_90a"，否则 "sm_{capability}"
    triple = 'nvptx64-nvidia-cuda'
    proc = 'sm_90a' if capability == 90 else f'sm_{capability}'
    features = get_features(opt)

    # 3. 将 LLVM IR 翻译成 PTX 汇编
    #    - 这里传入包括 triple, proc, features, 以及一些编译选项 (如 fp_fusion)
    #    - 返回值 ret 是 PTX 字符串
    ret = llvm.translate_to_asm(
        src,                 # LLVM IR
        triple,             # Target triple
        proc,               # GPU 架构
        features,           # 特性字符串
        ['nvptx-short-ptr'],# 一些附加选项
        opt.enable_fp_fusion,
        False               # 是否以 debug 方式生成
    )

    # 4. 从生成的 PTX 中找到 kernel 名 ('.visible .entry kernelName')
    #    - 预期只有一个 kernel
    names = re.findall(r".visible .entry ([a-zA-Z_][a-zA-Z0-9_]*)", ret)
    assert len(names) == 1, f"Expected exactly 1 kernel, got {len(names)}"
    metadata["name"] = names[0]  # 记录到 metadata

    # 5. 将 PTX 文件头中的 '.version x.x' 替换为实际 desired ptx_version
    #    - ptx_version 字符串形如 '6.3', '7.5' 等
    ptx_version = f'{ptx_version // 10}.{ptx_version % 10}'
    ret = re.sub(r'\.version \d+\.\d+', f'.version {ptx_version}', ret, flags=re.MULTILINE)

    # 6. 移除会阻止 ptxas 优化的 "debug" 标记，以取得更高性能
    ret = re.sub(r",\s*debug|debug,\s*", "", ret)

    # 7. 若环境变量里设置 NVPTX_ENABLE_DUMP=1，则打印生成的 PTX 以供调试
    if os.environ.get("NVPTX_ENABLE_DUMP", "0") == "1":
        print("// -----// NVPTX Dump //----- //")
        print(ret)

    # 8. 返回处理好的 PTX 字符串
    return ret

输入输出示例:

输入：LLVM IR 字符串

复制
; LLVM IR 输入（已优化）
target triple = "nvptx64-nvidia-cuda"
define void @triton_kernel(ptr %arg0) #0 {
  %val = load float, ptr %arg0
  store float %val, ptr null  ; 示例简化存储操作
  ret void
}

attributes #0 = { "target-features"="+ptx80,+sm_80" }

输出：PTX 汇编代码

// 生成的 PTX 代码
.version 7.8        ; 根据选项调整后的版本
.target sm_80       ; GPU 计算能力
.address_size 64    ; 64 位地址

.visible .entry triton_kernel(      ; 元数据中记录内核名称
  .param .u64 %arg0
) {
  ld.param.u64 %rd1, [%arg0];
  ld.global.f32 %f1, [%rd1];
  st.global.f32 [0], %f1;           ; 存储操作
  ret;
}

4.5 make_cubin

make_cubin 函数负责将PTX（字符串形式）编译成Cubin（NVIDIA GPU 的二进制可执行文件）。它通过在本地创建临时文件并调用 ptxas 命令行工具完成最终编译输出。如果编译成功，就会返回 cubin（二进制）数据；若编译出错，则会抛出相应的异常。

@staticmethod
def make_cubin(src, metadata, opt, capability):
    """
    将 PTX 字符串编译成 NVIDIA Cubin 二进制。
    :param src:         要编译的 PTX 源码（字符串）
    :param metadata:    存放编译中产生的一些上下文信息（本函数暂未直接修改其内容）
    :param opt:         编译选项对象 (enable_fp_fusion 等)
    :param capability:  GPU 架构 (如 80 对应 sm_80, 90 对应 sm_90a)
    :return:            编译成功时返回一个 bytes 类型的 cubin（可执行二进制）
    """

    # 1. 获取 ptxas 二进制所在路径
    #    如果找不到，_path_to_binary(...) 会抛出相应异常
    ptxas, _ = _path_to_binary("ptxas")

    # 2. 创建临时文件，用于写入 PTX 源码并接受日志输出
    #    mode='w' 以可写方式创建 .ptx 文件
    #    mode='r' 用来读取日志 .log 文件
    with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.ptx') as fsrc, \
         tempfile.NamedTemporaryFile(delete=False, mode='r', suffix='.log') as flog:

        # 2.1 写入 PTX 源码到临时文件
        fsrc.write(src)
        fsrc.flush()

        # 2.2 输出的目标文件名 (.o) 存在于同一路径
        fbin = fsrc.name + '.o'

        # 3. 构建 ptxas 命令行选项
        #    - line_info: 是否启用行号信息
        line_info = [] if os.environ.get('TRITON_DISABLE_LINE_INFO') else ['-lineinfo']
        #    - fmad: 是否允许 FMA 指令融合 (由 enable_fp_fusion 决定)
        fmad = [] if opt.enable_fp_fusion else ['--fmad=false']
        #    - suffix: 如果capability=90, 则在后面加 'a' => sm_90a
        suffix = 'a' if capability == 90 else ''
        #    - opt_level: 是否禁用 ptxas 优化
        opt_level = ['--opt-level', '0'] if os.environ.get("DISABLE_PTXAS_OPT", "0") == "1" else []

        ptxas_cmd = [
            ptxas,
            *line_info,
            *fmad,
            '-v',  # 打印一些编译信息到 stderr
            *opt_level,
            f'--gpu-name=sm_{capability}{suffix}',
            fsrc.name,  # 输入 PTX 文件
            '-o', fbin  # 输出目标文件
        ]

        try:
            # 4. 调用 ptxas 进行编译
            #    - close_fds=False 可能是为了兼容某些环境
            #    - stderr=flog 将错误日志写入 flog 文件
            subprocess.run(ptxas_cmd, check=True, close_fds=False, stderr=flog)

            # 4.1 若编译成功，清理生成的中间日志文件
            if os.path.exists(fsrc.name):
                os.remove(fsrc.name)
            if os.path.exists(flog.name):
                os.remove(flog.name)

        except subprocess.CalledProcessError as e:
            # 5. 如果编译出错，读取错误日志并抛出 PTXASError
            with open(flog.name) as log_file:
                log = log_file.read()
            if os.path.exists(flog.name):
                os.remove(flog.name)

            # 根据 e.returncode 的不同，返回不同错误信息
            if e.returncode == 255:
                error = 'Internal Triton PTX codegen error'
            elif e.returncode == 128 + signal.SIGSEGV:
                error = '`ptxas` raised SIGSEGV'
            else:
                error = f'`ptxas` failed with error code {e.returncode}'

            raise PTXASError(
                f"{error}\n"
                f"`ptxas` stderr:\n{log}\n"
                f'Repro command: {" ".join(ptxas_cmd)}\n'
            )

        # 6. 若编译成功，则读取生成的 .o 文件即 cubin
        with open(fbin, 'rb') as f:
            cubin = f.read()

        # 6.1 清理 .o 文件
        if os.path.exists(fbin):
            os.remove(fbin)

    # 7. 返回最终的 cubin 二进制 (bytes)
    return cubin

至此，Triton 编译 kernel 的完整流程梳理总结完毕。

参考资料

vllm 推理流程剖析

2024-12-02

简单来讲，大模型推理一般会经历 4 个流程：Tokenizer -> Model.forward -> Sampler -> DeTokenizer，但对于集成了各种优化技术的 vllm 框架来说，大模型推理服务流程是非常复杂的

先看下 vLLM 代码整体架构：

vllm 模型执行流程调用关系可以通过以下代码打印：

print(vllm_model.model.llm_engine.model_executor.driver_worker.model_runner.model)

这里先直接总结 vllm 模型推理涉及到的类，及它们之间的调用关系如下，贯穿这些类的是 execute_model 函数。

LLMEngine —> ExecutorBase —> WorkerBase —> ModelRunnerBase

LLMEngine 类

LLMEngine 类: vllm/engine/llm_engine.py

（涉及代码 2000 行）：用于管理大语言模型（LLM）的推理和生成过程。其中：

step 方法：是引擎的核心方法，每调用一次执行一次解码步骤。包括调度、执行模型、处理输出和清理任务。
_initialize_kv_caches 方法：初始化 KV 缓存，动态调整 GPU 和 CPU 缓存块数量。

step() 方法中真正调用模型推理的代码只有一行，通过调用 model_executor 类实例的 execute_model 方法执行模型推理。

outputs = self.model_executor.execute_model(execute_model_req=execute_model_req)

ExecutorBase 类

ExecutorBase 类: vllm/executor/executor_base.py

模型推理执行器, 针对不同的硬件平台 CPU/GPU/XPU 等，使用工厂方法注册得到不同平台的模型推理执行器类，基类不实现 execute_model 函数，每个平台的执行器类都有各自的实现。

execute_model 函数内部通过调用下述代码，实现模型推理，输入参数只有一个 execute_model_req。类型为 ExecuteModelRequest。初始化和执行模型推理代码如下所示:

def __init__(self, ):
    self.driver_worker = self._create_worker()

def execute_model(self, ):
    output = self.driver_worker.execute_model(execute_model_req)

_create_worker 函数

_create_worker() 函数由 ExecutorBase 类的初始化函数调用: vllm/executor/gpu_executor.py

_create_worker 通过 local_rank 和 rank 参数配置工作器（Worker），其本质上是通过调用 create_worker 函数，create_worker 函数的输入参数如下，函数输入参数通过另外一个类函数 _get_worker_module_and_class 获取。

worker_module_name: 工作器模块名字
worker_class_name：工作起类名
worker_class_fn：工作器函数

def create_worker(worker_module_name: str, worker_class_name: str,
                  worker_class_fn: Optional[Callable[[], Type[WorkerBase]]],
                  **kwargs):
    wrapper = WorkerWrapperBase(
        worker_module_name=worker_module_name,
        worker_class_name=worker_class_name,
        worker_class_fn=worker_class_fn,
    )
    wrapper.init_worker(**kwargs)
    return wrapper.worker

_get_worker_module_and_class 函数，支持根据当前的调度器配置（scheduler_config）和推测性配置（speculative_config）动态地选择适当的 Worker 模块和类。这些 Worker 模块负责执行实际的模型推理工作。函数实现代码如下所示:

def _get_worker_module_and_class(
        self) -> Tuple[str, str, Optional[Callable[[], Type[WorkerBase]]]]:
    worker_class_fn = None
    if self.scheduler_config.is_multi_step:
        worker_module_name = "vllm.worker.multi_step_worker"
        worker_class_name = "MultiStepWorker"
    elif self.speculative_config:
        worker_module_name = "vllm.spec_decode.spec_decode_worker"
        worker_class_name = "create_spec_worker"
    else:
        worker_module_name = "vllm.worker.worker"
        worker_class_name = "Worker"
    return (worker_module_name, worker_class_name, worker_class_fn)

_get_worker_module_and_class 函数逻辑说明

多步调度 (is_multi_step)：如果启用了多步调度，函数返回：
- 模块路径：”vllm.worker.multi_step_worker”
- 类名：”MultiStepWorker”
推测性解码 (speculative_config)：如果启用了推测性解码，函数返回：
- 模块路径：”vllm.spec_decode.spec_decode_worker”
- 类名：”create_spec_worker”
默认配置：如果既未启用多步调度，也未启用推测性解码，则返回默认的 Worker：
- 模块路径：”vllm.worker.worker”
- 类名：”Worker”
函数指针：
- 默认返回 None，表示无需额外的类处理逻辑。

总结：根据 scheduler_config 或 speculative_config 的状态，选择合适的 Worker 模块和类。

WorkerWrapperBase 类

WorkerWrapperBase 类： vllm/worker/worker_base.py

WorkerBase: 定义了所有 Worker 的抽象基类（接口），明确了每个 Worker 必须实现的功能。抽象方法包括设备初始化、缓存管理、模型执行等。

WorkerWrapperBase 类的关键方法 init_worker()：完成初始化 Worker 的功能，输入参数是前面 create_worker 的输入参数，支持自定义逻辑。init_worker 方法支持两种方式获取实际 Worker 类：

使用 worker_class_fn 函数返回 Worker 类。
importlib 动态导入模块，并从模块中获取 Worker 类。

其中动态导入模块的实现代码如下所示：

mod = importlib.import_module(self.worker_module_name)
worker_class = getattr(mod, self.worker_class_name)

上述代码是 Python 的动态导入机制，可以在运行时根据字符串形式的模块名称加载模块，而不是在代码编译时导入。getattr() 函数的作用是从一个对象（在这里是模块 mod）中获取指定名称的属性（在这里是 self.worker_class_name）。在这里的代码中，用于动态获取模块中的类。

LocalOrDistributedWorkerBase 类

LocalOrDistributedWorkerBase：vllm/worker/worker_base.py

抽象基类，用于在本地或分布式环境中运行模型推理任务。它通过以下功能实现模型的执行和数据广播：

元数据广播：在分布式环境中，主工作器（Driver Worker）将元数据和输入数据广播到其他工作器。辅助工作器从广播数据中提取输入。
输入准备和分布式支持：主工作器准备输入并广播。辅助工作器接收广播并提取输入。
推理执行：负责调用模型运行器（model_runner）来执行模型推理步骤。支持单步或多步推理。包含 SPMD（单程序多数据）模式执行
抽象方法定义：定义了 prepare_worker_input 和 execute_worker 等抽象方法，需子类实现本地逻辑。

真正的模型推理执行是通过调用 self.model_runner.execute_model 代码实现。

output = self.model_runner.execute_model(
    model_input=model_input,
    kv_caches=self.kv_cache[worker_input.virtual_engine]
    if self.kv_cache is not None else None,
    intermediate_tensors=intermediate_tensors,
    num_steps=num_steps,
    **kwargs,
)

self.model_runner 的初始化在各个不同 Worker 子类中实现。如 GPU 的 Worker 类的 init 初始化函数实现了对 self.model_runner 的赋值。

ModelRunnerClass: Type[GPUModelRunnerBase] = ModelRunner
if model_runner_cls is not None:
    ModelRunnerClass = model_runner_cls
elif model_config.task == "embedding":
    ModelRunnerClass = EmbeddingModelRunner
elif self.model_config.is_encoder_decoder:
    ModelRunnerClass = EncoderDecoderModelRunner
self.model_runner: GPUModelRunnerBase = ModelRunnerClass(
    vllm_config=self.vllm_config,
    kv_cache_dtype=self.cache_config.cache_dtype,
    is_driver_worker=is_driver_worker,
    **speculative_args,
)

Worker 类

Worker 类：vllm/worker/worker.py

默认的 GPU 工作器类，继承自 LocalOrDistributedWorkerBase。在 GPU 上执行模型（分区）的 worker 类。每个 worker 都与单个 GPU 相关联。worker 负责维护 KV 缓存并在 GPU 上执行模型。在分布式推理的情况下，每个 worker 都会被分配一个模型分区。其核心函数作用总结如下：

determine_num_available_blocks：分析 GPU 内存使用情况，确定可用的 KV 缓存块数量。并返回 num_gpu_blocks 和 num_cpu_blocks。
initialize_cache：分配 GPU 和 CPU 上的 KV 缓存，准备模型的计算缓存。
execute_worker：根据输入执行缓存操作，例如数据块的交换、复制等。

ModelRunner 类

ModelRunner：vllm/worker/model_runner.py

类定义：ModelRunner 是基于 GPU 的模型运行器，支持推理和采样。类包含 2 个成员变量：

_model_input_cls：定义模型输入的类型，ModelInputForGPUWithSamplingMetadata。
_builder_cls：定义构造模型输入的工具类，ModelInputForGPUBuilder。

核心函数：

make_model_input_from_broadcasted_tensor_dict：从广播的张量字典创建 ModelInputForGPUWithSamplingMetadata 实例。通过调用ModelInputForGPUWithSamplingMetadata 类的 from_broadcasted_tensor_dict 方法解析输入张量，并与注意力后端 (attn_backend) 配置绑定。
prepare_model_input：为推理准备模型输入，包括元数据和采样配置。按 prefill 和 decode 的顺序对输入数据进行批处理。使用 dataclasses.replace 方法更新模型输入。
execute_model：作用执行一次模型推理。推理执行流程包括：使用 CUDA 图 (decode_meta.use_cuda_graph) 或常规执行 —> 模型计算 logits —> 如果是最后一级流水线，使用 logits 采样下一个 token。最终返回的是采样输出 (SamplerOutput)。

调用 model.forward 模型推理的代码如下所示，这里的 model_executable 是应用了 cuda graph 技术后捕获 graph_runners[virtual_engine][graph_batch_size]。

with set_forward_context(model_input.attn_metadata):
    hidden_or_intermediate_states = model_executable(
        input_ids=model_input.input_tokens,
        positions=model_input.input_positions,
        kv_caches=kv_caches,
        attn_metadata=model_input.attn_metadata,
        intermediate_tensors=intermediate_tensors,
        **MultiModalKwargs.as_kwargs(multi_modal_kwargs,
                                        device=self.device),
        **seqlen_agnostic_kwargs)

self.graph_runners 的赋值是通过 GPUModelRunnerBase 类的 capture_model 函数捕获一系列不同 batch_size 的 capture_models。

至此，vllm 推理的各个类的调用关系和流程的简单分析完毕了，后续会继续优化细节。

参考资料

The first vLLM meetup

LLaVA 系列模型结构详解

2024-11-28

前言
LLaVA1
- ViT-L/14 模型结构
LLaVA1.5
- LLaVA-1.5-HD
LLaVA1.6（LLaVA-NeXT）
参考资料

前言

视觉语言模型 VIA 或者说多模态大模型 MLLM 架构通常都是: LLM + 视觉编码器 + 映射层的组合。英伟达发布的视觉语言模型 VILA 架构和训练流程如下图所示：

可以看出视觉语言模型的架构是由视觉 encoder、映射层和语言 decoder 组成。常见的视觉语言模型如下所示，模型架构都很相似。

VILA-1.5 (B/8B/13B/40B)
LLaVA(1.5,1.6) (7B-34B)
InternLM-XComposer2 (7B, 4khd-7B)
QWen-VL (7B)
DeepSeek-VL (7B)

LLaVA1

Llava1 的模型结构很简洁，CLIP 模型的视觉编码器 + 映射层 + LLM（Vicuna、LLama） ，利用 CLIP 模型的 Vison Encoder 结构对输入图片提取视觉特征，即转换为形状为 [N=1, grid_H x grid_W, hidden_dim] 的 feature map，然后通过一个映射层（线性层）将图像特征对齐到文本特征维度，即得到形状为 [N=1, grid_H x grid_W, embedding_dim] 的 image tokens embedding 向量，再然后将图片 tokens 向量和输入文本 tokens 向量 concat 后作为 LLM 的输入，生成回答文本。

LLaVA 模型架构如下图所示吧:

具体来说，对于输入图像 $X_v$，采用预训练 CLIP 模型的视觉编码器 ViT-L/14(224²)，其生成的视觉特征为 $Z_v = g(X_v)$，在作者的实验中，只用最后一个 Transformer 层之前和之后的网格特征。并使用一个简单的线性层将图像特征连接（映射）到词嵌入空间，通过一个可训练的投影矩阵 $W$ 将 $Z_v$ 转换为语言嵌入标记 $H_v$，$Z_v$ 向量的最后一个维度就是 LLM 的词嵌入空间维度 embedding_dim。

\[H_v = W\cdot X_v, with Z_v = g(X_v)\]

ViT-L/14 模型结构

ViT-L/14 模型的 L 表示模型的规模，为 “Large”，ViT(Vision Transformer) 模型有不同规模的模型，例如：

ViT-B（Base）：通常有 12 层 Transformer。
ViT-L（Large）：通常有 24 层 Transformer。
ViT-H（Huge）：通常有 32 层 Transformer。

ViT 会将输入图像分割成固定大小的 patch（例如 14x14），ViT-L/14 即表示 patch 大小为 14，

CLIP 模型的视觉部分使用 ViT 来编码图像的特征，文本部分使用 Transformer 来编码文本的特征。

不同版本 ViT 模型的参数总结：

模型版本	Transformer 层数	隐藏维度	参数量	Patch 分辨率
ViT-B/16	12	768	86M	16x16
ViT-L/16	24	1024	307M	16x16
ViT-L/14	24	1024	307M	14x14
ViT-H/14	32	1280	632M	14x14

LLaVA1.5

模型结构上和前作相比，LLaVA1.5 将之前用于维度映射的的简单一层线性层替换为 2 层线性层的 MLP 结构，并将 clip-L/14 的输入分辨率从 224*224 提升到 336*336，因为作者发现提高输入图像分辨率能够增强模型性能，LLM 换成了 Vicuna1.5（在 LLama2 上微调的模型）

LLaVA-1.5-HD

目前开源的 CLIP 视觉编码器的分辨率上限为 336*336，这意味着无法简单地替换视觉编码器来支持更高分辨率的图像。为了解决这个问题，论文探索了一种方法，既能让多模态语言模型（LMM）处理高分辨率图像，又能保持 LLaVA-1.5 的高效数据使用。

将输入图像划分为若干小块，每块的分辨率与视觉编码器原始训练时一致，然后分别对这些块进行编码。编码完成后，我们将这些块的特征图合并为目标分辨率的大特征图，并将其输入到 LLM 中。同时，为了给 LLM 提供全局上下文信息并减少图像分割、编码和合并操作带来的不良影响，我们还将一个经过下采样（resize）的图像特征连接到合并后的特征图中。

这样的设计允许我们处理任意分辨率的图像，同时保持 LLaVA-1.5 的数据高效性。这一新模型被作者命名为 LLaVA-1.5-HD。

LLaVA1.6（LLaVA-NeXT）

模型推理层面新的升级点在于，Vision Encoder 分辨率支持更大的分辨率，包括 672x672, 336x1344, 1344x336 几种分辨率的输入，并且支持通过图片裁切，编码，合并来实现，和前作一样的方法。毕竟，当提供高分辨率图像和保留细节的表征时，模型感知图像中复杂细节的能力会显著提高。它减少了面对低分辨率图像时的模型幻觉，即猜测想象的视觉内容。

参考资料

温度系数与 top-p 采样策略详解

2024-11-24

一、Temperature 温度系数作用
二、解码策略介绍
三、top-p 采样算法
- 3.1 top-p 采样算法步骤:
- 3.2 top-p 采样代码
参考资料

一、Temperature 温度系数作用

Temperature 采样的温度系数意义、公式和知识蒸馏很相似，结合 softmax 的公式，都是如下形式:

\[q_i = \frac{exp(z_i/T)}{\sum_j^K exp(z_j/T)}\]

当 $T$ 趋于无穷大时，输出概率分布将趋于均匀分布，概率为 $1/K$, 此时信息熵是最大的。反过来，$T$ 趋于0时，正确类别的概率接近 $1$，输出结果就是确定的，信息熵为 0，softmax 的效果与 argmax 差不多.

应用代码如下所示：

# logits 是 llm 推理输出, 形状为 [batch_size, seq_len, vocab_size]
probs = torch.softmax(logits[:, -1, :] / temperature, dim=-1)

代码详解：

logits[:, -1] 表示选择的是最后一个位置（seq_len 维度的最后一项）对应的 logits，形状变为 [batch_size, vocab_size]。因为在生成模型中的 prefill 阶段，我们只关心当前生成的最后一个 token 的分布。
temperature 作用是调整 logits 的分布，用于控制采样的随机性。总结就是，温度系数 $T$ 越大输出越平滑，结果越不确定，越小则越确定。具体来说，当 temperature < 1.0，分布会变得更加陡峭，更倾向于选择高概率的 token。temperature > 1.0，分布会变得更加平坦，增加随机性。

二、解码策略介绍

首先需要知道，LLM 的输出结果只是下一个 token 的概率分布 logits，即对下一个单词的预测概率张量，形状为 [batch_size, seq_len, logits]。而如何从概率分布中选择下一个单词，就是我要介绍的解码策略，也叫采样策略

解码策略里，常见的方法是是贪心策略，Top-K 采样和 Top_p 采样，这几个方法的不同点在于候选集的选择策略不同。

贪心策略取的是概率最大的 Top1 的样本作为候选项，也就是永远取概率最大的样本作为一个候选项，但这样只能保证是局部最优，也就是当前步是最优的，达不到全局最优。
Top-K 采样取的是概率的前 TopK 的样本作为候选项, 也就是每一步都保留有 K 个候选项，能在一定程度上保证全局最优。但 top-k 有个问题就是 k 取多少，是最优的，这个难以确定。
Top-p 采样，针对的就是 K 值难确定的问题，通过设定阈值 p, 根据候选集累积概率之和达到阈值 p，来选择候选项的个数，也叫核采样。

三、top-p 采样算法

Top-p 采样（也称为核采样，Nucleus Sampling）是一种用于自然语言生成模型的解码策略，旨在平衡生成文本的多样性和质量。核心思想是：在每一步生成 next_token 时，都从累积概率超过阈值 p 的tokens 集合中进行随机采样。具体操作是，每个时间步，按照 token出现的概率由高到底排序，当概率之和大于 top-p 的时候，就不考虑后面的低概率 tokens。

上图很好的展示了 Top-p 采样（Nucleus Sampling）的过程，可以分为两个步骤：

1，确定候选集：左图显示如何根据累积概率选择候选集：

每个单词（或 token）都有一个概率，例如：“United”: 12%，“Netherlands”: 2.7%，按照概率降序排列，逐步累加概率，直到累积概率达到阈值（例如 15%）。
一旦达到阈值，忽略其他概率更低的词（如 “Czech” 和 “U” 被排除）。
因此，此例中，候选集包括：“United” (12%)、“Netherlands” (2.7%)。

2，从候选集中采样: 右图显示如何从候选集中基于归一化概率进行采样：

候选集中的概率重新归一化。例如：
- United”: 原概率 12% 占候选集的 82%（12% / 15%）。
- “Netherlands”: 原概率 2.7% 占候选集的 18%（2.7% / 15%）。
根据归一化后的概率进行随机采样。最终生成的词可能是：“United”（较高概率）或 “Netherlands”（较低概率）。

很明显，top-p 采样方法可以动态调整候选词的数量，避免了固定数量候选词可能带来的问题。另外，可以发现，top_p 越小，则过滤掉的小概率 token 越多，采样时的可选项目就越少，生成结果的多样性也就越小。

3.1 top-p 采样算法步骤:

Top-p 采样的详细步骤：

概率排序：对模型在当前时间步生成的所有词汇的概率进行降序排序。
确定候选集：从排序后的词汇中，选择累积概率达到或超过设定阈值 p 的最小集合，记为 V_p。例如，若 p=0.9，则选择前几个词，使其概率之和至少为 0.9。
归一化概率：对候选集 V_p 中的词汇的概率进行重新归一化，使其和为 1。
随机采样：根据归一化后的概率分布，从候选集 V_p 中随机选择下一个生成的词。
token 索引映射：使用 torch.gather 函数将采样的索引映射回原始词汇表索引。

3.2 top-p 采样代码

top-p 采样代码详细解释：

def sample_top_p(probs, p):
    """
    执行 Top-p (Nucleus) 采样, 从概率分布中采样下一个词。
    参数：
        probs (torch.Tensor): 概率分布张量，形状为 `[batch_size, vocab_size]`。
        p (float): 累积概率阈值，取值范围在 0 到 1 之间。
    返回：
        torch.Tensor: 采样得到的词索引，形状为 `[batch_size, 1]`。

    说明：
        Top-p 采样算法: 选择概率累积和超过阈值 p 的最小集合，将这些词的概率重新归一化后进行采样。
    """
    # 对概率分布进行降序排序。probs_sort: 排序后的概率值，形状与 probs 相同。probs_idx: 排序后的索引，用于映射回原始词汇表。
    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
    # 计算排序后概率的累积和. 返回的 probs_sum 是累积概率分布。
    probs_sum = torch.cumsum(probs_sort, dim=-1)
    # 保留累积概率未超过阈值 p 的词汇的概率，其余词汇的概率被置为 0.0。
    mask = probs_sum - probs_sort > p # 创建掩码，对于每个位置，计算累积概率（不包括当前词）是否超过阈值 p。
    probs_sort[mask] = 0.0 # 将累积概率超过阈值 p 的词的概率置零。

    # 对剩余的概率重新归一化, 确保总和为 1。
    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
    # 从重新归一化的概率分布中采样下一个词. 返回的 next_token 是采样得到的词在排序后概率分布中的索引。
    next_token_sorted_idx = torch.multinomial(probs_sort, num_samples=1)
    # 在 probs_idx 的最后一维（dim=-1）中，使用 next_token_sorted_idx 作为索引，提取对应的值。沿着 dim=1（列）进行索引提取
    # NOTE: torch.gather 函数按照给定的索引张量 index，从输入张量中收集 (获取) 数据，并返回一个与索引张量形状一致的张量。
    next_token = torch.gather(probs_idx, -1, index = next_token_sorted_idx)
    
    return next_token # 返回采样得到的下一个词的索引

代码运行示例：

输入概率分布： probs = [0.1, 0.3, 0.4, 0.15, 0.05]
按降序排序： probs_sort = [0.4, 0.3, 0.15, 0.1, 0.05]
   原始索引：   probs_idx  = [2, 1, 3, 0, 4]
计算累积概率： probs_sum = [0.4, 0.7, 0.85, 0.95, 1.0]
根据 p=0.8 标记掩码： mask = [False, False, True, True, True]
将超出范围的概率置零： probs_sort = [0.4, 0.3, 0, 0, 0]
重新归一化： probs_sort = [0.5714, 0.4286, 0, 0, 0]
根据概率采样： next_token_index = 0
从原始索引还原： next_token = probs_idx[0] = 2

重点函数解释:

1，torch.gather 函数按照给定的索引张量 index，从输入张量中收集 (获取) 数据，并返回一个与索引张量形状一致的张量。

示例代码：

import torch

# 创建一个 3x4 的输入张量
input_tensor = torch.tensor([[10, 20, 30, 40],
                             [50, 60, 70, 80],
                             [90, 100, 110, 120]])

# 创建一个包含索引的张量
index_tensor = torch.tensor([[3, 2, 1, 0],
                             [0, 1, 2, 3],
                             [1, 0, 3, 2]])

# 沿着 dim=1（列）进行索引提取
output_tensor = torch.gather(input_tensor, dim=1, index=index_tensor)

print(output_tensor)

"""
程序运行后输出:
tensor([[ 40,  30,  20,  10],
        [ 50,  60,  70,  80],
        [100,  90, 120, 110]])
对于 input_tensor 的第二行 [50, 60, 70, 80]，index_tensor 的第二行 [0, 1, 2, 3] 指示提取顺序为第二行的 0 列、第二行的第 1 列、第二行的第 2 列、第二行的第 3 列，结果为 [50, 60, 70, 80]。
"""

2，torch.multinomial 用于从概率分布中抽取样本，支持带放回和不带放回两种方式。具体来说，它的功能是基于输入的概率权重进行采样。

函数签名如下:

torch.multinomial(input, num_samples, replacement=False, *, generator=None) -> LongTensor

参数解释：

input: 1D 或 2D 的张量，表示概率分布。它的值不需要是标准化的概率，但必须是非负的。如果是 2D 张量，每一行会被视为一个单独的分布。
num_samples: 需要采样的样本数量。
replacement: 是否是有放回采样。
- 如果为 True，可以多次采样同一个索引。
- 如果为 False，采样后不会重复选择。
generator（可选）: 控制采样随机性的生成器。

代码示例:

import torch

# 定义一个概率分布
probs = torch.tensor([0.1, 0.3, 0.6])

# 从分布中采样 1 个样本（不放回）
sample_idx = torch.multinomial(probs, num_samples=1)
print("采样的索引:", sample_idx.item())

# 输出采样的索引可能为 2，因为其对应概率最大

参考资料

vllm 优化之 PagedAttention 源码解读

2024-11-17

一 PagedAttention 内核
二 Paged(页表)原理分析
参考资料

PagedAttention 算法的原理可以参考我前面写的文章vllm优化技术速览。从源码的角度来看 PagedAttention，其实可以分为两部分:

PagedAttention Kernel 的实现，这里 v1 算法计算逻辑部分和标准 attention 差不多（v2 计算逻辑和 flashattentionv2 一致），但是 kv cache 的分配和管理使用了 kv cache 动态管理、存取优化技术。
PagedAttention 的对 kv cache 的内存分配管理技术（KV cache manager），之前的 kv cache 在 seq 这个维度都是固定为最大输入尺寸的 max_seq_len, 但实际单个请求可能消耗这么多内存，且在请求推理结束后也没有释放对应内存的功能，这必然会造成大量的内存浪费和碎片化。由此 PagedAttention 算法基于操作系统的 page table 思想构建了 block table 来动态分配管理 kv cache 内存，下图很好的展示了相比早前的其他推理框架，LLM 在显存利用率上极高。

这种 kv cache 内存管理分配的算法（思想）是可以应用到其他 llm 推理服务框架中。

一 PagedAttention 内核

1.1 主要函数

看一个文件代码之前，先快速过一下这个文件有哪些主要（模板）类或者函数，vllm 中 pagedattention 内核的实现 csrc/attention/attention_kernels.cu 文件中，其主要有以下模板函数。

template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
          bool IS_BLOCK_SPARSE,
          int PARTITION_SIZE = 0>  // Zero means no partitioning.
__device__ void paged_attention_kernel()

// v2 内核的算法实现逻辑对应的就是 Flashattentionv1
// Grid: (num_heads, num_seqs, 1).
template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
          bool IS_BLOCK_SPARSE>
__global__ void paged_attention_v1_kernel()

// v2 内核的算法实现逻辑对应的就是 Flashattentionv2, 因此并行度多了一个 kv cache seq！并行数量为 max_num_partitions.
// Grid: (num_heads, num_seqs, max_num_partitions).
template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE,
          int NUM_THREADS, vllm::Fp8KVCacheDataType KV_DTYPE,
          bool IS_BLOCK_SPARSE,
          int PARTITION_SIZE>
__global__ void paged_attention_v2_kernel()

// paged_attention_v1 内核的包装函数
template <typename T, typename CACHE_T, int BLOCK_SIZE,
          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
          int NUM_THREADS = 128>
void paged_attention_v1_launcher()

// paged_attention_v2 内核的包装函数
template <typename T, typename CACHE_T, int BLOCK_SIZE,
          vllm::Fp8KVCacheDataType KV_DTYPE, bool IS_BLOCK_SPARSE,
          int NUM_THREADS = 128, int PARTITION_SIZE = 512>
void paged_attention_v2_launcher()

// paged_attention_v1 对外提供的接口函数，也是生成 python 调用接口的函数形式，部分参数我省略了
void paged_attention_v1(
    torch::Tensor& out,    // [num_seqs, num_heads, head_size]
    torch::Tensor& query,  // [num_seqs, num_heads, head_size]
    torch::Tensor& key_cache,  // [num_blocks, num_heads, head_size/x, block_size, x]
    torch::Tensor& value_cache,       // [num_blocks, num_heads, head_size, block_size]
    int64_t num_kv_heads,  // [num_heads]
    double scale,
    torch::Tensor& block_tables,  // [num_seqs, max_num_blocks_per_seq]
    torch::Tensor& seq_lens,      // [num_seqs]
    int64_t block_size, int64_t max_seq_len,
)
// paged_attention_v2 对外提供的接口函数，参数和 v1 类似
void paged_attention_v2()

PagedAttention 本质上是集合了 kv cache 动态管理技术的优化版 flashattention。PagedAttention 的内核实现有两个版本 paged_attention_v1_kernel 和 paged_attention_v2_kernel，v1 改编自FasterTransformers 的 MHA 实现，适合长度小于 8192 或者 num_seqs * num_heads > 512 的情况，v2 是参考 FlashDecoding 方式进行实现，对 sequence 维度进行切分以增加并行粒度。

这两个版本的内核都是基于 paged_attention_kernel 内核通过输入不同参数来实现，v2 内核版本多了一个 kv cache seq 维度分区数量的参数，并行度层面多了 kv cache seq 分区的并行度！

flashattention 两种算法实现集成在一个内核里，这还是很考验作者工程功底的！

PagedAttention 内核的实现函数和常规 Attention 的实现相比最明显的就是多了 blocks 相关参数，以及 k_cache 的尺寸变成了 [num_blocks, num_kv_heads, head_size/x, block_size, x]，很明显了多了 num_blocks 和 block_size 维度（v_cache 变量也是），用于表示一个 seq 用多少个 blocks 存储，以及每个 block 存储多少个 tokens。

PagedAttention kernel 模板函数签名如下所示:

// Grid: (num_heads, num_seqs, 1).
// 这里为了方便阅读我删除了块稀疏的参数
template <typename scalar_t, typename cache_t, int HEAD_SIZE, int BLOCK_SIZE, int NUM_THREADS>
__global__ void paged_attention_v1_kernel(
    scalar_t* __restrict__ out,           // [num_seqs, num_heads, head_size]
    const scalar_t* __restrict__ q,       // [num_seqs, num_heads, head_size]
    const cache_t* __restrict__ k_cache,  // [num_blocks, num_kv_heads, head_size/x, block_size, x], 最后一个x 是 vectorize，一个thread fetch一个vector
    const cache_t* __restrict__ v_cache,  // [num_blocks, num_kv_heads, head_size, block_size]
    const int num_kv_heads,               // [num_heads]
    const float scale,
    const int* __restrict__ block_tables,  // [num_seqs, max_num_blocks_per_seq]
    const int* __restrict__ seq_lens,      // [num_seqs]
    const int max_num_blocks_per_seq,
    const float* __restrict__ alibi_slopes,  // [num_heads]
    const int q_stride, const int kv_block_stride, const int kv_head_stride,
    const float k_scale, const float v_scale) 

1.2 内核配置定义

先阅读 paged_attention_v1_kernel() 内核的调用（包装）函数 paged_attention_v1_launcher() 的内容来看 kernel 的配置如何。

可以看出 kernel 的 grid 和 block 配置如下所示，即分别定义了二维 grid 和一维 block 配置，其中每个 BLOCKS_X 处理一个 head，每个 BLOCKS_Y 处理一个 seq，每个 thread 处理最后一个维度 hidden_size 的计算。

dim3 grid(num_heads, num_seqs); // dim3 grid(BLOCKS_X, BLOCKS_Y)
dim3 block(NUM_THREADS);

知道了 kernel 的配置，我们再回过头去看 kernel 源码-paged_attention_kernel() 模板函数，按照 kernel 编写惯例，开头的代码依然是先计算全局线程 id 和偏移，只保留 v1 内核相关且注释后的代码如下所示：

// 用于对整数除法结果进行向上取整。
#define DIVIDE_ROUND_UP(a, b) (((a) + (b) - 1) / (b))

void paged_attention_kernel()
{
    /****** seq 维度的线程索引和 blocks 数量计算 ********/
    // 获取当前请求序列 seq 的索引，基于网格的 y 维度
    const int seq_idx = blockIdx.y;
    // 获取当前请求的实际长度，即 tokens 数量
    const int seq_len = seq_lens[seq_idx];
    // 计算序列被分割成多少个块，每块大小为 BLOCK_SIZE
    const int num_seq_blocks = DIVIDE_ROUND_UP(seq_len, BLOCK_SIZE);
    // 未启用分区处理（USE_PARTITIONING 为 false），起始块索引为 0。
    const int num_blocks = num_seq_blocks
    
    // 计算当前 block 块的起始和结束索引
    const start_block_idx = 0
    const end_block_idx = num_blocks
    // 计算当前 block 处理的令牌（token）范围
    const int start_token_idx = start_block_idx * BLOCK_SIZE
    const int end_token_idx = end_block_idx * BLOCK_SIZE

    // 计算 warp 数量
    constexpr int NUM_WARPS = NUM_THREADS / WARP_SIZE;
    // 获取当前线程在块内的索引
    const int thread_idx = threadIdx.x;
    // 计算当前线程所在的 warp 索引
    const int warp_idx = thread_idx / WARP_SIZE;
    // 计算当前线程在 warp 内的 lane 索引
    const int lane = thread_idx % WARP_SIZE;

    /****** num_heads 维度的 kv head 索引计算 ********/
    // 兼容了 GQA 技术的 kv head 地址计算
    // 获取当前查询头的索引，基于网格的 x 维度
    const int head_idx = blockIdx.x;
    const int num_heads = gridim.x;
    // 计算每个 Key/Value 头对应的查询头数, 看到没，GQA 是可以集成到 attention 内核里面的！
    const int num_queries_per_kv = num_heads / num_kv_heads;
    // 计算当前 Key/Value 头的索引
    const int kv_head_idx = head_idx / num_queries_per_kv

    /****** thread group 向量化load&store 相关代码, 不太理解 *****/
    // 定义线程组大小和数量
    constexpr int THREAD_GROUP_SIZE = MAX(WARP_SIZE / BLOCK_SIZE, 1);
    constexpr int NUM_THREAD_GROUPS =
        NUM_THREADS / THREAD_GROUP_SIZE;  // 确保 THREAD_GROUP_SIZE 能整除 NUM_THREADS
    assert(NUM_THREADS % THREAD_GROUP_SIZE == 0);

    // 定义向量类型，用于存储部分 Key 或 Query
    // 向量大小配置为线程组中的线程数 * 向量大小保证16字节对齐
    constexpr int VEC_SIZE = MAX(16 / (THREAD_GROUP_SIZE * sizeof(scalar_t)), 1);
    using K_vec = typename Vec<scalar_t, VEC_SIZE>::Type;
    using Q_vec = typename Vec<scalar_t, VEC_SIZE>::Type;
    using Quant_vec = typename Vec<cache_t, VEC_SIZE>::Type;

    // 每个线程处理的元素数
    constexpr int NUM_ELEMS_PER_THREAD = HEAD_SIZE / THREAD_GROUP_SIZE;
    
    // 每个线程处理的向量数
    constexpr int NUM_VECS_PER_THREAD = NUM_ELEMS_PER_THREAD / VEC_SIZE;

    // 计算线程组索引和偏移
    const int thread_group_idx = thread_idx / THREAD_GROUP_SIZE;
    const int thread_group_offset = thread_idx % THREAD_GROUP_SIZE;

    // 查询地址计算
    const scalar_t* q_ptr = q + seq_idx * q_stride + head_idx * HEAD_SIZE;
}

通过注释我们可以发现最前面代码的核心就是计算 seq、num_heads 维度的索引以及线程组索引和偏移。

1.3 基于 block_tables 读取 kv cache

这部分代码是真正属于 pagedattention 原创性的设计，即如何基于 block_tables 去 token 的 offset。

首先就是：获取当前序列 (seq_idx) 的块表。block_tables 是一个二维数组，每个序列有最多 max_num_blocks_per_seq 个块。通过 seq_idx * max_num_blocks_per_seq 计算当前序列的块表起始地址。
遍历 Key 块。
获取物理块号: block_table[block_idx]。
循环加载 Key 向量、计算点积和更新 qk_max。
- 最外层循环: 每个 warp 负责计算一个 block key，而每个 block key shape 为 [block_size, num_head, head_size]
- 第二层循环: 每个thread_group取一个key，即num_head个元素，计算QK dot

// x == THREAD_GROUP_SIZE * VEC_SIZE
// Each thread group fetches x elements from the key at a time.
constexpr int x = 16 / sizeof(cache_t);
// 获取当前序列的块表，确定要迭代的 Key 块。
// block_tables 是函数参数，形状为 [num_seqs, max_num_blocks_per_seq] 的二维数组，每个序列有最多 max_num_blocks_per_seq 个块。
// 通过 seq_idx * max_num_blocks_per_seq 计算当前序列的块表起始地址。
const int* block_table = block_tables + seq_idx * max_num_blocks_per_seq;

// 每个 warp 负责 blocksize * headsize 个元素
for (int block_idx = start_block_idx + warp_idx; block_idx < end_block_idx; block_idx += NUM_WARPS) {
    
    // 获取当前块的物理块号，并将其转换为 int64_t 以避免溢出。
    const int64_t physical_block_number = static_cast<int64_t>(block_table[block_idx]);

    // 加载一个 Key 向量到寄存器。
    // 每个线程组中的每个线程处理 Key 的不同部分。
    // 例如，如果线程组大小为 4，则组中的第一个线程处理 0, 4, 8,... 向量，第二个线程处理 1, 5, 9,... 向量，依此类推。
    for (int i = 0; i < NUM_TOKENS_PER_THREAD_GROUP; i++) {
        // 在当前 physical block 中找到当前 thread group 负责的局部 token id
        const int physical_block_offset =
            (thread_group_idx + i * WARP_SIZE) % BLOCK_SIZE;
        
        // 计算 token 在当前 seq 的所有 block 中的全局索引。
        const int token_idx = block_idx * BLOCK_SIZE + physical_block_offset;
        
        // 声明一个用于存储多个 Key 向量的数组。
        K_vec k_vecs[NUM_VECS_PER_THREAD];
        // 遍历每个向量，加载 Key 向量到 k_vecs 数组。
        // 根据上述 shape 算出当前 seq 的具体 k cache 的 block size 这一维度的 offset
        for (int j = 0; j < NUM_VECS_PER_THREAD; j++) {
            // k_cache.shape=[num_blocks, num_kv_heads, head_size/x, block_size, x]
            const cache_t* k_ptr =
                k_cache + physical_block_number * kv_block_stride +
                kv_head_idx * kv_head_stride + physical_block_offset * x;

            // 因为是向量化 LOAD，还需要计算出 vec 的全局id，和 vec 内元素的局部 offset
            const int vec_idx = thread_group_offset + j * THREAD_GROUP_SIZE;
            const int offset1 = (vec_idx * VEC_SIZE) / x;
            const int offset2 = (vec_idx * VEC_SIZE) % x;

            // 根据 Key 缓存的数据类型，加载并转换 Key 向量。
            if constexpr (KV_DTYPE == Fp8KVCacheDataType::kAuto) {
                // 直接加载 Key 向量。
                k_vecs[j] = *reinterpret_cast<const K_vec*>(
                    k_ptr + offset1 * BLOCK_SIZE * x + offset2);
            } else {
                // 从量化向量转换为 Key 向量。
                Quant_vec k_vec_quant = *reinterpret_cast<const Quant_vec*>(
                    k_ptr + offset1 * BLOCK_SIZE * x + offset2);
                k_vecs[j] = fp8::scaled_convert<K_vec, Quant_vec, KV_DTYPE>(
                    k_vec_quant, k_scale);
            }
        }

        // 计算查询与 Key 的点积。
        // 这包括线程组内的归约操作。
        float qk = scale * Qk_dot<scalar_t, THREAD_GROUP_SIZE>::dot(
            q_vecs[thread_group_offset], k_vecs);
        
        // 如果提供了 ALiBi 斜率，则添加偏置。
        qk += (alibi_slope != 0) ? alibi_slope * (token_idx - seq_len + 1) : 0;

        // 如果当前线程组偏移量为 0，则进行以下操作：
        if (thread_group_offset == 0) {
            // 计算当前令牌是否超出序列长度，用于掩码。
            const bool mask = token_idx >= seq_len;
            
            // 如果掩码为真，则将 logits 设为 0；否则，设为计算的 qk。
            logits[token_idx - start_token_idx] = mask ? 0.f : qk;
            
            // 更新查询与 Key 的最大值，用于后续的 softmax 计算。
            qk_max = mask ? qk_max : fmaxf(qk_max, qk);
        }
    }
}

后续的代码就是去更新 softmax 和姨同样的操作去加载 v, 然后再做 gemv（softmax(qk^t) * v），最终得到 attention 输出，这里不再分析具体算法逻辑。

二 Paged(页表)原理分析

这里的算法分析重点在于分析如何创建 block table、实现逻辑 table 和物理 table 的映射，以及如何针对每个 seq 动态分配相应数量的 block 用于存储 kv cache。

这块代码的实现是在 vllm 的请求调度层模块里，请求调度模块的作用是将服务接收到的请求进行状态管理，包括入队出队操作，并且将请求解析成推理引擎的输入，由 Worker 模块完成模型推理。

请求调度的核心源码在：vllm/core/scheduler.py 文件中。核心实现在 Scheduler 类中，这个类中的 _schedule 函数调用内部的相关完成具体的请求调度。另外，调度器中的下发 batch 请求到 Worker 模块中的相关函数，如 _schedule_prefills 函数会先调用 block_manager.can_allocate 函数判断是否有足够内存分配。而在初始化方法 __init__ 函数有 kv cache 显存块状态管理器的初始化。具体代码如下所示：

def __init__(
    self,
    scheduler_config: SchedulerConfig,
    cache_config: CacheConfig,
    lora_config: Optional[LoRAConfig],
    pipeline_parallel_size: int = 1,
    output_proc_callback: Optional[Callable] = None,
) -> None:
    self.scheduler_config = scheduler_config
    self.cache_config = cache_config
    # Note for LoRA scheduling: the current policy is extremely
    # simple and NOT fair. It can lead to starvation of some
    # LoRAs. This should be improved in the future.
    self.lora_config = lora_config

    version = "selfattn"
    if (self.scheduler_config.task == "embedding"
            or self.cache_config.is_attention_free):
        version = "placeholder"

    BlockSpaceManagerImpl = BlockSpaceManager.get_block_space_manager_class(
        version)

    # Create the block space manager.
    self.block_manager = BlockSpaceManagerImpl(
        block_size=self.cache_config.block_size,
        num_gpu_blocks=cache_config.num_gpu_blocks,
        num_cpu_blocks=cache_config.num_cpu_blocks,
        sliding_window=self.cache_config.sliding_window,
        enable_caching=self.cache_config.enable_prefix_caching)

初始化函数中定义的块管理器 block_manager 就是我们关心的，它是 KVCache 显存块状态管理器。用于分配、释放 KVCache 显存块以及状态更新，分配显存块时会返回显存块 id，用于 PagedAttention 计算时获取 KVCache 块显存地址。

值得注意的是，BlockManager（和调度器）实际上只负责管理页表（即管理逻辑块和每个 seq 到物理块的映射关系），实际的物理块中的数据不由它管理。这个实际上和 os 中的页表也差不多，BlockManager中的一个物理块就相当于页表中的一个PTE，而不是真实存放数据的物理块，实际进行内存分配的是 CacheEngine。

2.1 Block 管理相关类

BlockManager 相关类的包装关系对应文件: block_manager.py -> block_table.py -> naive_block.py

BlockTable

block_table.py 文件的 BlockTable 类将 tokens 序列映射到块列表 blocks 中，其中每个 block 代表序列一部分的连续内存分配。这些块由 DeviceAwareBlockAllocator 管理，它负责分配和释放这些逻辑块。

其中 SelfAttnBlockSpaceManager 类继承自 BlockSpaceManager，父类只负责定义接口，子类才负责具体的实现。

BlockTable 类最主要的函数是 allocate 用于将 tokens 序列映射相应物理内存块列表 blocks 中，具体物理内存块的分配是通过设备内存分配器的分配函数 ` self._allocator.allocate_immutable_blocks` 实现的。

class BlockTable:
    """管理特定序列的内存块的类。
    
    BlockTable 将一系列令牌映射到一组块中，每个块代表序列的一部分连续内存分配。这些块由 DeviceAwareBlockAllocator 管理，负责块的分配和释放。
    
    参数：
        block_size (int): 每个内存块可以存储的最大令牌数量。
        block_allocator (DeviceAwareBlockAllocator): 用于管理物理块内存的分配器。
        _blocks (Optional[List[Block]], optional): 可选的现有块列表，用于初始化 BlockTable。如果未提供，则创建一个空的 BlockTable。
        max_block_sliding_window (Optional[int], optional): 每个序列需要保留的最大块数。如果为 None，则保留所有块（例如，当不使用滑动窗口时）。至少应满足模型的滑动窗口大小。
    
    属性：
        _block_size (int): 每个内存块可以存储的最大令牌数量。
        _allocator (DeviceAwareBlockAllocator): 用于管理物理块内存的分配器。
        _blocks (Optional[List[Block]]): 由此 BlockTable 管理的逻辑块列表。
        _num_full_slots (int): 当前存储在块中的令牌数量。
    """

    def __init__(
        self,
        block_size: int,
        block_allocator: DeviceAwareBlockAllocator,
        _blocks: Optional[List[Block]] = None,
        max_block_sliding_window: Optional[int] = None,
    ):
        self._block_size = block_size  # 设置每个块的大小
        self._allocator = block_allocator  # 设置内存分配器
        if _blocks is None:
            _blocks = []
        self._blocks: BlockList = BlockList(_blocks)  # 初始化块列表

        self._max_block_sliding_window = max_block_sliding_window  # 设置滑动窗口的最大块数
        self._num_full_slots = self._get_num_token_ids()  # 获取当前存储的令牌数量


    def allocate(self,
                 token_ids: List[int],
                 device: Device = Device.GPU) -> None:
        """为给定的令牌序列分配内存块。
        
        此方法分配所需数量的块以存储给定的令牌序列。
        
        参数：
            token_ids (List[int]): 要存储的令牌 ID 序列。
            device (Device, optional): 要分配块的设备。默认为 GPU。
        """
        assert not self._is_allocated  # 确保尚未分配块
        assert token_ids  # 确保有令牌需要分配
        blocks = self._allocate_blocks_for_token_ids(prev_block=None,
                                                     token_ids=token_ids,
                                                     device=device)  # 分配块
        self.update(blocks)  # 更新块表
        self._num_full_slots = len(token_ids)  # 更新存储的令牌数量

    def _allocate_blocks_for_token_ids(self, prev_block: Optional[Block],
                                        token_ids: List[int],
                                        device: Device) -> List[Block]:
        """为给定的令牌 ID 分配内存块。
        
        参数：
            prev_block (Optional[Block]): 前一个块。如果没有，则为 None。
            token_ids (List[int]): 要存储的令牌 ID 列表。
            device (Device): 要分配块的设备。
        
        返回：
            List[Block]: 分配的块列表。
        """
        blocks: List[Block] = []

        block_token_ids = []
        tail_token_ids = []
        for cur_token_ids in chunk_list(token_ids, self._block_size):
            if len(cur_token_ids) == self._block_size:
                block_token_ids.append(cur_token_ids)
            else:
                tail_token_ids.append(cur_token_ids)

        if block_token_ids:
            blocks.extend(
                self._allocator.allocate_immutable_blocks(
                    prev_block, block_token_ids=block_token_ids,
                    device=device))  # 分配不可变块
            prev_block = blocks[-1]

        if tail_token_ids:
            assert len(tail_token_ids) == 1  # 仅有一个尾块
            cur_token_ids = tail_token_ids[0]

            block = self._allocator.allocate_mutable_block(
                prev_block=prev_block, device=device)  # 分配可变块
            block.append_token_ids(cur_token_ids)  # 追加令牌 ID

            blocks.append(block)

        return blocks
    
    def update(self, blocks: List[Block]) -> None:
        """重置块表为新提供的块（包括其对应的块 ID）。
        
        参数：
            blocks (List[Block]): 新分配的块列表。
        """
        self._blocks.update(blocks)

_allocate_blocks_for_token_ids 函数会据块大小 (block_size) 将令牌序列分割成多个块，并使用 DeviceAwareBlockAllocator 来分配这些块。具体函数流程总结如下:

创建空块列表: blocks: List[Block] = []，也是函数返回的结果
分配逻辑块列表，通过分割 token ids 实现，每个子逻辑 block 包含的内容实际是 token ids。
分配实际物理块列表: 调用 _ allocator.allocate_immutable_blocks 函数返回实际块列表，并拓展到逻辑块列表 blocks 中。
返回分配的块列表 blocks。

下述是一个测试示例，展示如何使用 _allocate_blocks_for_token_ids 方法：

def test_allocate_blocks_for_token_ids():
    # 初始化分配器和 BlockTable
    allocator = DeviceAwareBlockAllocator()
    block_size = 4
    block_table = BlockTable(block_size=block_size, block_allocator=allocator)

    # 定义测试令牌 ID
    token_ids = [1, 2, 3, 4, 5, 6, 7, 8, 9]

    # 分配块
    allocated_blocks = block_table._allocate_blocks_for_token_ids(
        prev_block=None,
        token_ids=token_ids,
        device=Device.GPU
    )

    # 更新块表
    block_table._blocks.update(allocated_blocks)
    block_table._num_full_slots = len(token_ids)

    # 打印分配结果
    print("分配的块数量:", len(allocated_blocks))
    for i, block in enumerate(allocated_blocks):
        print(f"块 {i}: Token IDs = {block.token_ids}, Mutable = {block.mutable}")

    # 预期：
    # - 前两个块完全填满（4个令牌）
    # - 最后一个块不完全填满（1个令牌，5）
    # 所以应有 3 个块
    assert len(allocated_blocks) == 3, "应分配 3 个块"
    assert allocated_blocks[0].token_ids == [1, 2, 3, 4], "第一个块令牌 ID 错误"
    assert allocated_blocks[1].token_ids == [5, 6, 7, 8], "第二个块令牌 ID 错误"
    assert allocated_blocks[2].token_ids == [9], "第三个块令牌 ID 错误"
    assert not allocated_blocks[0].mutable, "第一个块应为不可变块"
    assert not allocated_blocks[1].mutable, "第二个块应为不可变块"
    assert allocated_blocks[2].mutable, "第三个块应为可变块"

    print("测试通过！")

if __name__ == "__main__":
    test_allocate_blocks_for_token_ids()

运行上述测试代码，结果如下所示:

分配的块数量: 3
块 0: Token IDs = [1, 2, 3, 4], Mutable = False
块 1: Token IDs = [5, 6, 7, 8], Mutable = False
块 2: Token IDs = [9], Mutable = True
测试通过！

CpuGpuBlockAllocator 类

前面的分析内容我们知道，请求到逻辑表的分配是通过 CpuGpuBlockAllocator 类实现的。CpuGpuBlockAllocator 类是一个内存块分配器，能够在 CPU 和 GPU 内存中分配和管理内存块。它实现了 DeviceAwareBlockAllocator 基类的接口，提供了在多个设备（如 CPU 和 GPU）之间分配、释放、分叉（forking）和交换（swapping）内存块的功能。

class CpuGpuBlockAllocator(DeviceAwareBlockAllocator):
    """一个能够在 CPU 和 GPU 内存中分配块的块分配器。
    
    该类实现了 `DeviceAwareBlockAllocator` 接口，提供了在 CPU 和 GPU 设备上分配和管理内存块的功能。
    
    `CpuGpuBlockAllocator` 维护了独立的 CPU 和 GPU 内存块池，并允许在这些内存池之间进行分配、释放、分叉和交换操作。
    """
    
    @staticmethod
    def create(
        allocator_type: str,
        num_gpu_blocks: int,
        num_cpu_blocks: int,
        block_size: int,
    ) -> DeviceAwareBlockAllocator:
        """创建一个具有指定配置的 CpuGpuBlockAllocator 实例。
        
        这个静态方法根据提供的参数创建并返回一个 CpuGpuBlockAllocator 实例。它初始化了 CPU 和 GPU 块分配器，指定块的数量、块大小和分配器类型。
        
        参数：
            allocator_type (str): 用于 CPU 和 GPU 块的块分配器类型。目前支持的值为 "naive" 和 "prefix_caching"。
            num_gpu_blocks (int): 要为 GPU 内存分配的块数量。
            num_cpu_blocks (int): 要为 CPU 内存分配的块数量。
            block_size (int): 每个块的大小，以令牌数量表示。
        
        返回：
            DeviceAwareBlockAllocator: 一个具有指定配置的 CpuGpuBlockAllocator 实例。
        
        注意：
            - 块 ID 是连续分配的，GPU 块 ID 在前，CPU 块 ID 在后。
        """
        # 对于 HPU，块 ID 0 仅用于填充
        reserved_blocks = 1 if current_platform.is_hpu() else 0
        block_ids = list(
            range(reserved_blocks, num_gpu_blocks + num_cpu_blocks))
        num_gpu_blocks -= reserved_blocks
        gpu_block_ids = block_ids[:num_gpu_blocks]
        cpu_block_ids = block_ids[num_gpu_blocks:]
    
        # 根据 allocator_type 创建不同类型的块分配器
        if allocator_type == "naive":
            gpu_allocator: BlockAllocator = NaiveBlockAllocator(
                create_block=NaiveBlock,  # 创建不可变块的函数
                num_blocks=num_gpu_blocks,
                block_size=block_size,
                block_ids=gpu_block_ids,
            )
    
            cpu_allocator: BlockAllocator = NaiveBlockAllocator(
                create_block=NaiveBlock,  # 创建不可变块的函数
                num_blocks=num_cpu_blocks,
                block_size=block_size,
                block_ids=cpu_block_ids,
            )
        elif allocator_type == "prefix_caching":
            gpu_allocator = PrefixCachingBlockAllocator(
                num_blocks=num_gpu_blocks,
                block_size=block_size,
                block_ids=gpu_block_ids,
            )
    
            cpu_allocator = PrefixCachingBlockAllocator(
                num_blocks=num_cpu_blocks,
                block_size=block_size,
                block_ids=cpu_block_ids,
            )
        else:
            raise ValueError(f"未知的分配器类型 {allocator_type=}")
    
        return CpuGpuBlockAllocator(
            cpu_block_allocator=cpu_allocator,
            gpu_block_allocator=gpu_allocator,
        )
    
    def __init__(self, cpu_block_allocator: BlockAllocator,
                 gpu_block_allocator: BlockAllocator):
        """初始化 CpuGpuBlockAllocator 实例。
        
        参数：
            cpu_block_allocator (BlockAllocator): 用于管理 CPU 内存块的分配器。
            gpu_block_allocator (BlockAllocator): 用于管理 GPU 内存块的分配器。
        """
        # 确保 CPU 和 GPU 分配器的块 ID 没有交集
        assert not (
            cpu_block_allocator.all_block_ids
            & gpu_block_allocator.all_block_ids
        ), "CPU 和 GPU 块分配器的块 ID 不能有交集"
    
        # 将 CPU 和 GPU 分配器存储在字典中
        self._allocators = {
            Device.CPU: cpu_block_allocator,
            Device.GPU: gpu_block_allocator,
        }
    
        self._swap_mapping: Dict[int, int] = {}  # 记录交换操作的块 ID 映射
        self._null_block: Optional[Block] = None  # 用于存储空块
    
        # 记录每个块 ID 对应的分配器
        self._block_ids_to_allocator: Dict[int, BlockAllocator] = {}
        for _, allocator in self._allocators.items():
            for block_id in allocator.all_block_ids:
                self._block_ids_to_allocator[block_id] = allocator
    
    def allocate_mutable_block(self, prev_block: Optional[Block],
                               device: Device) -> Block:
        """在指定设备上分配一个新的可变块。
        
        参数：
            prev_block (Optional[Block]): 序列中的前一个块。用于前缀哈希。
            device (Device): 要分配新块的设备。
        
        返回：
            Block: 新分配的可变块。
        """
        return self._allocators[device].allocate_mutable_block(prev_block)
    
    def allocate_immutable_blocks(self, prev_block: Optional[Block],
                                  block_token_ids: List[List[int]],
                                  device: Device) -> List[Block]:
        """在指定设备上分配一组包含提供的块令牌 ID 的不可变块。
        
        参数：
            prev_block (Optional[Block]): 序列中的前一个块。用于前缀哈希。
            block_token_ids (List[int]): 要存储在新块中的块令牌 ID 列表。
            device (Device): 要分配新块的设备。
        
        返回：
            List[Block]: 新分配的包含提供的块令牌 ID 的不可变块列表。
        """
        return self._allocators[device].allocate_immutable_blocks(
            prev_block, block_token_ids)

NaiveBlockAllocator

上述类其实也还是一层包装，在看针对不同场景的类，以简单的 NaiveBlockAllocator 类为例分析。

class NaiveBlockAllocator(BlockAllocator):
    """一个简单的块分配器，不支持前缀缓存。
    
    该类实现了 `BlockAllocator` 接口，提供了基本的内存块分配和释放功能。
    
    参数：
        create_block (Block.Factory): 用于创建新块的工厂函数。当 NaiveBlockAllocator 被前缀缓存分配器组合使用时，必须能够创建前缀缓存块（但不应了解其余细节）。
        num_blocks (int): 要管理的块的总数量。
        block_size (int): 每个块的大小，以令牌数量表示。
        block_ids (Optional[Iterable[int]], optional): 可选的块 ID 可迭代对象。如果未提供，块 ID 将从 0 到 `num_blocks - 1` 顺序分配。
    """
    
    def __init__(
        self,
        create_block,
        num_blocks: int,
        block_size: int,
        block_ids: Optional[Iterable[int]] = None,
        block_pool: Optional[BlockPool] = None,
    ):
        if block_ids is None:
            block_ids = range(num_blocks)  # 如果未提供块 ID，则顺序分配
        
        self._free_block_indices: deque = deque(block_ids)  # 初始化自由块队列
        self._all_block_indices = frozenset(block_ids)  # 所有块 ID 的集合
        assert len(self._all_block_indices) == num_blocks, "块 ID 数量应与 num_blocks 相等"
    
        self._refcounter = RefCounter(
            all_block_indices=self._free_block_indices)
        self._block_size = block_size
    
        self._cow_tracker = CopyOnWriteTracker(
            refcounter=self._refcounter)
    
        if block_pool is None:
            extra_factor = 4
            # 预分配 "num_blocks * extra_factor" 个块对象。
            # "* extra_factor" 是为了允许分配比物理块更多的块对象
            self._block_pool = BlockPool(self._block_size, create_block, self,
                                         num_blocks * extra_factor)
        else:
            # 在这种情况下，块池由调用者提供，意味着可能需要在分配器之间共享块池
            self._block_pool = block_pool
    
    def allocate_immutable_blocks(
            self,
            prev_block: Optional[Block],
            block_token_ids: List[List[int]],
            device: Optional[str] = None) -> List[Block]:
        """分配一组新的不可变块，包含提供的块令牌 ID，并链接到前一个块。
        
        参数：
            prev_block (Optional[Block]): 序列中的前一个块。如果为 None，则分配的块为序列中的第一个块。
            block_token_ids (List[List[int]]): 要存储在新块中的块令牌 ID 列表。
            device (Optional[str], optional): 分配块的设备。对于 NaiveBlockAllocator，通常为 None。
        
        返回：
            List[Block]: 新分配的不可变块列表。
        """
        assert device is None, "NaiveBlockAllocator 不支持设备参数"
        num_blocks = len(block_token_ids)  # 需要分配的块数量
    
        block_ids = []
        for i in range(num_blocks):
            block_ids.append(self._allocate_block_id())  # 分配块 ID
    
        blocks = []
        for i in range(num_blocks):
            # 初始化块，设置前一个块、令牌 ID、块大小和物理块 ID
            prev_block = self._block_pool.init_block(
                prev_block=prev_block,
                token_ids=block_token_ids[i],
                block_size=self._block_size,
                physical_block_id=block_ids[i])
            blocks.append(prev_block)
    
        return blocks
    
    def allocate_mutable_block(self,
                               prev_block: Optional[Block],
                               device: Optional[str] = None) -> Block:
        """分配一个新的可变块，并链接到前一个块。
        
        参数：
            prev_block (Optional[Block]): 序列中的前一个块。如果为 None，则分配的块为序列中的第一个块。
            device (Optional[str], optional): 分配块的设备。对于 NaiveBlockAllocator，通常为 None。
        
        返回：
            Block: 新分配的可变块。
        """
        assert device is None, "NaiveBlockAllocator 不支持设备参数"
        block_id = self._allocate_block_id()  # 分配一个块 ID
        block = self._block_pool.init_block(prev_block=prev_block,
                                            token_ids=[],  # 初始化为空令牌 ID
                                            block_size=self._block_size,
                                            physical_block_id=block_id)
        block.mutable = True  # 设置为可变块
        return block

NaiveBlockAllocator 类的块函数也是调用 self._block_pool.init_block 接口，再跳转到 vllm/core/block/common.py 文件 BlockPool 类的 init_block 函数，其是通过 create_block: Block.Factory 类的相关接口来实现的。

也就是说 NaiveBlockAllocator、CpuGpuBlockAllocator 和 PrefixCachingBlockAllocator 三个 block 分配器类是为了适应不同设备和场景（简单场景和 PrefixCaching）而设计出来的，但真正的 block 定义类是通过块工厂创建函数 Block.Factory 创建得到。

BlockList 类

BlockList 类通过维护块及其对应的 ID 列表，优化对物理块 ID 的访问。提供方法来更新列表、向块添加令牌 ID，以及检索块或其 ID，避免在每次迭代块管理器时重新构建块 ID 列表。

class BlockList:
    """This class is an optimization to allow fast-access to physical 
    block ids. It maintains a block id list that is updated with the 
    block list and this avoids the need to reconstruct the block id 
    list on every iteration of the block manager
    """

    def __init__(self, blocks: List[Block]):
        self._blocks: List[Block] = []
        self._block_ids: List[int] = []

        self.update(blocks)

    def _add_block_id(self, block_id: Optional[BlockId]) -> None:
        assert block_id is not None
        self._block_ids.append(block_id)

    def _update_block_id(self, block_index: int,
                         new_block_id: Optional[BlockId]) -> None:
        assert new_block_id is not None
        self._block_ids[block_index] = new_block_id
    #####省略代码######

逻辑 block 管理类-SelfAttnBlockSpaceManager

SelfAttnBlockSpaceManager 类用于管理注意力机制中 KV（Key-Value）缓存块，主要负责逻辑内存块的分配、交换、以及其他高级功能如前缀缓存（即请求共享前缀systemp prompt）、分叉/写时复制（Forking/Copy-on-Write）和滑动窗口内存分配。

和前面几个是 block 模块内部类不同，它是对外部模块提供的类，但 BlockManager（和调度器）实际上只负责管理页表（即管理每个 seq 到物理块的映射关系），实际的物理块中的数据不由它管理。

先看构造函数 __init__()，函数中维护了一个逻辑 block_tables，它是一个字典，形式如 block_tables: Dict[SeqId, BlockTable] = {}，这个字典维护着整个 vllm 系统中每个 Sequence 实例到它的 block_table 之间的映射关系，方便快速查找到当前文本序列(Sequence) 对应的 PhysicalTokenBlock。构造函数的输入参数比较多，这里重点看三个参数的意义：

block_size: 每个内存块的大小，表示可以存储多少个令牌的 KV 数据。
num_gpu_blocks: 分配在 GPU 上的内存块数量。
num_cpu_blocks: 分配在 CPU 上的内存块数量。

__init__() 的部分代码如下所示:

SelfAttnBlockSpaceManager 类中内存块分配相关有 allocate 和 _allocate_sequence 函数，分别用于为给定的序列组分配所需的内存块和为单个序列分配块表。

块的分配策略由 CpuGpuBlockAllocator 来实现，分配策略有 naive 和 prefix_caching 两种。CpuGpuBlockAllocator 是一个管理 CPU 和 GPU 内存块的分配器类，继承自 DeviceAwareBlockAllocator，提供了在 CPU 和 GPU 设备上分配和管理内存块的功能。

块的分配策略类涉及到 4 个类，它们的类图关系如下所示：

classDiagram
    DeviceAwareBlockAllocator <-- CpuGpuBlockAllocator
    BlockAllocator <-- NaiveBlockAllocator
    class DeviceAwareBlockAllocator {
        +allocate_mutable_block()
        +allocate_immutable_block()
        +free()
    }
    class BlockAllocator {
        +allocate_mutable_block()
        +allocate_immutable_block()
        +free()
    }
    class CpuGpuBlockAllocator {
        -_allocators: Dict
        +create()
        +swap()
    }
    class NaiveBlockAllocator {
        -_free_blocks: List
        +allocate_block()
        +free()
    }

类的关系说明：

DeviceAwareBlockAllocator
- 顶层接口
- 定义设备相关的内存块分配方法
- CpuGpuBlockAllocator 实现此接口
- BlockAllocator
基础抽象类
- 定义基本的内存块分配操作
- NaiveBlockAllocator 继承此类
CpuGpuBlockAllocator
- 实现 DeviceAwareBlockAllocator 接口
- 内部使用 BlockAllocator 的具体实现(如 NaiveBlockAllocator)
- 管理 CPU 和 GPU 两种设备上的内存块
NaiveBlockAllocator
- BlockAllocator 的简单实现
- 被 CpuGpuBlockAllocator 使用来管理具体设备上的内存块

2.2 slot mapping

上一节讲的 block_tables 是逻辑层面的，而传给实际计算 kernel 的 block_tables 是形状为 [batch_size, max_blocks_per_seq] 的 torch.Tensor 表示每个序列的块地址列表，第一维表示序列 ID，第二维是物理块列表。

例如，[0, 1, 2] 表示 tokens 存储在 kv cache 的第 0、1 和 2 个块中。
每个块最多可容纳 block_size 个 tokens。
如果启用了 cuda-graph 捕获，则第二维将填充至 max_blocks_per_seq 的大小。

block_tables: Optional[torch.Tensor]

另外 vllm/attention/backends/utils.py 文件中提供了一些函数用于计算“槽映射”（slot mapping），并将序列中的 token 索引映射到内存块中的槽索引。

主函数 compute_slot_mapping，根据是否进行性能分析、是否需要填充以及使用哪种实现方式（Python 或 NumPy），计算序列的槽映射。

在模型 forward 过程中调用 flash_attention 做注意力分值计算时会按照 slot_mapping 指引位置将本层的 kv cache 存储到 vllm 初始化过程中分配的全零张量中，这在 cuda 函数中实现。

FlashAttentionMetadata 数据类的定义如下：

@dataclass
class FlashAttentionMetadata:
    # NOTE(sang): Definition of context_len, query_len, and seq_len.
    # |---------- N-1 iteration --------|
    # |---------------- N iteration ---------------------|
    # |- tokenA -|......................|-- newTokens ---|
    # |---------- context_len ----------|
    # |-------------------- seq_len ---------------------|
    #                                   |-- query_len ---|

    num_actual_tokens: int  # Number of tokens excluding padding.
    max_query_len: int
    query_start_loc: torch.Tensor
    max_seq_len: int
    seq_start_loc: torch.Tensor
    block_table: torch.Tensor
    slot_mapping: torch.Tensor

2.3 物理 block 分配类-CacheEngine

CacheEngine 给GPU分配空间的方式，本质上通过 pytorch 的接口在 gpu 上分配 num_blocks 大小的零 tensor 来作为物理块的空间的，而不是直接使用 cudaMalloc 进行操作的。

和 lightllm 的 tokenattention 直接提前分配一个可用最大形状的 Tensor，且后续 kv cache 的获取和释放都从这里操作不同。PagedAttention 为 transformer 模型的每个 layer 都分配一个可用最大尺寸的 tensor，并组合成列表的形式。也就是如果模型为 layer 为 16，对应的 kv cache 就是一个拥有 16 个 tensor 的列表。

代码通过循环遍历 self.num_attention_layers，为每个层分配独立的 KV 缓存张量，确保每层的 kv 张量能够被单独存储和访问，避免不同层之间的干扰。num_gpu_blocks 会通过 model_executor.determine_num_available_blocks 函数获取当前模型在指定设备上的每个 layer 的最大可用物理 blocks 数目。

绝大部分后端的 kv_cache_shape 形状都是 [2, num_blocks, block_size, num_kv_heads, head_size]。

# vllm/attention/ops/paged_attn.py
class PagedAttention:
    def get_kv_cache_shape(
        num_blocks: int,
        block_size: int,
        num_kv_heads: int,
        head_size: int,
    ) -> Tuple[int, ...]:
        return (2, num_blocks, block_size * num_kv_heads * head_size)

其中 block_size 由 llm 服务启动参数设定，而 num_gpu_blocks 既可以在服务启动的时候通过预热自动跑出来，也可通过服务启动参数 num_gpu_blocks_override 由用户自行设定并覆盖 num_gpu_blocks。

2.3.1 num_gpu_blocks 获取-determine_num_available_blocks 函数

determine_num_available_blocks 函数的具体实现是在 worker 目录下的各个设备的 work.py 实现，先以简单 cpu_work.py 的实现为例分析，cpu 中的 num_gpu_blocks（实际是 cpu 的可用内存块数量）计算是通过理论计算得到的，通过 cpu/gpu 设备可用的内存空间除以相关 kv_cache_block_size 得到可用 blocks 数量。

cache_block_size = self.get_cache_block_size_bytes()
# self.cache_config.cpu_kvcache_space_bytes 可通过环境变量 `VLLM_CPU_KVCACHE_SPACE` 定义
num_cpu_blocks = int(self.cache_config.cpu_kvcache_space_bytes // cache_block_size)

上述代码中的 get_cache_block_size_bytes 函数实际上是先计算对应模型 kv 每个 token 占用的空间，又因为 block_size 表示一个 block 对应 block_size 个 tokens，自然需要再乘以 block_size。get_cache_block_size 具体实现代码如下所示：

值得一提的是，指定设备的 kv cache 可用空间，以及可分配的 tokens 数量，也可以参考我的 llm_counts 工具理论计算得到，代码更优雅，使用更简单。

@staticmethod
def get_cache_block_size(
    block_size: int,
    cache_dtype: str,
    model_config: ModelConfig,
    parallel_config: ParallelConfig,
) -> int:
    head_size = model_config.get_head_size()
    num_heads = model_config.get_num_kv_heads(parallel_config)
    num_layers = model_config.get_num_layers(parallel_config)

    key_cache_block = block_size * num_heads * head_size
    value_cache_block = key_cache_block
    
    # 每层 layer 的 block size
    total = num_layers * (key_cache_block + value_cache_block)
    if cache_dtype == "auto":
        dtype = model_config.dtype
    else:
        dtype = STR_DTYPE_TO_TORCH_DTYPE[cache_dtype]
    dtype_size = torch.tensor([], dtype=dtype).element_size()
    return dtype_size * total

gpu_work.py 中的实现的计算逻辑和 cpu_work.py 一样，不同的是，这里借助了 torch.cuda.mem_get_info() 函数直接获取 gpu 总内存和在加载完模型之后的剩余显存。另外，用户在启动 llm 服务时，--block-size 参数可设置的取值范围是 {8,16,32,64,128}，默认值是 16。

@torch.inference_mode()
def determine_num_available_blocks(model_config, gpu_memory_utilization = 0.9) -> Tuple[int, int]:
    """
    评估模型的峰值内存使用情况，以确定在不发生内存溢出的情况下可以分配的 KV（键值）缓存块的数量。

    该方法首先清理 CUDA 缓存，然后使用虚拟输入执行一次前向传播，以评估模型的内存使用情况。
    接着，计算在剩余可用内存下，最多可以分配的 GPU 和 CPU 缓存块数量。

    提示：
        可以通过调整 `gpu_memory_utilization` 参数来限制 GPU 内存的使用。
    """
    # 清理 CUDA 缓存，以确保获取准确的内存使用信息
    torch.cuda.empty_cache()

    # 使用虚拟输入执行一次前向传播，以评估模型的内存使用情况
    self.model_runner.profile_run()

    # 同步 CUDA 操作，确保内存信息准确
    torch.cuda.synchronize()
    
    # 获取当前 GPU 的空闲内存和总内存（单位：字节）
    free_memory_pre_profile, total_gpu_memory = torch.cuda.mem_get_info()
    # 计算模型加载后的峰值内存使用量
    # Get the peak memory allocation recorded by torch
    peak_memory = torch.cuda.memory_stats()["allocated_bytes.all.peak"]
    
    # 清理未使用的缓存，计算非Torch分配的内存
    torch.cuda.empty_cache()
    torch_allocated_bytes = torch.cuda.memory_stats()["allocated_bytes.all.current"]

    total_allocated_bytes = torch.cuda.mem_get_info()[1] - torch.cuda.mem_get_info()[0]
    non_torch_allocations = total_allocated_bytes - torch_allocated_bytes
    
    if non_torch_allocations > 0:
        peak_memory += non_torch_allocations

    available_kv_cache_memory = (
        total_gpu_memory * gpu_memory_utilization -
        peak_memory)
    
    # 计算每个缓存块的大小
    cache_block_size = _get_cache_block_size(model_config)
    # 计算在剩余可用内存下，最多可以分配的 GPU 缓存块数量
    num_gpu_blocks = int(
        (total_gpu_memory * gpu_memory_utilization -
         peak_memory) // cache_block_size
    )
    # 确保缓存块数量不为负数
    num_gpu_blocks = max(num_gpu_blocks, 0)

    logger.info(
            "Memory profiling results: total_gpu_memory=%.2fGiB \n"
            " initial_memory_usage=%.2fGiB peak_torch_memory=%.2fGiB \n"
            " memory_usage_post_profile=%.2fGib \n"
            " non_torch_memory=%.2fGiB kv_cache_size=%.2fGiB \n"
            " gpu_memory_utilization=%.2f", total_gpu_memory / (1024**3),
            (total_gpu_memory - free_memory_pre_profile) / (1024**3),
            (peak_memory - non_torch_allocations) / (1024**3),
            total_allocated_bytes / (1024**3),
            non_torch_allocations / (1024**3),
            available_kv_cache_memory / (1024**3),
            gpu_memory_utilization)

    # 进行垃圾回收，释放未使用的内存
    gc.collect()
    # 再次清理 CUDA 缓存
    torch.cuda.empty_cache()
    # 返回可分配的 GPU 和 CPU 缓存块数量（此处 CPU 块数量为 0）

    return num_gpu_blocks, 0

总结：至此 vllm 的 pagedattention 内核设计和动态分配、管理 kv cache 内存的模块分析完毕，难点主要有三个：

cpu/gpu 设备在指定模型上的可分配的内存 blocks 的计算;
(不同设备的) block_tables 的创建和管理，以及基于 Scheduler 的调度策略去分配、释放和回收 blocks。
最后就是 pagedattention 内核代码中相关线程索引和偏移的计算怎么改成基于 block_tables 的形式，多了三个参数：k_cache, v_cache, block_table，其中 block_tables 尺寸为 [num_seqs, max_num_blocks_per_seq]，v_cache 尺寸为 [num_blocks, num_heads, head_size, block_size]。

以上，这都需要反复阅读理解代码才能得到清晰的理解。

harleyszhang Blog

Pytorch显存管理机制与显存占用分析方法

VGG

ResNet

Inceptionv3

Resnetv2

ResNeXt

Darknet53

DenseNet

CSPNet

VoVNet

一些结论

参考资料

DeepseekMoE 结构详解和代码实现

1. 基础 MOE 结构介绍

2. DeepseekMOE 结构介绍

2.1 Gate 网络与 DeepseekMOE 计算流程

3. DeepseekMOE 结构代码实现

3.1 DeepseekV2MLP 实现

3.2 门控/路由网络实现

3.3 DeepseekMOE 实现

参考资料

MLA 结构代码实现及优化

1. MLA 实现拆解

1.1 Q 向量计算

1.2 KV 向量计算

1.3 Self-Attention 计算

2. 标准 MLA 模块的代码实现

3 MLA 模块的代码优化-Projection Absorption

3.1 CC (CacheCompressed）

3.2 A_CC（AbsorbCacheCompressed）

参考资料

DeepSeekV2 论文解读

1. 介绍

2. 架构

2.1 多头潜变量注意力（MLA）：提升推理效率

2.1.1 Standard Multi-Head Attention

2.1.2 Low-Rank Key-Value Joint Compression

2.1.3 Decoupled Rotary Position Embedding

2.1.4 kv cache 大小的比较

2.1.5 总结

2.2 DeepSeekMoE：以经济成本训练强大的模型

2.2.2. 设备受限路由（Device-Limited Routing）

2.2.3. 负载均衡的辅助损失（Auxiliary Loss for Load Balance）

2.2.4. Token-Dropping 策略

参考资料

DeepSeekV3 简单概述

1. 介绍

2. 模型总结

参考资料

triton 内核编译流程解析

一 Triton 概述

二 Triton 编译（JIT）入口

2.1 jit() 函数

2.2 JITFunction 类

2.3 ASTSource 类

三 内核编译函数 compile

AST -> TritonIR

四 内核 compile 流程解析

Backend 系统

4.1 make_ttir

4.2 make_ttgir

4.3 make_llir

4.4 make_ptx

4.5 make_cubin

参考资料

vllm 推理流程剖析

LLMEngine 类

ExecutorBase 类

_create_worker 函数

WorkerWrapperBase 类

LocalOrDistributedWorkerBase 类

Worker 类

ModelRunner 类

参考资料

LLaVA 系列模型结构详解

前言

LLaVA1

ViT-L/14 模型结构

LLaVA1.5

三内核编译函数 compile

四内核 compile 流程解析