fix(pd): role chunk before decode errors; count only consumed MTP steps

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.

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",

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

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 = {

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"]}
@@ -303,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",
@@ -333,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
@@ -345,6 +384,8 @@ async def _gen():
                             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:
@@ -357,13 +398,15 @@ async def _gen():
                         yield _chunk(_event_delta(ev))
                 if saw_tool:
                     finish_reason = "tool_calls"
-                yield _chunk(
-                    {},
-                    finish=finish_reason,
-                    usage={
+                if not role_sent:
+                    yield _role_once()
+                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

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):