Building an LLM Inference Engine from ScratchPart 0

教程首页

教程导论：我们要从零构建什么？

项目式教材：从“LLM 是什么”走到一个能运行 Qwen2.5 的 Rust + DirectML 推理引擎。

学习目标

理解本教程最终产物：一个可读、可测试、可优化的 LLM 推理引擎 xinfer。
理解“无黑盒”哲学：为什么不用 ONNX Runtime / cuBLAS，而要亲手写 kernel。
能解释仓库的 crate 边界、依赖方向，以及各目录与教程模块的对应关系。
完成 Lab 0：构建工作区、检测 DirectML 设备、运行测试。

0.1我们正在构建什么，为什么值得构建？

大语言模型（Large Language Model, LLM）的基本行为可以精确描述：输入一串 token，输出一个长度等于词表大小的向量，其中每个分量表示对应 token 作为下一个 token 的得分。将选出的 token 追加到输入末尾并重新计算，即可逐 token 生成文本。

现成的推理 API 能够直接完成这一过程，但同时隐藏了其内部实现。本教程从零实现一个完整的推理引擎，覆盖 tokenizer 分词、Qwen2.5 权重加载、矩阵乘法在 GPU 上的计算，以及 KV cache 的作用；实现深入到 HLSL compute shader 一层，并使同一份核心代码既可运行于 Windows PC，也可运行于 Xbox 的 C++ host。该引擎即本教程的最终产物，命名为 xinfer。

定义：推理（Inference）

在本教程中，推理特指以下过程：加载训练完成的权重，对输入 token 执行一次前向计算，得到下一个 token 的分布。推理不涉及训练、梯度计算与优化器，关注点在于如何正确、高效且可移植地运行模型。

用一个公式概括语言模型：

P(x_{t+1}\mid x_1,\ldots,x_t) = \operatorname{softmax}\left(f_\theta(x_1,\ldots,x_t)\right)

其中 $x_i$ 是第 $i$ 个 token， $f_\theta$ 是由模型权重 $\theta$ 定义的 transformer 网络。推理引擎的任务就是高效计算这个 $f_\theta$ ，然后根据 logits 选择下一个 token。

图 0-1：LLM 推理不是“一次函数调用”，而是不断循环的数值计算流水线。

0.2“无黑盒”哲学：为什么不用现成运行时？

ONNX Runtime、TensorRT、llama.cpp、vLLM 等成熟运行时在生产环境中表现良好，但它们将最具学习价值的实现细节封装在 API 之后。调用一句 session.run() 即可得到结果，却难以回答以下问题：

一个 attention head 在计算 $QK^\top$ 时如何访问显存？
为什么 decode 阶段的矩阵乘是 GEMV，而 prefill 阶段是 GEMM？
为什么每个 token 都需要将约 608 KB 的 logits 从显存读回内存，且这一步足以拖慢整条流水线？
同一个 softmax 公式在 GPU 上为何涉及 threadgroup、barrier 与资源状态？

xinfer 采用相反的原则：核心计算路径不依赖现成库。矩阵乘法不调用 cuBLAS，算子不调用 ONNX Runtime，而是以 HLSL kernel 逐一实现；每个 kernel 均以一份 CPU 实现作为正确性基准（correctness oracle）进行对拍，数值一致后再继续。

工程现实

“无黑盒”并不等于“永远不用库”。本教程仍然使用 Rust、Direct3D 12、DirectML、HuggingFace tokenizer 等工具。区别在于：我们把 transformer 的核心计算路径 暴露出来，自己掌握它的正确性和性能。

0.3工具与环境：从第一天能跑起来开始

主 GPU backend 选用 Windows 上的 Direct3D 12 + DirectML，原因如下：D3D12 是 Windows 与 Xbox 共用的底层图形/计算 API，同一份代码可在两个平台复用；DirectML 的算子可直接使用 D3D12 资源；Rust 通过官方 windows crate 即可调用这些 COM 接口，无需额外的 C++ 胶水层。

组件	用途	本教程中的位置
Rust	系统语言、模块化、FFI	`crates/*`
D3D12	GPU 设备、buffer、command list、fence	`xinfer-dml`
HLSL	编写 compute shader	`shaders/*.hlsl`
DirectML	Windows ML operator API；也作为对比案例	`DmlOperator`
Qwen2.5	真实 decoder-only transformer checkpoint	`models/`（gitignored）
Xbox GDK	原生 host、device injection、部署路径	`xbox/`

你会反复看到两个量：吞吐量和显存占用。吞吐量通常写作：

\text{tokens/sec} = \frac{\text{生成的 token 数}}{\text{耗时（秒）}}

一个 f32 权重矩阵 $W\in\mathbb{R}^{m\times n}$ 的显存大约是：

\operatorname{bytes}(W) = m\cdot n\cdot 4

