From 24d9912c8d826757899d26bd5bcb9d9d56bad7b9 Mon Sep 17 00:00:00 2001 From: Shuang Wu Date: Wed, 5 Aug 2026 16:05:46 +0800 Subject: [PATCH 1/4] fix(pd_vllm): prefill request hygiene, ignore_eos support, usage in its own chunk (#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 --- pyproject.toml | 4 +- requirements.txt | 4 +- tilert/pd_vllm/pd_router.py | 67 ++++++++++++++++++++++-------- tilert/pd_vllm/profiles/mla_nsa.py | 6 ++- 4 files changed, 57 insertions(+), 24 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 407c497..1d1600e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,8 +20,8 @@ dependencies = [ # https://download.pytorch.org/whl/cu130``); installing from PyPI yields a # CUDA build that does not match the cu130-linked tilert binary. "torch==2.11.0", - "transformers==4.46.3", - "tokenizers==0.20.3", + "transformers>=4.46.3", + "tokenizers>=0.20.3", "numpy", "scipy", "einops", diff --git a/requirements.txt b/requirements.txt index c22551d..ade8509 100644 --- a/requirements.txt +++ b/requirements.txt @@ -8,8 +8,8 @@ # # The recommended path remains the prebuilt Docker image (see README). torch==2.11.0 -transformers==4.46.3 -tokenizers==0.20.3 +transformers>=4.46.3 +tokenizers>=0.20.3 numpy scipy einops diff --git a/tilert/pd_vllm/pd_router.py b/tilert/pd_vllm/pd_router.py index 87a61e5..e599827 100644 --- a/tilert/pd_vllm/pd_router.py +++ b/tilert/pd_vllm/pd_router.py @@ -97,6 +97,37 @@ def _thinking_enabled(body: dict) -> bool: return bool(ctk.get("enable_thinking", True)) +# Client fields that must not survive into the prefill request, which is +# forwarded verbatim apart from the fields we set: stream_options contradicts +# the stream=False we force (vLLM rejects the pair with a 400 during body +# parsing), and max_completion_tokens takes precedence over max_tokens, so it +# would override our max_tokens=1. Streaming clients send both. +_PREFILL_DROP_FIELDS = ("stream_options", "max_completion_tokens") + + +def build_prefill_body(path: str, body: dict, node: DecodeNode) -> dict: + """The vLLM request that prefills only and hands the KV state to ``node``. + + Lives outside ``build_app`` so the rewrite can be exercised without a + router process, a vLLM instance or a decode node. + """ + prefill_body = dict(body) + prefill_body["max_tokens"] = 1 + prefill_body["stream"] = False + for field in _PREFILL_DROP_FIELDS: + prefill_body.pop(field, None) + if path.endswith("chat/completions"): + prefill_body["logprobs"] = True + prefill_body["top_logprobs"] = 1 + else: + prefill_body["logprobs"] = 1 + prefill_body["kv_transfer_params"] = { + "tilert_host": node.host, + "tilert_ctrl_port": node.ctrl_port, + } + return prefill_body + + class RouterCtx: """Immutable per-process context (tokenizer, parser factory, config).""" @@ -133,24 +164,13 @@ def pool_status(): # ── shared prefill step ────────────────────────────────────────────── def _prefill(path, body, node): - prefill_body = dict(body) - prefill_body["max_tokens"] = 1 - prefill_body["stream"] = False - if path.endswith("chat/completions"): - prefill_body["logprobs"] = True - prefill_body["top_logprobs"] = 1 - else: - prefill_body["logprobs"] = 1 - prefill_body["kv_transfer_params"] = { - "tilert_host": node.host, - "tilert_ctrl_port": node.ctrl_port, - } + prefill_body = build_prefill_body(path, body, node) r = requests.post(f"{ctx.vllm_url}{path}", json=prefill_body, timeout=600) r.raise_for_status() return r.json() def _sampling_of(body): - return {k: body[k] for k in ("temperature", "top_p", "top_k") if k in body} + return {k: body[k] for k in ("temperature", "top_p", "top_k", "ignore_eos") if k in body} def _max_tokens_of(body): return int(body.get("max_tokens") or body.get("max_completion_tokens") or 256) @@ -267,6 +287,17 @@ def _chunk(delta: dict, finish=None, usage=None) -> str: payload["usage"] = usage return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n" + def _usage_chunk(usage: dict) -> str: + payload = { + "id": chunk_id, + "object": "chat.completion.chunk", + "created": int(time.time()), + "model": model, + "choices": [], + "usage": usage, + } + return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n" + def _event_delta(ev: dict) -> dict: if ev["kind"] == "reasoning": return {"reasoning_content": ev["text"]} @@ -357,13 +388,13 @@ async def _gen(): yield _chunk(_event_delta(ev)) if saw_tool: finish_reason = "tool_calls" - yield _chunk( - {}, - finish=finish_reason, - usage={ + yield _chunk({}, finish=finish_reason) + yield _usage_chunk( + { "prompt_tokens": prompt_tokens, "completion_tokens": n_tokens, - }, + "total_tokens": (prompt_tokens or 0) + n_tokens, + } ) yield "data: [DONE]\n\n" completed_ok = True diff --git a/tilert/pd_vllm/profiles/mla_nsa.py b/tilert/pd_vllm/profiles/mla_nsa.py index a270f7c..ed3c079 100644 --- a/tilert/pd_vllm/profiles/mla_nsa.py +++ b/tilert/pd_vllm/profiles/mla_nsa.py @@ -364,6 +364,7 @@ def __init__(self, generator, with_mtp: bool): self.max_seq_len = getattr(generator.decode_layer, "max_seq_len", 200000) self.last_stats: dict = {} self.stop_ids = self._resolve_stop_ids(generator) + self._ignore_eos = False @staticmethod def _resolve_stop_ids(generator) -> set: @@ -392,6 +393,7 @@ def decode(self, first_token_id, max_tokens, sampling, on_token=None, cancel_eve top_k=int(sampling.get("top_k", 256)), use_topp=True, ) + self._ignore_eos = bool(sampling.get("ignore_eos")) budget = min(int(max_tokens), self.max_seq_len - self._seq_len - 1) if budget <= 0: self.last_stats = {"finish_reason": "length"} @@ -403,7 +405,7 @@ def decode(self, first_token_id, max_tokens, sampling, on_token=None, cancel_eve def _decode_mtp(self, first_token_id, budget, on_token, cancel_event): dl = self.gen.decode_layer T = self.mtp_seq_len - stop_ids = self.stop_ids + stop_ids = set() if self._ignore_eos else self.stop_ids torch = self._torch tokens = [int(first_token_id)] if on_token: @@ -453,7 +455,7 @@ def _decode_standard(self, first_token_id, budget, on_token, cancel_event): from tilert.models.deepseek_v3_2.temp_var_indices import Idx dl = self.gen.decode_layer - stop_ids = self.stop_ids + stop_ids = set() if self._ignore_eos else self.stop_ids torch = self._torch tokens = [int(first_token_id)] if on_token: From 43d130d6007cc3e9db961d506d1d79e677ec7175 Mon Sep 17 00:00:00 2001 From: mzmssg Date: Thu, 6 Aug 2026 15:55:43 +0800 Subject: [PATCH 2/4] docs(readme): update install references to v0.1.5.post2 (#56) 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 --- README.md | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index fdd123e..463cb6c 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ ______________________________________________________________________ ## 📰 News -- 🔀 **2026-07-14 · [v0.1.5](https://github.com/tile-ai/TileRT/releases/tag/v0.1.5) Released**. Introduce [**PD (prefill–decode) disaggregation**](https://www.tilert.ai/blog/tilert-vllm-disaggregation.html) — vLLM prefill + TileRT decode, behind an OpenAI-compatible endpoint. Supported on GLM-5/5.1 and DeepSeek-V3.2. +- 🔀 **2026-07-14 · [v0.1.5](https://github.com/tile-ai/TileRT/releases/tag/v0.1.5.post2) Released**. Introduce [**PD (prefill–decode) disaggregation**](https://www.tilert.ai/blog/tilert-vllm-disaggregation.html) — vLLM prefill + TileRT decode, behind an OpenAI-compatible endpoint. Supported on GLM-5/5.1 and DeepSeek-V3.2. - 💥 **2026-06-08 · [Breaking 1000 TPS on a 1T Model](https://www.tilert.ai/blog/breaking-1000-tps.html)**. In collaboration with [Xiaomi MiMo](https://mimo.xiaomi.com/blog/mimo-tilert-1000tps), TileRT pushes [**MiMo-V2.5-Pro-UltraSpeed**](https://platform.xiaomimimo.com/docs/en-US/model-intro/mimo-v2.5-pro-ultraspeed) past **1000 tokens/s** on a **1-trillion-parameter** model through extreme model–system co-design — a first without custom silicon, all on a single 8-GPU node. @@ -70,7 +70,7 @@ ______________________________________________________________________ ### Build environment of the v0.1.5 wheel -The official `tilert==0.1.5.post1` wheel on PyPI was compiled against the following stack. Treat these as **hard requirements**, not lower bounds. +The official `tilert==0.1.5.post2` wheel on PyPI was compiled against the following stack. Treat these as **hard requirements**, not lower bounds (`transformers` / `tokenizers` are lower bounds since v0.1.5.post2). | Component | Pinned version | | ---------------- | --------------------------------------------------- | @@ -79,8 +79,8 @@ The official `tilert==0.1.5.post1` wheel on PyPI was compiled against the follow | Operating System | Linux **x86_64**, glibc **≥ 2.28** (manylinux_2_28) | | Python | **3.12** | | PyTorch | **`torch==2.11.0+cu130`** | -| `transformers` | **`4.46.3`** | -| `tokenizers` | **`0.20.3`** | +| `transformers` | **`>= 4.46.3`** | +| `tokenizers` | **`>= 0.20.3`** | ### Recommended: pre-built Docker image @@ -106,18 +106,18 @@ docker run --rm -it --gpus all --ipc=host \ ghcr.io/tile-ai/tilert:cu132-latest # Inside the container — install from PyPI: -pip install tilert==0.1.5.post1 +pip install tilert==0.1.5.post2 # Or pin the exact wheel from the GitHub Release page directly # (same artifact, useful when PyPI is unreachable): -pip install https://github.com/tile-ai/TileRT/releases/download/v0.1.5/tilert-0.1.5.post1-cp312-cp312-manylinux_2_28_x86_64.whl +pip install https://github.com/tile-ai/TileRT/releases/download/v0.1.5.post2/tilert-0.1.5.post2-cp312-cp312-manylinux_2_28_x86_64.whl ``` Verify the install: ```bash python -c "import tilert, torch; print('tilert', tilert.__version__, '/ torch', torch.__version__, '/ cuda', torch.version.cuda)" -# Expected: tilert 0.1.5.post1 / torch 2.11.0+cu130 / cuda 13.0 +# Expected: tilert 0.1.5.post2 / torch 2.11.0+cu130 / cuda 13.0 ``` Proceed to [Getting Started](#getting-started) to download and convert model weights. From d4c6ec1c5d9346cec5b7666708de9905d4a9e292 Mon Sep 17 00:00:00 2001 From: XieBaijie Date: Thu, 13 Aug 2026 11:00:22 +0800 Subject: [PATCH 3/4] perf(pd): close the PD decode-path gap to the non-PD engine reference (#59) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit # 摘要 优化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，随时可以摘掉。 --- tilert/pd_vllm/decode_server.py | 32 ++++++++++++++- tilert/pd_vllm/pd_router.py | 12 +++++- tilert/pd_vllm/profiles/glm5.py | 12 +++++- tilert/pd_vllm/profiles/mla_nsa.py | 63 +++++++++++++++++++----------- 4 files changed, 93 insertions(+), 26 deletions(-) diff --git a/tilert/pd_vllm/decode_server.py b/tilert/pd_vllm/decode_server.py index 3372694..4ebd777 100644 --- a/tilert/pd_vllm/decode_server.py +++ b/tilert/pd_vllm/decode_server.py @@ -17,6 +17,7 @@ import contextlib import json import logging +import os import queue as queue_mod import socket import threading @@ -33,6 +34,9 @@ logger = logging.getLogger("pd_vllm.decode_server") +DECODE_POLL_S = max(0.0, float(os.environ.get("TILERT_DECODE_POLL_MS") or "200")) / 1000.0 + + class DecodeBody(BaseModel): rid: str first_token_id: int @@ -176,6 +180,12 @@ def pd_decode(body: DecodeBody): # streaming: ndjson lines {"t":[ids...]}* then {"done":true,...}; # lock/engine ownership transfers to the generator. q: queue_mod.Queue = queue_mod.Queue() + fin: dict = {"loop": None, "ev": None} + + def _signal_done() -> None: + loop, ev = fin["loop"], fin["ev"] + if loop is not None and ev is not None: + loop.call_soon_threadsafe(ev.set) def _run(): try: @@ -187,9 +197,11 @@ def _run(): cancel_event=cancel, ) q.put(("done", tokens)) + _signal_done() except Exception as e: # pragma: no cover logger.exception("stream decode failed for %s", body.rid) q.put(("error", str(e))) + _signal_done() worker = threading.Thread(target=_run, name="pd-decode", daemon=True) @@ -204,11 +216,28 @@ async def _gen(): import anyio from starlette.concurrency import run_in_threadpool + fin["loop"] = asyncio.get_running_loop() + fin["ev"] = asyncio.Event() worker.start() try: batch: list[int] = [] done_msg = None last_activity = time.time() + while done_msg is None: + try: + first = q.get_nowait() + except queue_mod.Empty: + if time.time() - last_activity > 600: + yield json.dumps({"error": "decode stalled"}) + "\n" + return + await asyncio.sleep(0.001) + continue + if isinstance(first, int): + yield json.dumps({"t": [first]}) + "\n" + else: + done_msg = first + last_activity = time.time() + break while done_msg is None: drained = False while True: @@ -232,7 +261,8 @@ async def _gen(): yield json.dumps({"error": "decode stalled"}) + "\n" return else: - await asyncio.sleep(0.005) + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(fin["ev"].wait(), timeout=DECODE_POLL_S) kind, payload = done_msg if kind == "done": timing = { diff --git a/tilert/pd_vllm/pd_router.py b/tilert/pd_vllm/pd_router.py index e599827..896a98d 100644 --- a/tilert/pd_vllm/pd_router.py +++ b/tilert/pd_vllm/pd_router.py @@ -334,8 +334,14 @@ async def _gen(): detok = IncrementalDetok(ctx.tokenizer) sess = parser.stream() if parser else None client = httpx.AsyncClient(timeout=httpx.Timeout(600, read=600)) + role_sent = False + + def _role_once(): + nonlocal role_sent + role_sent = True + return _chunk({"role": "assistant"}) + try: - yield _chunk({"role": "assistant"}) async with client.stream( "POST", f"{node.http_base}/pd/decode", @@ -364,6 +370,8 @@ async def _gen(): text = detok.push(msg["t"]) if not text: continue + if not role_sent: + yield _role_once() if sess is None: yield _chunk({"content": text}) continue @@ -388,6 +396,8 @@ async def _gen(): yield _chunk(_event_delta(ev)) if saw_tool: finish_reason = "tool_calls" + if not role_sent: + yield _role_once() yield _chunk({}, finish=finish_reason) yield _usage_chunk( { diff --git a/tilert/pd_vllm/profiles/glm5.py b/tilert/pd_vllm/profiles/glm5.py index c98d7e9..ec77122 100644 --- a/tilert/pd_vllm/profiles/glm5.py +++ b/tilert/pd_vllm/profiles/glm5.py @@ -7,14 +7,22 @@ from __future__ import annotations +import os + from tilert.pd_vllm.profiles import base from tilert.pd_vllm.profiles.mla_nsa import ( MlaNsaEngineAdapter, MlaNsaProfile, ) -NUM_LAYERS = 79 # 78 main + 1 MTP draft -LAYOUT_VERSION = 10 # glm5 wire family +_NO_MTP = (os.environ.get("TILERT_PD_NO_MTP") or "0").strip().lower() not in ( + "0", + "false", + "no", + "off", +) +NUM_LAYERS = 78 if _NO_MTP else 79 # 78 main + 1 MTP draft +LAYOUT_VERSION = 1010 if _NO_MTP else 10 # glm5 wire family def _build_engine(model_weights_dir, max_seq_len, with_mtp, ar_steps): diff --git a/tilert/pd_vllm/profiles/mla_nsa.py b/tilert/pd_vllm/profiles/mla_nsa.py index ed3c079..f8747bd 100644 --- a/tilert/pd_vllm/profiles/mla_nsa.py +++ b/tilert/pd_vllm/profiles/mla_nsa.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import os import re from dataclasses import dataclass @@ -414,6 +415,7 @@ def _decode_mtp(self, first_token_id, budget, on_token, cancel_event): self.last_stats = {"finish_reason": "stop"} return [] dl.set_prefill_valid_tokens(0) + ar_steps = max(1, min(1024, int(os.environ.get("GLM5_AR_N", "8")))) draft = torch.full((1, T), int(self._last_prompt_token), dtype=torch.int32, device="cuda:0") accepted, finish, fwd, finished = [], "length", 0, False while not finished and len(tokens) < budget: @@ -424,18 +426,27 @@ def _decode_mtp(self, first_token_id, budget, on_token, cancel_event): draft = torch.full((1, T), int(first_token_id), dtype=torch.int32, device="cuda:0") elif fwd > 1: draft = dl.get_next_draft_tokens(0).reshape(1, T) - dl.forward(draft) - n_acc = dl.get_num_accepted(0) - pred = dl.get_predicted_tokens(0).flatten() + if fwd == 0: + steps = 1 + else: + rem = budget - len(tokens) + steps = max(1, min(ar_steps, -(-rem // T))) + dl.show_hands(draft, steps) + acc = dl.ar_accepted_tokens(0).cpu() + num = dl.ar_num_accepted(0).cpu() + n_tokens = int(acc[0].item()) + n_steps = int(num[0].item()) + emitted = acc[1 : 1 + n_tokens].tolist() + per_step = num[1 : 1 + n_steps].tolist() if fwd == 0: fwd += 1 continue - accepted.append(n_acc) + accepted.extend(per_step) fwd += 1 - for i in range(n_acc): + for tok in emitted: if len(tokens) >= budget: break - tok = int(pred[i].item()) + tok = int(tok) if tok in stop_ids: finished = True finish = "stop" @@ -452,8 +463,6 @@ def _decode_mtp(self, first_token_id, budget, on_token, cancel_event): return tokens def _decode_standard(self, first_token_id, budget, on_token, cancel_event): - from tilert.models.deepseek_v3_2.temp_var_indices import Idx - dl = self.gen.decode_layer stop_ids = set() if self._ignore_eos else self.stop_ids torch = self._torch @@ -463,23 +472,33 @@ def _decode_standard(self, first_token_id, budget, on_token, cancel_event): if int(first_token_id) in stop_ids: self.last_stats = {"finish_reason": "stop"} return [] - finish = "length" - cur = torch.tensor(int(first_token_id), dtype=torch.long, device="cuda:0") - while len(tokens) < budget: + dl.set_prefill_valid_tokens(0, with_mtp=False) + ar_steps = max(1, min(1024, int(os.environ.get("GLM5_AR_N", "8")))) + finish, finished = "length", False + last_tok = int(first_token_id) + prev = torch.tensor([last_tok], dtype=torch.int32, device="cuda:0") + while not finished and len(tokens) < budget: if cancel_event is not None and cancel_event.is_set(): finish = "cancelled" break - res = dl.forward(cur) - intermediates, *_ = res[0] - nxt = intermediates[Idx.TOKEN_OUT][0][0] - tok = int(nxt.item()) - if tok in stop_ids: - finish = "stop" - break - tokens.append(tok) - if on_token: - on_token(tok) - cur = nxt + steps = max(1, min(ar_steps, budget - len(tokens))) + dl.show_hands_no_mtp(prev, steps) + acc = dl.ar_accepted_tokens_no_mtp(0).cpu() + n_tokens = int(acc[0].item()) + emitted = acc[1 : 1 + n_tokens].tolist() + for tok in emitted: + if len(tokens) >= budget: + break + tok = int(tok) + if tok in stop_ids: + finished = True + finish = "stop" + break + tokens.append(tok) + last_tok = tok + if on_token: + on_token(tok) + prev = torch.tensor([last_tok], dtype=torch.int32, device="cuda:0") dl.reset_sequence() self.last_stats = {"finish_reason": finish} return tokens From 3dff7098075f4d2866f055554cc38661515199c3 Mon Sep 17 00:00:00 2001 From: CrimsonDump Date: Thu, 13 Aug 2026 15:14:33 +0800 Subject: [PATCH 4/4] fix(pd): role chunk before decode errors; count only consumed MTP steps --- tilert/pd_vllm/pd_router.py | 2 ++ tilert/pd_vllm/profiles/mla_nsa.py | 28 +++++++++++++++++----------- 2 files changed, 19 insertions(+), 11 deletions(-) diff --git a/tilert/pd_vllm/pd_router.py b/tilert/pd_vllm/pd_router.py index 896a98d..768c643 100644 --- a/tilert/pd_vllm/pd_router.py +++ b/tilert/pd_vllm/pd_router.py @@ -384,6 +384,8 @@ def _role_once(): if finish_reason == "cancelled": finish_reason = "stop" elif "error" in msg: + if not role_sent: + yield _role_once() yield _chunk({"content": f"\n[decode error: {msg['error']}]"}) finish_reason = "stop" if client_gone: diff --git a/tilert/pd_vllm/profiles/mla_nsa.py b/tilert/pd_vllm/profiles/mla_nsa.py index f8747bd..29c1b35 100644 --- a/tilert/pd_vllm/profiles/mla_nsa.py +++ b/tilert/pd_vllm/profiles/mla_nsa.py @@ -441,19 +441,25 @@ def _decode_mtp(self, first_token_id, budget, on_token, cancel_event): if fwd == 0: fwd += 1 continue - accepted.extend(per_step) fwd += 1 - for tok in emitted: - if len(tokens) >= budget: - break - tok = int(tok) - if tok in stop_ids: - finished = True - finish = "stop" + offset = 0 + for na in per_step: + step_emit = emitted[offset : offset + na] + offset += na + for tok in step_emit: + if len(tokens) >= budget: + break + tok = int(tok) + if tok in stop_ids: + finished = True + finish = "stop" + break + tokens.append(tok) + if on_token: + on_token(tok) + accepted.append(na) + if finished or len(tokens) >= budget: break - tokens.append(tok) - if on_token: - on_token(tok) dl.reset_sequence() self.last_stats = { "finish_reason": finish,