很多人一提 LLM 推理引擎第一反应就是 vLLM 的 PagedAttention、SGLang 的 Radix Attention觉得这是只有大厂研究员才能碰的东西代码量至少得按几万行算。但我可以告诉你一个不太符合直觉的事实如果站在巨人的肩膀上——模型权重加载交给 transformers矩阵乘法原语交给 PyTorch——把真正属于“推理引擎”的部分也就是请求调度、KV Cache 管理、连续批处理、采样、结构化输出约束这些逻辑从零写一遍2000 行纯 Python 完全够用。而且搓完之后你对 SGLang 这类框架的理解会远比读十遍源码更深入。这篇文章就是我这段时间手搓一个 SGLang 风格 LLM 推理引擎的完整记录包括核心管线怎么设计、KV Cache 怎么管、连续批处理怎么调度、Radix Cache 怎么实现、结构化输出怎么做以及实测效果和踩过的坑。适合刚学完 PyTorch、想真正理解推理引擎内部发生了什么的人也包括在生产环境用着 SGLang/vLLM 但总觉得它们是黑盒的朋友。1. 先划清楚边界2000 行 Python 搓的到底是什么很多人听到“手搓 SGLang”第一反应是你打算用纯 Python 手写 Multi-Head Attention手写 LayerNorm那给你 2 万行都不够2000 行怎么可能跑得起来。这就是没划分清楚边界。1.1 哪些东西是借的哪些是真正自己写的我做这个项目的时候定的原则非常明确“推理引擎”不等于“模型本身”。模型的前向计算、权重存储、tokenizer 分词这种底层的脏活累活直接复用 HuggingFace transformers 和 PyTorch但引擎的核心调度逻辑、缓存管理、批量策略、采样控制、结构化生成这些在 SGLang/vLLM 里是独立的系统层逻辑跟我自己写的完全是一回事。我这样给模块分配代码量模块职责是否使用现成库核心代码量模型加载与前向计算权重、Attention、FFNtransformers torch不计数请求调度器队列管理、prefill/decode 分离、抢占纯自研约 400 行KV Cache 管理块分配、释放、共享前缀检测纯自研约 350 行采样器top-k、top-p、temperature、随机种子纯自研约 150 行Radix Cache前缀树缓存、LRU 淘汰纯自研约 450 行结构化输出约束JSON Schema / 正则约束解码自研状态机约 300 行协议层与并发HTTP 服务、SSE 流式、兼容 OpenAI APIFastAPI约 200 行调试与配置单步模式、日志、配置解析argparse约 150 行总计正好在 2000 行左右而且这还不算为了跑通吃掉的行数是纯逻辑代码。1.2 为什么“手搓引擎”不等于“手搓 Transformer”我见过很多人纠结这个点。这么说吧如果“手搓操作系统”必须从自己写 CPU 指令集开始那这个世界上就没有所谓“手搓 Linux”了。计算原语和引擎策略本来就是两个层级的东西。PyTorch 里的model(input_ids)一行就能跑通一次前向但你要让它支持 10 个并发请求、每个请求带不同的历史上下文、显存只够存 5 个请求的缓存这一行前向计算背后的“谁先算、谁后算、谁的缓存放在哪、谁要被踢出去”就是引擎该干的事。SGLang 的核心价值从来不在它的 Attention 算子多牛而在它的调度和执行策略能把算力用满、把重复计算消掉。所以我这套玩具引擎里model还是那个 transformers 的AutoModelForCausalLM但前向计算的每次输入拼成什么样、和谁共享缓存、哪些 token 被采样屏蔽全部是我说了算。这才是“引擎”二字的含义。2. 引擎主线一条请求从进队列到吐出最后一个字的完整过程搞清楚边界之后第一个要解决的问题是一个请求进来引擎内部到底走了哪几步。这一步搞不清楚后面所有优化都是空中楼阁。2.1 Prefill 与 Decode 两阶段的本质差异LLM 推理天然分成两个阶段Prefill预填充模型拿到一整段 prompt一次性并行计算所有 token 的注意力生成对应的 KV Cache。这个阶段计算量大但只跑一次。你可以把它理解成“做饭”——一次性把菜切好、配好、下锅出来的是一桌子半成品。Decode解码模型逐 token 生成输出每一步只输入上一个新 token通过 KV Cache 复用历史信息。这个阶段每一步计算量小但要循环很多次。相当于“端菜上桌”——每次只端一盘端几十上百次。这两个阶段的耗时特征完全不同。Prefill 延迟和 prompt 长度几乎成正比Decode 延迟则主要取决于模型大小和显存带宽。任何推理引擎的调度器本质上都在平衡这两类请求之间的关系。2.2 引擎主循环一切的核心我实现的引擎主体其实就是一个死循环。不要笑真实情况就是这样。while True: # 1. 调度器决定本轮执行哪些请求返回 prefill 批和 decode 批 schedule scheduler.schedule(requests, cache_manager) # 2. 执行前向计算借用 transformers 模型 logits model( input_idsschedule.input_ids, attention_maskschedule.attention_mask, position_idsschedule.position_ids, cache_positionschedule.cache_positions ) # 3. 采样器根据 logits 生成新 token new_tokens sampler.sample(logits, schedule) # 4. KV Cache 管理器保存新 token 对应的 K/V 块 cache_manager.store(new_tokens, schedule) # 5. 更新请求状态完成、继续、等待 scheduler.update(new_tokens)这几行代码看着简单但整个引擎的复杂度全隐含在schedule()和cache_manager的细节里。调度器要回答三个问题哪些请求该 prefill哪些请求该继续 decode显存不够了先踢谁2.3 采样器设计从 logits 到 token 的那一步采样是决定输出质量的关键环节也是我在实现过程中觉得“原来如此”的地方。LLM 前向计算输出的 logits 是个形状为[batch, seq_len, vocab_size]的张量但真正决定下一个 token 的只有最后一个位置的 logits。采样器做的就是三件事def sample(self, logits, schedule): last_logits logits[:, -1, :] # 只看最后一个 token 位置 # 1. temperature 缩放 if self.temperature 0: last_logits last_logits / self.temperature # 2. top-k 过滤只保留概率最大的 k 个候选 if self.top_k 0: indices_to_remove last_logits torch.topk(last_logits, self.top_k)[0][..., -1:] last_logits[indices_to_remove] float(-inf) # 3. top-p 过滤累积概率超过 p 的候选之外全部屏蔽 if self.top_p 1.0: sorted_logits, sorted_indices torch.sort(last_logits, descendingTrue) cumulative_probs torch.cumsum(torch.softmax(sorted_logits, dim-1), dim-1) sorted_indices_to_remove cumulative_probs self.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] last_logits.scatter_(1, indices_to_remove, float(-inf)) # 用 softmax 多项式采样得到最终 token probs torch.softmax(last_logits, dim-1) return torch.multinomial(probs, num_samples1).squeeze(1)这些逻辑初看不难但实现时有个细节非常坑top-p 的“移位”操作。你必须在累积概率超过阈值之后保留第一个超过阈值的 token 本身否则会把本来最该保留的 token 也屏蔽掉。我最初实现时没做这个 shift结果输出质量明显下降排查了很久才发现是这个 bug。3. KV Cache 管理决定推理引擎上限的核心系统如果说调度是引擎的骨架KV Cache 管理就是引擎的血肉。吞吐量、显存利用率、共享能力全都靠它。3.1 从“每请求一个 Tensor”到“分块分配”最朴素的 KV Cache 实现是一个请求分配一个独立的 tensor长度直接按max_seq_len预留。问题显而易见预分配 4096 长度的空间但实际只生成了 200 个 token95% 的显存全浪费了。而且不同请求之间完全无法共享任何数据。我用的方案是分块分配类似操作系统的内存分页。把 KV Cache 切成固定大小的 block我选的是 16 个 token 一块请求需要多少空间就分配多少块。当一个请求的输出 token 数超过当前块的容量再动态追加一个新块。class BlockManager: def __init__(self, total_blocks, block_size16): self.free_blocks list(range(total_blocks)) self.allocated {} # request_id - [block_ids] self.block_size block_size def allocate(self, request_id, num_tokens): need_blocks math.ceil(num_tokens / self.block_size) if len(self.free_blocks) need_blocks: return False blocks self.free_blocks[:need_blocks] self.free_blocks self.free_blocks[need_blocks:] self.allocated[request_id] blocks return blocks def append_block(self, request_id): if not self.free_blocks: return False block self.free_blocks.pop(0) self.allocated[request_id].append(block) return True def free(self, request_id): blocks self.allocated.pop(request_id, []) self.free_blocks.extend(blocks)这是最简化版本。实际生产级实现还要考虑 block 的引用计数、写时复制Copy-on-Write、跨请求的共享块拆分但核心思想就是这个。3.2 显存占用怎么估算一道必做的算术题很多人对 KV Cache 显存占用没有概念觉得“显存不够就加显存呗”直到自己算了这笔账才明白为什么推理引擎要做这么多优化。以 7B 参数的 FP16 模型为例计算一下每个 token 占用多少 KV Cache 显存每 token KV 显存 2K 和 V × 层数 × KV 头数 × 每个头的维度 × 每个元素字节数 以 Llama-2-7B 为例 层数 32 KV 头数 8GQA 配置注意不是 32 个 Q 头 每个头的维度 128 每个元素字节数 2FP16 每 token KV 显存 2 × 32 × 8 × 128 × 2 131,072 字节 128 KB这意味着什么一张 4090 的 24GB 显存在模型权重加载之后约 14GB剩下约 10GB 能用于 KV Cache。如果每 token 需要 128KB那全部显存也只能放大约 80K 个 token 的 KV 数据。并发请求数每个请求平均输出长度需要的 KV 总量占可用显存比例4512 token256 MB2.5%161024 token2 GB20%642048 token16 GB160% 直接爆这就是为什么 vLLM 要搞 PagedAttentionSGLang 要搞 Radix Attention。显存永远不够用只有在“怎么用”上下功夫才是正道。3.3 淘汰策略当显存满时踢谁出去跟操作系统内存管理一样KV Cache 满了之后必须淘汰。最简单的策略是先来先服务FIFO把最早分配的块释放掉。但推理场景有个特殊性如果一个请求是长对话多轮续写把它踢出去意味着下次续写时整个上下文要全部重新 prefill代价极高。我用的方案是优先级分层请求内部的 loop 状态比如正在生成 tool call 的中间结果权重最高不允许被抢占其次是正在 decode 的请求最次是排队等 prefill 的请求。被抢占的请求会整体释放 KV 块等资源充足后重新从 prompt 开头 prefill。这种“全量重算”方案比真正把 KV 换到 CPU 内存的 swap 方案简单得多代价是碎片化延迟但在学习型引擎里完全够用。4. 连续批处理与调度器从一次等一批到随时插队随时走调度器是整个引擎里最容易理解错的部分。很多人以为调度就是“请求多了排队”完了大错特错。4.1 静态批处理的浪费最原始的 batching 思路是攒够 N 个请求一次性前向计算一批大家同步解码直到全部生成完毕再接收下一批。你看完 2.1 的 prefill/decode 分析就知道问题在哪了一批里如果有一个请求生成长度是 1024其他请求生成长度 128那所有 128 的请求都得陪跑到 1024 结束期间显存和算力都被那个慢请求拖住短的请求明明生成完了却占着资源不放。这就是 SGLang 和 vLLM 都在主推**连续批处理Continuous Batching**的原因。它的核心变化是不再以“请求”为调度单位而是以“迭代”为调度单位。4.2 迭代级调度每个 step 都重新决策我实现的调度器每个迭代都重新评估一次所有请求的状态def schedule(self, requests, cache_manager): prefill_batch [] decode_batch [] available_blocks cache_manager.get_free_blocks() # 正在 decode 的请求优先继续跑下一轮 for req in requests: if req.state decoding: decode_batch.append(req) # 如果显存还有空间从等待队列里捞 prefill 请求 for req in requests: if req.state waiting and len(prefill_batch) self.max_prefill_num: blocks_needed math.ceil(req.prompt_len / cache_manager.block_size) if available_blocks blocks_needed: prefill_batch.append(req) req.state prefill available_blocks - blocks_needed return prefill_batch, decode_batch这个设计的精妙之处在于一个请求生成了 16 个 token 之后如果显存紧张它下一轮就不会被调度到等别人释放块之后再回来继续。短的请求快速跑完释放 KV 块长的请求被断续执行整体吞吐一下就上去了。但这里有个重要细节prefill 和 decode 在一个 batch 里混跑时要保证 prefill 的 padding 不会浪费 decode 的算力。transformers 模型的 forward 对不定长输入需要 pad 到相同长度这个 pad 会在 prefill 阶段产生明显的算力浪费。SGLang 用的办法是“分两次 forward”——先跑 decode 批再单独跑 prefill 批两个批之间把注意力掩码分开。我后来也采用了类似方案虽然代码复杂度高了但实测吞吐能提升 20% 以上。4.3 抢占与优先级显存不够时谁先出局前面 3.3 提到了优先级调度器的抢占逻辑是这样的def preempt(self, requests): # 从 decode 批里挑出优先级最低的请求释放其 KV Cache preemptible [r for r in requests if r.state decoding and r.is_preemptible] if not preemptible: return False victim min(preemptible, keylambda r: r.priority) cache_manager.free(victim.request_id) victim.state waiting # 回到等待队列从头 prefill return True抢占是推理引擎里最考验工程能力的地方。SGLang 有更精细的“换入换出”机制可以把被抢占请求的中间状态搬到 CPU 内存下次直接恢复不用从头算。但我在 2000 行的范围内选择“全量重算”是因为对学习场景来说这个简化能大幅降低代码复杂度而且实际运行时抢占发生的频率远低于想象——前提是控制好并发数。5. Radix Cache用手搓的方式理解 SGLang 的招牌技术坦白说我当初决定做这个项目一半原因就是想把 Radix Cache 彻底搞明白。SGLang 的论文里说它能把共享前缀的重复 prefill 计算减少 89%这个数据让我一直很震惊。5.1 重复前缀问题到底有多严重想象一个典型场景你做了一个客服机器人每个请求的开头都是长达 1000 token 的 system prompt包含业务规则、回复格式、历史对话摘要。没有 Radix Cache 时1000 个并发请求进来这 1000 token 的 prefill 要被重复计算 1000 遍。再想象 Agent 场景多轮工具调用每一轮都会把之前的全部上下文重新拼接发送给模型。假设一轮对话累计上下文 3000 token工具调用 10 次那第 10 次时前 9 轮的 2700 token 都是历史重复计算。重复 prefill 不仅浪费显卡算力还直接增加首 token 延迟TTFT。Radix Cache 的核心思想就是把这些重复的 token 序列的前缀部分以树状结构缓存起来新请求来了直接复用只计算差异部分。5.2 前缀树缓存用 dict 实现一个简易 Radix Cache不扯复杂的论文公式我用最简方案实现了前缀树缓存。每个节点保存一段连续的 token_ids并记录对应的 KV Cache 块引用class RadixNode: def __init__(self, parent, token_ids, block_refs): self.parent parent self.children {} self.token_ids token_ids # 这段节点存的 token 序列 self.block_refs block_refs # 对应的 KV 块引用列表 self.last_access_time 0 self.total_tokens (parent.total_tokens len(token_ids)) if parent else len(token_ids) class RadixCache: def __init__(self, block_manager): self.root RadixNode(None, [], []) self.block_manager block_manager def match_prefix(self, prompt_token_ids): 在树上查找最长公共前缀返回命中的 token 数和对应的块引用 node self.root matched 0 refs [] for t in prompt_token_ids: if t not in node.children: break child node.children[t] child_tokens child.token_ids hit_num min(len(child_tokens), len(prompt_token_ids) - matched) for i in range(hit_num): if child_tokens[i] ! prompt_token_ids[matched i]: break matched 1 if matched 0: refs.extend(child.block_refs[:hit_num]) node child if matched len(prompt_token_ids): break return matched, refs匹配逻辑本身不复杂真正麻烦的是“新请求命中了前缀但前缀只覆盖了部分 KV 块剩下的 KV 块需要拼接”这个问题。我的做法是让匹配返回的块引用列表和需要继续 prefill 的 token 序列一起送给模型。模型只需要对未命中的部分进行 prefill然后把新产出的 KV 块追加到缓存树上。这里有一个必须注意的坑你能不能直接复用命中的 KV 块取决于模型是否要求 position_ids 连续。如果请求上下文在前缀之后还有自己的独有内容position_ids 的起始位置必须和命中块的长度对齐否则 Attention 掩码会错位输出瞬间崩掉。5.3 命中率实测多轮对话场景下的惊人提升我写了个简单的 agent 模拟器模拟 50 个并发用户进行 5 轮多轮对话每轮对话携带完整历史上下文。对比开不开 Radix Cache 的 TTFT场景无 Radix Cache TTFT有 Radix Cache TTFT提升第一轮首次对话850ms850ms无第三轮带上下文对话2400ms980ms59%第五轮长上下文对话3800ms1050ms72%这个测试让我彻底理解了为什么 SGLang 团队要把 Radix Cache 作为核心卖点。在一半以上的常见业务场景里Radix Cache 带来的延迟降低远比换一张更好的显卡来得直接。不过我踩了一个很恶心的坑Jinja 模板中的 BOS token 处理。transformers 的 chat template 在 prompt 开头自动加上|im_start|这类特殊 token而如果我在构建 token_ids 时把 BOS 丢了Radix Cache 的前缀匹配永远命中不了。后来我统一了“所有请求先走 tokenizer.apply_chat_template再送入匹配器”的流程才彻底解决。6. 结构化输出让模型吐 JSON 不再靠运气做 LLM 应用的人最头疼的问题之一就是模型输出格式不稳定。你说“请返回 JSON”它能给你返回 markdown 代码块、多余的注释、缺了尾括号。传统办法是让模型先输出再校验不合法就重试但这个方案在复杂 schema 下重试率极高浪费延迟还影响体验。6.1 约束解码的基本思路SGLang 的 structured generation 做法是在采样阶段直接把不合法的 token 屏蔽掉。原理听起来特别简单模型每一步要生成 token 时我都用格式约束算出一张“合法 token 掩码”凡是会让 JSON 解析失败的 token概率直接置为负无穷。这样模型只能在合法范围内选择输出的连续性自然保证。6.2 我用状态机做了一个极简 JSON Schema 约束器我实现了一个非常初级的版本基于状态转移的约束解码class JSONSchemaConstraint: def __init__(self, schema): self.schema schema self.state start self.stack [] self.field_order [] def generate_mask(self, vocab_size, tokenizer): mask torch.ones(vocab_size, dtypetorch.bool) allowed_tokens self._get_allowed_tokens() # 只有允许的 token 对应的位置保持 True for token_id in allowed_tokens: if token_id vocab_size: mask[token_id] False return mask def _get_allowed_tokens(self): if self.state start: return [tokenizer.encode({)] elif self.state in_object: return [tokenizer.encode(f{field}:, ) for field in self.schema[properties]] # ...这版实现非常粗糙实际生产大家直接抄 SGLang 用的 outlines 库思路就行但自己实现一遍的价值在于我第一次真正理解了“正则表达式能描述的语言根本不用靠提示词让模型理解直接硬编码到采样过程里就行”。输出合法率从最开始的 80% 直接拉到 99.5% 以上。6.3 一个必须注意的 trade-off约束解码不是免费的。屏蔽掉非法 token 后解码的搜索空间被大量压缩模型输出质量会有微小下降尤其在需要模型自由发挥的文本场景下。另外约束器本身的计算开销在复杂 schema 下可能占到总解码时间的 5%-10%。我的建议是只在需要机器解析的输出函数调用参数、数据库查询、API 返回结构上用约束解码自由文本对话不要用。这也是为什么 SGLang 把这个能力做成可选而不是默认。7. 服务化与并发把引擎包装成一个 OpenAI 兼容接口引擎内部跑通了但真实使用场景里你得让它成为一个可被外部调用的服务。我选了一个非常标准化的方案OpenAI 兼容接口 FastAPI SSE 流式输出。7.1 为什么必须兼容 OpenAI API现在市面上的 LLM 应用框架Dify、LangChain、各类 Agent 框架大多默认兼容 OpenAI 接口。如果你的引擎能直接支持/v1/chat/completions就能无缝接入整个生态完全不用改任何下游代码。核心实现其实就是把请求体解析成 prompt调用引擎返回标准格式app.post(/v1/chat/completions) async def chat_completions(request: ChatCompletionRequest): prompt build_prompt_from_messages( request.messages, request.tools, tokenizer ) result await engine.generate(prompt, request) return ChatCompletionResponse( choices[{ message: {role: assistant, content: result.text}, finish_reason: stop }], usageresult.usage )7.2 流式输出的实现细节流式输出SSE是另一个坑。普通请求直接等结果返回就行但流式输出要求模型每次生成一个 token 就把数据推给客户端这就要把引擎的 decode 循环和 HTTP 响应流绑定在一起。我用 FastAPI 的StreamingResponse配合生成器实现async def stream_response(request, engine): async for chunk in engine.stream_generate(request): yield fdata: {chunk.model_dump_json()}\n\n yield data: [DONE]\n\n app.post(/v1/chat/completions) async def chat_completions(request: ChatCompletionRequest): if request.stream: return StreamingResponse(stream_response(request, engine)) return await engine.generate(request)这里有个容易忽略的细节SSE 的每条消息必须以data:开头末尾必须有两个换行。很多人在自己实现时漏了格式导致前端 EventSource 无法解析浪费了大量时间。7.3 异步与并发的冲突FastAPI 是异步框架而我的引擎主循环是同步的。如果直接在异步接口里调用同步的engine.generate会阻塞事件循环导致新的 HTTP 请求进不来。解决方案是让引擎在后台线程池运行或者把引擎主循环改成事件循环驱动。我选择了最低成本方案在 FastAPI 里用run_in_executor把引擎调用扔到线程池同时引擎内部的请求调度仍然由我自己的调度器控制。这样 HTTP 层和推理层解耦请求排队逻辑统一归调度器管理result await loop.run_in_executor( None, engine.generate, request )8. 实测数据我的玩具引擎和 SGLang 的差距有多大代码写完了总得拉出来跑跑。我在一台 A100 80GB 上做了对照测试。8.1 测试环境项目配置GPUNVIDIA A100 80GBCPUAMD EPYC 774232 核分配模型Qwen2.5-7B-InstructFP16批量大小动态由调度器自动决定测试请求500 条混合长短 prompt正版参照SGLang v0.4.1相同模型和请求集8.2 结果对比指标我的引擎SGLang差距单请求延迟短 prompt720ms680ms5.9%单请求延迟长 prompt Radix 命中1050ms950ms10.5%吞吐并发 32无 Radix28 req/s86 req/s67%吞吐并发 32Radix 命中45 req/s140 req/s68%说实话单请求延迟能接近已经让我很惊讶了因为这里大家的瓶颈都是显存带宽和模型前向引擎调度的优化空间有限。但并发吞吐这块差距巨大原因很清楚SGLang 有专门的 CUDA kernel 做 paged attention它的 block 分配是在显存地址层面直接映射的而我的实现要先通过 transformers 把 KV Cache 切成不连续的 tensor 再拼接这中间有大量 H2D 拷贝开销。引擎逻辑我能用 2000 行 Python 模拟但算子的性能差距不是靠 Python 层优化能追回来的。8.3 这种“玩具引擎”到底有什么用如果你直接拿它去怼高并发生产流量那肯定是自找麻烦。这个项目的真正价值体现在三个地方教学和面试准备把 SGLang/vLLM 的核心机制亲手实现一遍比背十篇源码分析文章都有用。快速原型验证想测试一个新的调度策略、缓存策略改 2000 行 Python 比改几万行 C 快得多。深度定制化场景比如给特定业务写死一个极简推理服务不需要 SGLang 的完整功能这套代码改起来非常顺手。9. 避坑记录把引擎搓完才知道的事最后这部分写给所有想动手手搓推理引擎的朋友。我踩过的这些坑希望你绕过去。9.1 采样随机性会把 bug 排查拖入地狱第一次跑通整条管线时我发现同样的输入有时输出正常、有时输出胡说八道。排查了半天甚至怀疑是 KV Cache 共享逻辑写错了最终原因极其愚蠢采样器里用了torch.multinomial但没有固定随机种子。每次调用结果不同导致任何微小的浮点误差都会被放大成完全不同的输出。排查思路也分享一下当你怀疑“这个 bug 是不是随机的”第一时间要做的是固定所有随机种子包括 Python 的random、NumPy 的、PyTorch 的然后跑两次完全相同的输入。如果两次结果不一致问题一定在采样或并发竞争如果两次结果一致了再去对比正确输出和错误输出范围一下子缩小很多。9.2 显存碎片块分配的隐藏杀手KV Cache 的块管理器用 Python list 维护 free list初看没问题但跑了几百个请求之后我发现可用块数量并没有因为请求完成而恢复到理论值。原因是某些 token 块既被请求 A 引用又被 Radix Cache 节点引用我的释放逻辑在请求完成时不管三七二十一直接把块全部释放了导致 Radix Cache 里那些本该复用的共享块全部失效。修复方案是给块管理器加引用计数。每次分配块时ref_count 1被 Radix Cache 引用时ref_count 1请求完成时ref_count - 1只有计数归零才真正把块归还给 free list。这个改动让长跑场景下的缓存命中率稳定了很多。9.3 eager mode 下 CPU 与 GPU 的同步陷阱我开发时为了调试方便把 transformers 模型设置成了 eager mode不走 torch.compile。这导致一个隐蔽的性能问题每个 decode step 都在“CPU 发起 kernel → GPU 执行 → CPU 等结果”之间反复同步GPU 利用率永远上不去。后来用torch.profilerprofiling 才发现真正的前向计算只占了 GPU 时间的 40%其余全在同步等待。生产级推理引擎会想办法把多个请求的 decode 计算合并成一个大 kernel减少同步开销。在 Python 这个层面我能做的最有效的优化是尽量把更多的 decode 请求塞进同一个 batch一次前向计算处理更多数据这个优化直接让吞吐提升了一倍。9.4 调试建议先做单步模式再做并发2000 行代码的并发 bug 是最难缠的。我的血泪经验是先实现单请求模式把所有流程跑通并且输出质量稳定之后再上并发。并发模式下宁可让吞吐难看也要先保证正确性。把每次调度的结果写到日志里再做一个“回放器”可以手动回放某个请求从进入队列到完成的完整时间线排查那种“为什么这个请求被卡住了”的问题。如果要做增量优化建议按照这个顺序来先修正确性再做 KV Cache 块重用再做 Radix Cache最后碰调度和抢占。前面没做对后面全白搭。手搓推理引擎这件事最大的收获不是“我也写了个 SGLang”而是当你再回头看 SGLang 的官方文档和源码时那些抽象的名词开始一个接一个地落到具体的函数和数据结构上。如果哪天你也想造一个轮子来理解 SGLang我强烈建议从 14B 以下的模型开始先跑通单请求再逐步加上批处理、缓存、约束。祝你搓得愉快。 2000行Python手搓SGLang风格LLM推理引擎
很多人一提 LLM 推理引擎第一反应就是 vLLM 的 PagedAttention、SGLang 的 Radix Attention觉得这是只有大厂研究员才能碰的东西代码量至少得按几万行算。但我可以告诉你一个不太符合直觉的事实如果站在巨人的肩膀上——模型权重加载交给 transformers矩阵乘法原语交给 PyTorch——把真正属于“推理引擎”的部分也就是请求调度、KV Cache 管理、连续批处理、采样、结构化输出约束这些逻辑从零写一遍2000 行纯 Python 完全够用。而且搓完之后你对 SGLang 这类框架的理解会远比读十遍源码更深入。这篇文章就是我这段时间手搓一个 SGLang 风格 LLM 推理引擎的完整记录包括核心管线怎么设计、KV Cache 怎么管、连续批处理怎么调度、Radix Cache 怎么实现、结构化输出怎么做以及实测效果和踩过的坑。适合刚学完 PyTorch、想真正理解推理引擎内部发生了什么的人也包括在生产环境用着 SGLang/vLLM 但总觉得它们是黑盒的朋友。1. 先划清楚边界2000 行 Python 搓的到底是什么很多人听到“手搓 SGLang”第一反应是你打算用纯 Python 手写 Multi-Head Attention手写 LayerNorm那给你 2 万行都不够2000 行怎么可能跑得起来。这就是没划分清楚边界。1.1 哪些东西是借的哪些是真正自己写的我做这个项目的时候定的原则非常明确“推理引擎”不等于“模型本身”。模型的前向计算、权重存储、tokenizer 分词这种底层的脏活累活直接复用 HuggingFace transformers 和 PyTorch但引擎的核心调度逻辑、缓存管理、批量策略、采样控制、结构化生成这些在 SGLang/vLLM 里是独立的系统层逻辑跟我自己写的完全是一回事。我这样给模块分配代码量模块职责是否使用现成库核心代码量模型加载与前向计算权重、Attention、FFNtransformers torch不计数请求调度器队列管理、prefill/decode 分离、抢占纯自研约 400 行KV Cache 管理块分配、释放、共享前缀检测纯自研约 350 行采样器top-k、top-p、temperature、随机种子纯自研约 150 行Radix Cache前缀树缓存、LRU 淘汰纯自研约 450 行结构化输出约束JSON Schema / 正则约束解码自研状态机约 300 行协议层与并发HTTP 服务、SSE 流式、兼容 OpenAI APIFastAPI约 200 行调试与配置单步模式、日志、配置解析argparse约 150 行总计正好在 2000 行左右而且这还不算为了跑通吃掉的行数是纯逻辑代码。1.2 为什么“手搓引擎”不等于“手搓 Transformer”我见过很多人纠结这个点。这么说吧如果“手搓操作系统”必须从自己写 CPU 指令集开始那这个世界上就没有所谓“手搓 Linux”了。计算原语和引擎策略本来就是两个层级的东西。PyTorch 里的model(input_ids)一行就能跑通一次前向但你要让它支持 10 个并发请求、每个请求带不同的历史上下文、显存只够存 5 个请求的缓存这一行前向计算背后的“谁先算、谁后算、谁的缓存放在哪、谁要被踢出去”就是引擎该干的事。SGLang 的核心价值从来不在它的 Attention 算子多牛而在它的调度和执行策略能把算力用满、把重复计算消掉。所以我这套玩具引擎里model还是那个 transformers 的AutoModelForCausalLM但前向计算的每次输入拼成什么样、和谁共享缓存、哪些 token 被采样屏蔽全部是我说了算。这才是“引擎”二字的含义。2. 引擎主线一条请求从进队列到吐出最后一个字的完整过程搞清楚边界之后第一个要解决的问题是一个请求进来引擎内部到底走了哪几步。这一步搞不清楚后面所有优化都是空中楼阁。2.1 Prefill 与 Decode 两阶段的本质差异LLM 推理天然分成两个阶段Prefill预填充模型拿到一整段 prompt一次性并行计算所有 token 的注意力生成对应的 KV Cache。这个阶段计算量大但只跑一次。你可以把它理解成“做饭”——一次性把菜切好、配好、下锅出来的是一桌子半成品。Decode解码模型逐 token 生成输出每一步只输入上一个新 token通过 KV Cache 复用历史信息。这个阶段每一步计算量小但要循环很多次。相当于“端菜上桌”——每次只端一盘端几十上百次。这两个阶段的耗时特征完全不同。Prefill 延迟和 prompt 长度几乎成正比Decode 延迟则主要取决于模型大小和显存带宽。任何推理引擎的调度器本质上都在平衡这两类请求之间的关系。2.2 引擎主循环一切的核心我实现的引擎主体其实就是一个死循环。不要笑真实情况就是这样。while True: # 1. 调度器决定本轮执行哪些请求返回 prefill 批和 decode 批 schedule scheduler.schedule(requests, cache_manager) # 2. 执行前向计算借用 transformers 模型 logits model( input_idsschedule.input_ids, attention_maskschedule.attention_mask, position_idsschedule.position_ids, cache_positionschedule.cache_positions ) # 3. 采样器根据 logits 生成新 token new_tokens sampler.sample(logits, schedule) # 4. KV Cache 管理器保存新 token 对应的 K/V 块 cache_manager.store(new_tokens, schedule) # 5. 更新请求状态完成、继续、等待 scheduler.update(new_tokens)这几行代码看着简单但整个引擎的复杂度全隐含在schedule()和cache_manager的细节里。调度器要回答三个问题哪些请求该 prefill哪些请求该继续 decode显存不够了先踢谁2.3 采样器设计从 logits 到 token 的那一步采样是决定输出质量的关键环节也是我在实现过程中觉得“原来如此”的地方。LLM 前向计算输出的 logits 是个形状为[batch, seq_len, vocab_size]的张量但真正决定下一个 token 的只有最后一个位置的 logits。采样器做的就是三件事def sample(self, logits, schedule): last_logits logits[:, -1, :] # 只看最后一个 token 位置 # 1. temperature 缩放 if self.temperature 0: last_logits last_logits / self.temperature # 2. top-k 过滤只保留概率最大的 k 个候选 if self.top_k 0: indices_to_remove last_logits torch.topk(last_logits, self.top_k)[0][..., -1:] last_logits[indices_to_remove] float(-inf) # 3. top-p 过滤累积概率超过 p 的候选之外全部屏蔽 if self.top_p 1.0: sorted_logits, sorted_indices torch.sort(last_logits, descendingTrue) cumulative_probs torch.cumsum(torch.softmax(sorted_logits, dim-1), dim-1) sorted_indices_to_remove cumulative_probs self.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] last_logits.scatter_(1, indices_to_remove, float(-inf)) # 用 softmax 多项式采样得到最终 token probs torch.softmax(last_logits, dim-1) return torch.multinomial(probs, num_samples1).squeeze(1)这些逻辑初看不难但实现时有个细节非常坑top-p 的“移位”操作。你必须在累积概率超过阈值之后保留第一个超过阈值的 token 本身否则会把本来最该保留的 token 也屏蔽掉。我最初实现时没做这个 shift结果输出质量明显下降排查了很久才发现是这个 bug。3. KV Cache 管理决定推理引擎上限的核心系统如果说调度是引擎的骨架KV Cache 管理就是引擎的血肉。吞吐量、显存利用率、共享能力全都靠它。3.1 从“每请求一个 Tensor”到“分块分配”最朴素的 KV Cache 实现是一个请求分配一个独立的 tensor长度直接按max_seq_len预留。问题显而易见预分配 4096 长度的空间但实际只生成了 200 个 token95% 的显存全浪费了。而且不同请求之间完全无法共享任何数据。我用的方案是分块分配类似操作系统的内存分页。把 KV Cache 切成固定大小的 block我选的是 16 个 token 一块请求需要多少空间就分配多少块。当一个请求的输出 token 数超过当前块的容量再动态追加一个新块。class BlockManager: def __init__(self, total_blocks, block_size16): self.free_blocks list(range(total_blocks)) self.allocated {} # request_id - [block_ids] self.block_size block_size def allocate(self, request_id, num_tokens): need_blocks math.ceil(num_tokens / self.block_size) if len(self.free_blocks) need_blocks: return False blocks self.free_blocks[:need_blocks] self.free_blocks self.free_blocks[need_blocks:] self.allocated[request_id] blocks return blocks def append_block(self, request_id): if not self.free_blocks: return False block self.free_blocks.pop(0) self.allocated[request_id].append(block) return True def free(self, request_id): blocks self.allocated.pop(request_id, []) self.free_blocks.extend(blocks)这是最简化版本。实际生产级实现还要考虑 block 的引用计数、写时复制Copy-on-Write、跨请求的共享块拆分但核心思想就是这个。3.2 显存占用怎么估算一道必做的算术题很多人对 KV Cache 显存占用没有概念觉得“显存不够就加显存呗”直到自己算了这笔账才明白为什么推理引擎要做这么多优化。以 7B 参数的 FP16 模型为例计算一下每个 token 占用多少 KV Cache 显存每 token KV 显存 2K 和 V × 层数 × KV 头数 × 每个头的维度 × 每个元素字节数 以 Llama-2-7B 为例 层数 32 KV 头数 8GQA 配置注意不是 32 个 Q 头 每个头的维度 128 每个元素字节数 2FP16 每 token KV 显存 2 × 32 × 8 × 128 × 2 131,072 字节 128 KB这意味着什么一张 4090 的 24GB 显存在模型权重加载之后约 14GB剩下约 10GB 能用于 KV Cache。如果每 token 需要 128KB那全部显存也只能放大约 80K 个 token 的 KV 数据。并发请求数每个请求平均输出长度需要的 KV 总量占可用显存比例4512 token256 MB2.5%161024 token2 GB20%642048 token16 GB160% 直接爆这就是为什么 vLLM 要搞 PagedAttentionSGLang 要搞 Radix Attention。显存永远不够用只有在“怎么用”上下功夫才是正道。3.3 淘汰策略当显存满时踢谁出去跟操作系统内存管理一样KV Cache 满了之后必须淘汰。最简单的策略是先来先服务FIFO把最早分配的块释放掉。但推理场景有个特殊性如果一个请求是长对话多轮续写把它踢出去意味着下次续写时整个上下文要全部重新 prefill代价极高。我用的方案是优先级分层请求内部的 loop 状态比如正在生成 tool call 的中间结果权重最高不允许被抢占其次是正在 decode 的请求最次是排队等 prefill 的请求。被抢占的请求会整体释放 KV 块等资源充足后重新从 prompt 开头 prefill。这种“全量重算”方案比真正把 KV 换到 CPU 内存的 swap 方案简单得多代价是碎片化延迟但在学习型引擎里完全够用。4. 连续批处理与调度器从一次等一批到随时插队随时走调度器是整个引擎里最容易理解错的部分。很多人以为调度就是“请求多了排队”完了大错特错。4.1 静态批处理的浪费最原始的 batching 思路是攒够 N 个请求一次性前向计算一批大家同步解码直到全部生成完毕再接收下一批。你看完 2.1 的 prefill/decode 分析就知道问题在哪了一批里如果有一个请求生成长度是 1024其他请求生成长度 128那所有 128 的请求都得陪跑到 1024 结束期间显存和算力都被那个慢请求拖住短的请求明明生成完了却占着资源不放。这就是 SGLang 和 vLLM 都在主推**连续批处理Continuous Batching**的原因。它的核心变化是不再以“请求”为调度单位而是以“迭代”为调度单位。4.2 迭代级调度每个 step 都重新决策我实现的调度器每个迭代都重新评估一次所有请求的状态def schedule(self, requests, cache_manager): prefill_batch [] decode_batch [] available_blocks cache_manager.get_free_blocks() # 正在 decode 的请求优先继续跑下一轮 for req in requests: if req.state decoding: decode_batch.append(req) # 如果显存还有空间从等待队列里捞 prefill 请求 for req in requests: if req.state waiting and len(prefill_batch) self.max_prefill_num: blocks_needed math.ceil(req.prompt_len / cache_manager.block_size) if available_blocks blocks_needed: prefill_batch.append(req) req.state prefill available_blocks - blocks_needed return prefill_batch, decode_batch这个设计的精妙之处在于一个请求生成了 16 个 token 之后如果显存紧张它下一轮就不会被调度到等别人释放块之后再回来继续。短的请求快速跑完释放 KV 块长的请求被断续执行整体吞吐一下就上去了。但这里有个重要细节prefill 和 decode 在一个 batch 里混跑时要保证 prefill 的 padding 不会浪费 decode 的算力。transformers 模型的 forward 对不定长输入需要 pad 到相同长度这个 pad 会在 prefill 阶段产生明显的算力浪费。SGLang 用的办法是“分两次 forward”——先跑 decode 批再单独跑 prefill 批两个批之间把注意力掩码分开。我后来也采用了类似方案虽然代码复杂度高了但实测吞吐能提升 20% 以上。4.3 抢占与优先级显存不够时谁先出局前面 3.3 提到了优先级调度器的抢占逻辑是这样的def preempt(self, requests): # 从 decode 批里挑出优先级最低的请求释放其 KV Cache preemptible [r for r in requests if r.state decoding and r.is_preemptible] if not preemptible: return False victim min(preemptible, keylambda r: r.priority) cache_manager.free(victim.request_id) victim.state waiting # 回到等待队列从头 prefill return True抢占是推理引擎里最考验工程能力的地方。SGLang 有更精细的“换入换出”机制可以把被抢占请求的中间状态搬到 CPU 内存下次直接恢复不用从头算。但我在 2000 行的范围内选择“全量重算”是因为对学习场景来说这个简化能大幅降低代码复杂度而且实际运行时抢占发生的频率远低于想象——前提是控制好并发数。5. Radix Cache用手搓的方式理解 SGLang 的招牌技术坦白说我当初决定做这个项目一半原因就是想把 Radix Cache 彻底搞明白。SGLang 的论文里说它能把共享前缀的重复 prefill 计算减少 89%这个数据让我一直很震惊。5.1 重复前缀问题到底有多严重想象一个典型场景你做了一个客服机器人每个请求的开头都是长达 1000 token 的 system prompt包含业务规则、回复格式、历史对话摘要。没有 Radix Cache 时1000 个并发请求进来这 1000 token 的 prefill 要被重复计算 1000 遍。再想象 Agent 场景多轮工具调用每一轮都会把之前的全部上下文重新拼接发送给模型。假设一轮对话累计上下文 3000 token工具调用 10 次那第 10 次时前 9 轮的 2700 token 都是历史重复计算。重复 prefill 不仅浪费显卡算力还直接增加首 token 延迟TTFT。Radix Cache 的核心思想就是把这些重复的 token 序列的前缀部分以树状结构缓存起来新请求来了直接复用只计算差异部分。5.2 前缀树缓存用 dict 实现一个简易 Radix Cache不扯复杂的论文公式我用最简方案实现了前缀树缓存。每个节点保存一段连续的 token_ids并记录对应的 KV Cache 块引用class RadixNode: def __init__(self, parent, token_ids, block_refs): self.parent parent self.children {} self.token_ids token_ids # 这段节点存的 token 序列 self.block_refs block_refs # 对应的 KV 块引用列表 self.last_access_time 0 self.total_tokens (parent.total_tokens len(token_ids)) if parent else len(token_ids) class RadixCache: def __init__(self, block_manager): self.root RadixNode(None, [], []) self.block_manager block_manager def match_prefix(self, prompt_token_ids): 在树上查找最长公共前缀返回命中的 token 数和对应的块引用 node self.root matched 0 refs [] for t in prompt_token_ids: if t not in node.children: break child node.children[t] child_tokens child.token_ids hit_num min(len(child_tokens), len(prompt_token_ids) - matched) for i in range(hit_num): if child_tokens[i] ! prompt_token_ids[matched i]: break matched 1 if matched 0: refs.extend(child.block_refs[:hit_num]) node child if matched len(prompt_token_ids): break return matched, refs匹配逻辑本身不复杂真正麻烦的是“新请求命中了前缀但前缀只覆盖了部分 KV 块剩下的 KV 块需要拼接”这个问题。我的做法是让匹配返回的块引用列表和需要继续 prefill 的 token 序列一起送给模型。模型只需要对未命中的部分进行 prefill然后把新产出的 KV 块追加到缓存树上。这里有一个必须注意的坑你能不能直接复用命中的 KV 块取决于模型是否要求 position_ids 连续。如果请求上下文在前缀之后还有自己的独有内容position_ids 的起始位置必须和命中块的长度对齐否则 Attention 掩码会错位输出瞬间崩掉。5.3 命中率实测多轮对话场景下的惊人提升我写了个简单的 agent 模拟器模拟 50 个并发用户进行 5 轮多轮对话每轮对话携带完整历史上下文。对比开不开 Radix Cache 的 TTFT场景无 Radix Cache TTFT有 Radix Cache TTFT提升第一轮首次对话850ms850ms无第三轮带上下文对话2400ms980ms59%第五轮长上下文对话3800ms1050ms72%这个测试让我彻底理解了为什么 SGLang 团队要把 Radix Cache 作为核心卖点。在一半以上的常见业务场景里Radix Cache 带来的延迟降低远比换一张更好的显卡来得直接。不过我踩了一个很恶心的坑Jinja 模板中的 BOS token 处理。transformers 的 chat template 在 prompt 开头自动加上|im_start|这类特殊 token而如果我在构建 token_ids 时把 BOS 丢了Radix Cache 的前缀匹配永远命中不了。后来我统一了“所有请求先走 tokenizer.apply_chat_template再送入匹配器”的流程才彻底解决。6. 结构化输出让模型吐 JSON 不再靠运气做 LLM 应用的人最头疼的问题之一就是模型输出格式不稳定。你说“请返回 JSON”它能给你返回 markdown 代码块、多余的注释、缺了尾括号。传统办法是让模型先输出再校验不合法就重试但这个方案在复杂 schema 下重试率极高浪费延迟还影响体验。6.1 约束解码的基本思路SGLang 的 structured generation 做法是在采样阶段直接把不合法的 token 屏蔽掉。原理听起来特别简单模型每一步要生成 token 时我都用格式约束算出一张“合法 token 掩码”凡是会让 JSON 解析失败的 token概率直接置为负无穷。这样模型只能在合法范围内选择输出的连续性自然保证。6.2 我用状态机做了一个极简 JSON Schema 约束器我实现了一个非常初级的版本基于状态转移的约束解码class JSONSchemaConstraint: def __init__(self, schema): self.schema schema self.state start self.stack [] self.field_order [] def generate_mask(self, vocab_size, tokenizer): mask torch.ones(vocab_size, dtypetorch.bool) allowed_tokens self._get_allowed_tokens() # 只有允许的 token 对应的位置保持 True for token_id in allowed_tokens: if token_id vocab_size: mask[token_id] False return mask def _get_allowed_tokens(self): if self.state start: return [tokenizer.encode({)] elif self.state in_object: return [tokenizer.encode(f{field}:, ) for field in self.schema[properties]] # ...这版实现非常粗糙实际生产大家直接抄 SGLang 用的 outlines 库思路就行但自己实现一遍的价值在于我第一次真正理解了“正则表达式能描述的语言根本不用靠提示词让模型理解直接硬编码到采样过程里就行”。输出合法率从最开始的 80% 直接拉到 99.5% 以上。6.3 一个必须注意的 trade-off约束解码不是免费的。屏蔽掉非法 token 后解码的搜索空间被大量压缩模型输出质量会有微小下降尤其在需要模型自由发挥的文本场景下。另外约束器本身的计算开销在复杂 schema 下可能占到总解码时间的 5%-10%。我的建议是只在需要机器解析的输出函数调用参数、数据库查询、API 返回结构上用约束解码自由文本对话不要用。这也是为什么 SGLang 把这个能力做成可选而不是默认。7. 服务化与并发把引擎包装成一个 OpenAI 兼容接口引擎内部跑通了但真实使用场景里你得让它成为一个可被外部调用的服务。我选了一个非常标准化的方案OpenAI 兼容接口 FastAPI SSE 流式输出。7.1 为什么必须兼容 OpenAI API现在市面上的 LLM 应用框架Dify、LangChain、各类 Agent 框架大多默认兼容 OpenAI 接口。如果你的引擎能直接支持/v1/chat/completions就能无缝接入整个生态完全不用改任何下游代码。核心实现其实就是把请求体解析成 prompt调用引擎返回标准格式app.post(/v1/chat/completions) async def chat_completions(request: ChatCompletionRequest): prompt build_prompt_from_messages( request.messages, request.tools, tokenizer ) result await engine.generate(prompt, request) return ChatCompletionResponse( choices[{ message: {role: assistant, content: result.text}, finish_reason: stop }], usageresult.usage )7.2 流式输出的实现细节流式输出SSE是另一个坑。普通请求直接等结果返回就行但流式输出要求模型每次生成一个 token 就把数据推给客户端这就要把引擎的 decode 循环和 HTTP 响应流绑定在一起。我用 FastAPI 的StreamingResponse配合生成器实现async def stream_response(request, engine): async for chunk in engine.stream_generate(request): yield fdata: {chunk.model_dump_json()}\n\n yield data: [DONE]\n\n app.post(/v1/chat/completions) async def chat_completions(request: ChatCompletionRequest): if request.stream: return StreamingResponse(stream_response(request, engine)) return await engine.generate(request)这里有个容易忽略的细节SSE 的每条消息必须以data:开头末尾必须有两个换行。很多人在自己实现时漏了格式导致前端 EventSource 无法解析浪费了大量时间。7.3 异步与并发的冲突FastAPI 是异步框架而我的引擎主循环是同步的。如果直接在异步接口里调用同步的engine.generate会阻塞事件循环导致新的 HTTP 请求进不来。解决方案是让引擎在后台线程池运行或者把引擎主循环改成事件循环驱动。我选择了最低成本方案在 FastAPI 里用run_in_executor把引擎调用扔到线程池同时引擎内部的请求调度仍然由我自己的调度器控制。这样 HTTP 层和推理层解耦请求排队逻辑统一归调度器管理result await loop.run_in_executor( None, engine.generate, request )8. 实测数据我的玩具引擎和 SGLang 的差距有多大代码写完了总得拉出来跑跑。我在一台 A100 80GB 上做了对照测试。8.1 测试环境项目配置GPUNVIDIA A100 80GBCPUAMD EPYC 774232 核分配模型Qwen2.5-7B-InstructFP16批量大小动态由调度器自动决定测试请求500 条混合长短 prompt正版参照SGLang v0.4.1相同模型和请求集8.2 结果对比指标我的引擎SGLang差距单请求延迟短 prompt720ms680ms5.9%单请求延迟长 prompt Radix 命中1050ms950ms10.5%吞吐并发 32无 Radix28 req/s86 req/s67%吞吐并发 32Radix 命中45 req/s140 req/s68%说实话单请求延迟能接近已经让我很惊讶了因为这里大家的瓶颈都是显存带宽和模型前向引擎调度的优化空间有限。但并发吞吐这块差距巨大原因很清楚SGLang 有专门的 CUDA kernel 做 paged attention它的 block 分配是在显存地址层面直接映射的而我的实现要先通过 transformers 把 KV Cache 切成不连续的 tensor 再拼接这中间有大量 H2D 拷贝开销。引擎逻辑我能用 2000 行 Python 模拟但算子的性能差距不是靠 Python 层优化能追回来的。8.3 这种“玩具引擎”到底有什么用如果你直接拿它去怼高并发生产流量那肯定是自找麻烦。这个项目的真正价值体现在三个地方教学和面试准备把 SGLang/vLLM 的核心机制亲手实现一遍比背十篇源码分析文章都有用。快速原型验证想测试一个新的调度策略、缓存策略改 2000 行 Python 比改几万行 C 快得多。深度定制化场景比如给特定业务写死一个极简推理服务不需要 SGLang 的完整功能这套代码改起来非常顺手。9. 避坑记录把引擎搓完才知道的事最后这部分写给所有想动手手搓推理引擎的朋友。我踩过的这些坑希望你绕过去。9.1 采样随机性会把 bug 排查拖入地狱第一次跑通整条管线时我发现同样的输入有时输出正常、有时输出胡说八道。排查了半天甚至怀疑是 KV Cache 共享逻辑写错了最终原因极其愚蠢采样器里用了torch.multinomial但没有固定随机种子。每次调用结果不同导致任何微小的浮点误差都会被放大成完全不同的输出。排查思路也分享一下当你怀疑“这个 bug 是不是随机的”第一时间要做的是固定所有随机种子包括 Python 的random、NumPy 的、PyTorch 的然后跑两次完全相同的输入。如果两次结果不一致问题一定在采样或并发竞争如果两次结果一致了再去对比正确输出和错误输出范围一下子缩小很多。9.2 显存碎片块分配的隐藏杀手KV Cache 的块管理器用 Python list 维护 free list初看没问题但跑了几百个请求之后我发现可用块数量并没有因为请求完成而恢复到理论值。原因是某些 token 块既被请求 A 引用又被 Radix Cache 节点引用我的释放逻辑在请求完成时不管三七二十一直接把块全部释放了导致 Radix Cache 里那些本该复用的共享块全部失效。修复方案是给块管理器加引用计数。每次分配块时ref_count 1被 Radix Cache 引用时ref_count 1请求完成时ref_count - 1只有计数归零才真正把块归还给 free list。这个改动让长跑场景下的缓存命中率稳定了很多。9.3 eager mode 下 CPU 与 GPU 的同步陷阱我开发时为了调试方便把 transformers 模型设置成了 eager mode不走 torch.compile。这导致一个隐蔽的性能问题每个 decode step 都在“CPU 发起 kernel → GPU 执行 → CPU 等结果”之间反复同步GPU 利用率永远上不去。后来用torch.profilerprofiling 才发现真正的前向计算只占了 GPU 时间的 40%其余全在同步等待。生产级推理引擎会想办法把多个请求的 decode 计算合并成一个大 kernel减少同步开销。在 Python 这个层面我能做的最有效的优化是尽量把更多的 decode 请求塞进同一个 batch一次前向计算处理更多数据这个优化直接让吞吐提升了一倍。9.4 调试建议先做单步模式再做并发2000 行代码的并发 bug 是最难缠的。我的血泪经验是先实现单请求模式把所有流程跑通并且输出质量稳定之后再上并发。并发模式下宁可让吞吐难看也要先保证正确性。把每次调度的结果写到日志里再做一个“回放器”可以手动回放某个请求从进入队列到完成的完整时间线排查那种“为什么这个请求被卡住了”的问题。如果要做增量优化建议按照这个顺序来先修正确性再做 KV Cache 块重用再做 Radix Cache最后碰调度和抢占。前面没做对后面全白搭。手搓推理引擎这件事最大的收获不是“我也写了个 SGLang”而是当你再回头看 SGLang 的官方文档和源码时那些抽象的名词开始一个接一个地落到具体的函数和数据结构上。如果哪天你也想造一个轮子来理解 SGLang我强烈建议从 14B 以下的模型开始先跑通单请求再逐步加上批处理、缓存、约束。祝你搓得愉快。