SGLang

§ 00 · PrologueWhat we're after.

SGLang 这次精读的目标，是把"高吞吐、低延迟、跨硬件"这三件事——分别由 RadixAttention、Continuous Batching with Overlap Scheduling、可插拔 Attention Backend 三大子系统承担——拆开来看个明白。它的源码量比 vLLM 还大；而且不像 PyTorch 那种"按 op/module 切片"的结构，SGLang 是按系统组件切片（tokenizer / scheduler / runner / cache / kernel），所以读起来更像读一个分布式系统。

精读分 10 个模块（M0 → M9）。本档案是读后的总览：Tokenizer/Scheduler/Workers 的三进程拓扑、ZMQ 消息总线、Overlap 调度器的时间线、RadixCache 的树形结构、27 个 attention backend 中 aiter_backend.py 给 MI300X 提供的入口，以及 sgl-kernel 这个独立 wheel 包的多平台构建系统。

§ M0 · 20 minThe lab, laid out.

SGLang 不是单包仓库 — 顶层就有 python/（Python 主包）+ sgl-kernel/（独立的 kernel wheel）+ rust/（一些 Rust 组件）+ sgl-model-gateway/（gateway 服务）+ proto/（gRPC 协议定义）。这种"主包 + 多 sub-package"结构反映了不同组件不同节奏发布的工程现实。

Python 主包 python/sglang/ 分两大块：

sglang/srt/ 进一步切分（这是真正的"系统"）：

§ M1 · 30 minThe Engine, started.

用户最常用的入口有两条路径：(1) HTTP 服务模式 — python -m sglang.launch_server --model X --tp 8；(2) Python in-process 模式 — from sglang import Engine; engine = Engine(model="X")。两条都最终落到 sky/entrypoints/engine.py 的 Engine 类（1349 行，line 174）。

class SchedulerInitResult:
    """Engine 启动期间收集每个 scheduler 子进程的初始化结果。"""

class Engine(EngineScoreMixin, EngineBase):
    """The main engine class. Holds the tokenizer manager,
    spawns scheduler procs, manages their lifecycle."""

Engine 启动序列（line 174-1158 大概结构）：

1. _set_envs_and_config（line 1159）— 设环境变量（CUDA_VISIBLE_DEVICES、ROCm 等），GC 调优
2. init_tokenizer_manager（line 126）— 起 TokenizerManager 进程
3. _compute_parallelism_ranks（line 1332）— 算 TP × DP × EP 拓扑下的 rank 分配
4. spawn N 个 scheduler 子进程（multiprocessing），每个进程跑一个 Scheduler.event_loop()
5. _wait_for_scheduler_ready（line 1265）— 等所有 scheduler 通过 ZMQ ready barrier
6. HTTP / gRPC server 接管（如果是 server 模式）

关键细节：每个 TP rank 是一个独立 scheduler 进程（不是 thread），用 multiprocessing spawn。Tokenizer 是单独一个进程，Detokenizer 又是单独一个。所以一个 TP=8 的部署有 1 + 1 + 8 = 10 个 Python 进程 共同协作。

§ M2 · 25 minZMQ — the wires.

四个进程之间的消息总线是 ZeroMQ。SGLang 用 zmq 库的 PUSH/PULL 模式（不是 PUB/SUB）— 因为请求需要严格的 producer→consumer 单向流动，PUSH/PULL 比 PUB/SUB 多一层 fan-out / load-balancing 语义。

主要 socket 拓扑：

注意是环形拓扑：Tokenizer → Scheduler → Worker → Scheduler → Detokenizer → Tokenizer。Tokenizer 既是入口也是出口，负责把 streamed 输出聚合后通过 HTTP 推给客户端。

§ M3 · 60 minThe Scheduler — 4006 lines in one class.

managers/scheduler.py 是 SGLang 整个 codebase 里最重的一个文件。class Scheduler 从 line 324 开始，__init__ 在 line 340，光构造函数就有几百行 — 它要初始化 ZMQ sockets、tokenizer、tp_group、model_runner、KV cache pool、radix cache、sampling、speculative decoding 等几乎所有子系统。

核心方法（按调用频率排）：

Scheduler 类用了几个 Mixin 模式来切分代码（line 248 / 324 等）：

• SchedulerMlxOverlapMixin — Apple MLX 后端特化
• scheduler_dp_attn_mixin — DP attention 特化
• scheduler_output_processor_mixin — 输出后处理
• scheduler_input_blocker — 输入背压

