环境准备#

来到 DeepSeek-V3.2-Exp 的仓库，简单浏览了一下目录结构，非常整洁清爽，没有多余的 submodule 负担。

快速翻阅 README：DeepSeek-V3.2-Exp 基于 V3.1-Terminus 打造，进一步优化了 token 的效率，直觉是应该又改进了稀疏化的效率，细读文字材料发现确实如此：推出了 DeepSeek Sparse Attention，旨在以更低的成本获得不逊色于 V3.1-Termius 的性能。

翻看 inference 目录，看到非常熟悉的几个源文件，这份源码此前在 TileLang 仓库中以例子的形式呈现：tilelang/examples/deepseek_v32/inference。说明未来 DeepSeek infra 的算子应该都会基于 TileLang 来编写，进一步精进 tiled-based HPC 是 AI infra 的必备技能。

查看 inference/requirements.txt，依旧是非常清爽的依赖库，没有指定具体 torch 版本，需要注意的是引入了一个额外优化的算子 Dao-AILab/fast-hadamard-transform，进一步移步检查，该算子主要实现三种精度下的 Multiply each row of x by the Hadamard transform matrix. 查看了该库的 wheel release，发现没有 torch>2.2、cu>12.2、python>=3.12 的分发包，因此无法避免 CUDA 编译。

使用之前创建的 TileLang 环境做基底：torch2.9.1cu12.8、tilelang 0.1.7、ninja 1.13.0、nvcc v12.8.93， 尝试直接 pip install -r requirements.txt 或 pip install fast-hadamard-transform 是会收到报错的，大概的意图好像是 fast-hadamard-transform 会现在 release wheels 里面找到预编译的轮子，找不到就自动下载源码进行编译，编译过程大概率会报错一个 import torch 但是找不到 torch 的错误，原因好像是不会使用（不会检查？）当前的 venv 进行插件编译。

我的解决方式是直接下载 fast-hadamard-transform 源码，pip install --no-build-isolation . 使用这个进行安装（不知道什么时候开始的，后续建议所有插件编译都带上 --no-build-isolation 参数），这样在我的 Debian 12 和 服务器的 22.04 LTS，可以直接安装 built cache。

初步研读#

自己研究神经网络模型的时候，跑通 forward pass 对我理解整个模型非常重要，单步调试 + debug console 可以非常清晰的了解整个模型的运作。 由于没有多卡的分布式设备，我要设法在单张 GPU 上完成这个过程，同时也不需要真的推理 671B 级别的 DeepSeek 模型，只需要把模型框架搞明白即可，因此不需要准备权重，要设法传入参数构造最小可运行模型。

所有的 Transformer 模块全都在一个 model.py 文件中，构造模型的参数类 ModelArgs 也均含有推荐的默认值，没有随意置空超参数，还标注好了分类， 加速算子全部在 kernel.py 中使用 TileLang 实现， 完整模型参数也全部在 config_671B_v3.2.json 可读， 甚至没有多余的 utils 代码， 非常清爽、整洁、干净。 更棒的是，而且在 main 函数入口就提供了可以即刻启动的 forward pass 的推理测试流程： inference/model.py

1
if __name__ == "__main__":
2
    torch.set_default_dtype(torch.bfloat16)
3
    torch.set_default_device("cuda")
4
    torch.manual_seed(0)
5
    args = ModelArgs()
6
    x = torch.randint(0, args.vocab_size, (2, 128))
7
    model = Transformer(args)
8
    print(model(x).size())

DeepSeek 这份代码是一个模型实现很棒的例子，不管给初学者、老手、高手，都有使用代码的自信和欲望，甚至可以在没有运行环境下都可以阅读和学习这份代码。

顺便骂一嘴。多年经验积累，我觉得某些自诩灵活的、模块化的网络模型工具箱，都有面向对象、设计模式的软件工程大病。 这些模型库喜欢把并不复杂的模型拆散成各种模块类，然后放到不同的源文件，模块一多自然就很细碎。 然后用所谓的工厂模式，把所有的模块参数全部聚合到配置文件里，有时还妄图在配置文件中去组织连通性，用不明所以的脚本去产生模型的各个模块。 静态阅读和调试起来非常灾难，尤其对我这种笨蛋和新手而言非常不友好。

