编写 LLM 推理引擎

一份分步教程：加载一个量化 GGUF 模型，在 GPU 上运行它，并生成文本 — 大约只用四十行 SuperJ。本教程驱动的 sj.gpu 栈就是同一套引擎，它在同一台 Apple M4 Max 上运行 TinyLlama 比 llama.cpp 更快（prefill 5,753 对 5,624 tok/s，decode 快 1.24 倍，输出字节相同）。你以库的形式得到那个引擎；本页展示如何使用它，以及如何下探到它之下。

目录

1. 你需要什么
2. 四十行推理器
3. 五个类，逐一讲解
4. 诚实地基准测试
5. 下探：Tensor 层
6. 为什么它快
7. 旋钮与 kill switch
8. 正确性：CPU 预言机

1. 你需要什么

• 一台 Apple Silicon 的 Mac，用于 Metal 后端 — 或者任何机器都行：sj.gpu 附带一个 CPU 参考后端，运行完全相同的代码（慢，但逐位有用 — 见 §8）。Linux/NVIDIA 上存在一个 CUDA 后端，但它只运行 llama 家族模型 — 本教程的 TinyLlama 可以在上面工作；下文两段描述的 Gemma 家族则不行。那些模型所需的量化 matmul（matmul_q4_0/q5_0/q8_0）、gemma3 的 QK-norm（rmsnorm_heads）、gemma4 的比例 RoPE（rope_dual_ff_f32）以及整条 MoE 路径（moe_topk、逐专家 matmul、weighted_row_sum、带索引的 gather/scatter）都存在于 Metal 后端，而 sdk/gpu-kernels/cuda/kernels.cu 中没有它们，因此一个 Gemma 模型在那里无法派发哪怕一次前向过程。移植它们由 #4239 跟踪。
• 一个 GGUF 模型文件。 这里全程使用 TinyLlama 1.1B Q4_K_M（636 MB）：

curl -L -o tinyllama-q4k.gguf \
  https://huggingface.co/TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF/resolve/main/tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf

目前支持：llama 家族（Llama/Mistral/Qwen 形状的 GGUF）、稠密 Gemma 家族（gemma/gemma2/gemma3 — GeGLU、嵌入缩放、QK-norm、三明治归一化、带逐层 RoPE theta 的滑动窗口/全局交替注意力，全部从 general.architecture 与每个模型自己的 GGUF 头检测得到；#4041），以及 gemma4（#4043/#4045/#4046/#4047 — gemma2/3 的一切，外加逐层注意力形状、local/SWA 层上缩减的 nKvHead、经由共享 rope_freqs.weight 的比例 RoPE、无权重的 V-norm、完全省略 attn_v.weight 的“交替注意力”层（V 复用 K 的 pre-norm 投影），以及一个 128 专家 top-8 的 MoE FFN（融合的逐专家 gate_up/down、逐专家输出缩放）与一个稠密/共享 GeGLU 专家并存，通过一个 5-norm 三明治与逐层输出缩放组合起来）。gemma4 的跨层 KV-cache 共享与 per-layer-embeddings（AltUp）不受支持 — 若 checkpoint 头把其中任一项设为非零值，Model.load 会大声失败（对本次验证所针对的 gemma4:26b checkpoint，两者均为 0）。张量 dtype：F32/F16/Q4_K/Q6_K/Q8_0（Q4_K_M 与 Q4_K_S 文件是 llama 的甜点），外加 Q4_0/Q5_0 — 小型与 QAT 量化的 Gemma 发行版通常把它们用于整个权重矩阵，而不只是嵌入。

声明了非常大的 context_length 的模型（gemma4:26b：262144）若 KvCache.create 预分配其完整窗口，可能超出可用内存 — 参见 KvCache.create(dev, m, maxContext) 的 3 参重载，用它来限制可用上下文（对应 llama.cpp 自己的 --ctx-size）。

2. 四十行推理器

import sj.gpu.GpuDevice;
import sj.gpu.Model;
import sj.gpu.KvCache;
import sj.gpu.TextTokenizer;
import sj.gpu.Tokenizers;

public class tinychat {
    public static void main(String[] args) {
        if (args.length < 2) {
            System.out.println("usage: tinychat <model.gguf> <prompt> [maxTokens]");
            System.exit(1);
        }
        int maxTokens = args.length > 2 ? Integer.parseInt(args[2]) : 200;

        // 1. 打开一个 GPU 设备（macOS 上用 Metal；在任何别处回退到 CPU
        //    参考后端 — 相同代码，相同结果）。
        GpuDevice dev = GpuDevice.open();

        // 2. 加载模型。GGUF 文件以 mmap 零拷贝映射；量化
        //    权重在几毫秒内成为 GPU 张量。
        Model m = Model.load(dev, args[0]);
        TextTokenizer tk = Tokenizers.fromGguf(m.gguf());
        KvCache kv = KvCache.create(dev, m);

        // 3. Prefill：在一次批处理中将整个 prompt 过一遍模型。
        //    返回第一个生成的 token。
        int[] promptTokens = tk.encode(args[1]);
        int token = m.prefill(promptTokens, kv);

        // 4. Decode：一次一个 token。decodeStep 在位置 `pos` 处理 `token`
        //    并返回下一个 token（贪心 argmax）。
        for (int pos = promptTokens.length; pos < maxTokens; pos++) {
            token = m.decodeStep(token, pos, kv);
            if (token == tk.eos()) break;
            System.out.print(tk.decode(token));
        }
        System.out.println();
    }
}

