Building an LLM Inference Engine from ScratchPart IV / Module 9

Part IV · Building xinfer

Module 9：Phase 0 — 项目脚手架与 Backend Bring-up

从空目录开始：搭建 Cargo workspace，创建 D3D12 + DirectML device，跑通第一个 DML GEMM 和第一个 HLSL compute shader。

学习目标

理解为什么大型系统需要清晰的 workspace / crate 边界。
能解释 Phase 0 的最小成功标准：device 可创建、GEMM 可验证、custom kernel 可 dispatch。
理解 D3D12 + DirectML device creation 的基本步骤。
能写出第一个 smoke test：CPU reference vs GPU output。
完成 Lab 9：搭建脚手架并通过 GEMM / compute smoke tests。

9.1为什么先搭 Cargo workspace？

LLM 推理引擎同时包含数值计算、GPU backend、权重加载、tokenizer、runtime、CLI 与平台接入等职责。Phase 0 的首要任务是为这些职责建立稳定的 Cargo workspace 边界，而非扩大功能面。这样做可以使模型层、backend 层和用户入口各自演进，避免后续在同一组文件中交叉修改。

xinfer 的实际 workspace 包含 9 个 crate：xinfer-core、xinfer-backend、xinfer-dml、xinfer-loader、xinfer-model、xinfer-tokenizer、xinfer-runtime、xinfer-ffi、xinfer-cli。下表列出 Phase 0 直接涉及的主要边界及其后续职责：

Crate	Phase 0 中的角色	未来扩展
`xinfer-core`	dtype、shape、config、错误类型	CPU reference ops
`xinfer-backend`	HAL trait 草图	Device / Buffer abstraction
`xinfer-dml`	D3D12 + DirectML bring-up	kernels、copy queue、timer
`xinfer-model`	空壳	Qwen2 forward
`xinfer-loader`	空壳	safetensors / config
`xinfer-runtime`	空壳	KV cache + generation
`xinfer-cli`	`xinfer device`	generate 命令

工程原则

Phase 0 的验收重点是边界清楚、最小链路可运行。后续 phase 应在既有 crate 内扩展，而不是因早期边界混乱而重写目录结构。

图 9-1：Phase 0 的 workspace 不是“过度设计”，而是后续可维护性的基础。

9.2创建 D3D12 + DirectML device

GPU backend 的入口是 device creation。xinfer-dml 中的 DmlDevice::new(false) 通过 DXGI 创建 factory，枚举非 software adapter，用 D3D12CreateDevice 检查并创建 ID3D12Device，再创建 direct command queue。DirectML 的 IDMLDevice 由 DMLCreateDevice 建立在同一个 D3D12 device 之上。

Phase 0 的关键步骤可以概括为：

创建 DXGI factory；
枚举 hardware adapter，跳过 software adapter；
用 D3D12CreateDevice 测试并创建 D3D12 device；
创建 direct command queue；
调用 DMLCreateDevice 创建 DirectML device；
在 CLI 中打印 adapter 名称，作为 smoke test。

图 9-2：D3D12 + DirectML device creation 的对象链。

# Phase 0 的第一个用户可见 smoke test
cargo run -p xinfer-cli -- device

# 期望输出
D3D12 + DirectML device created.
Adapter: <name>

9.3第一个 DirectML GEMM 与第一个 HLSL kernel

Device 能创建只说明对象链可用，不能说明命令提交、资源绑定或数值计算正确。Phase 0 还需要两个计算 smoke tests：

DirectML GEMM：验证 operator create、compile、initialize、bind、dispatch 全链路。
Custom HLSL affine：验证 FXC 编译、root signature、root constants、SRV/UAV root descriptors 与 dispatch 路径。

GEMM 的 CPU reference

对矩阵 $A\in\mathbb{R}^{M\times K}$ 、 $B\in\mathbb{R}^{K\times N}$ ，GEMM 输出：

C_{i,j}=\sum_{p=0}^{K-1}A_{i,p}B_{p,j}

crates\xinfer-dml\tests\gemm_smoke.rs 使用 $M=4,K=3,N=5$ 的小矩阵，把 gemm_f32 的输出与三重循环 CPU reference 比较。GPU 输出必须与 CPU reference 在容差内一致：

\max_{i,j}|C^{\text{gpu}}_{i,j}-C^{\text{cpu}}_{i,j}| < \varepsilon

Affine kernel

自写 HLSL kernel 使用 shaders\affine.hlsl，计算下面的逐元素函数：

y_i = 2x_i + 1

该 kernel 的价值在于覆盖完整的 custom-kernel 路径：ComputeKernel::new 编译 shader 并创建 root signature 与 PSO；affine_kernel 上传输入、分配输出、记录 dispatch、等待 fence、下载结果；compute_smoke.rs 逐元素验证 out[i]=in[i]*2+1。后续 RMSNorm、RoPE、SwiGLU 与 attention 都沿用这条执行路径。

