diff --git a/tilert/pd_vllm/pd_router.py b/tilert/pd_vllm/pd_router.py index 87a61e5..7114a9d 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,18 +164,7 @@ 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()