框架Learning

1. SGLang：高性能大语言模型推理框架

SGLang 是一个专为低延迟、高吞吐推理设计的先进、高性能 LLM 服务框架。其架构经过精细划分，能够从容应对现代大规模语言模型及多模态模型的复杂挑战。

1.1 核心架构组件

1. SGLang 运行时 (SRT)

SRT 是系统的核心，专为极致效率而设计，由以下几个关键子模块组成：

• **TokenizerManager**：基于 FastAPI 的异步进程，是请求的主要入口。它负责反序列化、分词（tokenization）和多模态数据处理，将这些 CPU 密集型任务从核心调度逻辑中分离出来。

• Scheduler（调度器）：运行时的“大脑”。负责管理请求批处理、调度策略（如连续批处理/continuous batching），并与 GPU 工作线程协同。RadixCache 也驻留在此处，实现在具有相同前缀的不同请求间高效共享 KV 缓存。

• ModelExecutor（模型执行器）：负责实际的模型前向传播。它利用 CUDA Graphs 来最小化 CPU 开销，并集成了针对注意力机制和量化等性能关键操作的自定义算子（kernels）。

• **RadixCache**：一种特殊的内存管理系统，以基数树（radix tree）结构组织 KV 缓存。这使得前缀匹配和重用几乎在瞬时完成，显著减少了常见提示词（prompts）的重复计算。

2. SGLang 语言 (DSL)

一种高级编程接口，允许开发者使用 Python 原生风格的编程原语来表达复杂的 LLM 工作流（如多步推理、结构化输出生成）：

• **api.py**：提供面向用户的函数，如 sgl.gen、sgl.select 和 sgl.assistant。

• **Interpreter & Tracer**：这些组件管理 SGLang 程序的执行流，既可以直接解释执行，也可以将其追踪（trace）为中间表示，以便在 SRT 后端进行优化执行。

3. SGLang 算子库 (Kernel)

一个专门的自定义 CUDA 和 ROCm 算子库（sgl-kernel/），针对底层进行了深度优化：

• 注意力机制：针对 PagedAttention、FlashAttention 和 DeepSeek 的 MLA 架构定制的算子。

• 量化支持：支持高性能的 FP8、FP4、INT4 和 AWQ 操作。

• 采样：优化的 Logit 处理和受限采样（使用状态机 FSM 进行 JSON/语法约束）。

4. SGLang 模型网关

一个由 Rust 实现的高性能代理（sgl-model-gateway/），用于：

• 路由与负载均衡：在多个 SGLang 实例之间分配请求。

• 监控与可观测性：为大规模部署提供集中的指标收集和追踪。

• gRPC 客户端：为网关提供健壮且低开销的通信通道。

---

1.2 请求生命周期

1. 接入：请求（如文本、图像）通过 HTTP (FastAPI) 或 gRPC 到达 TokenizerManager。

2. 分词：请求文本被转换为 token，多模态数据进行预处理（如图像缩放、归一化）。

3. 调度：分词后的请求通过 ZeroMQ 发送到 Scheduler。调度器检查 RadixCache 中的 KV 缓存命中情况，并将请求加入内部队列。

4. 批处理：调度器将请求打包成动态批次，优先处理缓存重用率高的请求或根据配置的调度策略进行处理。

5. 执行：批次被分发给 ModelExecutor 工作线程（可能通过张量并行或流水线并行分布在多个 GPU 上）。

6. 采样：模型输出的 Logit 由采样算子处理（如贪婪采样、Top-p 采样或通过 FSM 进行的约束采样）。

7. 反分词：生成的 token 发回 DetokenizerManager，将其转换为文本并将响应流式传输给用户。

---

1.3 性能创新总结

• RadixAttention：利用 Prompt 中的公共前缀重用 KV 缓存，带来显著的吞吐量提升。

• 零开销 CPU 调度器：最小化调度循环的开销，确保 GPU 不会因为 CPU 处理速度而产生瓶颈。

• Prefill-Decode 分离（计算解耦）：将预填充（Prefill）和解码（Decode）阶段分离到不同的实例或节点上。这种解耦可以针对计算密集型的 Prefill 和访存密集型的 Decode 进行独立优化，防止长时间的 Prefill 任务导致 Decode 任务出现高抖动或停顿。