改用 f16 后，每个元素由 4 字节降为 2 字节，显存与带宽需求约减半。这也是 xinfer 将权重与 LM head 均存为 f16 的原因：Qwen2.5-0.5B 的全部权重因此约为 942 MiB，可完整驻留显存。

图 0-2：教程项目推进路线：先保证正确，再测量，再优化。

0.4如何阅读这个代码库：crate 图与依赖方向

xinfer 是一个 Cargo workspace。设计原则是：越靠近模型数学，越不应该依赖具体 GPU API；越靠近平台，越可以直接处理 D3D12、DirectML、Xbox GDK 这些细节。

图 0-3：依赖方向从应用层指向模型/运行时，再指向 backend；核心数学层不依赖具体平台。

阅读建议

第一次阅读代码库时，建议按教程顺序：xinfer-core::cpu → shaders/*.hlsl → xinfer-dml → xinfer-model → xinfer-runtime → xinfer-cli。

0.5Lab 0：把项目跑起来

Lab 0 的目标非常朴素：确认你的开发环境可以构建项目、找到 DirectML GPU、运行测试。

步骤 1：构建工作区

# 在仓库根目录执行
cargo build

步骤 2：确认 GPU backend 可用

cargo run -p xinfer-cli -- device

期望看到类似输出：

D3D12 + DirectML device created.
Adapter: AMD Radeon RX 9070 XT

步骤 3：运行测试

cargo test

本教程的基本工程纪律是：

\text{每次优化前后都必须满足：}\quad \lVert y_{\mathrm{gpu}} - y_{\mathrm{cpu}}\rVert_\infty < \varepsilon

图 0-4：真实项目历史中的优化结果：同一个模型、同一个问题，decode 从 3.6 tok/s 到 30.9 tok/s。

小结

本章建立了全局认识：最终目标是一个可运行 Qwen2.5 的推理引擎 xinfer；选择自行实现 kernel 而非依赖现成运行时，是为了暴露并掌握核心计算路径；crate 的划分与依赖方向遵循“上层依赖抽象、底层处理平台”的原则；Lab 0 通过 cargo build 与 cargo test 验证了开发环境。Transformer 内部算子的细节将从下一章 Part I 开始逐一展开。

思考与练习

基础用自己的话解释：为什么 LLM 推理可以被看作“反复计算下一 token 概率”？

语言模型本质上建模条件分布 $P(x_{t+1}\mid x_1,\ldots,x_t)$ ：给定已有 token，输出下一个 token 的概率。生成时，我们采样或取 argmax 得到一个新 token，把它接到序列末尾，再把更长的序列重新喂回模型预测下一个。如此循环，就把“一次预测”变成了“持续生成文本”。因此推理引擎的工作就是高效、反复地计算这个分布并选 token。

基础运行 cargo run -p xinfer-cli -- device，记录你的 GPU 名称。

命令会创建 D3D12 + DirectML device 并打印所选 adapter，例如：

D3D12 + DirectML device created.
Adapter: AMD Radeon RX 9070 XT

如果失败，通常说明没有 DirectML-capable GPU、驱动过旧，或缺少图形运行时——这正是 Lab 0 要先排除的环境问题。

进阶在仓库中找到 shaders/linear_f16.hlsl，读注释并写下它为什么使用 2D dispatch grid。

该 kernel 让一个 threadgroup 计算一个输出元素，因此需要的 group 数等于输出元素总数。LM head 的输出维度是词表大小 151936，远超过 D3D12 单个 dispatch 维度 65535 的上限。把 grid 拆成二维（idx = gy*grid_x + gx）就能覆盖超过 65535 的输出，而不会超出单轴限制。

进阶从图 0-3 中选择两个 crate，说明它们之间为什么应该是这个依赖方向，而不是反过来。

例如 xinfer-model 依赖 xinfer-backend，而不是反过来。模型层表达“需要什么能力”（分配 buffer、执行 kernel），backend 层表达“某个平台怎样实现”。让模型依赖抽象接口，可以在不改模型的情况下替换 backend；如果反过来让 backend 依赖具体模型，backend 就被绑死在 Qwen2 上，失去通用性。这就是“依赖指向稳定的抽象”。

挑战如果要把 backend 从 DirectML 换成 Vulkan，你认为哪些 crate 应该改？哪些 crate 不应该改？

应该改：xinfer-dml（替换为一个新的 xinfer-vulkan backend，实现同样的 HAL traits），以及 HLSL shader（需要改写或编译成 SPIR-V）。

不应该改：xinfer-core（纯数学/类型）、xinfer-backend（trait 定义本身）、xinfer-model、xinfer-runtime、xinfer-tokenizer。如果 HAL 设计得当，模型与 runtime 只依赖抽象 trait，因此理想情况下完全不需要修改。这正是引入 HAL 的价值。