源代码
class ForwardMode(IntEnum):
# Extend a sequence. The KV cache of the beginning part of the sequence is already computed (e.g., system prompt).
# It is also called "prefill" in common terminology.
EXTEND = auto()
# Decode one token.
DECODE = auto()
# Contains both EXTEND and DECODE when doing chunked prefill.
MIXED = auto()
# No sequence to forward. For data parallel attention, some workers will be IDLE if no sequence are allocated.
IDLE = auto()
# Used in speculative decoding: verify a batch in the target model.
TARGET_VERIFY = auto()
# Used in speculative decoding: extend a batch in the draft model.
DRAFT_EXTEND = auto()
# A dummy first batch to start the pipeline for overlap scheduler.
# It is now used for triggering the sampling_info_done event for the first prefill batch.
DUMMY_FIRST = auto()
def is_prefill(self):
return self.is_extend()
def is_extend(self):
return (
self == ForwardMode.EXTEND
or self == ForwardMode.MIXED
or self == ForwardMode.DRAFT_EXTEND
or self == ForwardMode.TARGET_VERIFY
)
def is_decode(self):
return self == ForwardMode.DECODE
def is_mixed(self):
return self == ForwardMode.MIXED
def is_idle(self):
return self == ForwardMode.IDLE
def is_target_verify(self):
return self == ForwardMode.TARGET_VERIFY
def is_draft_extend(self):
return self == ForwardMode.DRAFT_EXTEND
def is_extend_or_draft_extend_or_mixed(self):
return (
self == ForwardMode.EXTEND
or self == ForwardMode.DRAFT_EXTEND
or self == ForwardMode.MIXED
)
def is_cuda_graph(self):
return (
self == ForwardMode.DECODE
or self == ForwardMode.TARGET_VERIFY
or self == ForwardMode.IDLE
)
def is_dummy_first(self):
return self == ForwardMode.DUMMY_FIRST
def is_decode_or_idle(self):
return self == ForwardMode.DECODE or self == ForwardMode.IDLE
解析
SGLang 推理框架中的 ForwardMode 深度解析
在现代大语言模型（LLM）推理服务中，为了极致地优化性能，推理框架需要在底层对请求的处理流程进行精细的划分和调度。SGLang 作为一个高性能的 LLM 推理框架，其内部通过一个名为 ForwardMode 的枚举类来清晰地定义和区分不同的前向计算模式。理解 ForwardMode 中各个成员的含义和作用，对于深入了解 SGLang 的工作机制至关重要。
ForwardMode 是一个继承自 IntEnum 的枚举类，它定义了模型在处理不同类型的计算批次（batch）时所处的多种状态。这些状态直接关联到 SGLang 的核心优化特性，如前缀缓存（KV Cache）、连续批处理（Continuous Batching）、分块预填充（Chunked Prefill）以及投机解码（Speculative Decoding）。
以下将对 ForwardMode 中的每个成员进行详细的讲解。
核心推理模式
EXTEND: 序列扩展（预填充）
# Extend a sequence. The KV cache of the beginning part of the sequence is already computed (e.g., system prompt).
# It is also called "prefill" in common terminology.
EXTEND = auto()
EXTEND 是 LLM 推理中最核心的模式之一，通常被称为“预填充”（Prefill）。当一个新的请求（例如，用户的提问）进入系统时，模型需要一次性处理整个输入序列（prompt），并计算出每个 token 对应的键值缓存（KV Cache）。这些计算出的 KV Cache 会被存储起来，用于后续的解码阶段。
例子：
假设一个请求的输入是“中国的首都是哪里？”。在 EXTEND 模式下，SGLang 会将这整个句子作为一个批次送入模型，并行计算所有汉字的 KV Cache。这个过程是计算密集型的。值得注意的是，EXTEND 模式也包含了对已有序列进行扩展的场景，比如在处理一个很长的文档时，可以分块进行 EXTEND 计算。
在早期的 SGLang版本中，曾有一个独立的 PREFILL 模式，但后来被 EXTEND 统一，因为 EXTEND 的概念更具通用性，既能处理新序列的首次计算，也能处理已有序列的后续扩展。
DECODE: 单令牌解码
# Decode one token.
DECODE = auto()
DECODE 模式紧随 EXTEND 之后，是生成模型输出的阶段。在此模式下，模型利用 EXTEND 阶段生成的 KV Cache，自回归地一次生成一个 token。这个过程是内存带宽密集型的，因为它需要频繁地从内存中读取庞大的 KV Cache。
例子：
在上一个例子的基础上，模型在 EXTEND 计算完输入“中国的首都是哪里？”的 KV Cache 后，会进入 DECODE 模式。它会生成第一个 token “北”，然后将“北”的 KV Cache 添加到已有的缓存中，接着生成“京”，循环往复，直到生成完整的答案并遇到终止符。
SGLang 等现代推理框架的一个关键优化就是将计算密集型的 EXTEND 批次和内存密集型的 DECODE 批次进行分离调度，以最大化 GPU 的利用率。
MIXED: 混合模式（分块预填充）
# Contains both EXTEND and DECODE when doing chunked prefill.
MIXED = auto()
当系统采用“分块预填充”（Chunked Prefill）技术时，一个批次中可能同时包含需要进行 EXTEND 计算的新序列部分和需要进行 DECODE 计算的已有序列。MIXED 模式正是用来描述这种混合状态的。 这对于处理超长输入序列非常有效，可以避免一次性处理过长的序列导致的显存溢出。
例子：
假设一个用户的输入是一个数万字的超长文档。SGLang 不会一次性对整个文档进行 EXTEND。它可能会先 EXTEND 前 4096 个 token，然后当处理后续部分时，一个批次里可能既包含了对新一块 4096 个 token 的 EXTEND 操作，也包含了对之前已经处理过的序列的 DECODE 操作。
投机解码（Speculative Decoding）相关模式
投机解码是一种先进的优化技术，它使用一个较小的“草稿模型”（Draft Model）来快速生成一个 token 序列草稿，然后由主模型（Target Model）一次性地对整个草稿序列进行验证，从而加速解码过程。
DRAFT_EXTEND: 草稿模型扩展
# Used in speculative decoding: extend a batch in the draft model.
DRAFT_EXTEND = auto()
此模式专用于投机解码。当使用草稿模型生成候选 token 序列时，前向计算就处于 DRAFT_EXTEND 模式。
例子：
在投机解码中，一个轻量的草稿模型会快速地预测出主模型接下来可能生成的 5 个 token。这个生成过程就发生在 DRAFT_EXTEND 模式下。
TARGET_VERIFY: 目标模型验证
# Used in speculative decoding: verify a batch in the target model.
TARGET_VERIFY = auto()
草稿生成后，SGLang 会切换到 TARGET_VERIFY 模式。在这个模式下，主模型（目标模型）会接收草稿模型生成的 token 序列，并一次性地、并行地对它们进行计算和验证，判断哪些 token 是可以接受的。
例子：
主模型接收到草稿模型生成的 5 个候选 token 后，会进入 TARGET_VERIFY 模式。它可能会发现前 3 个 token 是正确的，第 4 个是错误的。于是，它会接受前 3 个 token，并从第 4 个 token 的位置开始重新生成。通过这种方式，原本需要 3 次 DECODE 迭代的操作，在一次 TARGET_VERIFY 前向计算中就完成了，从而实现了加速。
系统与调度相关模式
IDLE: 空闲状态
# No sequence to forward. For data parallel attention, some workers will be IDLE if no sequence are allocated.
IDLE = auto()
IDLE 模式表示当前计算单元（例如一个 GPU Worker）没有分配到任何需要处理的序列。这是一个调度概念，在数据并行或张量并行场景下，由于负载不均，某些计算资源可能会暂时处于空闲状态。
例子：
在一个拥有 4 个 GPU 的张量并行推理服务中，如果当前的请求批次大小较小，调度器可能只将任务分配给了其中 2 个 GPU，那么另外 2 个 GPU 在这个计算周期内就处于 IDLE 模式。
DUMMY_FIRST: 虚拟初始批次
# A dummy first batch to start the pipeline for overlap scheduler.
# It is now used for triggering the sampling_info_done event for the first prefill batch.
DUMMY_FIRST = auto()
DUMMY_FIRST 是一个非常特殊的模式，它代表一个虚拟的、不包含实际计算任务的初始批次。它的主要作用是“预热”或初始化整个推理流水线，特别是在使用重叠调度器（Overlap Scheduler）时，用于触发一些初始事件，确保后续的真实计算批次能够顺畅地执行。
例子：
在服务启动后的第一个请求到来之前，调度器可能会先发送一个 DUMMY_FIRST 批次，这个批次会流经整个系统，使得 CUDA Graph 等需要预先捕获计算图的机制能够被正确初始化，或者触发某些同步事件，为第一个真正的 EXTEND 批次做好准备。
补充
Chunked Prefill 是一种精妙的调度机制，它将宏观的预填充任务微观化，通过与解码任务混合批处理，有效地平衡了LLM推理中的计算和内存瓶颈。
传统的连续批处理（Continuous Batching）在处理新请求时，会执行一个完整的预填充操作，这个操作的耗时由最长的输入序列决定。在预填充期间，已经在进行解码的旧请求要么暂停，要么只能生成一个词元，这导致了解码请求的有效吞吐量下降，特别是遇到那些特别长的Prefill任务，那将会大大增加其他请求解码速度。
Chunked Prefill 的核心思想是将一个长序列的预填充任务分解成多个较小的“块”（Chunks）。 例如，一个包含 4096 个词元的长提示，可以被分解为 8 个 512 词元的块。推理引擎每次只处理一个预填充块，而不是整个序列。
这种机制带来了两大优势：
- 减少流水线气泡，提升GPU利用率：预填充块的计算量相对较小且固定，更容易与许多小的解码任务打包在一起，形成一个计算量饱和的批次。这使得计算密集型的预填充块可以充分利用GPU的计算单元，而内存密集型的解码任务则“搭便车”，利用剩余的计算周期完成，从而掩盖了各自的瓶颈。
- 改善解码请求的延迟：对于正在进行解码的用户来说，他们感受到的不再是长时间的完全停滞，而是在每个混合批次中均匀地、轻微地“减速”。这大大降低了输出每个词元的平均时间，改善了流式生成的体验。
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