• Structured Output Steering（结构化输出引导）：通过内置的有限状态机（FSM）引导采样过程。与在外部进行后处理不同，SGLang 在算子层级集成了格式约束（如 JSON Schema 或正则表达式），极大地提升了生成结构化数据的速度和准确性。

• Multi-GPU Tensor Parallelism（多 GPU 张量并行）：原生支持跨多显卡的模型并行。利用经过优化的通信算子（如基于 NCCL 的全归约），将大型模型拆分并分发，从而支持超大规模参数模型的高效推理。

• Data-Dependent Control Flow（数据依赖型控制流）：借助 SGLang 的 DSL 接口，系统能够根据模型中间生成的 Token 实时决定后续的推理路径（如 if-else 分支），而无需频繁在 Python 后端和 GPU 执行引擎之间切换上下文。

1.4 RadixCache

RadixCache 是 SGLang 的核心创新之一，它将 LLM 服务的 KV Cache 管理从传统的“线性队列”演进为“前缀树（Radix
Tree）”。这种结构允许系统在多个请求之间高效地复用、共享和回收 KV Cache。

1. RadixCache 的核心概念

在传统的 LLM 推理（如 vLLM）中，KV Cache 通常以请求为单位进行管理。即使两个请求有共同的前缀（例如相同的
System Prompt 或 Few-shot 示例），它们也可能重复计算并存储相同的前缀 KV Cache。

RadixCache 的做法：

• 前缀树结构：将 Token 序列看作路径，节点存储对应的 KV Cache 块。

• 多级复用：如果请求 A 的前缀是 [1, 2, 3]，请求 B 的前缀是 [1, 2, 4]，它们会共享节点 [1, 2] 的 KV
  Cache，只在 3 和 4 处产生分支。

• 动态演进：随着推理的进行，树会不断生长。当生成结束时，这些节点不会立即释放，而是保留在树中供后续请求匹
  配。

2. 显存管理 (VRAM Management)

SGLang 的显存管理是基于 分页内存管理 (PagedAttention) 的思想，但增加了一层逻辑索引（Radix Tree）。

A. 逻辑与物理分离

• 物理池 (Physical Pool)：在初始化时，SGLang 会预分配一大块 GPU 显存，切分为固定大小的“块（Blocks）”。

• 逻辑节点 (Logical Nodes)：RadixCache 中的每个节点都记录了它对应的物理块索引（Token IDs -> Physical
  Block Indices）。
  B. 引用计数 (Reference Counting)
  每个 Radix 节点维护一个 ref_count：

• ref_count > 0：表示当前有正在运行的推理请求正在读取或写入该节点。这些节点是“锁定”的，绝对不能被置换。

• ref_count = 0：表示该节点当前没有被任务占用，但它仍然保留在显存中。这些节点是“可回收”的候选者。

  3. 置换机制 (Eviction Policy)

  当显存不足以容纳新请求的 KV Cache 时，RadixCache 会启动置换流程。
  A. LRU (Least Recently Used) 策略
  RadixCache 维护了一个 LRU 队列，用于管理所有 ref_count == 0 的节点。

• 每次节点被匹配成功（Cache Hit）或新创建时，都会更新其 last_accessed_time。

• 当需要释放空间时，系统从 LRU 队列中选择最久未使用的节点进行删除。
  B. 递归删除与剪枝

• 从叶到根：置换通常从叶子节点开始。如果一个叶子节点被删除，其父节点若也满足 ref_count == 0
  且没有其他子节点，也可能被列入回收范围。

• 原子性：删除操作会将物理块索引归还给物理内存池的 free_list。
  C. 置换流程示例

1. 请求进入：系统搜索 RadixTree，找到最长匹配前缀。
2. 空间检查：计算新 Token 需要的 Block 数量。如果 free_list 长度不足。
3. 触发置换：
   • 查找 LRU 队列中 last_accessed_time 最早且 ref_count == 0 的节点。
   • 释放该节点占用的物理块。
   • 重复此过程直到满足新请求的空间需求。
