AI-Inference-Lab

一个面向学习的大模型推理实验室，包含三个相互衔接的子项目：

• Mini-LLaMA Inference Engine：简化版 LLaMA decoder-only 推理引擎，包含 Embedding、RMSNorm、RoPE、Self-Attention、SwiGLU MLP、KV Cache 和自回归生成。
• CUDA-Transformer-Kernels：面向 Transformer 的 CUDA 算子库示例，包含 RMSNorm 与 Softmax 的 C++/CUDA 扩展，以及 PyTorch fallback。
• Mini-vLLM：简化版大模型推理服务系统，演示 prefill/decode 两阶段、continuous batching、请求调度和分页 KV cache block 管理。

安装

pip install -r requirements.txt

CUDA 扩展是可选的。没有 CUDA 环境时，cuda_kernels/python/ops.py 会自动使用 PyTorch 实现。

cd cuda_kernels
python setup.py install

快速运行

python examples/run_mini_llama.py
python examples/run_cuda_ops.py
python examples/run_mini_vllm.py

测试

pytest cuda_kernels/tests

Benchmark

python benchmark/bench_rmsnorm.py
python benchmark/bench_softmax.py
python benchmark/bench_decode.py

这些 benchmark 会输出 latency、tokens/s、TTFT（Time To First Token）和 TPOT（Time Per Output Token）等指标，便于对比 PyTorch 原生实现与自定义 kernel。