这种"主类 + Mixin"是 SGLang 切分 4006 行的方式。但这也是个学习成本 — 跳到一个方法时不一定知道它属于哪个 Mixin。读源码时记得 grep "def method_name" 而不是 grep "class .*:"。

# 简化后的 event_loop_overlap 核心结构
def event_loop_overlap(self):
    """Overlap: while GPU runs forward(t), CPU plans forward(t+1)"""
    last_batch = None
    while True:
        # 1) 从 ZMQ 收新请求，加入 waiting queue
        recv_reqs = self.recv_requests()
        self.process_input_requests(recv_reqs)

        # 2) 决定下一个 batch (CPU work)
        batch = self.get_next_batch_to_run()

        # 3) 启动 forward (GPU work, async via cuda stream)
        if batch is not None:
            future = self.forward_async(batch)  # ← non-blocking

        # 4) 等上一个 batch 的 GPU 结果 (overlapped with #2/#3)
        if last_batch is not None:
            result = last_batch.future.wait()
            self.process_batch_result(last_batch, result)

        last_batch = batch

§ M5 · 40 minRadixAttention, the prefix tree.

SGLang 在论文里的最大卖点是 RadixAttention — 把多个请求共享的 prompt prefix 自动复用 KV cache。实现在 mem_cache/radix_cache.py（828 行），三个核心类：

class RadixKey:        # 66  — token sequence → hashable key
class TreeNode:        # 206 — radix tree 的节点 (children, key, value, ref_count)
class RadixCache(KVCacheEventMixin, BasePrefixCache):  # 269

Tree 上每个节点：

• key: 一段 token 序列（RadixKey，可哈希）
• value: 这段 tokens 对应的 KV cache 的物理 slot indices
• children: 后续 token 派生出的子节点
• ref_count: 多少个活跃请求正在引用这段 prefix
• last_access: LRU 时间戳

核心操作（来自 BasePrefixCache 抽象基类）：

mem_cache/ 还有 18 个其他文件 — 不同形态的 cache：hiradix_cache.py（多级 hierarchical）、chunk_cache.py（不用 radix 的简单 chunk）、mamba_radix_cache.py（Mamba/SSM 模型的特化）、hisparse_memory_pool.py（稀疏激活）等。还有 cpp_radix_tree/ — 一个 C++ 实现的 radix tree，给极端性能场景用。

§ M6 · 35 minThe ModelRunner — 3607 lines.

model_executor/model_runner.py 是另一个怪兽。class ModelRunner(ModelRunnerKVCacheMixin) 在 line 323，ModelRunnerOutput dataclass 在 315。它负责的事：

1. 从 HF 或自定义 loader 加载模型权重
2. 分配 KV cache memory pool（巨大 tensor，按 layer × block × head 切片）
3. 选择并初始化 attention backend（从 27 个候选里挑）
4. build / capture CUDA Graph（line 2339-2347 的关键分支）
5. 每次 forward(batch) 时把 ScheduleBatch 翻译成 ForwardBatch（含位置编码、cache indices、mask 等），调 self.model.forward(...)
6. 给 attention backend 调 init_forward_metadata(forward_batch) 准备元数据（page table / sliding window mask）
7. 对接 sampling（top_k / top_p / 结构化生成）

CUDA Graph 这一段值得单独说 — SGLang 提供多种 capture 策略：

• cuda_graph_runner.py — 普通整张图 capture
• piecewise_cuda_graph_runner.py — 分段 capture，更灵活
• breakable_cuda_graph_runner.py — 可中断的 capture（用于 disagg P/D）
• cpu_graph_runner.py — CPU 版本（用于 host-side 模拟）
• mindspore_runner.py — MindSpore 平台特化

§ M7-M8 · 50 minTwenty-seven backends, one interface.

layers/attention/ 是 SGLang 跨硬件能力的核心 — 27 个 attention 实现共享一个 AttentionBackend 抽象基类（base_attn_backend.py），attention_registry.py 维护"backend 名字 → class"映射。

AITER 后端 — 给 MI300X 的入口

aiter_backend.py 一个文件就 3284 行。class AiterAttnBackend(AttentionBackend) 在 line 117。它在 SGLang 里的位置：用户跑 sky launch --infra k8s ... sgl --attention-backend aiter，scheduler 用 AITER 的 forward / metadata 接口，背后调 aiter 这个独立 Python 包（github.com/ROCm/aiter）的 CK / hipBLASLt kernel。

