perf(pd): close the PD decode-path gap to the non-PD engine reference

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

@@ -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(
                     {

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

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