编译并运行：

superj compile tinychat.sj --sdk-path "$SJ_HOME/sdk" --link --output tinychat
./tinychat tinyllama-q4k.gguf "The capital of France is" 100

这就是整个引擎。下面的一切都是理解与控制。

3. 五个类，逐一讲解

GpuDevice — 后端句柄。GpuDevice.open() 为这台机器挑选最佳后端（Metal → CUDA → CPU）；GpuDevice.open("cpu") 强制指定一个。它也是通往原始计算的入口：arena、缓冲区、kernel（§5）。

Model — 一个已加载的 transformer。Model.load(dev, path) 解析 GGUF 头、映射文件，并把每个权重包装为零拷贝张量视图 — 没有 600 MB 的读取；加载只需毫秒。内省：nLayers()、dim()、nHead()、nKvHead()、vocabSize()、maxSeq()。两个动词：

• prefill(int[] tokens, KvCache kv) -> int — 在一次批处理 GPU 过程中处理整个 prompt（仅当 prompt 超出临时内存预算时内部才分块）。填充 KV cache，返回第一个生成的 token。
• decodeStep(int token, int pos, KvCache kv) -> int — 一次自回归步骤：在位置 pos 处理 token，追加到 cache，返回贪心的下一个 token。

Tokenizers.fromGguf — 返回该文件真正需要的 TextTokenizer，然后是 encode(string) -> int[]、decode(int) -> string、isEos(int)、eos()、bos()。

不要用手挑实现。GGUF 要么是 SentencePiece（sj.gpu.Tokenizer — Gemma、Llama 1/2、Mistral），要么是字节级 BPE（sj.gpu.HfTokenizer — Llama 3、Qwen、2024 年之后的大多数发行版），而用错了并不会失败：它会把每个单词边界拆成字节回退 token，模型回答出流利的胡话（#4240）。Tokenizers.fromGguf 读取 tokenizer.ggml.model 并据此选择。

一旦超出裸补全 prompt 的范围，encode 有三件事就变得要紧（#4241）：

• 它会把词表中的控制 token — <start_of_turn>、<|turn>、<|eot_id|> — 解析成各自单独的 id，正是这一点让聊天模板成为可表达的。encode(text, false) 会关掉这一行为；对用户消息请用它，这样包含字面 turn 标记的内容就无法伪造出一个 turn。
• 它只在文件要求时才添加 BOS（tokenizer.ggml.add_bos_token）。gemma4 要求不加，因为它的模板自己写入 <bos>。
• 用 isEos(token)，永远不要用 token == eos()。生成结束是一个集合：gemma4 的是 [<eos>, <turn|>, <|tool_response>]，而一个聊天 turn 以 <turn|> 结束。只跟 eos() 比较会径直冲过 turn 边界。

demo/sj/demo/ChatDemo.sj 是这一切的一个约 100 行的完整示例 — 分词器选择、按家族的聊天组帧，以及在 turn 边界处停止。已针对 gemma2:2b 与 gemma3:1b 验证：

$ ./chatdemo gemma3-1b.gguf "What is the capital of France? Answer in one sentence."
backend=metal arch=gemma3 layers=26 vocab=262144
promptTokens=23

> What is the capital of France? Answer in one sentence.
Paris.

replyTokens=2 endedOnTurnBoundary=true

如今有它不能在上面工作的两个模型，原因都在分词器之外：gemma4（#4248 — 前向过程在数值上是错的）与任何 Llama-3.x GGUF（#4258 — 在 gemma4 之外 rope_freqs.weight 被忽略，因此 RoPE 以未缩放方式运行）。

KvCache — 注意力的 key/value cache，默认按模型的 maxSeq() 定尺寸。每个对话创建一个；要重新开始就创建一个新的。两个 cache 共用一个 Model = 两个共享权重的独立会话。KvCache.create(dev, m, maxContext) 则把可用上下文限制在模型声明的最大值之下 — 对 gemma4:26b（context_length=262144）这类模型是必需的：预分配完整窗口所需的内存远超大多数机器的拥有量（#4047）。

Gguf — 原始文件读取器（m.gguf()），如果你需要 Model 未暴露的头或张量。