它们都忽视了一个最重要的准则，神经网络本质就是清晰的张量数据流、计算图。 请时刻谨记：一个模型永远对应一个完整的输入到输出链路，它就应该是少量/单个源码上完整可读的、可调试的，而不应该使用配置文件去保证模型之间的连通性，配置文件分类组织好参数、管理好实验就够了。 多余的软工病，都是蠢笨的行为。

简言之，网络模型本来就是一进一出的事，那为何不把单个模型全部放到单个文件，形成一个模型对应一个模型文件、一个模型文件一个配置文件的超清晰模式呢？ 人脑这点空间本来就不擅长处理长记忆长连接，为什么非要组织结构上故弄玄虚呢？嘶……感觉好像想到某些图形渲染的库也有同样的毛病，也该骂。

继续阅读代码。

直接使用 python model.py 跑起来，果然 24GB 的显存直接炸掉了（好像需要使用 32.66GB 的显存），意料之中的错误。

进入 ModelArgs 类中找到参数说明 n_layers (int): Number of transformer layers.，默认值为27层，将层数降低至2层，保证重复又足够小。 此外我还将最大的单词表大小，从 102400 缩小到 2048，只是为了提升速度，后经实测不修改也可以。

1
args = ModelArgs()
2
args.n_layers = 2
3
args.vocab_size = 2048

继续运行起来，根据 ModelArgs 构造完 Transformer 模型本体后，直接进入模型类的 inference forward：

1
@torch.inference_mode()
2
def forward(self, tokens: torch.Tensor, start_pos: int = 0):
3
    """
4
    Forward pass for the Transformer model.
5
     Args:
6
        tokens (torch.Tensor): Input tensor of token IDs with shape (batch_size, seq_len).
7
        start_pos (int, optional): Starting position in the sequence for rotary embeddings. Defaults to 0.
8
     Returns:
9
        torch.Tensor: Logits tensor of shape (batch_size, vocab_size).
10
    """
11

12
    seqlen = tokens.size(1)
13
    freqs_cis = self.freqs_cis[start_pos:start_pos+seqlen]
14
    mask = torch.full((seqlen, seqlen), float("-inf"), device=tokens.device).triu_(1) if seqlen > 1 else None
15
    h, residual = self.embed(tokens), None
16

17

18
    for layer in self.layers:
19
        h, residual = layer(h, residual, start_pos, freqs_cis, mask)
20
    h, _ = self.norm(h, residual)
21

22

23
    logits = self.head(h[:, -1].float())
24

25

26
    if world_size > 1:
27
        all_logits = [torch.empty_like(logits) for _ in range(world_size)]
28
        dist.all_gather(all_logits, logits)
29
        logits = torch.cat(all_logits, dim=-1)
30

31
    return logits

继续执行，抛出了在执行到 MLA (Multi-head Latent Attention) 模块里面的 Indexer 模块时出现运行时错误 RuntimeError: kernel fp8_index_kernel_ input q_s ndim expected 3, but got 4, 这是第一个指向 TileLang fp8_index 算子的错误。

简单翻阅了 MLA 模块的 forward，看起来像是做完 KV cache 之后有一个对 cache 的评分，然后筛选出评分最高的部分进行后续计算。 后续询问 AI 得知，这是在实现 low-rank 压缩之后，动态筛选高相关度 tokens 的做法，以达到更稀疏、高效的相关度标记，这就是传说中 DeepSeek Sparse Attention 的关键算子。 此外，注释里也说明 # we use fp8 kv cache in actual deployment, so here we simulate the precision by casting kv to fp8 and then back to bf16. 实际采用 FP8 低精度部署的 KV cache，对应后续的 fp8_indexer TileLang kernel，或者说这是一个 index score applier。

算子调用位于：index_score = fp8_index(...)

算子错误陷于：fp8_index_kernel(q.shape[2], q.shape[3])(q, q_s, k, k_s)

