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/pd_router.py b/tilert/pd_vllm/pd_router.py index 87a61e5..e599827 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"]} @@ -357,13 +388,13 @@ async def _gen(): yield _chunk(_event_delta(ev)) if saw_tool: finish_reason = "tool_calls" - yield _chunk( - {}, - finish=finish_reason, - usage={ + 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/mla_nsa.py b/tilert/pd_vllm/profiles/mla_nsa.py index a270f7c..ed3c079 100644 --- a/tilert/pd_vllm/profiles/mla_nsa.py +++ b/tilert/pd_vllm/profiles/mla_nsa.py @@ -364,6 +364,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 +393,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 +405,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: @@ -453,7 +455,7 @@ 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: