源码精读：io_struct.py —— 进程间消息体大全

谁该读这一篇？ 想看懂 SGLang 进程间到底传些什么字段、要给 SGLang 加新功能（比如自定义采样参数）的开发者。 前置阅读： 01-overview/02-architecture.md、02-tokenizer-manager.md、07-detokenizer.md。 耗时： 25 分钟（更像字典查阅） 学完能： 1. 列出 SGLang 进程间通信的主要消息体； 2. 看懂 GenerateReqInput 的字段含义； 3. 知道在 batch 和单请求 dataclass 之间怎么转换； 4. 给消息体加新字段时知道注意什么（向后兼容、pickle 大小）。

1. 文件总览

源码：io_struct.py（1300+ 行）。 所有 ZMQ 消息体 + 内部数据 contract 都在这里。按方向分组：

客户端 → TM:
    GenerateReqInput           # 文本生成
    EmbeddingReqInput          # embedding

TM → Scheduler:
    TokenizedGenerateReqInput
    TokenizedEmbeddingReqInput
    BatchTokenizedGenerateReqInput    # 批
    BatchTokenizedEmbeddingReqInput

Scheduler → Detokenizer:
    BatchTokenIDOutput
    BatchEmbeddingOutput

Detokenizer → TM:
    BatchStrOutput

控制 / 维护:
    ClearHiCacheReqInput / Output
    FreezeGCReq
    AbortReq
    UpdateWeightFrom*Input
    ...

2. `GenerateReqInput`（138+）

最重要的入口消息。客户端 OpenAI 兼容 API 经过 TM 的 normalization 后变成这个 dataclass。

@dataclass
class GenerateReqInput(BaseReq):
    text: Optional[Union[str, List[str]]] = None
    input_ids: Optional[Union[List[int], List[List[int]]]] = None

    sampling_params: Union[Dict, List[Dict]] = None     # temperature, top_p, ...
    rid: Optional[Union[str, List[str]]] = None         # request id
    stream: bool = False
    log_metrics: bool = True

    # 受约束解码
    regex: Optional[str] = None
    json_schema: Optional[str] = None
    grammar: Optional[str] = None

    # 多模态
    image_data: Optional[List[ImageInputItem]] = None
    audio_data: Optional[List[AudioInputItem]] = None
    video_data: Optional[List[VideoInputItem]] = None

    # 输出元信息
    return_logprob: bool = False
    top_logprobs_num: int = 0
    return_text_in_logprobs: bool = False

    # 控制
    no_stop_trim: bool = False
    priority: Optional[int] = None
    bootstrap_room: Optional[int] = None                # PD 分离用

    # session
    session_params: Optional[SessionParams] = None

    # LoRA
    lora_path: Optional[Union[str, List[str]]] = None

    # ...还有十几个字段

字段类别：

输入内容：text / input_ids / image_data / audio_data / video_data。
采样：sampling_params。
约束：regex / json_schema / grammar。
元信息：return_logprob 等。
调度提示：priority / lora_path / bootstrap_room。

注意：所有字段都允许 List[...]（一次提交一批），TM 内部按 batch 处理。

3. `TokenizedGenerateReqInput`（720+）

TM 把 GenerateReqInput tokenize 后变成它：

@dataclass
class TokenizedGenerateReqInput(BaseReq):
    rid: str
    input_text: str                               # 原文（log 用）
    input_ids: List[int]                          # tokenize 后的
    image_inputs: Optional[Dict[str, Any]]        # 多模态预处理后
    sampling_params: SamplingParams
    return_logprob: bool
    logprob_start_len: int
    top_logprobs_num: int
    stream: bool
    lora_id: Optional[str]
    session_params: Optional[SessionParams]
    custom_logit_processor: Optional[bytes]       # pickle 后的自定义 processor
    bootstrap_host: Optional[str]
    bootstrap_room: Optional[int]
    # ...

差异：

输入从字符串 → token id。
多模态从原始字节 → embedding placeholder + 预处理元信息。
受约束解码从 string → grammar / regex 编译后的对象。
LoRA path 解析成 LoRA id。

4. `BatchTokenizedGenerateReqInput`（814+）

批量版本：

@dataclass
class BatchTokenizedGenerateReqInput(BaseBatchReq):
    reqs: List[TokenizedGenerateReqInput]

TM 在高 QPS 下会把短时间收到的多个请求批起来发，省 ZMQ 调用次数。

5. `BatchTokenIDOutput`（1083+）

Scheduler → Detokenizer：

