mojo-dllm - 用 Mojo 写的扩散语言模型本地推理运行时

mojo-dllm 是一个用 Mojo 编写的扩散语言模型本地推理项目。据项目说明，它可以从量化 GGUF 文件运行 LLaDA-8B 和 Dream-7B 两个扩散语言模型，既能在 CPU 上跑，也能选用 NVIDIA GPU；GGUF 读取器、分词器、量化矩阵乘法、GPU kernel、transformer 与去噪循环全部用 Mojo 实现，运行时不需要 Python。它适合希望在本地跑扩散式（非自回归）语言模型、又不想引入 Python 运行时依赖的开发者，也适合关注 Mojo 语言在推理场景中实际落地情况的人。

项目自述解释了这类模型的运行特征：扩散模型不从左到右逐字生成，而是从一块被掩码的 token 画布出发，在固定步数内逐步确认最有把握的 token，直到没有剩余掩码；每一步都要在整个画布上重跑一遍模型，因此运行时问题看起来像一串 GEMM。mojo-dllm 的设计正是围绕这一点展开的。

维护者与状态

项目由 vladimirrott 维护，仓库创建于 2026-10-04，最后一次提交在 2026-10-29。最新 release 为 v0.1.0，发布于 2026-10-04。截至 2026-11-01，仓库星标数为 206。项目另有一个官网页面，用于承载文档与说明。

主要能力

以下内容均据项目自述整理，未经本站独立验证：

• 纯 Mojo 实现：从 GGUF 解析、分词、量化矩阵乘法到 GPU kernel、transformer 与去噪循环，全部由 Mojo 编写，运行时不需要 Python，GPU 为可选项。
• 两个模型支持：可运行 LLaDA-8B 与 Dream-7B，读取量化 GGUF 文件；-m 参数接受 .gguf 文件路径，或形如 owner/repo[:QUANT] 的 Hugging Face 仓库标识。
• 为窄 GEMM 重排权重：加载时把每个 Q4_K 与 Q6_K 矩阵按每 256 位寄存器 8 行的方式交错排列，使一条 AVX-VNNI vpdpbusd 指令即可对 int8 激活完成 8 行 × 4 次乘加；Q4_K 仍保持 4 位。
• 适配混合架构 CPU 的线程调度：任务通过原子计数器逐项分发，性能核自然承担更多工作。
• GPU 上的 int8 张量核心：使用 --device gpu 时，各投影从 GPU 显存读取同一批 Q4_K / Q6_K 块并在张量核心上相乘；注意力在 f16 张量核心上运行，每个 warp 处理 16 个 query，并通过共享内存与在线 softmax 流式读取 key。
• 提前停止：每一步只填充被掩码的位置，因此当完整答案后出现文本结束 token 后，后续步骤无法再改变已输出的文本，运行时便在此停止，输出保持一致。
• 只在采样器需要的位置计算 logits：一步只需要当前块中被掩码位置的 logits，因此 126 464 维的输出投影只在这些行上执行。
• 各模型各自的采样器：LLaDA 的分块与低置信度重掩码遵循其 generate.py，Dream 的时间步调度与熵排序遵循其 diffusion_generate。

项目还自述了一套正确性校验方式：解码器与 GEMM 对照 gguf 包的反量化结果；2 层模型的前向对照 numpy；两个分词器各用 36 个多语言用例与 llama-tokenize 逐 id 比对；两个真实模型上的 logits 与生成 token 对照 llama.cpp；GPU kernel 也逐项校验。

使用入口

仓库地址：vladimirrott/mojo-dllm；文档与说明见项目官网。

据 README 的快速开始，先安装 uv（已有可跳过），克隆仓库后执行 setup 与 run：

curl -LsSf https://astral.sh/uv/install.sh | sh   # skip if you have uv
git clone https://github.com/vladimirrott/mojo-dllm && cd mojo-dllm
./mojo-dllm setup
./mojo-dllm run "Explain speculative decoding in two sentences."

据项目说明，setup 会检查机器、把固定版本的 Mojo 工具链装进 .venv、构建 build/mojo-dllm，并询问是否下载 LLaDA-8B（4.9 GB，带 sha256 校验），随后跑一次简短冒烟测试；若构建产物与下载内容仍是最新，再次执行 setup 会跳过这两步。./mojo-dllm doctor 用于报告缺失项及对应的修复命令。run 在发现计算能力 8.0 及以上的 NVIDIA GPU 时优先使用 GPU，否则回退到 CPU，可用 --device cpu 或 --device gpu 强制指定。快速开始文档还覆盖了模型选择、离线使用、inspect、tokenize 及相关参数。

许可与限制

许可证为 Apache-2.0。

据项目说明，运行环境要求 Linux on x86-64、支持 AVX2、glibc 2.34 或更新版本，并需要约 6 GB 空闲内存。GPU 路径目前只在 RTX 3050 Laptop GPU 上做过测量；若显卡空闲显存不足，程序会提示所需显存并退出。项目自述的基准结果显示，在 CPU 上 mojo-dllm 是所列表格中最快的运行时，而在 GPU 上 llama.cpp 的 CUDA 构建更快，mojo-dllm 的张量核心 kernel 尚属第一版，当前主要优势是较小的主机内存占用（权重驻留显存）。这些数字来自项目自己的 bench/results/ 结果文件，属于维护者自述，本站未复现。