采样：decodeStep 目前是贪心（argmax）。要做 temperature/top-k，你需要自己从 Tensor 层取 logits（§5）— 一个不错的第一个扩展项目。

4. 诚实地基准测试

两条来之不易的规则：

1. 计时前先预热。 Apple GPU 以低功耗状态启动，需要约 100 ms 的持续工作才能达到满频。一次冷 prefill 测的是时钟爬坡，而不是你的代码 — 读数会慢约 40%。先向一个临时 KvCache 做一次丢弃式 prefill，再计时真正的那个；如果首次调用延迟对你的应用重要，单独报告冷启动数字。
2. 同会话比较。 温度会改变每一个数字。把你的引擎和任何基线背靠背地、在同一台机器、在同一次运行中基准测试 — 一小时前在凉机器上测的基线不是比较。

要测逐 kernel 的 GPU 时间，有两个内置 profiler（环境变量，无需重新构建）：

• SUPERJ_GPU_PROFILE2=1 — 在满频下把每个 dispatch 放进各自的命令缓冲区，退出时打印逐 kernel 表。对大 kernel（matmul、attention）可信；它对大量微小 kernel 会高估，因为逐缓冲区开销主导了一个 13 µs 的 dispatch。
• SUPERJ_GPU_CBSTATS=1 — dispatch/barrier 计数，加上来自硬件时间戳的真实 GPU 忙时与占用率。

5. 下探：Tensor 层

Model 是公共原语之上的几百行普通 SuperJ — sj.gpu 源码随 SDK 发布，通读 Model 是学习整个前向过程的最快方式（用 superj doc --list / superj doc sj.gpu.Tensor 浏览 API）。要构建不同的架构或自定义采样，使用相同的构件：

GpuArena a = dev.scratch();                       // bump 分配器，每步重置
Tensor w  = Tensor.view(dev, buf, DType.Q4_K, rows, cols);  // 零拷贝权重
Tensor x  = Tensor.alloc(dev, DType.F32, n, dim);           // 激活

Tensor xn  = x.rmsNorm(normWeight);               // 融合归一化
Tensor q   = wq.matmul(xn);                       // GEMV（n==1）或分块 GEMM（n>1）
Tensor.ropeDualInPlace(q, k, ropeSpec, pos);      // 旋转位置编码，q+k 一次 dispatch
Tensor att = Tensor.attention(q, keys, vals, attnSpec, seqLen);  // flash attention
wo.matmulAccumInto(att, x);                       // x += Wo·att，融合残差
int next   = logits.argmax();                     // GPU argmax，4 字节回读

一切都按 batch size 自动 dispatch：一行激活走内存受限的 GEMV 路径，多行走分块 GEMM 路径。一个会咬每个人的坑：arena 在 SuperJ 中是保留字 — 不要用它命名变量。

6. 为什么它快

下面每一项都经过测量，并在落地时验证为输出相同：

• 零拷贝权重 — 在映射文件之上一个无拷贝的 GPU 缓冲区；加载约 5 ms。
• 批处理命令编码 + 滚动提交 — 每 token 一条命令流；GPU 执行早层的同时 CPU 编码后续层。
• GEMV/GEMM 分流 — decode 流式读取量化权重（内存受限，630 MB/token 墙）；prefill 跑一个分块 GEMM，对所有 prompt 行复用每个权重块（计算受限）。
• f16 权重孪生 — prefill 权重在加载时预反量化为 f16（TinyLlama 常驻约 2 GB），从 GEMM 内循环删除反量化：达到测得 15.3 TFLOP/s 硬件上限的 14.3。Decode 保留量化原件 — 它是带宽受限，是相反的权衡。
• 融合 epilogue — 旋转位置编码在 q/k GEMM 的 epilogue 内施加，而 k/v 投影直接写入 KV cache：整条流水线阶段（以及它们的同步停顿）消失了。
• 单遍 prefill + 因果跳过 — 一个 pp512 级 prompt 跑成一次批处理过程，attention kernel 跳过每个 query 块因果视界之外的 key/value 块，而不是计算后再掩码。
• Token replay — 稳态 decode 用两个打补丁的标量原生地重新编码每个 token 的整条 dispatch 流。

7. 旋钮与 kill switch

每个优化都有一个环境 kill switch（无需重新构建）— 二分一个疑似误算或回归时第一个该拿的工具：

8. 正确性：CPU 预言机

CPU 后端运行相同的 Model 代码，带标量 kernel，是真相基准：用 SUPERJ_GPU=cpu 运行任何生成并 diff 文本。GPU 路径被要求达到逐 token 相同的输出 — 其中有几项优化是构造上逐位相同：它们复制了所替换路径的精确浮点求值顺序。当你扩展引擎时，继承这条纪律：给每个改动一个 kill switch，对自身做 A/B，并与预言机 diff。快但错，只是错。