@dataclass
class BatchTokenIDOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin):
    rids: List[str]
    finished: List[bool]
    finished_reasons: List[Optional[Dict]]
    output_ids: List[List[int]]                   # 这一批每个请求新的 token id 列表
    decoded_texts: List[Optional[str]]            # 如果 Scheduler 已经 detokenize 一些
    vids: List[int]                               # 验证 id

    # logprob 相关
    input_token_logprobs_val: Optional[...]
    output_token_logprobs_val: Optional[...]
    top_logprobs_val: Optional[...]
    # ...

vids 这字段防 stale：Scheduler 一次发送对应一个 vid，Detokenizer 用来验证消息不是旧的（避免 abort 后还处理）。

6. `BatchStrOutput`（1155+）

Detokenizer → TM：

@dataclass
class BatchStrOutput(BaseBatchReq, SpeculativeDecodingMetricsMixin):
    rids: List[str]
    finished: List[bool]
    finished_reasons: List[Optional[Dict]]
    output_strs: List[str]                        # detokenize 后的字符串增量
    decoded_offsets: List[int]                    # 在完整 token 流中的位置

    # 继承自 spec mixin
    spec_verify_ct: int
    # ...

TM 收到后按 rids[i] 找 ReqState，把 output_strs[i] 塞队列。

7. `EmbeddingReqInput` / `BatchEmbeddingOutput`（829+ / 1221+）

非生成任务的等价物：

@dataclass
class EmbeddingReqInput(BaseReq):
    text: Optional[Union[str, List[str]]] = None
    input_ids: ...
    is_cross_encoder_request: bool = False
    rid: Optional[Union[str, List[str]]] = None
    ...

@dataclass
class BatchEmbeddingOutput(BaseBatchReq):
    rids: List[str]
    embeddings: List[List[float]]
    finished: List[bool]
    ...

Embedding 没有 detokenize 步骤，Scheduler 直接发回 TM。

8. 控制类消息

@dataclass
class AbortReq(BaseReq):
    rid: Optional[str] = None                     # None 表示 abort all
    abort_all: bool = False

@dataclass
class ClearHiCacheReqInput(BaseReq):
    """触发 HiCache flush."""

@dataclass
class FreezeGCReq(BaseReq):
    """让 Detokenizer 暂时不跑 gc."""

@dataclass
class UpdateWeightFromDiskReqInput / UpdateWeightFromDistributedReqInput / ...
    """权重热更新."""

这些都是带外控制：TM 收到客户端管理 API 后转发给 Scheduler / Detokenizer。

9. 加新字段的注意事项

向后兼容：用 Optional 默认 None，老版本 TM/Scheduler 收到新字段不报错。
pickle 大小：避免传大张量（用 shared memory 或 ZMQ 多帧）。
类型简单：保持原生类型 + dataclass，避免 enum / Path / 自定义 class（pickle 跨进程麻烦）。
__post_init__：可以加校验，但不要做重计算。
批 vs 单：新字段同时加到 batched 版本，否则批序列化会丢。

10. 序列化与传输

ZMQ 用 send_pyobj / recv_pyobj，内部是 pickle 协议 4：

self.send_to_scheduler.send_pyobj(req)
# 内部 ≈ self.send(pickle.dumps(req, protocol=4))

性能：

100 KB 以下消息 < 0.5 ms。
1 MB 级（多模态 embedding） ~5 ms。
10 MB+ 应该考虑共享内存（SharedMemoryManager）。

11. 关键源码索引

内容	行
`BaseReq` / `BaseBatchReq`	53 / 76
`SpeculativeDecodingMetricsMixin`	88
`SessionParams`	112
`GenerateReqInput`	138
`TokenizedGenerateReqInput`	720
`BatchTokenizedGenerateReqInput`	814
`EmbeddingReqInput`	829
`TokenizedEmbeddingReqInput`	1036
`BatchTokenizedEmbeddingReqInput`	1068
`BatchTokenIDOutput`	1083
`BatchStrOutput`	1155
`BatchEmbeddingOutput`	1221
`ClearHiCacheReqInput` / `Output`	1248 / 1253

12. 小结

io_struct.py 是 SGLang 进程间通信的"接口定义文件"。
客户端 → TM 用 GenerateReqInput；TM → Scheduler 转 TokenizedGenerateReqInput；Scheduler → Detokenizer 用 BatchTokenIDOutput；Detokenizer → TM 用 BatchStrOutput。
所有消息都 pickle 通过 ZMQ 传输。
加新字段要 Optional 默认 None 保持兼容，避免传大张量。

13. 自检

text 字段允许 Union[str, List[str]]，为什么？后端怎么区分？

