tiny-vllm正确性验证体系：PyTorch参考实现+softmax单元测试清单，让CUDA内核不跑偏-开发者社区

【免费下载链接】tiny-vllm

Build your own high performance LLM inference engine in C++ and CUDA - a smaller version of vLLM

项目地址：https://gitcode.com/gh_mirrors/ti/tiny-vllm

🧪 tiny-vllm 是一个用 C++ 和 CUDA 从零构建的高性能 LLM 推理引擎，是 vLLM 的"精简兄弟"。自研 CUDA 内核最怕的就是精度悄悄跑偏，而 tiny-vllm 的正确性验证体系给出的答案是：用 PyTorch 参考实现对标每一步、用 softmax 单元测试守住数值边界，再用端到端 token 校验做最终兜底。本文带你完整梳理这套验证方法，帮你给自己的 GPU 内核工程建立一张可照抄的"防跑偏清单"。

验证体系总览：5层防线各管什么

写推理引擎时，一次索引写错、一个 epsilon 漏加，都可能让输出从"Paris"变成乱码，而且中间没有任何报错。tiny-vllm 把验证手段直接沉淀在仓库里，分层覆盖：

验证层	关键文件	管什么
数值边界	tests/test_softmax.cu	softmax 求和为 1、无 NaN/Inf、平局/尖峰等边界
参考输出	python/reference.py	用 PyTorch 跑同模型，打印各阶段"标准答案"
层间对标	python/rms_norm.py、reference.txt	embedding → RMSNorm → 逐层 hidden state 逐值对比
端到端	full_test.sh、python/decode_test.py	固定 prompt 跑完整推理，核对最终 token
内存安全	check.sh	compute-sanitizer 抓越界读写、非法访问

上面这张图来自 README.md 对 cublasGemmEx 转置技巧的讲解：cuBLAS 默认按列优先读矩阵，而 LLM 权重按行优先存储，两者错位是矩阵乘法出错的经典源头——也正是验证体系需要重点盯防的地方。

softmax 单元测试：5个用例守住数值边界

tests/test_softmax.cu 针对的是 online softmax 中"分块部分结果合并"（fused）的逻辑——它只改变计算顺序、不改变数学。测试为 5 组精心挑选的输入行调用softmax()kernel，逐行验证两件事：求和 ≈ 1（容差 0.01）、没有 NaN/Inf：

用例	输入特征	压的是什么
`ordinary_8`	8 个普通数值	常规正确性
`tied_max_8`	8 个全相同的最大值	平局处理
`stability_spike_8`	一堆 0.001 里混一个 50.0	数值稳定性（防溢出）
`all_negative_8`	全负数	负值区间
`bigger_16`	16 元素混合行	不同行大小

文件头部注释还点出了验证标准：融合版与两遍法（先 max 后 sum）版本在相同输入下逐位一致，精确到最后一个 bfloat16 位（见 tests/test_softmax.cu）——这比"看起来差不多"严格得多。

PyTorch参考实现：用成熟工具链生成"标准答案"

核心思路很实用：与其手工推导每个算子该输出什么，不如用 PyTorch + Hugging Face Transformers 加载同一模型（Llama 3.2 1B Instruct），让成熟工具链跑出参考值，再和 C++ 引擎逐项对比。

一键生成全链路参考值

python python/reference.py "The capital of France is" --model meta-llama/Llama-3.2-1B-Instruct

python/reference.py 会依次打印：

Embedding（embed_tokens查表）
第 0 层input_layernorm的 RMSNorm 输出
Q / K / V 投影（q_proj、k_proj、v_proj）
逐层 hidden state（16 层每层入口的前 10 个值）
最终 RMSNorm 与预测 token ID

--num-values参数可调整每个 token 打印前多少个数值，方便对齐 C++ 侧的打印格式。

对标快照：不用重新跑 PyTorch 也能核对

仓库预存了两份"标准答案"文本：

reference.txt：reference.py对 "The capital of France is" 的完整输出，从 Token IDs 一路到Predicted token ID: 12366、Decoded: Paris（reference.txt），改完 C++ 后直接和它逐段比对即可。
python/rms_norm_crosscheck.txt：RMSNorm 专用对标，由 python/rms_norm.py 生成，打印每个 token 的 embedding 前 10 值、RMSNorm 输出前 10 值，以及input_layernorm的 norm 权重（python/rms_norm_crosscheck.txt）——权重也核对，才能排除"读错权重地址"这类错误。

端到端验证：一句"Paris"兜住整条链路

单元测试查得再细，也可能漏掉组合起来的误差。tiny-vllm 用固定 prompt 的完整推理做最终判定：

./full_test.sh

full_test.sh 把一段固定 token 序列（对应 "What is 2+2?" 的对话模板）管道输入引擎，期望最终生成 token4320，即答案 "4"。RMSNorm 的 epsilon、RoPE 的角度、causal mask、cublasGemmEx 的转置标志，任何一环错位都会让最后一个 token 变掉——这是最便宜的全链路守门员。

辅助工具还有：

python/decode_test.py：常用 prompt → token 序列的对照表（如 "Capital of France?" →791, 6864, 315, 9822, 374, 12366, 13），便于把引擎打印的 token ID 人工还原成可读文本。
python/batching_test_tokens.py：把 4 条 prompt 按 Llama 3.2 对话模板编码，生成可直接粘贴进 C++ 的push_back序列（含prompt_offsets、prompt_lengths），用于验证 continuous batching。
python/tokenizer.py：文本 ↔ token 编解码小工具，支持--decode模式，日常调试很顺手。

内存安全与性能检查：正确性的另一半

数值对不代表内存干净。check.sh 一行命令拉起 CUDA 官方内存检查器：

compute-sanitizer ./build/tiny-vllm

越界读写、未初始化访问、非法拷贝都会被精确到 kernel 和线程报出来——对自写 kernel 尤其重要。

性能侧的 ncu.sh（Nsight Compute）和 nsys.sh（Nsight Systems）虽属 profiling 工具，却常能顺带暴露"启动配置异常"这类隐性错误，排查时值得顺手跑一下。

正确性验证清单：照着做就行

把上面的手段串成一条可复现流程：

改完单个 kernel → 跑对应单元测试（如 tests/test_softmax.cu），先过数值边界
./check.sh→ compute-sanitizer 确认内存访问干净
python python/reference.py "The capital of France is"→ 与 reference.txt 逐段对标
./full_test.sh→ 端到端 token 校验收尾
需要性能数据时再上 ncu.sh / nsys.sh

单元测试（数值边界）→ 内存检查（访问安全）→ PyTorch 参考实现（逐值对标）→ 端到端（最终 token），四层防线层层兜底。对新手来说，这套组合的最大价值在于每一步都只需一条命令，成本极低却覆盖面完整。

写在最后

tiny-vllm 把"验证"当作和"实现"同等的工程内容来组织：python/ 目录放参考与对标脚本，tests/ 放单元测试，根目录的 shell 脚本承担各类验证入口。无论你要自研推理引擎、优化 GPU 内核，还是给深度学习项目补测试，这套"PyTorch 参考实现 + softmax 单元测试 + 端到端 token 校验"的清单都拿来即用。