在 kernel.py 源码中，看到了熟悉的 TileLang JIT kernel 的双括号单行启动：fp8_index_kernel(q.shape[2], q.shape[3]) 通过传入 tile shape 相关的参数构造 kernel，(q, q_s, k, k_s) 则是启动 kernel 传入的计算数据。通过字母猜测，q 和 q_s 可能是评分相关的张量，然后应用各种奇妙的矩阵乘和 top-k max 计算在 key 上的评分。

因为 tile-based 的 GPU 编程模型中 shape 是 layout 最关键的构造参数之一，TileLang 也是如此，后续参与计算的 tensor shape 肯定是会定义清楚的。 所以为了修复常见的 shape 不对的 runtime error，需要继续阅读 kernel 的实现。 从 fp8_index(...) 调用处继续跟踪，就看到了真正的实现，也就是装饰有 @T.prim_func 的函数。

1
@T.prim_func
2
def fp8_index_kernel_(
3
    q: T.Tensor[(b, m, h, d), FP8],
4
    q_s: T.Tensor[(b, m, h), FP32],
5
    k: T.Tensor[(b, n, d), FP8],
6
    k_s: T.Tensor[(b, n), FP32],
7
    o: T.Tensor[(b, m, n), FP32],
8
) -> None:

这样就清晰了，甚至根据这个 shape 都能大概猜到 GEMM 会怎么做，返回到 fp8_kernel 的调用，通过 debug console 检查所有 q、k 的形状。 q_s 和 k_s shape 第四维是 1，需要将其去掉才能对上 kernel 的 tensor shape 定义。

1
q.shape = torch.Size([2, 128, 64, 128])
2
q_s.shape = torch.Size([2, 128, 64, 1])
3
k.shape = torch.Size([2, 128, 128])
4
k_s.shape = torch.Size([2, 128, 1])

为了保证算子的源码的整洁，我不打算在 kernel.py 的调用处进行，而是将 model.py Indexer 模块部分做修改。

1
index_score = fp8_index(q_fp8.contiguous(), weights, self.k_cache[:bsz, :end_pos].contiguous(), self.k_scale_cache[:bsz, :end_pos].contiguous())
2
index_score = fp8_index(q=q_fp8.contiguous(),
3
                        q_s=weights.squeeze(-1).contiguous(),
4
                        k=self.k_cache[:bsz, :end_pos].contiguous(),
5
                        k_s=self.k_scale_cache[:bsz, :end_pos].squeeze(-1).contiguous())

重新运行，Indexer 模块又炸了 ValueError: Default process group has not been initialized, please make sure to call init_process_group.，这次是要将 topk indices broadcast 到分布式组里面。

1
import torch.distributed as dist
2
# ...
3
dist.broadcast(topk_indices_, src=0)

嗯，可以想到，这是一个在所有的多 GPU 切片上共享相同的评分索引（top-k indicies），保证后续模型在多 GPU 上计算的一致性， 这也是 Transformer 模块的 inference forward 最后一步要从各个分布式节点上 gather 最终结果的原因。

于是，用 DeepSeek 写了一个 init 的片段，运行之后警告没有 destroy，所以我自己补了一个，用这两个片段包住主函数内的代码。 其中 world_size、rank、block_size 的设置是源码默认提供的，不做修改，我曾尝试使用 torchrun 也跑不起来（好像我用 torchrun 就没跑起来过任何东西）。

1
# DeepSeek-V3.2-Exp pre-settings
2
world_size = 1
3
rank = 0
4
block_size = 128
5

6
# Torch distributed init
7
if not dist.is_initialized():
8
    dist.init_process_group(
9
        backend='nccl' if torch.cuda.is_available() else 'gloo',
10
        init_method='tcp://127.0.0.1:23456',
11
        world_size=world_size,
12
        rank=rank
13
    )
14

15
# Torch distributed destroy
16
if dist.is_initialized():
17
    dist.destroy_process_group()

终于跑通了，结果应该非常清爽，没有任何错误和警告提示，对应输入和输出的 tensor shape，如下：

1
torch.Size([2, 128])
2
torch.Size([2, 2048])

至此，就完成了 DeepSeek-V3.2-Exp 模型的修改，可以完美的进行 forward pass。若要继续研究和实验，这会是一个非常棒的开始。