4. 分配与锁定：从物理池获取新块，将新节点插入树中，并将 ref_count 加 1。

4. 为什么 RadixCache 更高效？

1. 自动前缀复用：无需用户手动指定，系统自动识别 System Prompt、Few-shot、Chat History 的重复部分。
2. 跨请求共享：在多轮对话或并行采样（如 Tree-of-Thought）中，公共路径的 KV Cache 只存一份。
3. 零拷贝匹配：匹配过程只是简单的树遍历，开销极低（CPU 端完成），而节省的是昂贵的 GPU
   计算（Prefill）和显存。

5. 总结

RadixCache 将 KV Cache的生命周期从“请求生命周期”延长到了“服务器运行生命周期”。它通过引用计数保证运行安全性，通过 LRU算法实现高效的显存置换，通过前缀树最大化了显存的复用率，是 SGLang 高吞吐量的核心保障。

2. Megatron-LM：大规模语言模型训练框架

Megatron-LM 是一个用于训练大规模语言模型的强力框架，其特点是具有高度可扩展性（Scalable）和模块化的架构。该代码库目前正从早期的单体式设计（Monolithic design）向名为 Megatron Core (MCore) 的库导向结构转型。

2.1 高层架构概览

该项目被划分为以下几个关键功能区域：

• **入口点 (pretrain_*.py)**：顶层脚本（如 pretrain_gpt.py, pretrain_bert.py），作为用户界面使用。它们定义了模型配置并调用训练编排逻辑。

• **Megatron Core (megatron/core/)**：该库现代化的模块化核心。它包含针对各种并行策略（张量并行 TP、流水线并行 PP、数据并行 DP、上下文并行 CP 和专家并行 EP）高度优化的算子原语。

• **遗留组件 (megatron/legacy/)**：为了向后兼容而保留的旧版单体实现。这包括旧版本的模型并行单元（MPU）和数据集加载器。

• **训练编排 (megatron/training/)**：将模型、数据和优化器绑定在一起的“粘合剂”。megatron/training/training.py 中的核心预训练函数负责管理初始化、训练循环、验证和检查点保存。

• **数据管理 (megatron/core/datasets/ & megatron/training/datasets/)**：处理复杂的数据加载、多数据集混合（Blending）以及在并行 Rank 之间的高效分片（Sharding）。

• **检查点 (megatron/training/checkpointing.py)**：管理模型状态的保存与加载，越来越多地利用 megatron/core/dist_checkpointing 来支持分布式环境。

---

2.2 核心组件与逻辑流

2.2.1 并行状态管理 (megatron/core/parallel_state.py)

这是 Megatron 分布式执行的“大脑”。它为不同的并行维度（TP, PP, DP, CP, EP）初始化进程组，确保每个 Rank 都能明确其在全球计算中的角色。

2.2.2 模型提供者模式 (Model Provider Pattern)

入口脚本定义了一个 model_provider 函数。该函数利用 megatron/core/models/（或 megatron/legacy/model/）中的构建器，根据请求的并行设置构造特定的架构（如 GPT, BERT, T5）。

2.2.3 训练循环 (megatron/training/training.py)

1. 初始化：设置分布式环境、随机种子和全局变量。

2. 模型/优化器设置：调用模型提供者并初始化优化器（通常使用分布式优化器以提高显存效率）。

3. 数据迭代器：构建高效的数据加载器，处理分片和数据混合。

4. 训练迭代：

   • 前向/反向传播：由 megatron/core/pipeline_parallel/schedules.py 处理，负责管理复杂的流水线调度策略（如 1F1B）。

   • 权重更新：优化器执行 Step 操作，并处理梯度裁剪与缩放。

5. 验证/评估：定期在验证集上运行推理。

6. 检查点保存：将模型和优化器状态保存到磁盘。

---

2.3 支持工具与示例

