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. 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/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 87a61e5..768c643 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"]} @@ -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 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 a270f7c..29c1b35 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 @@ -364,6 +365,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 +394,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 +406,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: @@ -412,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: @@ -422,25 +426,40 @@ 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) fwd += 1 - for i in range(n_acc): - if len(tokens) >= budget: + 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 - tok = int(pred[i].item()) - if tok in stop_ids: - finished = True - finish = "stop" - break - tokens.append(tok) - if on_token: - on_token(tok) dl.reset_sequence() self.last_stats = { "finish_reason": finish, @@ -450,10 +469,8 @@ 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 = 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: @@ -461,23 +478,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