答案

`str` 是单请求 prompt，`List[str]` 是一次提交一批 prompt（OpenAI API 的 `n` 字段或者 client 批量请求）。后端在 TM 里把 `List[str]` 拆成多个独立的 `TokenizedGenerateReqInput`，每个 rid 不同；如果客户端传了 `List[str]` 但只想要 1 个结果，TM 用 broadcast 语义复制。实际处理代码用 `isinstance(text, list)` 分支判断。这种 polymorphic 字段在 OpenAI 协议里很常见（如 `messages`、`stop` 也都可以是 str 或 list）。

vid 字段防的是什么 race？

答案

`vid` (version id) 防的是 **abort 之后 stale 消息回流**：场景：客户端 abort 请求 → TM 标记 abort → Scheduler 已经在跑这个 rid 的 batch → 跑完发 `BatchTokenIDOutput` 给 Detokenizer → Detokenizer 又发 `BatchStrOutput` 给 TM → 但 TM 这边 rid 已经 del 了，或者 rid 复用给了新请求（如果 rid 是用户自传的）。有 vid 后：每次 Scheduler 开始处理一个新版本（abort 后重发）就 ++vid，Detokenizer / TM 拿到带 stale vid 的消息直接丢弃。类似 etcd 的 revision、Kafka 的 offset，分布式系统经典 pattern。

我加一个新采样参数，应该在哪几个 dataclass 都加？

答案

至少四个层级： (a) **`SamplingParams`**（[`srt/sampling/sampling_params.py`](../sglang/python/sglang/srt/sampling/)）—— 字段定义 + 校验 + 默认值。 (b) **`GenerateReqInput`** —— 顶级入参（OpenAI 兼容路径），用 `Optional[T] = None` 保持向后兼容。 (c) **`TokenizedGenerateReqInput`** —— TM 转给 Scheduler 时透传。 (d) **`BatchTokenizedGenerateReqInput`** —— batch 版本同步加。还要在 [`srt/layers/sampler.py`](../sglang/python/sglang/srt/layers/sampler.py) 实现实际采样逻辑；可能在 `SamplingBatchInfo`（GPU 端的 batched 采样信息）里也加字段。最后 [`entrypoints/openai/`](../sglang/python/sglang/srt/entrypoints/openai/) 的请求 schema 加映射。

多模态 image data 在 GenerateReqInput 是 raw bytes，在 TokenizedGenerateReqInput 是什么？

答案

`GenerateReqInput.image_data`：`List[ImageInputItem]`，每个 item 可以是 base64 / URL / 文件路径 / PIL.Image。 TM 收到后过 multimodal processor（[`srt/managers/multimodal_processor.py`](../sglang/python/sglang/srt/managers/multimodal_processor.py)）：(1) 下载 / decode → PIL Image；(2) resize / normalize → 张量；(3) 跑 ViT encoder → image embeddings（或者把"待编码"先发给 Scheduler 让 GPU 跑）。 `TokenizedGenerateReqInput.image_inputs`：`Dict[str, Any]`，含 image embeddings、image placeholder token 位置、grid 信息等。Scheduler 在 forward 时把 placeholder token 替换成实际 image embedding。

pickle 序列化 10 MB 张量为什么慢？怎么绕？

答案

pickle 慢的原因：(a) 默认 protocol 把张量序列化为 numpy bytes，10 MB 数据要 memcpy + send + recv + memcpy 反序列化；(b) 跨进程发的是 host bytes，GPU tensor 还要先 .cpu()，多一次 device-host 拷贝。绕法： - **CUDA IPC**：发 IPC handle 而不是数据；接收端 `cudaIpcOpenMemHandle` 直接拿到同一 GPU 上的指针，零拷贝。SGLang 的 `update_weights_from_ipc` 就是这条。 - **共享内存**：用 `multiprocessing.shared_memory` 或 ZMQ 多帧（separate body），张量数据走 OS 共享段。 - **避免传张量**：消息体只带元数据（shape / dtype / offset），张量本身在共享的 KV pool / shared tensor pool 里。 SGLang 的常规 ZMQ 消息保持在 KB 级以下，10 MB+ 张量都走 IPC 或共享段。

14. 下一步

02-tokenizer-manager.md — 这些消息怎么被生产/消费。
07-detokenizer.md — BatchTokenIDOutput 怎么被消费。
03-scheduler.md — Scheduler 怎么用 TokenizedGenerateReqInput。
源码：io_struct.py。

上游源码：sglang/python/sglang/srt/managers/io_struct.py。