LLM Inference 🚀

从零实现大模型推理框架 — Qwen3 的本地部署与优化

由于好奇大模型推理的原理，学习了 tiny-llm 课程并将苹果芯片相关代码适配 Nvidia 芯片。侧重 GPU 优化技术，完成 Qwen3-4B-Instruct-2507 模型的本地部署和推理加速。显存不够可以考虑 Qwen3-0.5B，主要是学技术。

pre-requirements

• 需要 Nvidia 显卡，安装 cuda 环境，以及 pytorch
• 为测试代码准性，需要安装 torchao 和 torchtune
• 为实现 flashattention 算子，需要安装 triton
• 为了加载 huggingface 的 Qwen3 model，需要安装 transformers。如果网络失败，需要 export HF_ENDPOINT=https://hf-mirror.com

我的版本：

torch                     2.6.0+cu126
torchao                   0.15.0
torchtune                 0.6.1
transformers              4.57.6
triton                    3.2.0

QWen3 实现与本地部署

#	任务	测试命令
1	✅ 实现 scaled_dot_product_attention	python -m unittest llm.test.attention_test.TestScaleDotAttention
2	✅ 实现 MultiHeadAttention	python -m unittest llm.test.attention_test.TestMultiHeadAttention
3	✅ 实现 RoPE 旋转位置编码	python -m unittest llm.test.rope_test
4	✅ 实现 RMSNorm 标准化	python -m unittest llm.test.norm_test
5	✅ 实现千问的 MLP	python -m unittest llm.test.mlp_test
6.1	✅ scaled_dot_product_attention 添加 GQA 支持	python -m unittest llm.test.attention_test.TestScaleDotAttention
6.2	✅ MultiHeadAttention 添加 GQA 支持	python -m unittest llm.test.attention_test.TestGroupedMultiHeadAttention
7	✅ 实现 tied embedding	python -m unittest llm.test.tie_embedding_test.TestTieEmbedding
8	✅ 实现 Qwen3 TransformerBlock	—
9	✅ 加载 Qwen3 模型，简单推理	python -m llm.executor.run_model

工程优化

#	优化	测试命令
1.1	✅ 解决重复生成问题，实现 TopK 采样	python -m llm.executor.run_model --topk 100
1.2	✅ 解决重复生成问题，实现 TopN 采样	python -m llm.executor.run_model --topp 0.7
2	✅ 实现 PD 分离与 KV Cache	python -m llm.executor.run_model --kv_cache
3	✅ 学习 CUDA 或 triton(建议)，实现 FlashAttention V1	python -m unittest llm.test.flash_attn_test

• page attention
• continuous batching
• load cuda graph

服务设计

• scheduler

结语

大模型推理到最终落地还有艰难的一公里，但是这些内容足以单开一个仓库：

• 不同芯片内存布局不同，NCHW？NHWC？
• 子图匹配与计算图融合
• 多卡推理，多卡的通信与内存管理
• 模型低 bit 量化

芯片部署这是个不折不扣的糙活，难以想象各个厂的芯片会有什么大坑，暂时没有兴趣；子图匹配与计算图融合有时间再看；模型量化暂时没兴趣；工作相关，后面会去看多卡相关的项目。

此外，大模型对话还会涉及很多优化技术，但这些优化技术更侧重大模型的服务和 Agent，更像大模型时代后端工程师的工作内容

• Speculative Decoding：用一个小而快的草稿模型提前预测多个 token，再用大模型并行验证这些预测。
• Memory：长期记忆和短期记忆的管理工具。
• RAG(Retrieval-Augmented Generation)：生成前先检索外部知识，让回答更准确、最新。
• MCP(Model Context Protocol)：大模型与外部数据源和工具之间的通信协议。
• Skills：有了 MCP 就可以调用多个工具，skills 定义了多个工具该如何搭配使用，是 Agent 的操作指南。

RAG 涉及数据库和向量检索，Memory 的历史记忆涉及摘要生成，也许会用其他技术来记录用户的操作习惯。however，这些东西每家公司的实现方式都不同，没有什么统一的标准答案，和 GPU 优化关系不大。由于本人的工作和大模型服务/Agent相差甚远，所以对大模型推理框架的探索到此结束。

参考

• tiny-llm
• torch MultiHeadAttention 转置
• torch 实现 RoPE
• RoPE 公式推导
• torch 实现 RMSNorm
• Qwen 实现 MLP
• GQA 概念
• Tie Embedding
• vllm 中 Qwen3 实现
• 大模型推理为什么需要采样、惩罚
• 看图学 KV Cache
• triton 入门
• FlashAttention 公式推导