兄弟们我最近干了件有点“逆天”的事用 2000 行左右的纯 Python手搓了一个简易版的LLM 推理引擎起名就叫SGLang.py专门用来把SGLang这类高性能推理框架的核心机制给“扒光”了看。这玩意儿不是为了替代真正的 SGLang也不是为了跑赢 vLLM纯粹是为了搞清楚一件事当我们在用大模型的时候那些看起来“理所当然”的流式输出、并发调度、KV Cache、前缀复用在底层到底是怎么一步步跑起来的我一直觉得用大模型就像开自动挡汽车但一个合格的工程师得知道发动机内部是怎么点火的。所以我花了几个周末把 SGLang 的调度逻辑、RadixAttention 算法、Continuous Batching 这些硬核概念全部用 Python 原生的数据结构字典、列表、队列重写了一遍没依赖任何重量级的 CUDA 扩展。这篇文章我不讲空理论直接把我踩过的坑、设计时的纠结、以及最后跑起来的完整思路都倒给你。如果你是那种用着 vLLM/SGLang 但总觉得隔着一层纱的人或者你想深入理解 LLM serving 机制、想自己动手写点推理服务的 Python 开发者这篇内容应该能给你不少启发。1. 内容整体设计与思路拆解1.1 为什么要“手搓”一个推理引擎在开始写代码之前我先给自己定了三个规矩。第一绝对不用 CUDA C只用 Python 和 PyTorch 做前向推理。因为我的目标是搞懂调度逻辑而不是搞工业级性能优化。真正的 SGLang 里有一堆用 CUDA 写的融合算子比如 RadixAttention 的前缀缓存算子这些对我来说是黑盒我要做的是把黑盒变成白盒。第二必须吃透 Continuous Batching连续批处理和 RadixAttention基数树注意力这两个核心概念。前者解决的是 GPU 利用率问题后者解决的是多请求重复 prompt比如系统提示词的显存浪费问题。我把这俩在 Python 里用最朴素的方式实现了虽然慢但逻辑一目了然。第三支持 OpenAI 兼容的接口。因为只有这样我才能用 vLLM 的测试脚本直接怼上去看看它在功能层面除了速度和吞吐到底是不是缺胳膊少腿。说实话业界对推理引擎的印象往往是“高深莫测”。但你把它的外皮剥掉核心就是一个不断循环的“调度器”Scheduler和“执行器”Executor。调度器决定这一秒该处理哪些 token执行器负责跑模型的前向计算。1.2 从 SGLang 身上“抄”了哪些核心机制SGLang 之所以快主要靠三招RadixAttention前缀树缓存、Continuous Batching和高效的 Prefix Caching前缀缓存调度。我没有能力完全复刻它的工程细节但我把这三招的“思想”给化用了RadixAttention 的简化版实现我用一棵字典嵌套的树Trie来存储所有请求的 token 序列树的每条边代表一段 KV Cache 的引用。当新请求进来时先走一遍这棵树把自己最长的公共前缀找出来直接复用之前的 KV Cache只对剩下的部分计算新的 KV。这样如果大家共享一个长系统提示词显存占用能降一个量级。Continuous Batching 的迭代级调度传统的静态批处理要等一个 batch 的所有序列都生成完了才能释放显存跑新的 batch。而 SGLang 做到了“token 级”的动态调度——每一步前向序列都可以进来或出去GPU 永远在处理有意义的 token而不是等待那些已经生成完 EOS结束符的序列。Prefill 和 Decode 的本体分离我明确区分了“Prefill”预填充阶段处理整个输入 prompt生成 KV Cache和“Decode”解码阶段逐个预测下一个 token。这两个阶段的计算密度完全不同Prefill 是计算密集Decode 是访存密集工业级框架常用不同的 CUDA kernel 去优化。我在自己的 Python 实现里对这三块做了完整的复刻但用了一种“教学级”的方式我甚至会在每次调度循环里打印当前 GPU 上有多少个空闲 token slot这样你能肉眼看到连续批处理是怎么把显存“挤出”水来的。1.3 技术选型为什么是 Python 而不是 C/Rust写这种系统级的服务大部分团队会选 C 或 Rust因为 Python 的 GIL全局解释器锁是并发的心头痛。但作为一个 demoPython 有它的独特优势——表达能力强开发速度快。在做这个项目的时候我特意把线程和异步分开了推理执行器放在独立线程里跑避免阻塞主线程的调度。HTTP 服务层用异步框架FastAPI处理并发请求。队列通信用queue.Queue做请求缓冲调度器从队列里取请求塞进 batch执行器跑完一次前向再把新 token 塞回流水线。这样的架构虽然不如 C 精细但在理解算法层面它完全没问题。而且 Python 生态里有transformers和torch我可以直接加载一个小模型比如meta-llama/Llama-3.2-1B-Instruct立刻就能看到输出。提示如果你只是对调度逻辑感兴趣完全可以不依赖torch甚至自己写一个假的“伪模型”直接返回随机 token这样能更聚焦在算法本身。2. 核心细节解析与实操要点2.1 Scheduler 是如何“见缝插针”的先讲调度器Scheduler它是整个系统的心脏。我把它分成两个等级的操作加入Add和前进Step。Add 操作当一个new_request进入调度器系统会为它分配一个Sequence对象这个对象里装着它的 token IDs经过 tokenizer 后的整数列表、生成的 KV Cache 引用、采样参数temperature、top_p等元数据。此时这个序列还处于“未加入批处理”的状态它会被放进一个waiting_queue等待队列。Step 操作这是每次循环的核心动作。调度器会先处理所有running正在运行状态的序列检查它们是否已经生成完毕或者是否触发了max_new_tokens。如果它们的输出长度已经足够就会被移出 batch释放它的 KV cache 槽位。然后调度器会从waiting_queue里取出新请求填满刚刚空出来的槽位。这样做的好处是GPU 每一步都在处理“活”请求而不是陪着那些已经结束的请求空转。这里有个细节值得吹一下如何给序列分配显存在真正的 SGLang 里显存是以Memory Pool内存池的形式管理的。每一个 KV Cache 块通常包含 16 个或 32 个 token 的显存占用。在 Python 实现里我用一个简单的kv_pool字典来模拟这个过程初始状态available_kv_blocks [0, 1, 2, ..., NUM_BLOCKS - 1]。当一个序列需要更多 KV 空间时调度器从空闲列表里弹出一个块号分配给该序列。当序列结束时它的所有块号被回收重新放回空闲列表。这么做的直观好处是你可以通过打印空闲块数量直接看到“批处理”是如何压榨显存利用率的。比如有 8 个空闲块来了 10 个请求但每个请求都要 2 个块那么第 9、10 个请求只能被阻塞在等待队列直到前面的请求释放块资源。注意这个显存池的分配逻辑虽然不够精确真实框架会有内存碎片整理但用来理解preemption抢占机制足够了。当显存不足时你需要考虑是抢占swap 出旧的序列还是拒绝拒绝新请求。我的 demo 用了简单的先来先服务真遇到了瓶颈直接让新请求挂起等待空槽。2.2 RadixAttention 的“前缀树”到底怎么存这是SGLang 最引以为傲的创新也是我研究最久的部分。常规的注意力机制里每个请求都需要为它的所有历史 token 计算并保存 KVKey 和 Value向量。如果有 1000 个请求同时带着同一个长系统提示词比如 500 token 的限制说明那么这 500 token 的 KV 就会被重复存 1000 次浪费极其惊人。RadixAttention 的思路是把 KV 缓存看作一棵前缀树Prefix Tree的共享节点。因为自然语言天然有公共前缀比如“你是一个有用的助手”这句话在很多请求里都会出现。这句话的 token 序列对应的 KV Cache只需要在树上存一次新来的请求只需要一个指针引用它即可。我的实现方案class RadixTree: def __init__(self): # 根节点用一个空字典 self.root {_kv: None, _children: {}} # 记录整个树里的 token 数量 self.total_len 0 def find_prefix(self, tokens): 返回最长公共前缀的 token 数和对应的 kv 节点 node self.root prefix_len 0 kv_refs [] for i, token in enumerate(tokens): if token in node[_children]: node node[_children][token] if node[_kv] is not None: kv_refs.append(node[_kv]) prefix_len 1 else: break return prefix_len, kv_refs关键点解析树的每条路径是一个 token ID。假设某个请求的完整 token 序列是[你, 是, 个, 好, 人, 吗]那么在树上会有一条从根到叶子节点的对应路径。而每一个中间节点的_kv字段存放的是“从根到这个节点”所有 token 对应的 KV 张量的叠加引用。当一个新请求[你, 是, 个, 好, 朋, 友]进来时find_prefix会沿着树往下走直到找到[你, 是, 个, 好]这个公共前缀。此时这 4 个 token 的 KV 不用重新计算直接复用之前请求缓存下来的张量引用即可。新请求只需要为[朋, 友]这两个 token 计算新的 KV并把它们挂到树上。当所有使用某个 KV 节点的引用都被删除时比如所有相关请求都结束了这个节点需要被释放。这就是“引用计数”机制。工业级的 SGLang 用这种机制做到极致而我用 Python 的sys.getrefcount也能模拟一个不完整的版本。这个思路有点像缓存数据库里的 LRU最近最少使用缓存只不过这是用树形结构做的缓存能感知到 token 的序列模式。实操心得我刚开始实现时把 KV 映射在树的边上后来发现调试巨痛苦。建议你把 KV 存在“节点”上而不是“边”上。每个节点代表“从根到这个位置的前缀序列”这样在续写时可以直接从该节点对应的隐藏层状态开始 forward而不用管边的事。2.3 Continuous Batching 的“填坑”玩法现在讲讲 Continuous Batching 的实现这是提升吞吐的关键。传统的静态批处理Static Batching是这么干的攒够一个 batch 的请求比如 8 个然后同时跑前向直到这 8 个全部生成 EOS再释放资源跑下一批。假如其中一个请求生成了 10 个 token 后就结束了而另外 7 个请求还在生成 100 个 token那么这 10 个 token 计算完成后它们占用的 GPU slot槽位就浪费了一直空转到整个 batch 结束。而 Continuous Batching 做的是token 级别的动态调度每一轮前向结束后就检查是否有序列已经结束生成了 EOS token如果有立刻从运行列表里剔除。同时从等待队列里拉取新的请求填充到刚空出的位置。这样 GPU 永远在处理有意义的序列那个生成完的 slot 能立刻被新请求顶上。这种做法本质上就是“厕所蹲位管理”——一个人上完了第二个人马上进去而不是等整排厕所全空了才放人。在典型的多轮对话场景中这种方式能让吞吐量提升 2 到 10 倍不等。我在 SGLang.py 里实现这个逻辑时用了一个非常朴素的循环def step(self): # 1. 处理完成的序列 for seq in self.running: if seq.is_finished(): self.running.remove(seq) self.free_kv_blocks(seq.kv_blocks) # 2. 接受新请求直到 GPU 槽位满 while self.waiting_queue and len(self.running) self.max_batch_size: new_seq self.waiting_queue.popleft() new_seq.allocate_kv_blocks(self.kv_pool) self.running.append(new_seq) # 3. 执行预填充对新增的序列和解码对所有运行中序列 self.prefill_new_sequences() decoded_tokens self.decode_all_sequences()这段代码的逻辑很简单但它体现的是调度器的核心哲学每一轮都很短但每一轮都在做最有用的事情。2.4 Prefill 和 Decode 为什么必须分开SGLang以及所有高性能推理引擎都将 LLM 的生成过程严格分为两个阶段Prefill和Decode。Prefill预填充当模型第一次接收一个完整的 prompt 时它需要并行处理所有输入 token计算出每个 token 对应的隐藏层状态并构建完整的 KV Cache。这个阶段是 compute-bound计算密集型GPU 的算力利用率通常很高。Decode解码模型根据已有的 KV Cache每次只计算下一个 token 的概率分布。这个阶段极度依赖从 KV Cache 中读取历史信息是 memory-bound访存密集型瓶颈往往在 HBM 带宽而不是计算能力。在工业级推理引擎里这两个阶段通常会使用不同的 CUDA kernel 来优化。而在我的 Python 版本里我通过两个独立的方法来实现prefill(seq, input_ids)它会对一个序列做完整的前向计算并把 KV 缓存存到seq.kv_cache里。decode(seq)它只读取seq.kv_cache的最后一行预测下一个 token并追加新的 KV。我甚至在这个项目里加了“选择阶段”的日志这样你在跑的时候能看到每个请求当前处于什么状态理解 GPU 在哪个阶段会更忙碌。提示很多刚入门的朋友以为 Decode 阶段的“慢”是模型输出太慢其实不然。是因为 Decode 时 GPU 一直在做低效的访存操作吞吐瓶颈在带宽。所以你会发现 eager mode 下 decode 的 token 吞吐是很低的优化方式就是搞更大的 batch一次处理更多序列摊薄访存开销。3. 实操过程与核心环节实现3.1 工程结构一个 2000 行的 Python 项目我先放一下工程的目录结构你可以直接照着搭sglang-py/ ├── README.md ├── requirements.txt ├── server.py # HTTP 入口FastAPI 服务 ├── scheduler.py # 调度器等待队列/Running队列/Batch组装 ├── radix_tree.py # 前缀树缓存实现 ├── model_runner.py # 模型加载与推理的封装 ├── tokenizer_wrapper.py # 分词器封装 ├── sampling.py # 采样参数与实现temperature/top_p ├── sequence.py # 序列对象token状态/KV缓存引用 └── test_client.py # 测试脚本模拟并发请求我强烈建议你按这个模块化方式组织代码因为推理引擎的逻辑非常复杂如果全部塞在一个文件里几个小时后你自己都会看不懂。特别是sequence.py它定义了一个序列的整个生命周期是核心中的核心。Sequence类的核心属性class Sequence: def __init__(self, _id, prompt_ids, sampling_params): self.id _id self.prompt_ids prompt_ids # 输入的 token ids self.output_ids [] # 生成的 token ids self.kv_cache_refs [] # 引用的 radix tree 节点 self.kv_blocks [] # 占用显存池的块列表 self.sampling_params sampling_params # 采样超参 self.status Status.WAITING # 枚举状态 property def all_token_ids(self): return self.prompt_ids self.output_ids3.2 模型加载与前向计算的最小实现我使用transformers库来加载模型这样能减少不必要的代码量专注在调度上。import torch from transformers import AutoModelForCausalLM, AutoTokenizer class ModelRunner: def __init__(self, model_namemeta-llama/Llama-3.2-1B-Instruct): self.model AutoModelForCausalLM.from_pretrained( model_name, torch_dtypetorch.float16, device_mapcuda ) self.tokenizer AutoTokenizer.from_pretrained(model_name) self.pad_token_id self.tokenizer.pad_token_id self.model.eval() def prefill(self, input_ids, past_kvNone): with torch.no_grad(): output self.model( input_idsinput_ids, past_key_valuespast_kv, use_cacheTrue, ) # 返回 logits 和更新后的 KV cache return output.logits, output.past_key_values def decode_one_token(self, input_id, past_kv): with torch.no_grad(): output self.model( input_idsinput_id, past_key_valuespast_kv, use_cacheTrue, ) return output.logits[:, -1, :], output.past_key_values需要特别说明的“坑”past_key_values是transformers里 kv cache 的标准格式它是一个嵌套的 Tuple按层数、按 K/V。在我的实现里RadixTree 缓存的就是这个对象。因为是教学代码我没有优化past_key_values的 format。实际上高效的做法是用 GQA分组查询注意力的缓存格式或者用张量分块存储。我用float16而不是float32纯粹为了省显存但小模型用 float32 也没问题。运行时会发现加载 1B 模型一共消耗大约 4GB 左右的显存float16如果你在本地没显卡也可以用 MPS 或者 CPU。不过 CPU 上跑decode会慢到让人怀疑人生建议有 GPU 再试。3.3 FastAPI 外衣与 OpenAI 兼容流式输出为了让外界可以像调用 OpenAI 一样调用我们的引擎我实现了/v1/chat/completions接口关键是处理streamTrue的流式返回。这里用 FastAPI 的StreamingResponse是关键。代码逻辑如下app.post(/v1/chat/completions) async def chat_completions(request: ChatCompletionRequest): # 从请求里取出 prompt并做 tokenize prompt build_prompt_from_messages(request.messages, self.tokenizer) input_ids self.tokenizer.encode(prompt) # 创建一个 sequence 并加入调度器等待队列 seq Sequence(...) self.scheduler.add_sequence(seq) async def event_generator(): while True: # 用生成器不断拉取新 token new_token await self.scheduler.wait_for_next_token(seq.id) if new_token is None: break yield fdata: {json.dumps({choices: [{delta: {content: new_token}}]})}\n\n return StreamingResponse(event_generator(), media_typetext/event-stream)这里有一点很巧妙我把调度器和 HTTP 服务解耦了。调度器运行在自己的线程里不断执行step()HTTP 的生成器循环则阻塞在wait_for_next_token上等待新 token 产出。听上去简单但中间涉及 Python 的线程同步、队列操作还有状态检查第一次跑通的时候我差点激动得跳起来。3.4 采样与约束解码的基础实现采样模块是另一个容易被忽略但至关重要的部分。生成 token 时我们不仅需要模型输出的 logits还要应用温度temperature、Top-K取概率最高的K个、Top-P核采样等策略。我的sampling.py实现了这些基础算法def sample_next_token(logits, temperature1.0, top_k0, top_p0.9): # 温度缩放 if temperature ! 1.0: logits logits / temperature # Top-K 过滤 if top_k 0: indices_to_remove logits torch.topk(logits, top_k)[0][..., -1, None] logits[indices_to_remove] float(-inf) # Top-Pnucleus过滤 if top_p 1.0: sorted_logits, sorted_indices torch.sort(logits, descendingTrue) cumulative_probs torch.cumsum(torch.softmax(sorted_logits, dim-1), dim-1) # 移除累积概率超过 top_p 的 token sorted_indices_to_remove cumulative_probs top_p sorted_indices_to_remove[..., 1:] sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] 0 indices_to_remove sorted_indices[sorted_indices_to_remove] logits[indices_to_remove] float(-inf) # 从剩余 token 中采样 probs torch.softmax(logits, dim-1) next_token_id torch.multinomial(probs, num_samples1) return next_token_id经验之谈采样参数看着简单但调试起来很上头。第一次跑模型时我忘了加 temperature导致每个序列的输出每次都一样后来加上了但忘了 top_p导致偶尔会有不连贯的重复。真正跑出稳定的下一代对话经历了至少三次 bug 修正。注意如果你的模型生成效果特别“死板”大概率是采样参数设置问题不是模型坏了。我在实现里默认是temperature1.0, top_k0, top_p0.9但跑对话时建议调低温度到 0.7 左右出来的文字更自然。4. 常见问题与排查技巧实录手搓推理引擎的过程本质上就是一个不断“爆雷”又“排雷”的过程。下面这些问题都是我自己踩过的列出来供大家参考。4.1 一个 Python 线程就能把 GIL 卡死吗现象服务并发压测时发现 GPU 利用率只有 30%但 CPU 有一个核心跑满了HTTP 请求排队严重。原因我的调度器step()是一个死循环又是跑在 Python 主线程里导致 FastAPI 的异步事件循环被阻塞了新的 HTTP 请求根本进不来。解决我把调度循环挪到了独立线程用threading.Thread(daemonTrue)启动。然后在主线程的 FastAPI 里只做队列操作。这样即使调度线程持续占用 GIL异步 I/O 也能得到部分执行时间。深层经验如果项目以后要卷性能需要把调度器和执行器分别放到独立进程中用高效的 IPC比如共享内存或 ZeroMQ来通信。Python 做 demo 没问题上生产千万别这么搞。4.2 KV Cache 的内存泄漏现象跑了几轮对话后显存占用持续上涨但没有崩。原因RadixTree 里的节点在引用计数归零后没有被及时清理。尤其是当一个前缀被多个序列同时引用时如果其中一个序列提前结束它所占用的 KV 缓存还没来得及释放留下了一个“幽灵缓存”。解决我在RadixTree里加了一个dereference(node)方法每次序列结束都会调用它减少节点的引用计数。当引用计数为 0 时就递归地把空闲的 KV 块归还给显存池。4.3 前缀树命中率极低Radix Cache 形同虚设现象连续输入两个相同系统提示词的请求但前缀树总是匹配不上导致每次都要重新全量 prefill。原因我发现问题出在 tokenization 上。同一个句子在带结尾换行符时多了一个换行 token导致前缀树匹配中断。而且我忘了对 prompt 做特殊 token 处理比如|begin_of_text|前缀。解决在做缓存 key 时强制把系统的特殊 token 也纳入序列。另外我写了一个简单的“前缀规范化”函数将所有 prompt 的末尾统一去掉换行让公共前缀尽量长。清单总结问题核心原因解决办法GPU 利用率低CPU 卡满Python 调度线程阻塞事件循环把调度器放独立线程或进程显存持续上涨不回收前缀树节点引用未释放实现引用计数和递归回收前缀缓存命中率低Tokenization 不一致对 prompt 做规范化处理统一分隔符流式输出不流畅HTTP 和调度耦合wait 循环无超时用 asyncio.Queue 解耦加超时机制模型输出重复内容采样参数设置不当调整 temperature/top_p关闭 beam search4.4 并发时请求丢失与乱序现象一旦并发数超过max_batch_size部分请求会超时甚至返回空结果。原因我最初用queue.Queue存所有未处理的请求但在 FastAPI 的event_generator和调度线程之间对seq.output_ids的访问没有加锁导致多端读写冲突。解决给每个Sequence加一个专属的asyncio.Event对象新 token 产生后触发事件。生成器循环会阻塞等待该事件而不是手动轮询。同时用threading.Lock保护共享的 KV 池分配与释放操作。提示这也是工业级框架里一个很重要的点调度器是单线程的但服务于并发的 HTTP 请求。如果用 Python 做简化实现锁的粒度一定要小否则并发性能会很难看。4.5 Decode 阶段为何特别慢现象在 CPU 或者纯 Python 环境下decode 阶段的 token 生成速度极慢甚至每秒只有几个 token。原因decode 是访存密集型操作每次前向计算都要读取所有 KV Cache。而 Python 的torch在每次小 batch 计算时有很大的 kernel 启动开销。而且我没有使用任何量化手段导致带宽占用极高。应对如果是做 demo建议直接上 GPU如果实在没 GPU可以尝试把decode_all_sequences里的多个序列合并到同一个 batch 里处理这样能摊薄一些 CPU 的开销。我的代码默认是每步处理整个 running 列表这样会比一个序列一个序列地跑更有吞吐。不过说句实在话如果想真正让解码变快还是得靠编译优化、算子融合和并发推理框架。Python 版的核心价值在于教学如果真的在意速度建议去看一看 vLLM 的 PagedAttention 和 SGLang 的原生实现。5. 从手搓引擎到理解 LLM 推理全景做完这个小项目我对 LLM serving 的理解完全上了一个台阶甚至看官方文档都轻松了很多。这里说几个我个人最深刻的认知变化。第一流式输出不是模型的功能而是工程的设计。模型本身一次只产生一个 token你只要拿到那个 token 就能立刻返回给前端。真正的难点在于你怎么高效地组织调度循环让每个 token 产生后能马上被消费者收到而不是傻等一个完整句子生成完。第二显存管理是推理引擎的天花板。谁能更高效地管理 KV Cache谁就能在同样的 GPU 上吃下更多的并发请求。RadixAttention 的精髓在于“存在树里”PagedAttention 的精髓在于“存在表中”本质上都是解决显存碎片化、共享化的问题。你理解了这一点再去看 vLLM 和 SGLang 的官方博客时那些抽象的概念一下就落地了。第三调度器的核心是不断做转场。它一会儿把新请求从 waiting 拉到 running一会儿把完成的序列从 running 踢走一会儿又把显存块回收。每一次转场都在给 GPU 喂“有效工作”。所以优化推理引擎很大程度上是在优化“转场”策略——该抢占谁该给谁更多 compute slot这些都是策略问题。至于后续扩展我给自己列了几个方向如果你们有兴趣也可以顺着做接入 LoRA 微调支持、实现 PagedAttention 式的块管理、跑一个真实的 benchmark 对比 vLLM、或者试图加入 beam search。任何一个方向拎出来都够再写几篇长文了。最后说一句题外话如果你真的想系统学习大模型推理优化直接上手读 SGLang 的源码会有一些门槛但我这种“先自己写一个慢的再去读快的”的方法论反而能帮你快速定位到源码里的关键模块。先用 2000 行 Python 把引擎的逻辑骨架搭起来再去看工业级框架的细节那种豁然开朗的感觉绝对比死磕源码来得痛快。