class AiterAttnBackend(AttentionBackend):    # line 117
    def init_forward_metadata(self, forward_batch):
        # build page table, mask, sliding window indices
        ...
    def forward_extend(self, ...): ...    # prefill
    def forward_decode(self, ...): ...    # generation
    def forward_target_verify(self, ...): ...  # speculative decoding

class AiterIndicesUpdaterPrefill:            # line 2962
class AiterMlaIndicesUpdaterPrefill:         # line 3064  ← MLA variant
class AiterMultiStepDraftBackend:            # line 3146  ← speculative draft

三个 helper 类分别处理：

• Indices updater — 把 SGLang 的 RadixCache slot indices 翻译成 AITER 期望的 page table 布局（每个 kernel 都要求一种特定的内存布局）
• MLA indices updater — DeepSeek / Kimi-K2 的 Multi-head Latent Attention 特化（压缩 KV）
• MultiStep draft — 投机解码 draft model 的多步 forward 优化

§ M9 · 25 minsgl-kernel — five wheels, one tree.

sgl-kernel/ 是 SGLang 工程组织最值得偷的部分 — 把 kernel 代码做成独立 wheel 包，5 个平台各一份 pyproject_*.toml，但共享同一份 C++ / Triton 源码。

包结构：

sgl-kernel/
├── csrc/             # C++ / CUDA / HIP 源码
├── include/          # 头文件
├── python/           # Python binding（pybind11 / cppextension）
├── cmake/            # 各平台的 CMake module
├── CMakeLists.txt    # 主 CMake
├── build.sh          # 构建脚本
├── Makefile
├── pyproject.toml    ← CUDA
├── pyproject_rocm.toml
├── pyproject_cpu.toml
├── pyproject_musa.toml
├── setup_metal.py
└── setup_musa.py

§ Traps · five of themWhat I wish I'd known.

读 SGLang 时最容易栽的几个坑。每个坑都对应一类具体的失误模式。

§ Red lines · threeThe big questions.

§ AMD · for JhinWhere to start touching.

基于这次精读，给你 AMD 工作的具体 actionable 起手点。

Workstream 1 — AITER kernel 调试

你 CLAUDE.md 里描述的 agent team："Kernel Debugger" / "SGLang Integrator" / "Arch Tracer"。具体到代码位置：

Workstream 2 — Grid search for MI300X / MI355X

当你跑 Kimi-K2.5 / Qwen3-Coder-Next 的 throughput 实验时：

• Tier 1 attention backend sweep: --attention-backend ∈ {aiter, triton, fa3}。固定 TP=8，看哪个吞吐高。
• Tier 2 parallelism: 在最优 backend 下扫 TP / EP / DP（EP=2/4/8）。
• Tier 3 CUDA Graph + mem: --enable-torch-compile / --enable-cuda-graph / --mem-fraction-static。

每次实验都用 SGLang 自带的 bench_serving.py（你已经在用）。结果存进性能数据库 — 这正是你 CLAUDE.md 提到的"performance database is the moat"。

Workstream 3 — 长期：把你的 kernel 做成独立 wheel

如果 multi-agent kernel optimization 系统产出了若干"AMD 上更快的 kernel"，仿照 sgl-kernel/ 把它们做成 amd-kernels wheel：

• 独立 pyproject.toml，only ROCm platform tag
• CMake + hipcc 构建管线
• 暴露 Python API 给 SGLang / vLLM / 任何 PyTorch 程序直接调用
• 独立 release 节奏 — kernel 改进不用等 SGLang 主包发版

这是把你的工作 leverage 出 AMD 团队、影响整个开源 ecosystem 的最直接路径。SGLang 的 sgl-kernel、vLLM 的 vllm-flash-attn 都是这种独立 kernel 包，已经形成工业惯例。

§ EpilogueWhat we found, and what's next.

SGLang 不是一个 fast LLM server — 它是一个用四进程 + ZMQ 总线 + 27 个 attention backend + RadixAttention 共同实现的 inference orchestration system。它的工程价值在三个层面：(1) 多进程 IPC 的精巧解耦，(2) 跨硬件 attention 多态的优雅抽象，(3) 把"用户 prompt prefix"这种应用层语义带进系统层做共享的洞察。

下一份 File № 003 是 vLLM 的对应分析 — 它走的是不同路径（PagedAttention 而非 RadixCache、v1 重构后的 EngineCore vs Worker 拆分）。两份合起来看才能理解"现代 LLM serving 的设计空间"。

— Fin.