• **tools/**：包含用于数据预处理（preprocess_data.py）、检查点转换和专门性能测试的重要工具。

• **examples/**：提供参考 Shell 脚本和配置，用于训练 Llama、Mistral 和 GPT-3 等知名模型。

• **tasks/**：用于微调和评估等下游任务的工具。

---

2.4 设计哲学总结

Megatron-LM 将可扩展性视为最高优先级。通过将模型定义（Core）与训练运行器（Training）解耦，它允许相同的核心并行原语在不同的模型架构和训练范式中复用，同时在 NVIDIA GPU 集群上保持巅峰性能。

2.5 并行实现

Megatron-LM 实现并行的核心思想是将深度学习模型在多个维度上进行拆分，并通过高效的通信原语（Communication
Primitives）保持同步。所有并行策略的协调中心是 megatron/core/parallel_state.py。

以下是五种主要并行方式的实现机制：

1. 数据并行 (Data Parallelism, DP)

这是最基础的并行方式。

• 实现：将训练数据切分为 N 份，每个 DP 组内的 GPU 维护一份完整的模型副本（或分片的参数）。
• 分布式优化器 (Distributed Optimizer)：为了节省内存，Megatron 通常使用分布式优化器。它不让每个 GPU
  都存一份完整的优化器状态（如 Adam 的动量），而是将优化器状态和梯度均匀地分片到 DP 组中的所有 GPU
  上。在更新权重后，通过 all-gather 同步最新的参数。

2. 张量并行 (Tensor Parallelism, TP)

这是 Megatron 的“看家本领”，通过切分算子内部的矩阵来实现。

• 路径：megatron/core/tensor_parallel/
• ColumnParallelLinear：将权重矩阵按列切分。输入数据 X 在所有 GPU 上复制，每个 GPU 计算一部分输出 Y_i =
  X * W_i。
• RowParallelLinear：将权重矩阵按行切分。每个 GPU 接收上一层切分后的输入 X_i，计算 Y_i = X_i *
  W_i，最后通过 all-reduce 操作将结果汇总。
• 通信优化：TP 通常局限在单机 8 卡内，利用 NVLink 提供的高带宽。

3. 流水线并行 (Pipeline Parallelism, PP)

将模型的层（Layers）按顺序切分到不同的 GPU 组（Stages）上。

• 路径：megatron/core/pipeline_parallel/
• 调度策略 (Schedules)：为了减少“流水线气泡（Bubble）”，Megatron 使用 1F1B (One Forward, One Backward)
  调度。即在同一个 Stage 上，交替执行微批次（Micro-batch）的前向和后向计算。
• P2P 通信：Stage 之间通过 send 和 recv 操作传递隐藏状态（Hidden States）和梯度。

4. 上下文并行 (Context Parallelism, CP)

针对超长文本（Long Context）设计的并行。

• 实现：将 Sequence（序列）维度进行切分。与 TP 切分特征维度不同，CP 在计算 Attention 时，通过在 CP
  组内交换 Key/Value 数据来完成完整的注意力计算。
• 优势：解决了单张显存无法容纳极长序列 Attention 矩阵的问题。

5. 专家并行 (Expert Parallelism, EP)

专门用于混合专家模型 (Mixture of Experts, MoE)。

• 实现：在 MoE 层中，将不同的“专家（Experts）”分布在不同的 GPU 上。
• Router：根据输入 Token 的特征，将其路由到对应的 GPU 专家上。这涉及到复杂的 all-to-all 通信操作，将
  Token 发送到对应的专家，再将结果收集回来。

---

6. 总结：并行是如何协调的？

Megatron 通过一个 3D（甚至 5D）并行网格 来管理所有 GPU。

• 初始化：在启动时，parallel_state.py 会根据用户设置（如 –tensor-model-parallel-size 2
  –pipeline-model-parallel-size 4）将全局 Rank 划分到不同的正交进程组中。

• 嵌套关系：

  • 最内层：TP（通信最频繁，通常在单机内）。

  • 中间层：CP / EP。

  • 最外层：PP 和 DP（通信频率较低，跨机扩展性好）。

• 全局一致性：所有并行组通过 parallel_state 获取自己的 Rank 信息，确保在计算和通信时，每个 GPU
  都知道自己该处理哪一部分数据和权重。

参考

NVIDIA/Megatron-LM: Ongoing research training transformer models at scale

sgl-project/sglang: SGLang is a high-performance serving framework for large language models and multimodal models.

[细读经典]Megatron论文和代码详细分析(1) - 知乎