diff --git a/README.md b/README.md index fdd123e..bd88e06 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.post3) 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.post3` 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.post3 # 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.post3/tilert-0.1.5.post3-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.post3 / torch 2.11.0+cu130 / cuda 13.0 ``` Proceed to [Getting Started](#getting-started) to download and convert model weights. @@ -363,6 +363,8 @@ python -m tilert.pd_vllm.pd_router \ Send OpenAI requests to `http://:23333/v1/chat/completions`. The router runs the prefill on vLLM (first token), hands the attention state to the TileRT decode node over RDMA, and streams the completion back. +A decode engine serves one sequence at a time, so the router reserves a node per request and answers `429` while they are all busy. Add `--queue-timeout ` to make a request wait for a free node instead of failing: useful when a single client fans out into concurrent sub-conversations β€” an agentic session spawning sub-agents, say β€” and the burst is wider than the pool but short-lived. Waits longer than 0.1 s are logged. The default, `0`, keeps the fail-fast behaviour. + ### Topology B: shared prefill β†’ TileRT decode **and** native vLLM decode One prefill pool feeds two decode pools side by side, composed under vLLM's `MultiConnector`. Each request is claimed by exactly one connector β€” the TileRT connector claims requests marked with `tilert_host`, and vLLM's native connector handles the rest β€” so latency-critical traffic goes to TileRT while general traffic stays on native vLLM decode, behind the same OpenAI surface. 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..4c5ea85 100644 --- a/tilert/pd_vllm/pd_router.py +++ b/tilert/pd_vllm/pd_router.py @@ -4,7 +4,8 @@ non-streaming. Flow per request (phase-1 hybrid, see design doc): - 1. pick a free decode node (in-memory busy tracking; all busy -> 429) + 1. pick a free decode node (in-memory busy tracking; all busy -> wait up to + --queue-timeout, then 429) 2. forward to vLLM with max_tokens=1 + logprobs and inject kv_transfer_params {tilert_host, tilert_ctrl_port} β€” the connector claims the request and RDMA-sends state to the decode node @@ -40,6 +41,8 @@ logger = logging.getLogger("pd_vllm.router") +QUEUE_LOG_SECONDS = 0.1 + class DecodeNode: def __init__(self, host: str, ctrl_port: int, http_port: int): @@ -54,21 +57,31 @@ def http_base(self) -> str: class Pool: - def __init__(self, nodes: list[DecodeNode]): + """Decode-node reservation; ``queue_timeout`` > 0 waits instead of failing fast.""" + + def __init__(self, nodes: list[DecodeNode], queue_timeout: float = 0.0): self.nodes = nodes - self._lock = threading.Lock() + self.queue_timeout = queue_timeout + self._cv = threading.Condition() def acquire(self) -> DecodeNode | None: - with self._lock: - for n in self.nodes: - if not n.busy: - n.busy = True - return n - return None + """Reserve a node, or None once ``queue_timeout`` elapses.""" + deadline = time.monotonic() + self.queue_timeout + with self._cv: + while True: + for n in self.nodes: + if not n.busy: + n.busy = True + return n + remaining = deadline - time.monotonic() + if remaining <= 0: + return None + self._cv.wait(remaining) def release(self, node: DecodeNode) -> None: - with self._lock: + with self._cv: node.busy = False + self._cv.notify() def first_token_from_logprobs(resp: dict, is_chat: bool) -> int: @@ -97,6 +110,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).""" @@ -123,6 +167,24 @@ def build_app(ctx: RouterCtx) -> FastAPI: app = FastAPI() pool = ctx.pool + def _acquire_node() -> tuple[DecodeNode | None, float]: + """Reserve a node, and report how long the caller queued for it.""" + t0 = time.monotonic() + node = pool.acquire() + waited = time.monotonic() - t0 + if node is not None and waited >= QUEUE_LOG_SECONDS: + logger.info("queued %.1fs for decode node %s", waited, node.host) + return node, waited + + def _busy_response(waited: float) -> JSONResponse: + """429 body, distinguishing a full pool from an exhausted queue timeout.""" + detail = ( + f"no decode node free after waiting {waited:.1f}s" + if pool.queue_timeout > 0 + else "all decode nodes busy" + ) + return JSONResponse({"error": detail}, status_code=429) + @app.get("/health") def health(): return {"status": "ok", "decode_free": sum(1 for n in pool.nodes if not n.busy)} @@ -133,24 +195,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) @@ -158,9 +209,9 @@ def _max_tokens_of(body): # ── non-streaming ──────────────────────────────────────────────────── def _handle(path: str, body: dict): is_chat = path.endswith("chat/completions") - node = pool.acquire() + node, waited = _acquire_node() if node is None: - return JSONResponse({"error": "all decode nodes busy"}, status_code=429) + return _busy_response(waited) t0 = time.time() try: prefill = _prefill(path, body, node) @@ -235,20 +286,29 @@ def _handle(path: str, body: dict): # ── streaming (chat only) ──────────────────────────────────────────── async def _handle_stream(path: str, body: dict, request: Request): + import anyio from starlette.concurrency import run_in_threadpool - node = pool.acquire() - if node is None: - return JSONResponse({"error": "all decode nodes busy"}, status_code=429) - + node = None # reserved inside the try: both handlers below release it try: + # Shielded: a disconnect must not strand a reservation mid-acquire. + with anyio.CancelScope(shield=True): + node, waited = await run_in_threadpool(_acquire_node) + if node is None: + return _busy_response(waited) prefill = await run_in_threadpool(_prefill, path, body, node) rid = derive_rid(prefill["id"]) first_token_id = first_token_from_logprobs(prefill, True) except Exception as e: - pool.release(node) + if node is not None: + pool.release(node) logger.exception("pd stream request failed before streaming") return JSONResponse({"error": str(e)}, status_code=502) + except BaseException: + # CancelledError is not an Exception; without this the node stays busy. + if node is not None: + pool.release(node) + raise chunk_id = prefill["id"] model = prefill.get("model") @@ -267,6 +327,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 +374,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 +410,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 @@ -344,9 +423,13 @@ async def _gen(): finish_reason = msg.get("finish_reason", "stop") if finish_reason == "cancelled": finish_reason = "stop" + break elif "error" in msg: + if not role_sent: + yield _role_once() yield _chunk({"content": f"\n[decode error: {msg['error']}]"}) finish_reason = "stop" + break if client_gone: logger.info("client gone mid-stream for %s", rid) return # finally fires the cancel @@ -357,13 +440,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 @@ -425,6 +510,12 @@ def main() -> None: default="glm47", help="output parser (reasoning + tool calls)", ) + ap.add_argument( + "--queue-timeout", + type=float, + default=0.0, + help="seconds to wait for a free decode node before answering 429 (0: fail fast)", + ) args = ap.parse_args() nodes = [] @@ -440,7 +531,7 @@ def main() -> None: args.model_path, trust_remote_code=True ) # nosec B615 - ctx = RouterCtx(args.vllm_url, Pool(nodes), tokenizer, args.parser) + ctx = RouterCtx(args.vllm_url, Pool(nodes, args.queue_timeout), tokenizer, args.parser) app = build_app(ctx) logger.info( "router on :%d -> vllm=%s, %d decode node(s), parser=%s", 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..2b6d656 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 @@ -12,6 +13,14 @@ logger = logging.getLogger("pd_vllm.profile.mla_nsa") +_AR_MTP_API = ("show_hands", "ar_accepted_tokens", "ar_num_accepted") +_AR_PLAIN_API = ("show_hands_no_mtp", "ar_accepted_tokens_no_mtp") + + +def _has_api(dl, names: tuple[str, ...]) -> bool: + return all(hasattr(dl, n) for n in names) + + KV_LORA_RANK = 512 QK_ROPE_HEAD_DIM = 64 INDEX_HEAD_DIM = 128 @@ -364,6 +373,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 +402,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 +414,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 +423,8 @@ 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")))) + ar_ok = _has_api(dl, _AR_MTP_API) 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 +435,47 @@ 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 ar_ok: + 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() + else: + dl.forward(draft) + n_acc = int(dl.get_num_accepted(0)) + pred = dl.get_predicted_tokens(0).flatten() + emitted = [int(pred[i].item()) for i in range(n_acc)] + per_step = [n_acc] if fwd == 0: fwd += 1 continue - accepted.append(n_acc) fwd += 1 - for i in range(n_acc): - if len(tokens) >= budget: - break - tok = int(pred[i].item()) - 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, @@ -450,10 +485,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,8 +494,47 @@ 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 [] + if not _has_api(dl, _AR_PLAIN_API): + return self._decode_plain_per_step(tokens, budget, on_token, cancel_event) + 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 + 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 + + def _decode_plain_per_step(self, tokens, 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 finish = "length" - cur = torch.tensor(int(first_token_id), dtype=torch.long, device="cuda:0") + cur = torch.tensor(int(tokens[0]), dtype=torch.long, device="cuda:0") while len(tokens) < budget: if cancel_event is not None and cancel_event.is_set(): finish = "cancelled"