编写 LLM 推理引擎

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

1. 你需要什么

一台 Apple Silicon 的 Mac，用于 Metal 后端 — 或者任何机器：sj.gpu 附带一个 CPU 参考后端，运行相同的代码（慢，但逐位有用 — 见 §8）。Linux/NVIDIA 上有一个 CUDA 后端。
一个 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 架构家族，带 F32/F16/Q4_K/Q6_K/Q8_0 张量（Q4_K_M 与 Q4_K_S 文件是甜点）。

2. 四十行推理器

import sj.gpu.GpuDevice;
import sj.gpu.Model;
import sj.gpu.KvCache;
import sj.gpu.Tokenizer;

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]);
        Tokenizer tk = Tokenizer.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。

Tokenizer — 由模型自带词表构建：fromGguf(m.gguf())，然后 encode(string) -> int[]、decode(int) -> string、eos()、bos()。

KvCache — 注意力的 key/value cache，按模型的 maxSeq() 定尺寸。每个对话创建一个；要重新开始就创建一个新的。两个 cache 共用一个 Model = 两个共享权重的独立会话。

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

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

4. 诚实地基准测试

两条来之不易的规则：

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

要测逐 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（无需重新构建）— 二分一个疑似误算或回归时第一个该拿的工具：

变量	效果
`SUPERJ_GPU=cpu`	强制 CPU 参考后端
`SUPERJ_GPU_PREFILL_F16=0`	禁用 f16 权重孪生
`SUPERJ_GPU_ROPE_FUSED=0`	取消旋转位置编码 / KV 写入与 GEMM 的融合
`SUPERJ_GPU_EMBED_BATCH=0`	逐 token 嵌入查找
`SUPERJ_GPU_REPLAY=0`	禁用 decode token replay
`SUPERJ_GPU_PROFILE2=1`	退出时满频逐 kernel GPU profile
`SUPERJ_GPU_CBSTATS=1`	退出时 dispatch/barrier 计数 + GPU 占用率

8. 正确性：CPU 预言机

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