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