图 9-3：Phase 0 的三个门槛：能创建设备、能跑库算子、能跑自写 kernel。

9.4Smoke test 的设计原则

Phase 0 的测试规模应小，但每个测试必须覆盖一个明确风险：

测试	覆盖风险	失败时优先检查
`xinfer device`	D3D12/DML device creation	驱动、DXGI、DirectML runtime
`gemm_smoke`	DML operator 生命周期 + binding	tensor desc、binding table、resource state
`compute_smoke`	HLSL 编译 + root signature + dispatch	shader 编译、root 参数顺序、UAV 输出
CPU reference 对比	数值正确性	shape、row-major layout、转置约定

Phase 0 的核心心态

完整 transformer 依赖多个尚未验证的子系统。Phase 0 先建立最小 GPU 计算闭环：上传 → dispatch → readback → verify。这个闭环可靠后，复杂 kernel 的错误才能被定位到具体算子或数据布局。

Lab 9站起工作区，通过 GEMM 与 compute smoke tests

本实验要求从空目录或 starter branch 复现 Phase 0。验收标准是 workspace 可构建、device 命令可打印 adapter、GEMM 与 affine kernel 均能通过 CPU/GPU 对比。

创建 Cargo workspace 和主要 crate 目录。
实现 xinfer-core 中的 DType、Shape、错误类型。
在 xinfer-dml 中实现 D3D12 + DirectML device creation。
实现 CLI 命令 xinfer device。
实现一个 DirectML GEMM smoke test，与 CPU reference 对比。
实现一个自写 HLSL affine kernel smoke test。

# 最终验收命令
cargo build
cargo run -p xinfer-cli -- device
cargo test -p xinfer-dml --test gemm_smoke
cargo test -p xinfer-dml --test compute_smoke

小结

Module 9 建立 xinfer 的工程起点：Cargo workspace 划分 crate 职责，DmlDevice 完成 D3D12 + DirectML 对象链创建，DirectML GEMM 与 affine.hlsl 分别验证库算子路径和 custom compute path。Phase 0 的核心产物是可重复验证的 GPU 闭环；完整模型会在后续阶段建立在这个闭环之上。

思考与练习

基础为什么 Phase 0 不直接实现 transformer，而先实现 affine kernel？

因为要先打通最小可行管线：device 创建、buffer 分配/上传/下载、kernel 编译/dispatch、fence 同步和结果 readback。affine kernel（本阶段为 out[i]=2*in[i]+1）不引入复杂数值问题，适合确认整条 GPU 路径可用。若直接实现 transformer，错误来源会同时包含管线、shape、权重布局和模型逻辑。

基础GEMM 中 $M,K,N$ 分别代表什么？

$C_{M\times N}=A_{M\times K}\cdot B_{K\times N}$ 。 $M$ 是输出行数（如序列长度/batch）， $K$ 是被求和消去的内积维（contraction，如 in_features）， $N$ 是输出列数（如 out_features）。decode 时 $M=1$ （GEMV），prefill 时 $M=$ 序列长度。

进阶如果 CPU/GPU GEMM 结果不一致，列出三个可能原因。

① 行/列主序或转置约定不一致（A 或 B 的 layout、是否 $W^\top$ ）；② 累加顺序/精度差异（GPU 并行归约 + f16 输入导致与 CPU f32 串行结果有浮点误差——若超出容差才算 bug）；③ 索引/边界错误（越界、stride 算错、threadgroup 覆盖不全导致部分输出未写）。其他还有未同步（读了未完成的 UAV 写）、未初始化输出 buffer 等。

进阶解释为什么 readback 是 smoke test 的必要部分。

因为只有把 GPU 输出拷回 CPU 并和参考值比较，才能证明 kernel 真的算对了。dispatch 不报错不代表结果正确——可能根本没写、写错位置或数值错。readback（经 readback heap + fence 等待）是闭环验证的最后一环，没有它 smoke test 只验证了“能跑”而非“跑对”。

挑战把 affine kernel 改成 out[i]=a*in[i]+b，其中 a,b 来自 root constants。

HLSL：用 root constants 传 a、b（例如 cbuffer Params : register(b0) { float a; float b; } 或 32-bit root constants）。kernel：

uint i = DTid.x; if (i < count) outBuf[i] = a * inBuf[i] + b;

Rust 侧：root signature 增加一个 root constant 槽（2 个 float），dispatch 前用 SetComputeRoot32BitConstants 写入 a、b，再绑定 in/out UAV。相比把 a、b 放进 buffer，root constants 更轻量、无需额外 buffer 与 barrier，适合这种少量标量参数。