fix(pd): role chunk before decode errors; count only consumed MTP steps - #61
Draft
CrimsonDump wants to merge 4 commits into
Draft
fix(pd): role chunk before decode errors; count only consumed MTP steps#61CrimsonDump wants to merge 4 commits into
CrimsonDump wants to merge 4 commits into
Conversation
…ts own chunk (tile-ai#55) Two fixes for serving OpenAI-compatible clients over a PD deployment. The prefill request now drops `stream_options` and `max_completion_tokens` from the client body: the first contradicts the `stream=False` we force (vLLM rejects the pair with a 400 during body parsing) and the second overrides our `max_tokens=1`. Verified with `vllm bench serve --backend openai-chat`, which sends both unconditionally and previously failed every request. `ignore_eos` is now forwarded to the decode engine and honoured by the MLA/NSA adapter, streaming usage moves into its own trailing chunk carrying `total_tokens` (the shape vLLM and the OpenAI API emit), and the `transformers` / `tokenizers` pins are relaxed to `>=`. --------- Co-authored-by: CrimsonDump <56749892+CrimsonDump@users.noreply.github.com> Co-authored-by: Shuang Wu <leo@aigcsemi.com>
Update the README for the v0.1.5.post2 release: install commands and expected version now reference `0.1.5.post2`, the News and wheel download links point to the v0.1.5.post2 release (the previous `v0.1.5` release URL does not exist), and the `transformers` / `tokenizers` rows reflect the relaxed `>=` pins. Co-authored-by: Ziming Miao <max@aigcsemi.com>
…tile-ai#59) # 摘要 优化PD分离开销。 测量条件：8×B200（TP8）、GLM-5.1 FP8、`conc=1`、vLLM prefill + TileRT decode over NIXL、 `benchmark_serving.py --ignore-eos`、2 个 warm-up 之后 16 个请求。`fwd` 是前向/秒 —— 与非 PD 生成器上报的 `Effective TPS (AR, ar_steps=N)` 是同一个量，可直接对比： | 配置 | 非 PD 参照 | PD 改后 | 差 | PD 改前 | |---|---|---|---|---| | ISL 1k / OSL 1k, MTP on | 188.29 | **186.13** | −1.15% | 179.2 | | ISL 8k / OSL 1k, MTP on | 134.41 | **133.11** | −0.97% | 128.4 | | ISL 1k / OSL 1k, MTP off | 294.79 | **293.18** | −0.55% | — | | ISL 8k / OSL 1k, MTP off | 206.89 | **205.49** | −0.68% | — | MTP-off 两行没有"改前"：这个形态在当前 `main` 上**根本跑不起来**。 改完之后对解码循环做逐语句拆解，`show_hands` 占每执行步时间的 **99.76%** （5320.9 µs 里的 5309.9 µs），整个 host 侧降到 0.21%。这份预算与客户端实测值对得上， 残差在 0.08% 以内。 # 改了什么 ## `profiles/mla_nsa.py` —— 两条解码循环 两条循环原来都调 `decode_layer.forward()`，那是 **prefill 阶段**的入口：它把 `ar_steps` 写死为 1，并且给每个设备返回一个 `DeviceResult` 供调用方丢弃。decode 阶段的入口是 `show_hands()` / `show_hands_no_mtp()`，也正是 `models/glm_5/generator.py` 在它自己的 解码循环里用的。用 prefill 入口付了两笔代价，都实测过： - 没有 AR 链式，每步都付一次 `cudaStreamSynchronize` 加一次 host 往返 —— **35.6 µs/步**； - `show_hands()` 返回 `None`，于是单步 getter 成了唯一的收获路径：每次调用一次 `get_num_accepted()`，加上**每 token** 一次 `int(pred[i].item())` —— **40.6 µs/步**。 `_decode_standard` 还额外每 token 多付一次同步：它用 `int(nxt.item())` 把下一个 token 读回来，再把设备张量交给 `forward()`，后者又拷回主机一次。 两条循环现在都按 `generator.py` 的做法一次整块 D2H 取回 AR 平铺缓冲，并且每次调用链式执行 `ar_steps` 个设备端步骤。`ar_steps` 取自 `GLM5_AR_N` —— 与 `models/glm_5/generator.py` 读的是同一个旋钮、同一个默认值（8），两条路链式行为一致。 光链式会让请求超算：`show_hands` 把整条链一次入队，`len(tokens) < budget` 只能等它返回后 才能重查，无条件的 `ar_steps` 链平均浪费约 4 步（OSL=1024 占墙钟 1.2%，OSL=256 占 4.3%）。 所以链长按剩余预算收窄： - **MTP off** —— 一步一个 token，`min(ar_steps, rem)` 是精确的：**零超算**。 - **MTP on** —— `ceil(rem / mtp_seq_len)`，即按每步吃满整个 draft 排程。这可证明是最小步数 （一步最多产出 `mtp_seq_len` 个 token，凑 `rem` 至少要 `ceil(rem/mtp_seq_len)` 步）， 而且永不为没人要的 token 烧掉一整步。少排一步只多付一次 launch（约 40 µs）， 对比白烧一步的约 5.3 ms。 `_decode_standard` 还补上了显式的 `set_prefill_valid_tokens(0, with_mtp=False)`。 这不是装饰：这个调用决定的就是 prefill / decode 模式 （"Select prefill (num_valid_tokens > 0) vs decode (0) mode"），而链式和 AR 平铺缓冲 只在 decode 模式下存在。旧的循环体用 `forward()`，它把 `ar_steps` 钉死 1、两种模式下都 合法，所以从不需要这个调用；链式的循环体需要。`_decode_mtp` 一直都调它，非 PD 的 no-MTP 循环也调。 ## `decode_server.py` —— 流式 - **首 token 单独成行。** 它由 prefill 在请求体里给出，`engine.decode()` 在跑任何前向 之前就把它入队，但原来它要和第一次 drain 一起发。立刻发出去才是客户端计的 TTFT； 压着不发等于把 KV convert/inject 的开销（这里约 70 ms/请求）算进 TPOT 的分子。 SGLang 对它的 handoff token 就是这么做的 —— 从一个不跑前向的 batch 里直接 stream 出去。 这里等的是短轮询而不是事件，因为我们持有的唯一事件表示"完成"；实际上 `on_token()` 在 `decode` 一进来就触发，所以最多空转一次。 - **完成不再等 poll。** 队列一空生成器就睡固定间隔，**包括最后一个 token 之后**，于是 每个请求都把流多开着最多一个 poll 周期，而这段延迟落在 TPOT 里。现在 worker 会点一个 `asyncio.Event`，睡眠改成 `wait_for(ev.wait(), timeout=poll)`。poll=200 ms 下这一项是 **93 ms/请求**。事件循环和 Event 在 worker 启动**之前**创建 —— 快请求可能在生成器第一次 迭代之前就结束。 - **`TILERT_DECODE_POLL_MS`**（默认 200）让 drain poll 可配。有了上面那个修复，它对性能 是中性的：200 ms 与 5 ms 相差 ±0.6% 以内，而每行的 token 数差 30–190 倍。 ## `pd_router.py` —— role delta `{"role": "assistant"}` 那个 chunk 原来在发上游请求**之前**就 yield，于是它比任何 token 都先到客户端。而一个带 `choices` 但没有 token 的 chunk，在不要求内容非空的客户端里照样会 起 TTFT 计时 —— InferenceX 的 `backend_request_func.py` 不检查，SGLang 的 `benchmark/serving.py` 检查 —— 所以这把 `_drain_own_kv` + `convert` + `inject` 塞进了被 上报的解码窗口。vLLM 的 role delta 是在结果循环里面发的（`async for res in result_generator` 下的 `if first_iteration:`），即引擎已经产出之后；现在改成与第一个带文本 的 chunk 配对，与之对齐。 条件是"第一个非空的 detokenizer 输出"而不是"收到的第一行"：`IncrementalDetok.push` 对不 完整的 UTF-8 序列返回 `""`，所以回复以多字节字符开头时，role delta 又会落在一个没有文本的 行上、把 TTFT 标记重新交给空 chunk。最后还有一个兜底：如果整条流没有任何一行带过文本 （立即 stop，或 decode 节点报错），补发一次，保持流的形状不变。 ## `profiles/glm5.py` —— 78 层形态放到环境变量后面 `TILERT_PD_NO_MTP=1` 选 78 层。PD 的 cache 布局是两端共用的，而 prefill 进程没有自己的 `--with-mtp` 可读，所以只能用一个在 P 和 D 上设成一致的环境变量。`LAYOUT_VERSION` 跟着 一起动，所以配错的 P/D 组合会触发 `prefill_connector` 里既有的 `layout_version` assert， 而不是静默地选到 `index_topk` 之外的 key。不设或设 `0` 就保持当前的 79 层 MTP 形态， 所以默认状态下这项改动是惰性的。 这是本 PR 里唯一一项"增加形态"而非"去掉开销"的改动，而且**单独成一个 commit** —— 如果你们更希望它单独开一个 PR，随时可以摘掉。
CrimsonDump
force-pushed
the
fix/pd-role-and-step-accounting
branch
from
August 13, 2026 08:01
6631df4 to
3dff709
Compare
xysmlx
marked this pull request as draft
August 13, 2026 13:31
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
fix(pd): role chunk before decode errors; count only consumed MTP steps