naklecha/simple-llm

解决的问题

SimpleLLM 是一个轻量级、高性能的推理引擎，专门设计用于在单个 NVIDIA H100 GPU 上运行 OpenAI 的 gpt-oss-20b 和 gpt-oss-120b 模型。其目标是在特定工作负载下提供比 vLLM 更快的解码速度，同时保持原始模型的准确性。

工作原理

此引擎以约 1,351 行的 Python 实现，利用 Triton 进行 GPU kernel 运算。它直接加载官方的 MXFP4 检查点，无需重新量化。主要的性能优化包括：

• 自定义 Kernel： 一个处理 sinks、滑动窗口、GQA 和 KV append 的单一解码注意力 kernel，以及一个融合的 RoPE kernel。
• 内存管理： 针对全注意力层使用分页 KV cache，针对滑动窗口层使用环形缓冲区，以优化上下文处理。
• 执行效率： 使用 CUDA graphs 处理活动槽位，并直接在 CUDA graph 内计算贪婪 argmax/下一步输入。
• MoE 优化： 在 tensor cores 上实现 MXFP4 混合专家（MoE），并优化路由，为每个专家写入 tile 表。

适用对象

需要在单个 H100 上运行 OpenAI 的 gpt-oss 模型，并希望获得最大解码吞吐量的开发人员和研究人員，特别是在小到中等批次大小或长上下文长度的情况下。

亮点

• 性能： 在批次 1 和 64 时，解码速度比 vLLM 0.31.0 快 11-15%。
• 效率： 极度精简的代码库（少于 1,400 行 Python）。
• 准确性： 与 OpenAI 发布的 GPQA-Diamond 和 AIME 2025 分数相符。
• 硬件优化： 使用 MXFP4 权重，专为 NVIDIA H100（80 GB）调校。