diff --git a/.flake8 b/.flake8
index a247806..67c009c 100644
--- a/.flake8
+++ b/.flake8
@@ -34,6 +34,11 @@ ignore =
per-file-ignores =
__init__.py:F401
tests/*:D100,D101,D102,D103
+ # pd_vllm tests are built from test doubles that mirror the real signatures
+ # (U100/U101), stack ``with`` blocks for readability (SIM117), and carry
+ # rationale docstrings that open with an identifier or run one sentence
+ # over several lines (D205/D210/D403).
+ tests/pd_vllm/*:D205,D210,D403,U100,U101,SIM117,R504,VNE002
setup.py:D100,D101,D102,D103,B009
# pd_vllm implements vLLM connector / abstract-profile interfaces, so many
# method params are unused by design (U100); black collapses the ``...``
diff --git a/README.md b/README.md
index fdd123e..59e1c8c 100644
--- a/README.md
+++ b/README.md
@@ -315,11 +315,14 @@ TileRT can run as the **decode engine behind a vLLM prefill**, integrated throug
**Prerequisites**
- Convert the model weights for TileRT decode (see [Step 2](#step-2-shard-weights-with-weight_converter)).
-- On the **prefill** node, a vLLM build with V1 disaggregation and support for the GLM-5/5.1 / DeepSeek-V3.2 (DSA) model and the `fp8_ds_mla` KV-cache dtype. Install `tilert` in the same environment so the connector plugin is importable.
-- **The KV-cache dtype must match on both ends.** These examples use fp8: `--kv-cache-dtype fp8_ds_mla` on the vLLM prefill and `--kv-cache-dtype fp8` on the TileRT decode (a mismatch is rejected at the connector handshake).
+- On the **prefill** node, a vLLM build with V1 disaggregation and support for the GLM-5/5.1/5.2/5.3 / DeepSeek-V3.2 (DSA) model. Install `tilert` in the same environment so the connector plugin is importable.
+- **The KV-cache dtype must match on both ends.** These examples use fp8: `--kv-cache-dtype fp8_ds_mla` on the vLLM prefill and `--kv-cache-dtype fp8` on the TileRT decode (a mismatch is rejected at the connector handshake). On a vLLM build without `fp8_ds_mla` (the ROCm build, for example) use the bf16 path instead: `--kv-cache-dtype auto` on the prefill and `--kv-cache-dtype bf16` on the decode.
+- **Both ends must run the same `tilert` release.** The control plane is versioned (`PROTOCOL_VERSION`): a sender now waits to be *admitted* before it RDMA-writes, and a receiver refuses the handshake with a sender that would not. `tilert_sync_send` (sending inside the forward window) has been removed and is ignored with a warning; `tilert_admission_attempts` (default 5) bounds the admission retries.
- The examples use the **NIXL** transfer engine. On multi-NIC hosts, pin NIXL to the RDMA NICs via `UCX_NET_DEVICES` (otherwise UCX may pick the wrong interface). Mooncake is also supported (`--transport mooncake` on the decode, `"tilert_transport": "mooncake"` on the prefill).
-Commands below use GLM-5/5.1. For DeepSeek-V3.2, use `--model deepseek_v3_2`, the DeepSeek-V3.2-TileRT weights, and `--parser none`.
+Commands below use GLM-5/5.1. For DeepSeek-V3.2, use `--model deepseek_v3_2`, the DeepSeek-V3.2-TileRT weights, and `--parser none`. For GLM-5.2 / GLM-5.3 use `--model glm5_2` (or `glm5_3`; same profile): on a ROCm torch build the ROCm engine adapter is selected automatically (`TILERT_PD_ENGINE_BACKEND=rocm|cuda` forces it), and the sparse-indexer layer set is checked at connector registration.
+
+**Router behaviour.** The router reads each decode node's `GET /capabilities` and refuses, before any prefill runs, a request whose sampling fields no node can execute (penalties, `min_p`, `seed`, `n`, ... are refused when non-neutral; `stop`, `logprobs`/`top_logprobs`, `response_format` and `ignore_eos` are served). Sampling defaults are resolved once and written to both legs: `--model` names the profile the decode nodes serve, `--generation-config auto|vllm` mirrors vLLM's flag, and `--default-temperature/--default-top-p/--default-top-k/--default-repetition-penalty` override the resolved values. `--queue-timeout` waits for a free decode node before answering 429; `--force-include-usage` adds the trailing usage chunk to every stream.
### Topology A: vLLM prefill → TileRT decode
diff --git a/pyproject.toml b/pyproject.toml
index 407c497..e1b1bad 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -121,6 +121,17 @@ disallow_untyped_defs = false
disallow_incomplete_defs = false
warn_return_any = false
+# The pd_vllm tests drive the real modules through duck-typed fakes
+# (SimpleNamespace servers, recording engines), so def-level annotations and
+# attribute checks on those doubles are noise; syntax and import errors still
+# surface.
+[[tool.mypy.overrides]]
+module = "tests.pd_vllm.*"
+disallow_untyped_defs = false
+disallow_incomplete_defs = false
+check_untyped_defs = false
+warn_return_any = false
+
[tool.bandit]
exclude_dirs = ["tests", "3rd-party"]
skips = ["B101", "B311", "B404", "B603", "B607"]
@@ -164,6 +175,7 @@ ignore = [
[tool.flake8.per-file-ignores]
"__init__.py" = ["F401"]
"tests/*" = ["D100", "D101", "D102", "D103"]
+"tests/pd_vllm/*" = ["D205", "D210", "D403", "U100", "U101", "SIM117", "R504", "VNE002"]
"setup.py" = ["D100", "D101", "D102", "D103", "B009"]
[tool.flake8.bugbear]
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/pd_vllm/__init__.py b/tests/pd_vllm/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/pd_vllm/test_generation_defaults.py b/tests/pd_vllm/test_generation_defaults.py
new file mode 100644
index 0000000..f33c350
--- /dev/null
+++ b/tests/pd_vllm/test_generation_defaults.py
@@ -0,0 +1,405 @@
+"""Sampling defaults resolve once, the way vLLM resolves them, for both PD legs.
+
+vLLM works out its ``default_sampling_params`` at startup
+(``ModelConfig.get_diff_sampling_param``): read the model's
+``generation_config.json`` unless ``--generation-config vllm``, apply
+``--override-generation-config``, keep a six-key allowlist, then per request
+resolve ``client value > that > the neutral defaults``.
+
+The decode adapters used to carry literals of their own instead, which for a
+checkpoint recommending 0.6 / 20 put ``temperature`` at 0.6 for the first token
+and 1.0 for the rest, and ``top_k`` at 20 for the first token and uncapped for
+the rest. This module resolves once and
+the router writes the result into both requests, so the legs cannot drift.
+
+Guarded fields are the other half: ``min_p`` has no decode implementation on any
+member, and ``repetition_penalty`` only where the profile declares a penalty
+pre-pass (no public profile does today, so a stub stands in). A config asking for one
+this deployment cannot execute fails at STARTUP -- the client never sent it, so no
+per-request check would ever see it.
+
+CPU only -- no GPU, no tilert, no vLLM.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_generation_defaults.py -v
+"""
+
+import json
+import types
+
+import pytest
+
+from tilert.pd_vllm import generation_defaults as gd
+from tilert.pd_vllm.capabilities import (
+ CapabilityUnavailable,
+ NodeCapabilities,
+ validate_generation_request,
+)
+
+RECOMMENDED_CONFIG = {
+ # A shipped generation_config.json with recommended sampling, sampling keys only.
+ "temperature": 0.6,
+ "top_p": 0.95,
+ "top_k": 20,
+ # Keys vLLM's allowlist ignores; they must not reach the sampler.
+ "do_sample": True,
+ "bos_token_id": 248044,
+ "eos_token_id": [248046, 248044],
+}
+
+
+@pytest.fixture(autouse=True)
+def penalty_profile(monkeypatch):
+ """Register a profile whose decode runtime declares penalties.
+
+ ``penalties_supported_by`` reads the profile's static ``declares_penalties``;
+ none of the public profiles (GLM-5 / GLM-5.2 / DSV3.2) set it, so the
+ adopt-a-penalty path is exercised through a stub registered for the test.
+ """
+ from tilert.pd_vllm.profiles import base
+
+ stub = types.SimpleNamespace(name="stub_penalties", declares_penalties=True)
+ monkeypatch.setitem(base._REGISTRY, "stub_penalties", stub)
+
+
+@pytest.fixture
+def model_dir(tmp_path):
+ def _write(config):
+ (tmp_path / "generation_config.json").write_text(json.dumps(config))
+ return str(tmp_path)
+
+ return _write
+
+
+# --------------------------------------------------------------------------- #
+# The chain, and what it takes from the file
+# --------------------------------------------------------------------------- #
+def test_the_model_config_supplies_the_defaults(model_dir):
+ d = gd.load(model_dir(RECOMMENDED_CONFIG), "auto", model="stub_penalties")
+ assert (d.temperature, d.top_p, d.top_k) == (0.6, 0.95, 20)
+
+
+def test_keys_outside_the_sampling_allowlist_are_ignored(model_dir):
+ """``generation_config.json`` is a general HF file: it carries token ids and
+ ``do_sample`` too, and vLLM's allowlist is what keeps those out of the
+ sampler. Adopting the file wholesale would put ``eos_token_id`` on the wire
+ as a sampling parameter.
+ """
+ d = gd.load(model_dir(RECOMMENDED_CONFIG), "auto", model="stub_penalties")
+ resolved = d.resolve({})
+ # Exactly vLLM's six-key allowlist minus max_new_tokens, which max_tokens=1
+ # already overrides on the prefill leg.
+ assert set(resolved) == {"temperature", "top_p", "top_k", "repetition_penalty", "min_p"}
+
+
+def test_generation_config_vllm_ignores_the_file(model_dir):
+ """vLLM's own escape hatch, spelled the same way."""
+ d = gd.load(model_dir(RECOMMENDED_CONFIG), "vllm", model="stub_penalties")
+ assert (d.temperature, d.top_p, d.top_k) == (1.0, 1.0, 0)
+
+
+def test_a_missing_file_falls_back_to_the_neutral_defaults(tmp_path):
+ d = gd.load(str(tmp_path), "auto", model="stub_penalties")
+ assert (d.temperature, d.top_p, d.top_k) == (1.0, 1.0, 0)
+
+
+def test_no_model_path_falls_back_to_the_neutral_defaults():
+ d = gd.load("", "auto", model="stub_penalties")
+ assert (d.temperature, d.top_p, d.top_k) == (1.0, 1.0, 0)
+
+
+def test_a_non_object_file_is_refused(tmp_path):
+ (tmp_path / "generation_config.json").write_text("[1, 2, 3]")
+ with pytest.raises(gd.UnsupportedGenerationDefault):
+ gd.load(str(tmp_path), "auto", model="stub_penalties")
+
+
+@pytest.mark.parametrize(
+ "field,value",
+ [
+ ("temperature", 0.9),
+ ("top_p", 0.5),
+ ("top_k", 40),
+ ],
+)
+def test_a_command_line_override_wins_over_the_file(model_dir, field, value):
+ d = gd.load(model_dir(RECOMMENDED_CONFIG), "auto", model="stub_penalties", **{field: value})
+ assert getattr(d, field) == value
+
+
+def test_a_partial_config_only_displaces_what_it_states(model_dir):
+ d = gd.load(model_dir({"top_p": 0.8}), "auto", model="stub_penalties")
+ assert d.top_p == 0.8
+ assert (d.temperature, d.top_k) == (1.0, 0) # still neutral
+
+
+def test_the_source_is_named_for_the_startup_log(model_dir):
+ path = model_dir(RECOMMENDED_CONFIG)
+ assert "generation_config.json" in gd.load(path, "auto", model="stub_penalties").source
+ assert "neutral" in gd.load(path, "vllm", model="stub_penalties").source
+ assert "override" in gd.load(path, "auto", model="stub_penalties", top_p=0.5).source
+
+
+# --------------------------------------------------------------------------- #
+# Per-request resolution: the client always wins
+# --------------------------------------------------------------------------- #
+DEFAULTS = gd.GenerationDefaults(temperature=0.6, top_p=0.95, top_k=20)
+
+
+@pytest.mark.parametrize(
+ "field,sent,want",
+ [
+ ("temperature", 0.2, 0.2),
+ ("top_p", 0.5, 0.5),
+ ("top_k", 5, 5),
+ ],
+)
+def test_an_explicit_client_value_wins(field, sent, want):
+ assert DEFAULTS.resolve({field: sent})[field] == want
+
+
+@pytest.mark.parametrize(
+ "field,want",
+ [
+ ("temperature", 0.6),
+ ("top_p", 0.95),
+ ("top_k", 20),
+ ],
+)
+def test_an_absent_field_takes_the_deployment_default(field, want):
+ assert DEFAULTS.resolve({})[field] == want
+
+
+@pytest.mark.parametrize(
+ "field,want",
+ [
+ ("temperature", 0.6),
+ ("top_p", 0.95),
+ ("top_k", 20),
+ ],
+)
+def test_an_explicit_null_is_treated_as_absent(field, want):
+ """How SDKs spell "unset", and how vLLM resolves it."""
+ assert DEFAULTS.resolve({field: None})[field] == want
+
+
+def test_zero_is_a_value_not_an_absence():
+ """``temperature: 0`` is a greedy request, not a missing field.
+
+ Falling back to the default here would silently make it sample.
+ """
+ assert DEFAULTS.resolve({"temperature": 0})["temperature"] == 0.0
+ assert DEFAULTS.resolve({"top_k": 0})["top_k"] == 0
+
+
+@pytest.mark.parametrize("field", ["temperature", "top_p", "top_k"])
+def test_a_bool_is_not_a_number(field):
+ with pytest.raises(ValueError):
+ DEFAULTS.resolve({field: True})
+
+
+def test_top_k_travels_in_the_request_domain():
+ """0 is vLLM's "no rank cut" sentinel, which is what belongs on the wire.
+
+ The kernel's own disabled value (``TOP_K_DISABLED`` = the candidate-pool
+ bound) is an engine-side mapping and would be a real rank cut to vLLM.
+ """
+ from tilert.pd_vllm.sampling import TOP_K_DISABLED, resolve_top_k
+
+ wire_value = gd.GenerationDefaults().resolve({})["top_k"]
+ assert wire_value == 0
+ # ... and the engine maps that sentinel to "no cut" on its side.
+ assert resolve_top_k({"top_k": wire_value}) == TOP_K_DISABLED
+
+
+# --------------------------------------------------------------------------- #
+# Guarded fields: refused at startup, per model family
+# --------------------------------------------------------------------------- #
+@pytest.mark.parametrize(
+ "model,supported",
+ [
+ ("stub_penalties", True),
+ ("glm5", False),
+ ("glm5_2", False),
+ ("glm5_3", False),
+ ("dsv32", False),
+ ("", False),
+ ("not-a-model", False),
+ ],
+)
+def test_which_families_declare_penalties(model, supported):
+ assert gd.penalties_supported_by(model) is supported
+
+
+@pytest.mark.parametrize("model", ["stub_penalties"])
+def test_a_penalty_is_adopted_on_a_family_that_declares_it(model_dir, model):
+ d = gd.load(model_dir({**RECOMMENDED_CONFIG, "repetition_penalty": 1.05}), "auto", model=model)
+ assert d.repetition_penalty == 1.05
+
+
+@pytest.mark.parametrize("model", ["glm5", "glm5_2", "dsv32"])
+def test_a_penalty_refuses_startup_on_a_family_without_the_pre_pass(model_dir, model):
+ """Serving would mean penalising the first token and not the rest, on every
+ request, with no per-request check able to see it.
+ """
+ with pytest.raises(gd.UnsupportedGenerationDefault) as e:
+ gd.load(model_dir({**RECOMMENDED_CONFIG, "repetition_penalty": 1.05}), "auto", model=model)
+ assert "repetition_penalty" in str(e.value)
+ # The message must say how to proceed, not just that it stopped.
+ assert "--generation-config vllm" in str(e.value)
+
+
+def test_an_unnamed_model_refuses_a_penalty(model_dir):
+ """Unknown family is the conservative answer: the cost is refusing to adopt
+ a default, not serving half-penalised.
+ """
+ with pytest.raises(gd.UnsupportedGenerationDefault) as e:
+ gd.load(model_dir({**RECOMMENDED_CONFIG, "repetition_penalty": 1.05}), "auto")
+ assert "--model" in str(e.value)
+
+
+@pytest.mark.parametrize("model", ["stub_penalties", "glm5_2"])
+def test_min_p_refuses_startup_on_every_family(model_dir, model):
+ with pytest.raises(gd.UnsupportedGenerationDefault) as e:
+ gd.load(model_dir({**RECOMMENDED_CONFIG, "min_p": 0.05}), "auto", model=model)
+ assert "min_p" in str(e.value)
+
+
+@pytest.mark.parametrize("model", ["stub_penalties", "glm5_2"])
+def test_a_guarded_field_at_its_neutral_value_is_accepted(model_dir, model):
+ """Stating the no-op asks for nothing, so there is nothing to refuse."""
+ d = gd.load(
+ model_dir({**RECOMMENDED_CONFIG, "repetition_penalty": 1.0, "min_p": 0.0}),
+ "auto",
+ model=model,
+ )
+ assert d.repetition_penalty == 1.0
+
+
+def test_generation_config_vllm_clears_a_guarded_field(model_dir):
+ """The documented way out of the startup refusal."""
+ d = gd.load(model_dir({**RECOMMENDED_CONFIG, "min_p": 0.05}), "vllm", model="glm5_2")
+ assert d.repetition_penalty == 1.0
+
+
+def test_an_unadoptable_penalty_is_pinned_to_the_no_op(model_dir):
+ """On a family without the pre-pass, the penalty must not be left for vLLM
+ to resolve from the model config on the prefill leg alone -- both legs are
+ pinned to the runtime no-op instead.
+ """
+ d = gd.load(
+ model_dir({**RECOMMENDED_CONFIG, "repetition_penalty": 1.0}), "auto", model="glm5_2"
+ )
+ assert d.resolve({})["repetition_penalty"] == 1.0
+
+
+# --------------------------------------------------------------------------- #
+# An adopted default is gated like a client-sent one
+# --------------------------------------------------------------------------- #
+def test_an_adopted_penalty_is_refused_on_a_node_that_cannot_apply_it():
+ """The family declares support, but THIS node's engine demoted its claim (an
+ older tilert wheel without the pre-pass). The value never appears in the
+ request body, so the gate has to be told about it or it decodes unpenalised.
+ """
+ adopted = gd.GenerationDefaults(repetition_penalty=1.05).resolve({})
+ body = {"messages": [{"role": "user", "content": "hi"}]}
+
+ # A node that can apply it: served.
+ validate_generation_request(body, NodeCapabilities(penalties=True), adopted)
+
+ # A node that cannot: refused, and the message says where the value is from.
+ with pytest.raises(CapabilityUnavailable) as e:
+ validate_generation_request(body, NodeCapabilities(penalties=False), adopted)
+ assert "deployment" in str(e.value)
+
+
+def test_a_neutral_adopted_penalty_needs_no_capability():
+ adopted = gd.GenerationDefaults().resolve({})
+ validate_generation_request({}, NodeCapabilities(penalties=False), adopted)
+
+
+def test_a_client_value_is_reported_as_the_client_s():
+ """The two origins must be distinguishable in the error, or an operator
+ cannot tell "drop the field" from "fix the deployment".
+ """
+ with pytest.raises(CapabilityUnavailable) as e:
+ validate_generation_request(
+ {"repetition_penalty": 1.2},
+ NodeCapabilities(penalties=False),
+ gd.GenerationDefaults().resolve({}),
+ )
+ assert "the request" in str(e.value)
+
+
+# --------------------------------------------------------------------------- #
+# Review findings on PR #40 (codex)
+# --------------------------------------------------------------------------- #
+def test_min_p_is_pinned_even_when_the_router_ignores_the_model_config():
+ """`--generation-config vllm` governs the ROUTER's reading, not the vLLM
+ server's.
+
+ The prefill instance is launched separately and still defaults to loading the
+ checkpoint's generation_config.json, so a checkpoint carrying min_p would
+ have it applied to token 1 and ignored for the rest -- on the very path the
+ startup refusal documents as the way out. Pinning the no-op explicitly is
+ what closes it.
+ """
+ for source in ("auto", "vllm"):
+ assert gd.load("", source, model="stub_penalties").resolve({})["min_p"] == 0.0
+
+
+def test_min_p_is_pinned_regardless_of_what_the_client_sent():
+ """A non-neutral client min_p is refused by the gate, so the only value that
+ may reach the sampler is the no-op.
+ """
+ assert DEFAULTS.resolve({"min_p": 0.4})["min_p"] == 0.0
+
+
+def test_the_resolution_covers_vllms_whole_allowlist():
+ """vLLM takes six keys from generation_config.
+
+ Any one left unpinned is a field its own resolution can still move on the prefill leg alone.
+ """
+ vllm_allowlist = {
+ "repetition_penalty",
+ "temperature",
+ "top_k",
+ "top_p",
+ "min_p",
+ "max_new_tokens",
+ }
+ resolved = set(gd.GenerationDefaults().resolve({}))
+ # max_new_tokens is covered by max_tokens=1 on the prefill request instead.
+ assert vllm_allowlist - resolved == {"max_new_tokens"}
+
+
+@pytest.mark.parametrize("model", ["glm5", "glm5_2", "dsv32", ""])
+def test_a_command_line_penalty_override_is_guarded_too(model):
+ """The overrides win over the file, so they need the same guard.
+
+ Without it the router starts with an adopted default the gate refuses on EVERY request --
+ worse than refusing to start.
+ """
+ with pytest.raises(gd.UnsupportedGenerationDefault) as e:
+ gd.load("", "vllm", model=model, repetition_penalty=1.2)
+ assert "command-line overrides" in str(e.value)
+
+
+@pytest.mark.parametrize("model", ["stub_penalties"])
+def test_a_command_line_penalty_override_is_honoured_where_executable(model):
+ assert gd.load("", "vllm", model=model, repetition_penalty=1.2).repetition_penalty == 1.2
+
+
+def test_a_neutral_penalty_override_needs_no_capability():
+ assert gd.load("", "vllm", model="glm5_2", repetition_penalty=1.0).repetition_penalty == 1.0
+
+
+def test_a_min_p_override_is_refused_on_every_family():
+ """There is no --default-min-p, but the guard covers the override path
+ generically, so adding one later cannot bypass it.
+ """
+ with pytest.raises(gd.UnsupportedGenerationDefault):
+ gd._check_guarded(
+ {"min_p": 0.05}, "command-line overrides", model="stub_penalties", penalties_ok=True
+ )
diff --git a/tests/pd_vllm/test_glm5_rocm_engine.py b/tests/pd_vllm/test_glm5_rocm_engine.py
new file mode 100644
index 0000000..41629fb
--- /dev/null
+++ b/tests/pd_vllm/test_glm5_rocm_engine.py
@@ -0,0 +1,282 @@
+"""RocmGlm52EngineAdapter: cache slot mapping and decode loop, no GPU.
+
+A FakeShowHands stands in for the ROCm e2e: CPU cache tensors in the real
+per-rank layout, a scripted accepted-token stream, and a call log. The tests
+pin the contract the GPU run relies on: which slot each (ki, kv, pe) lands
+in, cur_pos == seq_len after inject, and the emitted stream honouring stop /
+budget / cancel in both MTP and plain modes.
+"""
+
+from __future__ import annotations
+
+import threading
+import types
+
+import pytest
+import torch
+
+from tilert.pd_vllm.grammar_spec import GrammarBackendUnavailable
+from tilert.pd_vllm.profiles.glm5_rocm_engine import RocmGlm52EngineAdapter
+
+
+def _layer_kind(i: int) -> int:
+ """Mirror of the engine's rule: 0 = dense, 1 = full (indexer), 2 = shared."""
+ if i < 3:
+ return 0
+ return 1 if (i - 2) % 4 == 0 else 2
+
+
+def _full_layer_ordinals(n: int) -> list[int]:
+ return [i for i in range(n) if _layer_kind(i) != 2]
+
+
+@pytest.fixture(autouse=True, scope="module")
+def _rocm_engine_helpers():
+ """Stand in for the two helpers the adapter imports from the ROCm tilert build.
+
+ ``tilert.models.glm_5.model_args.full_layer_ordinals`` and
+ ``tilert.models.glm_5.weight_converter.{fp8_ki_enabled,pure_tp8_enabled}``
+ ship with the ROCm engine only. When the installed ``tilert`` lacks them
+ (the CUDA tree, or no engine at all) stub modules are put in ``sys.modules``
+ for the duration of this module so the mapping logic is testable anywhere.
+ """
+ import sys
+
+ try:
+ from tilert.models.glm_5.model_args import full_layer_ordinals # noqa: F401
+ from tilert.models.glm_5.weight_converter import fp8_ki_enabled # noqa: F401
+ except ImportError:
+ pass
+ else:
+ yield
+ return
+ ma = types.ModuleType("tilert.models.glm_5.model_args")
+ ma.layer_kind, ma.full_layer_ordinals = _layer_kind, _full_layer_ordinals
+ wc = types.ModuleType("tilert.models.glm_5.weight_converter")
+ wc.fp8_ki_enabled = lambda: False
+ wc.pure_tp8_enabled = lambda: True
+ with pytest.MonkeyPatch.context() as mp:
+ mp.setitem(sys.modules, "tilert.models.glm_5.model_args", ma)
+ mp.setitem(sys.modules, "tilert.models.glm_5.weight_converter", wc)
+ yield
+
+
+N_LAYERS = 8 # 3 dense + 5 MoE -> full layers {0,1,2,6}
+L = 32 # max_seq_len
+NPES = 8
+KV, PE, KI = 512, 64, 128
+
+
+class FakeShowHands:
+ def __init__(self, num_mtp: int, pure_tp8: bool, fp8_ki: bool, stream: list[int]):
+ self.args = types.SimpleNamespace(max_seq_len=L, num_devices=NPES)
+ self.n_layers = N_LAYERS
+ self.num_mtp = num_mtp
+ self.npes = NPES
+ self.calls: list[tuple] = []
+ self._stream = list(stream) # tokens the "device" will emit
+ self._ar: list[int] = []
+ n_extra = 1 if num_mtp > 0 else 0
+ n_full = len(_full_layer_ordinals(N_LAYERS))
+ self._caches = []
+ for rank in range(NPES):
+ c = []
+ if rank == 0:
+ if pure_tp8:
+ for _ in range(N_LAYERS):
+ c.append(torch.zeros(1, L, KV, dtype=torch.bfloat16))
+ c.append(torch.zeros(1, L, PE, dtype=torch.bfloat16))
+ for _ in range(n_full + n_extra):
+ if fp8_ki:
+ c.append(torch.zeros(L * (KI + 4), dtype=torch.uint8))
+ else:
+ c.append(torch.zeros(1, L, KI, dtype=torch.bfloat16))
+ else:
+ for _ in range(N_LAYERS + n_extra):
+ c.append(torch.zeros(1, L, KV, dtype=torch.bfloat16))
+ c.append(torch.zeros(1, L, PE, dtype=torch.bfloat16))
+ self._caches.append(c)
+
+ # -- e2e API --
+ def reset_sequence(self):
+ self.calls.append(("reset",))
+ self._ar = []
+
+ def set_cur_pos(self, p):
+ self.calls.append(("set_cur_pos", p))
+
+ def update_sampling(self, use_topp, temperature, top_p):
+ self.calls.append(("sampling", use_topp, temperature, top_p))
+
+ def seed_draft(self, tok, draft):
+ self.calls.append(("seed_draft", tok, draft))
+
+ def _take(self, n):
+ out, self._stream = self._stream[:n], self._stream[n:]
+ self._ar.extend(out)
+ return out
+
+ def mtp_n(self, k):
+ self.calls.append(("mtp_n", k))
+ # every verify step accepts 2 tokens (or what is left)
+ return len(self._take(2 * k))
+
+ def step(self, tok):
+ self.calls.append(("step", tok))
+ return self._take(1)[0]
+
+ def decode_n(self, n):
+ self.calls.append(("decode_n", n))
+ self._take(n)
+
+ @property
+ def accepted_count(self):
+ return len(self._ar)
+
+ def accepted_tokens(self, start=0, end=None):
+ end = len(self._ar) if end is None else end
+ return list(self._ar[start:end])
+
+
+def _gen(dl, stop=(999,)):
+ return types.SimpleNamespace(decode_layer=dl, stop_token_ids=set(stop))
+
+
+def _req(seq, n_layers_sent, first=7):
+ layers = []
+ for lid in range(n_layers_sent):
+ ki = torch.full((seq, KI), float(lid) + 0.5, dtype=torch.bfloat16)
+ kv = torch.full((seq, KV), float(lid), dtype=torch.bfloat16)
+ pe = torch.full((seq, PE), -float(lid), dtype=torch.bfloat16)
+ layers.append((ki, kv, pe))
+ return types.SimpleNamespace(
+ seq_len=seq, layers=layers, first_token_id=first, last_prompt_token=3
+ )
+
+
+@pytest.fixture(autouse=True)
+def _no_gpu_sync(monkeypatch):
+ monkeypatch.setattr(torch.cuda, "synchronize", lambda *a, **k: None)
+
+
+@pytest.mark.parametrize("pure_tp8", [True, False])
+def test_inject_slot_mapping_bf16_ki(pure_tp8):
+ dl = FakeShowHands(num_mtp=3, pure_tp8=pure_tp8, fp8_ki=False, stream=[])
+ ad = RocmGlm52EngineAdapter(_gen(dl), with_mtp=True, pure_tp8=pure_tp8, fp8_ki=False)
+ seq = 5
+ ad.inject(_req(seq, N_LAYERS + 1))
+ assert ("reset",) in dl.calls and ("set_cur_pos", seq) == dl.calls[-1]
+ # ranks 1..7: pair 2*lid / 2*lid+1 holds layer lid, incl. the MTP block (lid=N_LAYERS)
+ for rank in range(1, NPES):
+ for lid in range(N_LAYERS + 1):
+ assert dl._caches[rank][2 * lid][0, :seq].float().unique().tolist() == [float(lid)]
+ assert dl._caches[rank][2 * lid + 1][0, :seq].float().unique().tolist() == [-float(lid)]
+ assert dl._caches[rank][2 * lid][0, seq:].abs().sum() == 0
+ # rank 0: ki slots follow the full-layer ordinals, MTP block last
+ ki_base = 2 * N_LAYERS if pure_tp8 else 0
+ full = _full_layer_ordinals(N_LAYERS)
+ for slot, lid in enumerate(full + [N_LAYERS]):
+ got = dl._caches[0][ki_base + slot][0, :seq].float().unique().tolist()
+ assert got == [float(lid) + 0.5], (slot, lid, got)
+ if pure_tp8:
+ # rank 0 also receives kv/pe for the main layers (its shared-layer MLA chain)
+ assert dl._caches[0][2 * 6][0, :seq].float().unique().tolist() == [6.0]
+ assert len(dl._caches[0]) == 2 * N_LAYERS + len(full) + 1
+
+
+def test_inject_fp8_ki_plane_roundtrips():
+ dl = FakeShowHands(num_mtp=0, pure_tp8=True, fp8_ki=True, stream=[])
+ ad = RocmGlm52EngineAdapter(_gen(dl), with_mtp=False, pure_tp8=True, fp8_ki=True)
+ seq = 4
+ req = _req(seq, N_LAYERS + 1)
+ # a non-constant row so the per-token scale is exercised
+ ki0 = torch.arange(seq * KI, dtype=torch.float32).reshape(seq, KI) / 37.0 - 3.0
+ req.layers[0] = (ki0.to(torch.bfloat16), req.layers[0][1], req.layers[0][2])
+ ad.inject(req)
+ plane = dl._caches[0][2 * N_LAYERS] # ki slot of layer 0
+ q = plane[: L * KI].view(torch.float8_e4m3fn).view(L, KI)[:seq].float()
+ s = plane[L * KI : L * KI + L * 4].view(torch.float32)[:seq]
+ deq = q * s.unsqueeze(-1)
+ assert torch.allclose(deq, ki0, rtol=0.13, atol=0.05) # e4m3 has 3 mantissa bits
+ assert plane[L * KI + seq * 4 : L * KI + L * 4].view(torch.float32).abs().sum() == 0
+
+
+def test_inject_without_mtp_drops_the_tail_layer():
+ dl = FakeShowHands(num_mtp=0, pure_tp8=True, fp8_ki=False, stream=[])
+ ad = RocmGlm52EngineAdapter(_gen(dl), with_mtp=False, pure_tp8=True, fp8_ki=False)
+ ad.inject(_req(3, N_LAYERS + 1)) # profile always ships the tail
+ assert len(dl._caches[1]) == 2 * N_LAYERS # no slot for it, silently dropped
+
+
+def test_inject_rejects_missing_mtp_layer_and_bad_seq():
+ dl = FakeShowHands(num_mtp=3, pure_tp8=True, fp8_ki=False, stream=[])
+ ad = RocmGlm52EngineAdapter(_gen(dl), with_mtp=True, pure_tp8=True, fp8_ki=False)
+ with pytest.raises(RuntimeError, match="speculative-config"):
+ ad.inject(_req(3, N_LAYERS))
+ with pytest.raises(RuntimeError, match="seq_len"):
+ ad.inject(_req(L + 1, N_LAYERS + 1))
+
+
+def test_decode_mtp_stream_stop_and_budget():
+ stream = [11, 12, 13, 14, 999, 15]
+ dl = FakeShowHands(num_mtp=3, pure_tp8=True, fp8_ki=False, stream=stream)
+ ad = RocmGlm52EngineAdapter(_gen(dl), with_mtp=True, ar_steps=1, pure_tp8=True, fp8_ki=False)
+ ad.inject(_req(4, N_LAYERS + 1))
+ seen = []
+ out = ad.decode(7, 100, {"temperature": 0.0}, on_token=seen.append)
+ assert out == [7, 11, 12, 13, 14] and seen == out
+ assert ad.last_stats["finish_reason"] == "stop"
+ assert ("seed_draft", 7, 7) in dl.calls
+ assert ("sampling", False, 1.0, 1.0) in dl.calls # greedy arm
+ # budget cut
+ dl2 = FakeShowHands(num_mtp=3, pure_tp8=True, fp8_ki=False, stream=list(range(100, 140)))
+ ad2 = RocmGlm52EngineAdapter(_gen(dl2), with_mtp=True, ar_steps=2, pure_tp8=True, fp8_ki=False)
+ ad2.inject(_req(4, N_LAYERS + 1))
+ out2 = ad2.decode(7, 5, {"temperature": 0.7, "top_p": 0.9})
+ assert out2 == [7, 100, 101, 102, 103] and ad2.last_stats["finish_reason"] == "length"
+ assert ("sampling", True, 0.7, 0.9) in dl2.calls
+
+
+def test_decode_mtp_respects_cache_edge():
+ # seq_len close to max_seq_len: no room for a verify chunk -> length, no mtp_n call
+ dl = FakeShowHands(num_mtp=3, pure_tp8=True, fp8_ki=False, stream=list(range(100, 120)))
+ ad = RocmGlm52EngineAdapter(_gen(dl), with_mtp=True, pure_tp8=True, fp8_ki=False)
+ ad.inject(_req(L - 4, N_LAYERS + 1)) # room = 4, needs mtp_seq(4)+slack(2)
+ out = ad.decode(7, 50, {})
+ assert out == [7] and ad.last_stats["finish_reason"] == "length"
+ assert not any(c[0] == "mtp_n" for c in dl.calls)
+
+
+def test_decode_plain_uses_step_then_decode_n_and_cancel():
+ stream = [21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31]
+ dl = FakeShowHands(num_mtp=0, pure_tp8=True, fp8_ki=False, stream=stream)
+ ad = RocmGlm52EngineAdapter(_gen(dl), with_mtp=False, pure_tp8=True, fp8_ki=False)
+ ad.inject(_req(2, N_LAYERS + 1))
+ out = ad.decode(7, 6, {"temperature": 1.0})
+ assert out == [7, 21, 22, 23, 24, 25]
+ kinds = [c[0] for c in dl.calls if c[0] in ("step", "decode_n")]
+ assert kinds[0] == "step" and "decode_n" in kinds
+ ev = threading.Event()
+ ev.set()
+ dl3 = FakeShowHands(num_mtp=0, pure_tp8=True, fp8_ki=False, stream=list(range(50, 90)))
+ ad3 = RocmGlm52EngineAdapter(_gen(dl3), with_mtp=False, pure_tp8=True, fp8_ki=False)
+ ad3.inject(_req(2, N_LAYERS + 1))
+ out3 = ad3.decode(7, 30, {}, cancel_event=ev)
+ assert out3 == [7, 50] and ad3.last_stats["finish_reason"] == "cancelled"
+
+
+def test_first_token_stop_ignore_eos_and_unsupported_features():
+ dl = FakeShowHands(num_mtp=0, pure_tp8=True, fp8_ki=False, stream=[999, 5])
+ ad = RocmGlm52EngineAdapter(_gen(dl), with_mtp=False, pure_tp8=True, fp8_ki=False)
+ ad.inject(_req(2, N_LAYERS + 1))
+ assert ad.decode(999, 10, {}) == [] and ad.last_stats["finish_reason"] == "stop"
+ dl2 = FakeShowHands(num_mtp=0, pure_tp8=True, fp8_ki=False, stream=[999, 5])
+ ad2 = RocmGlm52EngineAdapter(_gen(dl2), with_mtp=False, pure_tp8=True, fp8_ki=False)
+ ad2.inject(_req(2, N_LAYERS + 1))
+ assert ad2.decode(999, 3, {"ignore_eos": True}) == [999, 999, 5]
+ assert ad.prepare_grammar(None) is None
+ with pytest.raises(GrammarBackendUnavailable):
+ ad.prepare_grammar({"type": "regex", "value": "a+"})
+ with pytest.raises(NotImplementedError):
+ ad.decode(7, 3, {"repetition_penalty": 1.2})
+ assert not ad.supports_logprobs() and not ad.supports_penalties() and ad.supports_ignore_eos()
diff --git a/tests/pd_vllm/test_grammar_adapter.py b/tests/pd_vllm/test_grammar_adapter.py
new file mode 100644
index 0000000..4058773
--- /dev/null
+++ b/tests/pd_vllm/test_grammar_adapter.py
@@ -0,0 +1,94 @@
+"""Adapter tests for MlaNsaEngineAdapter.prepare_grammar (Stage 2/3).
+
+The full masked decode loops (_decode_standard / _decode_mtp) need the real GPU
+decode layer + a GLM tokenizer, so they are validated on-cluster (ladder rung
+2/3). Here we cover what IS reachable off-GPU:
+ * no spec -> None (unconstrained; no backend import attempted)
+ * prepare_grammar builds a GrammarSession with the right num_positions
+ (mtp_seq_len for MTP, 1 for AR) and the per-request think_end_id.
+
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_grammar_adapter.py -v
+"""
+
+import types
+
+from tilert.pd_vllm.profiles.mla_nsa import MlaNsaEngineAdapter
+
+
+def _fake_generator():
+ return types.SimpleNamespace(
+ mtp_seq_len=4,
+ decode_layer=types.SimpleNamespace(max_seq_len=4096),
+ stop_token_ids={2},
+ config=types.SimpleNamespace(vocab_size=288),
+ tokenizer=types.SimpleNamespace(convert_tokens_to_ids=lambda s: 257), # id
+ _grammar_engine=None,
+ )
+
+
+def test_prepare_none_returns_none():
+ for with_mtp in (True, False):
+ adapter = MlaNsaEngineAdapter(_fake_generator(), with_mtp=with_mtp)
+ assert adapter.prepare_grammar(None) is None
+
+
+def _install_fake_grammar(monkeypatch):
+ """Stub the backend lookup with recording doubles, so session construction
+ can be checked without a real tokenizer/GPU.
+
+ Patch ``load_grammar_backend`` — the indirection — rather than a module
+ path in ``sys.modules``. The stubs used to be installed at
+ ``tilert.models.glm_5_2.grammar``; the engine moved the wrapper to
+ ``tilert.grammar`` and left that path as a fallback the loader only reaches
+ when the new one is absent. With a real engine installed the new path
+ resolves, the stub was never consulted, and these tests ran the REAL
+ GrammarEngine against a SimpleNamespace tokenizer — xgrammar rejected it
+ and the tests failed. Patching the lookup keeps them independent of where
+ the engine happens to keep the wrapper, which is the whole point of
+ ``grammar_backend``.
+
+ The name is bound at import time in the profile module, so the patch has to
+ land on that module's attribute, not on ``grammar_backend``'s.
+ """
+ calls: dict = {}
+
+ class FakeEngine:
+ def __init__(self, tok, padded_vocab_size, stop_token_ids):
+ calls["engine"] = (padded_vocab_size, tuple(stop_token_ids))
+
+ class FakeSession:
+ def __init__(self, engine, spec, num_positions, think_end_id):
+ calls["session"] = {
+ "num_positions": num_positions,
+ "think_end_id": think_end_id,
+ "spec": spec,
+ }
+
+ monkeypatch.setattr(
+ "tilert.pd_vllm.profiles.mla_nsa.load_grammar_backend", lambda: (FakeEngine, FakeSession)
+ )
+ return calls
+
+
+def test_prepare_mtp_uses_mtp_seq_len_positions(monkeypatch):
+ calls = _install_fake_grammar(monkeypatch)
+ adapter = MlaNsaEngineAdapter(_fake_generator(), with_mtp=True)
+ adapter.prepare_grammar({"type": "regex", "value": "[0-9]"}, enable_thinking=True)
+ assert calls["session"]["num_positions"] == 4 # == mtp_seq_len
+ assert calls["session"]["think_end_id"] == 257 # resolved
+ assert calls["engine"][0] == 288 # padded vocab size
+
+
+def test_prepare_ar_uses_single_position(monkeypatch):
+ calls = _install_fake_grammar(monkeypatch)
+ adapter = MlaNsaEngineAdapter(_fake_generator(), with_mtp=False)
+ adapter.prepare_grammar({"type": "regex", "value": "[0-9]"}, enable_thinking=True)
+ assert calls["session"]["num_positions"] == 1 # non-MTP AR
+
+
+def test_prepare_no_thinking_gate_when_disabled(monkeypatch):
+ calls = _install_fake_grammar(monkeypatch)
+ adapter = MlaNsaEngineAdapter(_fake_generator(), with_mtp=True)
+ adapter.prepare_grammar({"type": "regex", "value": "[0-9]"}, enable_thinking=False)
+ assert calls["session"]["think_end_id"] is None
diff --git a/tests/pd_vllm/test_grammar_backend_path.py b/tests/pd_vllm/test_grammar_backend_path.py
new file mode 100644
index 0000000..17b6d75
--- /dev/null
+++ b/tests/pd_vllm/test_grammar_backend_path.py
@@ -0,0 +1,77 @@
+"""Which module the xgrammar host wrapper is loaded from.
+
+The wrapper moved from ``tilert.models.glm_5_2.grammar`` to ``tilert.grammar``
+in the engine. One serve version must work against engine wheels from
+either side of that move, so both paths are exercised here with stubbed
+modules — no engine install required.
+"""
+
+from __future__ import annotations
+
+import sys
+import types
+
+import pytest
+
+from tilert.pd_vllm.grammar_backend import load_grammar_backend
+
+_NEW = "tilert.grammar"
+_OLD = "tilert.models.glm_5_2.grammar"
+_PARENTS = ("tilert", "tilert.models", "tilert.models.glm_5_2")
+
+
+def _stub(monkeypatch, path: str, tag: str) -> None:
+ """Install a fake wrapper module at ``path`` whose classes carry ``tag``."""
+ for parent in _PARENTS:
+ monkeypatch.setitem(
+ sys.modules, parent, sys.modules.get(parent) or types.ModuleType(parent)
+ )
+ mod = types.ModuleType(path)
+ for cls_name in ("GrammarEngine", "GrammarSession"):
+ setattr(mod, cls_name, type(cls_name, (), {"came_from": tag}))
+ monkeypatch.setitem(sys.modules, path, mod)
+
+
+def _absent(monkeypatch, path: str) -> None:
+ """Make importing ``path`` raise, the way a wheel that omits it would."""
+
+ class _Blocker:
+ def find_module(self, name, path=None): # pragma: no cover - legacy hook
+ return None
+
+ def find_spec(self, name, path=None, target=None):
+ if name == globals()["_blocked"]:
+ raise ModuleNotFoundError(f"No module named {name!r}")
+ return None # noqa: R501 (finder protocol: None means "not mine")
+
+ globals()["_blocked"] = path
+ monkeypatch.delitem(sys.modules, path, raising=False)
+ monkeypatch.setattr(sys, "meta_path", [_Blocker()] + sys.meta_path)
+
+
+def test_the_new_path_is_preferred(monkeypatch):
+ _stub(monkeypatch, _NEW, "new")
+ _stub(monkeypatch, _OLD, "old")
+ engine, session = load_grammar_backend()
+ assert engine.came_from == "new"
+ assert session.came_from == "new"
+
+
+def test_an_older_engine_wheel_still_works(monkeypatch):
+ _absent(monkeypatch, _NEW)
+ _stub(monkeypatch, _OLD, "old")
+ engine, session = load_grammar_backend()
+ assert engine.came_from == "old"
+ assert session.came_from == "old"
+
+
+def test_with_neither_path_the_error_names_no_model(monkeypatch):
+ _absent(monkeypatch, _NEW)
+ monkeypatch.delitem(sys.modules, _OLD, raising=False)
+ for parent in _PARENTS:
+ monkeypatch.delitem(sys.modules, parent, raising=False)
+ with pytest.raises(ModuleNotFoundError) as e:
+ load_grammar_backend()
+ msg = str(e.value)
+ assert "xgrammar host backend" in msg
+ assert "glm" not in msg.lower(), "the client-visible message must not name another model"
diff --git a/tests/pd_vllm/test_grammar_plumbing.py b/tests/pd_vllm/test_grammar_plumbing.py
new file mode 100644
index 0000000..e576162
--- /dev/null
+++ b/tests/pd_vllm/test_grammar_plumbing.py
@@ -0,0 +1,223 @@
+"""Stage-1 plumbing tests: grammar_spec end-to-end through the HTTP seams with
+a StubEngine — no GPU / no tilert / no real vLLM.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_grammar_plumbing.py -v
+
+Covers:
+ * decode_server: compile-before-inject classification (400/500) returns
+ BEFORE the wire-wait; grammar_session threads into decode; runtime
+ violation -> 400.
+ * pd_router: bad spec -> 400 before any network; grammar_spec forwarded to
+ /pd/decode; a decode grammar error propagates its status (not masked 502).
+"""
+
+import queue
+import types
+
+from fastapi.testclient import TestClient
+
+from tilert.pd_vllm import pd_router
+from tilert.pd_vllm.decode_server import build_app
+from tilert.pd_vllm.engine_iface import StubEngine
+
+
+# --------------------------------------------------------------------------- #
+# decode_server via TestClient + a fake ReceiveServer
+# --------------------------------------------------------------------------- #
+class _FakeReq:
+ def __init__(self, rid):
+ self.rid = rid
+ self.seq_len = 8
+ self.last_prompt_token = 5
+
+
+class _FakeServer:
+ """Minimal ReceiveServer stand-in: hands back one matching req and
+ converts it to a sentinel the StubEngine happily injects.
+ """
+
+ def __init__(self, rid):
+ self.completed: queue.Queue = queue.Queue()
+ self.completed.put(_FakeReq(rid))
+ self.profile = types.SimpleNamespace(convert=lambda *a, **k: "converted", num_ranks=8)
+ self.buffer = None
+ self.base_ptr = 0
+ self.max_seq_len = 4096
+
+ def expect(self, rid=None):
+ # /pd/decode announces its rid so the real ReceiveServer can drop a
+ # tombstone left by a previous attempt at the same request. Recorded,
+ # so a test can assert the announcement happened.
+ self.expected_rids = getattr(self, "expected_rids", [])
+ self.expected_rids.append(rid)
+
+ def release(self, rid=None):
+ pass
+
+
+def _client(rid="rid-1"):
+ return TestClient(build_app(_FakeServer(rid), StubEngine()))
+
+
+def test_decode_invalid_grammar_400_before_inject():
+ r = _client().post(
+ "/pd/decode",
+ json={
+ "rid": "rid-1",
+ "first_token_id": 7,
+ "max_tokens": 8,
+ "grammar_spec": {"type": "not_a_real_type"},
+ },
+ )
+ assert r.status_code == 400
+ assert r.json()["error_type"] == "invalid_grammar"
+
+
+def test_decode_backend_missing_500():
+ r = _client().post(
+ "/pd/decode",
+ json={
+ "rid": "rid-1",
+ "first_token_id": 7,
+ "max_tokens": 8,
+ "grammar_spec": {"type": "__backend_missing__"},
+ },
+ )
+ assert r.status_code == 500
+ assert r.json()["error_type"] == "grammar_backend_unavailable"
+
+
+def test_decode_valid_grammar_threads_and_returns_200():
+ r = _client().post(
+ "/pd/decode",
+ json={
+ "rid": "rid-1",
+ "first_token_id": 7,
+ "max_tokens": 8,
+ "grammar_spec": {"type": "regex", "value": "[0-9]"},
+ "enable_thinking": False,
+ },
+ )
+ assert r.status_code == 200
+ assert r.json()["token_ids"][0] == 7
+
+
+def test_decode_runtime_violation_400():
+ r = _client().post(
+ "/pd/decode",
+ json={
+ "rid": "rid-1",
+ "first_token_id": 7,
+ "max_tokens": 8,
+ "grammar_spec": {"type": "regex", "value": "__violate__"},
+ },
+ )
+ assert r.status_code == 400
+ assert r.json()["error_type"] == "grammar_violation"
+
+
+def test_decode_unconstrained_still_works():
+ r = _client().post("/pd/decode", json={"rid": "rid-1", "first_token_id": 7, "max_tokens": 8})
+ assert r.status_code == 200
+ assert r.json()["token_ids"][0] == 7
+
+
+# --------------------------------------------------------------------------- #
+# pd_router: extraction + forwarding + status propagation (network mocked)
+# --------------------------------------------------------------------------- #
+class _Resp:
+ def __init__(self, payload, status=200):
+ self._payload = payload
+ self.status_code = status
+
+ def json(self):
+ return self._payload
+
+ def raise_for_status(self):
+ if self.status_code >= 400:
+ raise RuntimeError(f"HTTP {self.status_code}")
+
+
+def _router_client():
+ pool = pd_router.Pool([pd_router.DecodeNode("127.0.0.1", 5556, 5557)])
+ ctx = pd_router.RouterCtx("http://vllm.invalid", pool, tokenizer=None, parser_name="none")
+ return TestClient(pd_router.build_app(ctx))
+
+
+def test_router_bad_spec_400_no_network(monkeypatch):
+ # If extraction fails, no prefill/decode POST should ever be attempted.
+ def _boom(*a, **k):
+ raise AssertionError("network must not be touched on bad spec")
+
+ monkeypatch.setattr(pd_router.requests, "post", _boom)
+ r = _router_client().post(
+ "/v1/chat/completions",
+ json={
+ "messages": [{"role": "user", "content": "hi"}],
+ "response_format": {"type": "json_schema", "json_schema": {}},
+ },
+ )
+ assert r.status_code == 400
+ assert r.json()["error_type"] == "invalid_grammar"
+
+
+def test_router_forwards_grammar_spec_and_propagates_violation(monkeypatch):
+ captured = {}
+
+ def fake_post(url, json=None, timeout=None, **kw):
+ if url.endswith("/pd/decode"):
+ captured["decode_body"] = json
+ # simulate decode-side fail-closed grammar violation
+ return _Resp(
+ {"error": "first token violates the grammar", "error_type": "grammar_violation"},
+ status=400,
+ )
+ # vLLM prefill response
+ return _Resp(
+ {
+ "id": "cmpl-abc",
+ "choices": [{"logprobs": {"content": [{"token": "token_id:7"}]}}],
+ "usage": {"prompt_tokens": 3},
+ "model": "glm5p2-tilert",
+ }
+ )
+
+ monkeypatch.setattr(pd_router.requests, "post", fake_post)
+ r = _router_client().post(
+ "/v1/chat/completions",
+ json={"messages": [{"role": "user", "content": "hi"}], "regex": r"[0-9]{3}"},
+ )
+ # grammar_spec forwarded to the decode node
+ assert captured["decode_body"]["grammar_spec"] == {"type": "regex", "value": r"[0-9]{3}"}
+ assert "enable_thinking" in captured["decode_body"]
+ # decode's 400 grammar_violation propagates (NOT masked as 502)
+ assert r.status_code == 400
+ assert r.json()["error_type"] == "grammar_violation"
+
+
+def test_router_omits_grammar_spec_for_plain_request(monkeypatch):
+ captured = {}
+
+ def fake_post(url, json=None, timeout=None, **kw):
+ if url.endswith("/pd/decode"):
+ captured["decode_body"] = json
+ return _Resp(
+ {"rid": "x", "token_ids": [7], "seq_len": 8, "timing_ms": {"finish_reason": "stop"}}
+ )
+ return _Resp(
+ {
+ "id": "cmpl-abc",
+ "choices": [{"logprobs": {"content": [{"token": "token_id:7"}]}}],
+ "usage": {"prompt_tokens": 3},
+ "model": "m",
+ }
+ )
+
+ monkeypatch.setattr(pd_router.requests, "post", fake_post)
+ r = _router_client().post(
+ "/v1/chat/completions", json={"messages": [{"role": "user", "content": "hi"}]}
+ )
+ assert r.status_code == 200
+ assert "grammar_spec" not in captured["decode_body"]
diff --git a/tests/pd_vllm/test_grammar_spec.py b/tests/pd_vllm/test_grammar_spec.py
new file mode 100644
index 0000000..7705a34
--- /dev/null
+++ b/tests/pd_vllm/test_grammar_spec.py
@@ -0,0 +1,215 @@
+"""Stage-1 unit tests: request -> grammar_spec translation + error hierarchy +
+StubEngine classification. Pure CPU, no GPU / no vLLM / no tilert.
+
+ CUDA_VISIBLE_DEVICES= python -m pytest tests/pd_vllm/test_grammar_spec.py -v
+"""
+
+import pytest
+
+from tilert.pd_vllm.engine_iface import StubEngine
+from tilert.pd_vllm.grammar_spec import (
+ GrammarBackendUnavailable,
+ GrammarViolationError,
+ InvalidGrammarError,
+ extract_request_grammar_spec,
+)
+
+
+# --------------------------- spec extraction ------------------------------- #
+def test_regex_extension_field():
+ assert extract_request_grammar_spec({"regex": r"[0-9]{3}"}) == {
+ "type": "regex",
+ "value": r"[0-9]{3}",
+ }
+
+
+def test_ebnf_extension_field():
+ spec = extract_request_grammar_spec({"ebnf": 'root ::= "hi"'})
+ assert spec == {"type": "ebnf", "value": 'root ::= "hi"'}
+
+
+def test_json_object():
+ spec = extract_request_grammar_spec({"response_format": {"type": "json_object"}})
+ assert spec == {"type": "json_object", "value": None}
+
+
+def test_json_schema():
+ schema = {"type": "object", "properties": {"n": {"type": "integer"}}}
+ spec = extract_request_grammar_spec(
+ {"response_format": {"type": "json_schema", "json_schema": {"schema": schema}}}
+ )
+ assert spec == {"type": "json_schema", "value": schema}
+
+
+def test_structural_tag():
+ rf = {"type": "structural_tag", "structures": [], "triggers": []}
+ spec = extract_request_grammar_spec({"response_format": rf})
+ assert spec == {"type": "structural_tag", "value": rf}
+
+
+def test_priority_json_schema_over_regex_and_ebnf():
+ req = {
+ "regex": r"[0-9]+",
+ "ebnf": 'root ::= "x"',
+ "response_format": {"type": "json_schema", "json_schema": {"schema": {"type": "object"}}},
+ }
+ assert extract_request_grammar_spec(req)["type"] == "json_schema"
+
+
+def test_priority_regex_over_ebnf():
+ assert (
+ extract_request_grammar_spec({"regex": r"[0-9]+", "ebnf": 'root ::= "x"'})["type"]
+ == "regex"
+ )
+
+
+def test_json_object_returns_before_regex():
+ # json_object short-circuits even when regex is also present.
+ req = {"regex": r"[0-9]+", "response_format": {"type": "json_object"}}
+ assert extract_request_grammar_spec(req)["type"] == "json_object"
+
+
+def test_none_when_unconstrained():
+ assert extract_request_grammar_spec({"messages": [], "temperature": 0.7}) is None
+
+
+def test_json_schema_missing_schema_raises_invalid():
+ with pytest.raises(InvalidGrammarError):
+ extract_request_grammar_spec(
+ {"response_format": {"type": "json_schema", "json_schema": {}}}
+ )
+
+
+def test_response_format_not_object_raises_invalid():
+ with pytest.raises(InvalidGrammarError):
+ extract_request_grammar_spec({"response_format": "json"})
+
+
+# --------------------------- error hierarchy ------------------------------- #
+def test_error_status_and_payload():
+ assert InvalidGrammarError("x").http_status == 400
+ assert GrammarBackendUnavailable("x").http_status == 500
+ assert GrammarViolationError("x").http_status == 400
+ p = InvalidGrammarError("bad schema").to_payload()
+ assert p == {"error": "bad schema", "error_type": "invalid_grammar"}
+
+
+# --------------------------- StubEngine classification --------------------- #
+def test_stub_prepare_none():
+ assert StubEngine().prepare_grammar(None) is None
+
+
+def test_stub_prepare_valid_returns_session():
+ sess = StubEngine().prepare_grammar({"type": "regex", "value": "[0-9]"}, enable_thinking=False)
+ assert sess["spec"]["type"] == "regex" and sess["enable_thinking"] is False
+
+
+def test_stub_prepare_unknown_type_invalid():
+ with pytest.raises(InvalidGrammarError):
+ StubEngine().prepare_grammar({"type": "nope"})
+
+
+def test_stub_prepare_not_a_dict_invalid():
+ with pytest.raises(InvalidGrammarError):
+ StubEngine().prepare_grammar(["not", "a", "dict"])
+
+
+def test_stub_prepare_backend_missing_500():
+ with pytest.raises(GrammarBackendUnavailable):
+ StubEngine().prepare_grammar({"type": "__backend_missing__"})
+
+
+def test_stub_decode_violation_raises():
+ eng = StubEngine()
+ sess = eng.prepare_grammar({"type": "regex", "value": "__violate__"})
+ with pytest.raises(GrammarViolationError):
+ eng.decode(first_token_id=7, max_tokens=8, sampling=None, grammar_session=sess)
+
+
+def test_stub_decode_threads_session_ok():
+ eng = StubEngine()
+ sess = eng.prepare_grammar({"type": "regex", "value": "[0-9]"})
+ out = eng.decode(first_token_id=7, max_tokens=8, sampling=None, grammar_session=sess)
+ assert out[0] == 7 and eng.last_stats["finish_reason"] == "stop"
+
+
+# --------------------- spec validation: compile-cost caps ------------------ #
+# Compilation runs inside the decode node's single-slot lock, so an unbounded
+# schema is downtime for everyone, not just a slow request.
+def _rf(schema):
+ return {
+ "response_format": {"type": "json_schema", "json_schema": {"name": "t", "schema": schema}}
+ }
+
+
+def test_deep_nesting_rejected():
+ schema = {"type": "string"}
+ for _ in range(200):
+ schema = {"type": "object", "properties": {"x": schema}}
+ with pytest.raises(InvalidGrammarError, match="deeper than"):
+ extract_request_grammar_spec(_rf(schema))
+
+
+def test_many_properties_rejected():
+ schema = {"type": "object", "properties": {f"p{i}": {"type": "string"} for i in range(5000)}}
+ with pytest.raises(InvalidGrammarError, match="subschemas"):
+ extract_request_grammar_spec(_rf(schema))
+
+
+def test_huge_enum_rejected():
+ schema = {"type": "object", "properties": {"v": {"enum": [f"o{i}" for i in range(50000)]}}}
+ with pytest.raises(InvalidGrammarError, match="enum"):
+ extract_request_grammar_spec(_rf(schema))
+
+
+def test_large_enum_does_not_count_as_nodes():
+ """Enum entries are literal data, not subschemas."""
+ schema = {"type": "object", "properties": {"v": {"enum": [f"o{i}" for i in range(900)]}}}
+ assert extract_request_grammar_spec(_rf(schema))["type"] == "json_schema"
+
+
+def test_long_regex_allowed():
+ """regex/EBNF are not capped: 200 000 chars compiles in 0.7s, so a length
+ limit would only over-block.
+ """
+ spec = extract_request_grammar_spec({"regex": "a" * 50000})
+ assert spec == {"type": "regex", "value": "a" * 50000}
+
+
+def test_structural_tag_schemas_are_walked():
+ """The walk follows spec['value'] whatever its shape, so a tag's schemas are
+ covered by the same caps without a per-type branch.
+ """
+ deep = {"type": "string"}
+ for _ in range(200):
+ deep = {"type": "object", "properties": {"x": deep}}
+ with pytest.raises(InvalidGrammarError, match="deeper than"):
+ extract_request_grammar_spec(
+ {
+ "response_format": {
+ "type": "structural_tag",
+ "structures": [{"begin": "", "schema": deep, "end": ""}],
+ }
+ }
+ )
+
+
+def test_ordinary_schema_unaffected():
+ schema = {
+ "type": "object",
+ "properties": {
+ "city": {"type": "string"},
+ "temp_c": {"type": "integer", "minimum": -90, "maximum": 60},
+ "forecast": {
+ "type": "array",
+ "items": {
+ "type": "object",
+ "properties": {"day": {"type": "string"}, "high": {"type": "integer"}},
+ "required": ["day", "high"],
+ },
+ },
+ },
+ "required": ["city", "temp_c"],
+ "additionalProperties": False,
+ }
+ assert extract_request_grammar_spec(_rf(schema))["value"] == schema
diff --git a/tests/pd_vllm/test_ignore_eos.py b/tests/pd_vllm/test_ignore_eos.py
new file mode 100644
index 0000000..527ec5c
--- /dev/null
+++ b/tests/pd_vllm/test_ignore_eos.py
@@ -0,0 +1,97 @@
+"""`ignore_eos` must empty the MLA/NSA adapter's stop set for that request.
+
+The router half is pinned in tests/pd_vllm/test_stream_e2e.py (the flag reaches
+the decode node's `sampling` dict). This file covers the far end: the adapter
+reading it, and both decode loops binding the stop set through it.
+
+Wiring only one of the two loops is the failure this guards. It is silent --
+nothing errors, the request just stops at the first EOS, which is precisely
+what a fixed-output-length benchmark (`vllm bench serve --ignore-eos`) asked it
+not to do, so measured decode throughput becomes whatever the model happened to
+emit.
+
+Scope note: the MLA/NSA adapter honours the flag (tile-ai/TileRT#55); the
+ROCm GLM adapter does too, and is covered in test_glm5_rocm_engine.py.
+
+No GPU, no tilert: the adapter is built with ``object.__new__`` and handed a
+recording stand-in for ``self.gen``, with ``max_tokens=0`` so ``decode()``
+returns right after ``update_sampling_params``.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_ignore_eos.py -v
+"""
+
+from __future__ import annotations
+
+import pathlib
+
+import pytest
+
+from tilert.pd_vllm.profiles.mla_nsa import MlaNsaEngineAdapter
+
+
+class _StubGen:
+ """Enough generator for decode() to reach its budget short-circuit."""
+
+ def update_sampling_params(self, **kw):
+ pass
+
+
+def _adapter() -> MlaNsaEngineAdapter:
+ a = object.__new__(MlaNsaEngineAdapter)
+ a.gen = _StubGen()
+ a.with_mtp = False
+ a.mtp_seq_len = 4
+ a.max_seq_len = 4096
+ a._seq_len = 8
+ a.stop_ids = {7, 8}
+ a._ignore_eos = False
+ a.last_stats = {}
+ return a
+
+
+def _decode_once(sampling: dict) -> MlaNsaEngineAdapter:
+ """max_tokens=0 => budget <= 0 => returns before touching the model."""
+ a = _adapter()
+ assert a.decode(
+ first_token_id=5, max_tokens=0, sampling=sampling, cancel_event=None, grammar_session=None
+ ) == [5]
+ return a
+
+
+def test_the_adapter_records_the_flag() -> None:
+ assert _decode_once({"ignore_eos": True})._ignore_eos is True
+
+
+def test_absent_means_stop_at_eos() -> None:
+ """The default every other request has must not shift."""
+ assert _decode_once({"temperature": 0.0})._ignore_eos is False
+
+
+def test_explicit_false_means_stop_at_eos() -> None:
+ assert _decode_once({"ignore_eos": False})._ignore_eos is False
+
+
+def test_the_flag_does_not_persist_into_the_next_request() -> None:
+ """The adapter is reused across requests on a decode node, so a sticky flag
+ would leak EOS-less decoding into every request that followed.
+ """
+ a = _adapter()
+ a._ignore_eos = True
+ a.decode(first_token_id=5, max_tokens=0, sampling={}, cancel_event=None, grammar_session=None)
+ assert a._ignore_eos is False
+
+
+def test_both_decode_loops_read_the_guarded_stop_set() -> None:
+ """Neither loop may bind the raw stop set.
+
+ The adapter has two (`_decode_mtp` and `_decode_standard`) and the flag only
+ bites where the loop actually breaks. Source-level, in the shape
+ test_top_k_resolution.py uses, because reaching those loops needs a GPU and
+ weights.
+ """
+ src = pathlib.Path(pytest.importorskip("tilert.pd_vllm.profiles.mla_nsa").__file__).read_text()
+ binds = [ln.strip() for ln in src.splitlines() if ln.strip().startswith("stop_ids = ")]
+ assert len(binds) == 2, f"expected one binding per decode loop: {binds}"
+ assert all(b == "stop_ids = set() if self._ignore_eos else self.stop_ids" for b in binds), binds
diff --git a/tests/pd_vllm/test_logprobs.py b/tests/pd_vllm/test_logprobs.py
new file mode 100644
index 0000000..05d45ce
--- /dev/null
+++ b/tests/pd_vllm/test_logprobs.py
@@ -0,0 +1,346 @@
+"""The chat ``logprobs`` contract: what is accepted, and what comes back.
+
+Ranges and interdependencies come from the OpenAI chat-completions reference as
+narrowed to the stricter vendor reading:
+
+* ``logprobs``: boolean, default false.
+* ``top_logprobs``: integer, default 0, ``[0, 5]``, only valid with
+ ``logprobs: true``.
+* logprobs cover ``message.content``; a reasoning segment carries none.
+* ``-9999.0`` is the documented stand-in when a real log probability is absent.
+
+No GPU, no tilert, no vllm: request parsing and response assembly are pure.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest tests/pd_vllm/test_logprobs.py -v
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from tilert.pd_vllm.logprobs import (
+ LOGPROB_UNAVAILABLE,
+ TOP_LOGPROBS_MAX,
+ LogprobsRequest,
+ LogprobsUnsupported,
+ build_logprobs,
+ resolve_logprobs_request,
+)
+
+# A toy vocabulary; 300 decodes to a multi-byte character so `bytes` is
+# exercised on something other than ASCII.
+VOCAB = {1: "Hello", 2: " world", 3: "!", 300: "世", 400: ""}
+
+
+def decode_one(tid: int) -> str:
+ return VOCAB.get(tid, f"<{tid}>")
+
+
+# --------------------------------------------------------------------------- #
+# request parsing
+# --------------------------------------------------------------------------- #
+
+
+def test_absent_means_no_logprobs() -> None:
+ assert resolve_logprobs_request({"model": "m"}) is None
+
+
+def test_false_means_no_logprobs() -> None:
+ assert resolve_logprobs_request({"logprobs": False}) is None
+
+
+def test_true_alone_defaults_top_n_to_zero() -> None:
+ """`logprobs: true` on its own returns the chosen token's logprob only."""
+ assert resolve_logprobs_request({"logprobs": True}) == LogprobsRequest(0)
+
+
+@pytest.mark.parametrize("n", [0, 1, 3, 5])
+def test_accepted_top_logprobs_range(n) -> None:
+ req = resolve_logprobs_request({"logprobs": True, "top_logprobs": n})
+ assert req == LogprobsRequest(n)
+
+
+@pytest.mark.parametrize("n", [6, 20, 256, -1])
+def test_out_of_range_is_rejected_not_clamped(n) -> None:
+ """A client asking for more than we serve must be told, not quietly cut."""
+ with pytest.raises(LogprobsUnsupported) as e:
+ resolve_logprobs_request({"logprobs": True, "top_logprobs": n})
+ assert e.value.http_status == 400
+ assert str(TOP_LOGPROBS_MAX) in str(e.value)
+
+
+def test_top_logprobs_without_logprobs_is_rejected() -> None:
+ with pytest.raises(LogprobsUnsupported) as e:
+ resolve_logprobs_request({"top_logprobs": 3})
+ assert e.value.http_status == 400
+ assert "logprobs must be set to true" in str(e.value)
+
+
+def test_zero_top_logprobs_without_logprobs_is_not_an_error() -> None:
+ """0 is the default, so its presence alone does not request anything."""
+ assert resolve_logprobs_request({"top_logprobs": 0}) is None
+
+
+def test_top_logprobs_true_is_a_type_error_not_one() -> None:
+ """bool is an int subclass; `top_logprobs: true` must not mean 1."""
+ with pytest.raises(LogprobsUnsupported):
+ resolve_logprobs_request({"logprobs": True, "top_logprobs": True})
+
+
+@pytest.mark.parametrize("bad", ["3", 3.5, [3]])
+def test_non_integer_top_logprobs_is_rejected(bad) -> None:
+ with pytest.raises(LogprobsUnsupported):
+ resolve_logprobs_request({"logprobs": True, "top_logprobs": bad})
+
+
+def test_non_boolean_logprobs_is_rejected() -> None:
+ with pytest.raises(LogprobsUnsupported):
+ resolve_logprobs_request({"logprobs": "yes"})
+
+
+def test_error_payload_shape_matches_the_grammar_errors() -> None:
+ """The router returns these the same way it returns grammar errors."""
+ try:
+ resolve_logprobs_request({"logprobs": True, "top_logprobs": 9})
+ except LogprobsUnsupported as e:
+ assert e.to_payload()["error_type"] == "invalid_logprobs"
+ assert "error" in e.to_payload()
+
+
+# --------------------------------------------------------------------------- #
+# response assembly
+# --------------------------------------------------------------------------- #
+
+
+def test_shape_is_content_plus_nullable_refusal() -> None:
+ out = build_logprobs([1], [-0.5], [[(1, -0.5)]], LogprobsRequest(1), decode_one)
+ assert set(out) == {"content", "refusal"}
+ assert out["refusal"] is None
+ assert set(out["content"][0]) == {"token", "logprob", "bytes", "top_logprobs"}
+
+
+def test_one_entry_per_content_token() -> None:
+ out = build_logprobs([1, 2, 3], [-0.1, -0.2, -0.3], None, LogprobsRequest(0), decode_one)
+ assert [c["token"] for c in out["content"]] == ["Hello", " world", "!"]
+ assert [c["logprob"] for c in out["content"]] == [-0.1, -0.2, -0.3]
+
+
+def test_top_n_zero_gives_an_empty_candidate_list() -> None:
+ out = build_logprobs([1], [-0.1], [[(1, -0.1), (2, -2.0)]], LogprobsRequest(0), decode_one)
+ assert out["content"][0]["top_logprobs"] == []
+
+
+def test_candidates_are_truncated_to_top_n() -> None:
+ """A decode node may return more than asked; the response must not."""
+ alts = [(1, -0.1), (2, -2.0), (3, -3.0), (300, -4.0)]
+ out = build_logprobs([1], [-0.1], [alts], LogprobsRequest(2), decode_one)
+ assert len(out["content"][0]["top_logprobs"]) == 2
+ assert [a["token"] for a in out["content"][0]["top_logprobs"]] == ["Hello", " world"]
+
+
+def test_bytes_is_utf8_of_the_token() -> None:
+ out = build_logprobs([300], [-1.0], None, LogprobsRequest(0), decode_one)
+ assert out["content"][0]["bytes"] == list("世".encode())
+ assert len(out["content"][0]["bytes"]) == 3
+
+
+def test_missing_logprob_uses_the_documented_sentinel() -> None:
+ out = build_logprobs([1], [None], None, LogprobsRequest(0), decode_one)
+ assert out["content"][0]["logprob"] == LOGPROB_UNAVAILABLE
+
+
+def test_negative_infinity_becomes_the_sentinel() -> None:
+ """JSON cannot carry -inf; it must not reach the client as Infinity."""
+ out = build_logprobs([1], [float("-inf")], None, LogprobsRequest(0), decode_one)
+ assert out["content"][0]["logprob"] == LOGPROB_UNAVAILABLE
+
+
+def test_sentinel_also_applies_inside_candidates() -> None:
+ out = build_logprobs([1], [-0.1], [[(2, float("-inf"))]], LogprobsRequest(1), decode_one)
+ assert out["content"][0]["top_logprobs"][0]["logprob"] == LOGPROB_UNAVAILABLE
+
+
+def test_response_is_json_serialisable() -> None:
+ """The whole point of the sentinel: no inf/nan reaches json.dumps."""
+ import json
+
+ out = build_logprobs(
+ [1, 300],
+ [float("-inf"), -0.2],
+ [[(1, float("-inf"))], [(300, -0.2)]],
+ LogprobsRequest(1),
+ decode_one,
+ )
+ assert "Infinity" not in json.dumps(out, allow_nan=False)
+
+
+@pytest.mark.parametrize("n_lp,n_top", [(2, 1), (1, 2)])
+def test_length_mismatch_is_caught(n_lp, n_top) -> None:
+ """A decode node returning the wrong count is a bug, not a client error."""
+ with pytest.raises(ValueError):
+ build_logprobs([1], [-0.1] * n_lp, [[(1, -0.1)]] * n_top, LogprobsRequest(1), decode_one)
+
+
+# --------------------------------------------------------------------------- #
+# reasoning segment carries no logprobs
+# --------------------------------------------------------------------------- #
+
+
+# --------------------------------------------------------------------------- #
+# token 1: sourced from the prefill response
+#
+# The decode node echoes first_token_id without sampling it, so it sends null
+# for that position. The distribution that produced the token exists only at the
+# prompt's last position, which the vLLM prefill instance evaluated -- and the
+# router already reads that same entry to recover the token id.
+# --------------------------------------------------------------------------- #
+
+
+def _prefill_resp(logprob=-0.25, cands=((7, -0.25), (8, -1.5), (9, -2.5))):
+ """A vLLM prefill reply shaped as --return-tokens-as-token-ids produces."""
+ return {
+ "choices": [
+ {
+ "logprobs": {
+ "content": [
+ {
+ "token": "token_id:7",
+ "logprob": logprob,
+ "top_logprobs": [
+ {"token": f"token_id:{i}", "logprob": lp} for i, lp in cands
+ ],
+ }
+ ]
+ }
+ }
+ ]
+ }
+
+
+def test_prefill_entry_is_parsed_with_its_candidates() -> None:
+ from tilert.pd_vllm.pd_router import first_token_logprob_from_prefill
+
+ lp, cands = first_token_logprob_from_prefill(_prefill_resp(), top_n=3)
+ assert lp == -0.25
+ assert cands == [(7, -0.25), (8, -1.5), (9, -2.5)]
+
+
+def test_prefill_candidates_are_capped_at_the_requested_count() -> None:
+ from tilert.pd_vllm.pd_router import first_token_logprob_from_prefill
+
+ _, cands = first_token_logprob_from_prefill(_prefill_resp(), top_n=1)
+ assert cands == [(7, -0.25)]
+
+
+def test_prefill_entry_without_logprobs_is_absent_not_fatal() -> None:
+ """A prefill reply carrying no usable entry must not fail the request.
+
+ Every other position is still correct, so the documented sentinel for that
+ one entry beats a 500 for the whole completion.
+ """
+ from tilert.pd_vllm.pd_router import first_token_logprob_from_prefill
+
+ assert first_token_logprob_from_prefill({}, 3) == (None, [])
+ assert first_token_logprob_from_prefill({"choices": [{"logprobs": {"content": []}}]}, 3) == (
+ None,
+ [],
+ )
+
+
+def test_prefill_candidates_that_are_not_token_ids_are_dropped() -> None:
+ """Without --return-tokens-as-token-ids the alternatives are plain text.
+
+ There is no id to report then, so the entry is dropped rather than guessed
+ at -- the token's own logprob still comes through.
+ """
+ from tilert.pd_vllm.pd_router import first_token_logprob_from_prefill
+
+ resp = _prefill_resp()
+ resp["choices"][0]["logprobs"]["content"][0]["top_logprobs"][1]["token"] = "he"
+ lp, cands = first_token_logprob_from_prefill(resp, top_n=3)
+ assert lp == -0.25
+ assert cands == [(7, -0.25), (9, -2.5)]
+
+
+def test_prefill_body_raises_top_logprobs_to_the_requested_count() -> None:
+ """Token 1's candidate row can only come from the prefill instance."""
+ from tilert.pd_vllm.pd_router import DecodeNode, build_prefill_body
+
+ node = DecodeNode(host="h", ctrl_port=1, http_port=2)
+ body = {"messages": [], "logprobs": True, "top_logprobs": 4}
+ out = build_prefill_body("/v1/chat/completions", body, node, LogprobsRequest(4))
+ assert out["logprobs"] is True
+ assert out["top_logprobs"] == 4
+
+
+def test_prefill_body_still_asks_for_one_entry_without_logprobs() -> None:
+ """The id extraction needs an entry even when the client asked for none."""
+ from tilert.pd_vllm.pd_router import DecodeNode, build_prefill_body
+
+ node = DecodeNode(host="h", ctrl_port=1, http_port=2)
+ out = build_prefill_body("/v1/chat/completions", {"messages": []}, node)
+ assert out["top_logprobs"] == 1
+
+
+# --------------------------------------------------------------------------- #
+# temperature bound
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize("temp", [1e-5, 0.05, 0.1, 0.19])
+def test_middle_band_temperature_is_refused(temp) -> None:
+ """Between greedy and the bound there is no path that can serve it.
+
+ The temperature selects the distribution being sampled, so it cannot be
+ substituted the way greedy's can, and the top-p path's raw denominator is
+ imprecise there.
+ """
+ with pytest.raises(LogprobsUnsupported) as e:
+ resolve_logprobs_request({"logprobs": True, "temperature": temp})
+ assert "temperature" in str(e.value)
+
+
+@pytest.mark.parametrize("temp", [0.0, 1e-9, 9e-6])
+def test_greedy_is_served_without_candidates(temp) -> None:
+ """The top-1 kernel exports the chosen token's value, so this is servable."""
+ req = resolve_logprobs_request({"logprobs": True, "temperature": temp})
+ assert req == LogprobsRequest(0)
+
+
+@pytest.mark.parametrize("temp", [0.0, 1e-9])
+@pytest.mark.parametrize("n", [1, 3, 5])
+def test_greedy_with_candidates_is_served(temp, n) -> None:
+ """The greedy kernel exports a candidate row too, so the full range serves.
+
+ TOP_LOGPROBS_MAX is that row's width, so nothing this accepts can ask for
+ more entries than the kernel writes.
+ """
+ assert resolve_logprobs_request(
+ {"logprobs": True, "temperature": temp, "top_logprobs": n}
+ ) == LogprobsRequest(n)
+
+
+@pytest.mark.parametrize("temp", [0.0, 1e-9])
+def test_greedy_beyond_the_row_width_is_refused(temp) -> None:
+ with pytest.raises(LogprobsUnsupported) as e:
+ resolve_logprobs_request(
+ {"logprobs": True, "temperature": temp, "top_logprobs": TOP_LOGPROBS_MAX + 1}
+ )
+ assert "top_logprobs" in str(e.value)
+
+
+@pytest.mark.parametrize("temp", [0.2, 0.6, 1.0, 2.0])
+def test_supported_temperature_is_accepted(temp) -> None:
+ req = resolve_logprobs_request({"logprobs": True, "temperature": temp})
+ assert req is not None
+
+
+def test_absent_temperature_is_accepted() -> None:
+ """vLLM's default is 1.0, which is inside the bound."""
+ assert resolve_logprobs_request({"logprobs": True}) is not None
+
+
+def test_low_temperature_without_logprobs_is_not_an_error() -> None:
+ """The bound constrains the logprobs export, not sampling."""
+ assert resolve_logprobs_request({"temperature": 0.0}) is None
diff --git a/tests/pd_vllm/test_logprobs_decode_protocol.py b/tests/pd_vllm/test_logprobs_decode_protocol.py
new file mode 100644
index 0000000..d4c93b3
--- /dev/null
+++ b/tests/pd_vllm/test_logprobs_decode_protocol.py
@@ -0,0 +1,339 @@
+"""`/pd/decode` carries per-token logprobs, in both response branches.
+
+Drives the real ``decode_server`` app over HTTP with ``StubEngine``, so the
+protocol is pinned without a GPU: the request field, the non-streaming
+``logprobs`` object, the per-line ``lp``/``tp`` fields on the NDJSON stream, and
+the refusal to answer at all when the engine cannot produce logprobs.
+
+The alignment property that matters: ``lp[i]`` and ``tp[i]`` describe ``t[i]``.
+The engine writes its entry before putting the token on the queue, so a token the
+generator has dequeued already has its logprob visible; the stream batches
+tokens, so this is asserted across batch boundaries rather than assumed.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_logprobs_decode_protocol.py -v
+"""
+
+from __future__ import annotations
+
+import json
+import queue
+import time
+import types
+
+import pytest
+from fastapi.testclient import TestClient
+
+from tilert.pd_vllm.decode_server import build_app
+from tilert.pd_vllm.engine_iface import StubEngine
+
+RID = "cmpl-lp"
+FIXED = (11, 22, 33)
+FIRST = 7
+
+
+class _FakeReq:
+ def __init__(self, rid):
+ self.rid = rid
+ self.seq_len = 8
+ self.last_prompt_token = 5
+
+
+class _FakeServer:
+ """Minimal ReceiveServer stand-in (same shape as test_grammar_plumbing)."""
+
+ def __init__(self, rid):
+ self.completed: queue.Queue = queue.Queue()
+ self.completed.put(_FakeReq(rid))
+ self.profile = types.SimpleNamespace(convert=lambda *a, **k: "converted", num_ranks=8)
+ self.buffer = None
+ self.base_ptr = 0
+ self.max_seq_len = 4096
+
+ def expect(self, rid=None):
+ # /pd/decode announces its rid so the real ReceiveServer can drop a
+ # tombstone left by a previous attempt at the same request. Recorded,
+ # so a test can assert the announcement happened.
+ self.expected_rids = getattr(self, "expected_rids", [])
+ self.expected_rids.append(rid)
+
+ def release(self, rid=None):
+ pass
+
+
+def _client(engine=None, rid=RID):
+ return TestClient(build_app(_FakeServer(rid), engine or StubEngine(FIXED)))
+
+
+def _body(**kw):
+ return {"rid": RID, "first_token_id": FIRST, "max_tokens": 8, **kw}
+
+
+class _NoLogprobsEngine(StubEngine):
+ """An engine predating the capability."""
+
+ def supports_logprobs(self) -> bool:
+ return False
+
+
+class _SlowEngine(StubEngine):
+ """Emits with a gap so the stream generator drains one token per batch.
+
+ Needed to exercise the batch offset: when every token arrives in a single
+ batch, a wrong offset is indistinguishable from a right one.
+ """
+
+ def decode(
+ self,
+ first_token_id,
+ max_tokens,
+ sampling,
+ on_token=None,
+ cancel_event=None,
+ grammar_session=None,
+ top_logprobs=None,
+ ):
+ out = ([int(first_token_id)] + list(self._fixed))[:max_tokens]
+ for t in out:
+ if on_token:
+ if top_logprobs is None:
+ on_token(t)
+ else:
+ on_token(
+ t,
+ self.fake_logprob(t),
+ [(t + k, self.fake_logprob(t) - 0.5 * k) for k in range(top_logprobs)],
+ )
+ time.sleep(0.03) # > the generator's 5 ms idle poll
+ self.last_stats = {"finish_reason": "stop"}
+ return out
+
+
+class _ShortLogprobsEngine(StubEngine):
+ """Declares support but emits a bare token -- an engine bug, not a client one.
+
+ Omits at position 1, not 0. Position 0 is ``first_token_id``, which the
+ prefill instance sampled, so a bare emit there is the contract (the router
+ fills that entry in from the prefill response) rather than a fault.
+ """
+
+ def decode(
+ self,
+ first_token_id,
+ max_tokens,
+ sampling,
+ on_token=None,
+ cancel_event=None,
+ grammar_session=None,
+ top_logprobs=None,
+ ):
+ out = ([int(first_token_id)] + list(self._fixed))[:max_tokens]
+ for i, t in enumerate(out):
+ if on_token:
+ if top_logprobs is None or i <= 1:
+ on_token(t)
+ else:
+ on_token(t, self.fake_logprob(t), [])
+ self.last_stats = {"finish_reason": "stop"}
+ return out
+
+
+class _EchoFirstEngine(StubEngine):
+ """The real engines' shape: token 0 bare, every later token with a logprob.
+
+ A real adapter behaves this way because it cannot report a distribution
+ for ``first_token_id`` -- it echoed that token rather than sampling it.
+ """
+
+ def decode(
+ self,
+ first_token_id,
+ max_tokens,
+ sampling,
+ on_token=None,
+ cancel_event=None,
+ grammar_session=None,
+ top_logprobs=None,
+ ):
+ out = ([int(first_token_id)] + list(self._fixed))[:max_tokens]
+ for i, t in enumerate(out):
+ if on_token:
+ if top_logprobs is None or i == 0:
+ on_token(t)
+ else:
+ on_token(t, self.fake_logprob(t), [(t, self.fake_logprob(t))])
+ self.last_stats = {"finish_reason": "stop"}
+ return out
+
+
+# --------------------------------------------------------------------------- #
+# not requested -> byte-for-byte the old response
+# --------------------------------------------------------------------------- #
+
+
+def test_absent_field_leaves_the_response_unchanged() -> None:
+ r = _client().post("/pd/decode", json=_body())
+ assert r.status_code == 200
+ assert "logprobs" not in r.json()
+
+
+def test_absent_field_leaves_the_stream_unchanged() -> None:
+ with _client().stream("POST", "/pd/decode", json=_body(stream=True)) as resp:
+ lines = [json.loads(x) for x in resp.iter_lines() if x]
+ tok_lines = [x for x in lines if "t" in x]
+ assert tok_lines and all("lp" not in x and "tp" not in x for x in tok_lines)
+
+
+# --------------------------------------------------------------------------- #
+# non-streaming
+# --------------------------------------------------------------------------- #
+
+
+def test_non_streaming_returns_one_logprob_per_token() -> None:
+ r = _client().post("/pd/decode", json=_body(top_logprobs=0))
+ body = r.json()
+ ids = body["token_ids"]
+ assert body["logprobs"]["lp"] == [StubEngine.fake_logprob(t) for t in ids]
+
+
+def test_non_streaming_candidate_count_matches_the_request() -> None:
+ r = _client().post("/pd/decode", json=_body(top_logprobs=3))
+ tp = r.json()["logprobs"]["tp"]
+ assert all(len(row) == 3 for row in tp)
+
+
+def test_top_logprobs_zero_gives_empty_candidate_rows() -> None:
+ r = _client().post("/pd/decode", json=_body(top_logprobs=0))
+ assert all(row == [] for row in r.json()["logprobs"]["tp"])
+
+
+def test_candidates_are_id_logprob_pairs_chosen_token_first() -> None:
+ r = _client().post("/pd/decode", json=_body(top_logprobs=2))
+ body = r.json()
+ first_id = body["token_ids"][0]
+ row = body["logprobs"]["tp"][0]
+ assert row[0] == [first_id, StubEngine.fake_logprob(first_id)]
+ assert row[0][1] > row[1][1], "candidates must be descending"
+
+
+# --------------------------------------------------------------------------- #
+# streaming: lp[i] / tp[i] line up with t[i] across batch boundaries
+# --------------------------------------------------------------------------- #
+
+
+@pytest.fixture
+def streamed():
+ with _client().stream("POST", "/pd/decode", json=_body(stream=True, top_logprobs=2)) as resp:
+ return [json.loads(x) for x in resp.iter_lines() if x]
+
+
+def test_every_token_line_carries_aligned_logprobs(streamed) -> None:
+ for line in (x for x in streamed if "t" in x):
+ assert len(line["lp"]) == len(line["t"])
+ assert len(line["tp"]) == len(line["t"])
+
+
+def test_streamed_logprobs_match_the_token_they_describe(streamed) -> None:
+ """The alignment property, checked per token across all batches."""
+ for line in (x for x in streamed if "t" in x):
+ for tok, lp, cands in zip(line["t"], line["lp"], line["tp"]):
+ assert lp == StubEngine.fake_logprob(tok)
+ assert cands[0][0] == tok
+
+
+def test_stream_covers_exactly_the_emitted_tokens(streamed) -> None:
+ flat = [t for x in streamed if "t" in x for t in x["t"]]
+ n_lp = sum(len(x["lp"]) for x in streamed if "t" in x)
+ assert n_lp == len(flat)
+ assert flat[0] == FIRST
+
+
+# --------------------------------------------------------------------------- #
+# an engine that cannot do it must say so, not answer without the field
+# --------------------------------------------------------------------------- #
+
+
+def test_unsupported_engine_returns_501_not_a_silent_omission() -> None:
+ r = _client(_NoLogprobsEngine(FIXED)).post("/pd/decode", json=_body(top_logprobs=1))
+ assert r.status_code == 501
+ assert r.json()["error_type"] == "logprobs_unavailable"
+
+
+def test_unsupported_engine_still_serves_requests_without_logprobs() -> None:
+ r = _client(_NoLogprobsEngine(FIXED)).post("/pd/decode", json=_body())
+ assert r.status_code == 200
+ assert "logprobs" not in r.json()
+
+
+def test_multi_batch_stream_keeps_logprobs_aligned() -> None:
+ """One token per batch, so a wrong batch offset shows up."""
+ with _client(_SlowEngine(FIXED)).stream(
+ "POST", "/pd/decode", json=_body(stream=True, top_logprobs=2)
+ ) as resp:
+ lines = [json.loads(x) for x in resp.iter_lines() if x]
+ tok_lines = [x for x in lines if "t" in x]
+ assert len(tok_lines) > 1, "expected several batches"
+ for line in tok_lines:
+ for tok, lp, cands in zip(line["t"], line["lp"], line["tp"]):
+ assert lp == StubEngine.fake_logprob(tok)
+ assert cands[0][0] == tok
+
+
+def test_engine_omitting_a_logprob_is_a_501_not_a_sentinel() -> None:
+ r = _client(_ShortLogprobsEngine(FIXED)).post("/pd/decode", json=_body(top_logprobs=1))
+ assert r.status_code == 501
+ assert r.json()["error_type"] == "logprobs_unavailable"
+
+
+def test_first_token_may_be_bare_and_travels_as_null() -> None:
+ """Position 0 is the one legitimate omission, and it must not be a sentinel.
+
+ The decode node never sampled ``first_token_id``, so it has no value to
+ report there. Sending -9999.0 would be indistinguishable from a real
+ measurement, and refusing would make every logprobs request a 501, so the
+ slot is held and sent as null for the router to fill from prefill.
+ """
+ r = _client(_EchoFirstEngine(FIXED)).post("/pd/decode", json=_body(top_logprobs=1))
+ assert r.status_code == 200
+ body = r.json()
+ lp = body["logprobs"]["lp"]
+ assert len(lp) == len(body["token_ids"]), "one entry per token"
+ assert lp[0] is None, "token 1 carries no decode-side value"
+ assert body["logprobs"]["tp"][0] == [], "and no candidate row"
+ assert all(v is not None for v in lp[1:]), "every later token has one"
+
+
+def test_first_token_bare_is_still_refused_when_a_later_one_is_too() -> None:
+ """Tolerating position 0 must not weaken the rule for the rest.
+
+ Guards against the obvious over-correction: accepting any bare emit once the
+ first has been seen would let a real engine fault through as a sentinel.
+ """
+ r = _client(_ShortLogprobsEngine(FIXED)).post("/pd/decode", json=_body(top_logprobs=1))
+ assert r.status_code == 501
+
+
+def test_the_streaming_worker_keeps_the_logprobs_type():
+ """Both branches of the decode server report the same inability the same way.
+
+ The blocking branch answers 501 `logprobs_unavailable`. The streaming worker
+ caught `LogprobsUnavailable` in its generic handler and emitted only
+ `{"error": ...}`, so the type was lost and the router -- which cannot see a
+ status inside a 200 body -- reported 502. Adding a `stop` string is what puts
+ a non-streaming request on that protocol, so the same failure changed status
+ depending on an unrelated field.
+ """
+ import pathlib as _p
+
+ src = _p.Path(decode_server.__file__).read_text()
+ run = src[src.index(" def _run():") : src.index("worker = threading.Thread")]
+ assert "except LogprobsUnavailable" in run, (
+ "the streaming worker no longer classifies LogprobsUnavailable, so the "
+ "router sees an untyped error and answers 502"
+ )
+ typed = run[run.index("except LogprobsUnavailable") :]
+ assert '"error_type": "logprobs_unavailable"' in typed
+
+
+from tilert.pd_vllm import decode_server # noqa: E402
diff --git a/tests/pd_vllm/test_no_silent_degradation.py b/tests/pd_vllm/test_no_silent_degradation.py
new file mode 100644
index 0000000..4ab8be1
--- /dev/null
+++ b/tests/pd_vllm/test_no_silent_degradation.py
@@ -0,0 +1,305 @@
+"""Every request field the router owns is executed, refused, or declared inert.
+
+The capability gate already has this guard for its own 15 fields:
+`test_every_static_field_has_a_live_and_a_neutral_fixture` fails if one is added
+without a fixture. The fields the ROUTER handles -- `stop`, `logprobs`,
+`chat_template_kwargs` and the rest -- had no such enumeration, and that is where
+seven consecutive review findings landed. Each was the same defect:
+
+ 200 OK, and the field silently did not take effect.
+
+That is worse than a refusal, because the client is told the request succeeded.
+The codebase's stated rule is the opposite -- accept only what can be executed
+completely, refuse the rest immediately -- and nothing enforced it.
+
+Two guards here:
+
+* :data:`ROUTER_FIELDS` enumerates them. A field added to the router without an
+ entry fails :func:`test_every_router_field_is_accounted_for`, so the next one
+ cannot be forgotten rather than found.
+* Each entry declares what a request carrying it must produce: ``served`` with a
+ predicate proving the field took effect, ``refused`` with the status, or
+ ``inert`` with the reason it legitimately changes nothing.
+
+No GPU and no weights: the real `build_app` against stub backends.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_no_silent_degradation.py -v
+"""
+
+from __future__ import annotations
+
+import json
+import re
+
+import httpx
+import pytest
+
+from tests.pd_vllm.test_stop_strings import ( # noqa: E402
+ _SEEN,
+ _make_decode,
+ _make_vllm,
+ _serve,
+ _Tok,
+)
+from tilert.pd_vllm.decode_pool import DecodeNode, Pool
+from tilert.pd_vllm.pd_router import RouterCtx, build_app
+
+# ── the enumeration ─────────────────────────────────────────────────────────
+#
+# `served` -- the request is 200 AND `check(reply, sent)` proves the field took
+# effect. A predicate that only checks the status would pass for a
+# field that was silently dropped, which is the whole point.
+# `reply` is the response JSON; `sent` is the body the decode node
+# received, which is where a FORWARDED field's effect is visible --
+# the stub does not implement `max_tokens`, and asserting on the
+# reply length would pin the stub rather than the router.
+# `refused` -- the request must fail with this status and a typed error.
+# `inert` -- 200 with nothing observable, and a reason that says why that is
+# correct rather than a gap.
+
+SERVED, REFUSED, INERT = "served", "refused", "inert"
+
+
+def _reply(body: dict) -> str:
+ msg = body["choices"][0].get("message") or {}
+ return (msg.get("reasoning_content") or "") + (msg.get("content") or "")
+
+
+ROUTER_FIELDS: dict[str, list[tuple]] = {
+ "stop": [
+ (
+ SERVED,
+ {"stop": ["STOP"]},
+ lambda j, sent: j["choices"][0]["stop_reason"] == "STOP" and "STOP" not in _reply(j),
+ ),
+ (REFUSED, {"stop": [""]}, 400),
+ (REFUSED, {"stop": 5}, 400),
+ ],
+ "include_stop_str_in_output": [
+ (
+ SERVED,
+ {"stop": ["STOP"], "include_stop_str_in_output": True},
+ lambda j, sent: _reply(j).endswith("STOP"),
+ ),
+ (
+ SERVED,
+ {"stop": ["STOP"], "include_stop_str_in_output": 1},
+ lambda j, sent: _reply(j).endswith("STOP"),
+ ),
+ (REFUSED, {"stop": ["STOP"], "include_stop_str_in_output": None}, 400),
+ (REFUSED, {"stop": ["STOP"], "include_stop_str_in_output": "maybe"}, 400),
+ ],
+ "logprobs": [
+ (
+ SERVED,
+ {"logprobs": True, "temperature": 0.6},
+ lambda j, sent: (j["choices"][0]["logprobs"] or {}).get("content"),
+ ),
+ (SERVED, {"logprobs": False}, lambda j, sent: j["choices"][0]["logprobs"] is None),
+ ],
+ "top_logprobs": [
+ # Forwarded to the node, which is the only place it can act. Position 0's
+ # candidate row comes from the prefill reply and carries one entry, so
+ # the response alone cannot show the requested count.
+ (
+ SERVED,
+ {"logprobs": True, "top_logprobs": 2, "temperature": 0.6},
+ lambda j, sent: sent.get("top_logprobs") == 2,
+ ),
+ (REFUSED, {"logprobs": True, "top_logprobs": 99}, 400),
+ (REFUSED, {"top_logprobs": 1}, 400),
+ ],
+ "chat_template_kwargs": [
+ (
+ SERVED,
+ {"chat_template_kwargs": {"enable_thinking": False}},
+ lambda j, sent: j["choices"][0]["message"]["content"] is not None,
+ ),
+ (REFUSED, {"chat_template_kwargs": "bad"}, 400),
+ (REFUSED, {"chat_template_kwargs": 5}, 400),
+ ],
+ "max_tokens": [
+ (SERVED, {"max_tokens": 3}, lambda j, sent: sent.get("max_tokens") == 3),
+ # Coerced, not rejected: the request model takes these and so must we.
+ # Asserting on `sent` is the point: the coerced value has to reach
+ # the node, where re-reading the body used to 502 after prefill.
+ (SERVED, {"max_tokens": "20.0"}, lambda j, sent: sent.get("max_tokens") == 20),
+ (SERVED, {"max_tokens": True}, lambda j, sent: sent.get("max_tokens") == 1),
+ # Refused before the prefill, not after it.
+ (REFUSED, {"max_tokens": 0}, 400),
+ (REFUSED, {"max_tokens": 1.9}, 400),
+ ],
+ "max_completion_tokens": [
+ # vLLM resolves this ahead of `max_tokens`, so the router has to too --
+ # otherwise the prefill leg and the decode leg disagree about the length.
+ (SERVED, {"max_completion_tokens": 3}, lambda j, sent: sent.get("max_tokens") == 3),
+ # Precedence survives the coercion, and the shadowed name is checked
+ # for its type but not its range -- as on vLLM.
+ (
+ SERVED,
+ {"max_completion_tokens": "7.0", "max_tokens": 3},
+ lambda j, sent: sent.get("max_tokens") == 7,
+ ),
+ (
+ SERVED,
+ {"max_completion_tokens": 5, "max_tokens": 0},
+ lambda j, sent: sent.get("max_tokens") == 5,
+ ),
+ (REFUSED, {"max_completion_tokens": 5, "max_tokens": 1.9}, 400),
+ ],
+ "temperature": [
+ (SERVED, {"temperature": 0.0}, lambda j, sent: sent["sampling"].get("temperature") == 0.0),
+ ],
+ "stream": [
+ (INERT, {"stream": False}, "false is the default; the streaming path has its own tests"),
+ ],
+ "stream_options": [
+ (
+ INERT,
+ {"stream_options": {"include_usage": True}},
+ "meaningful only with stream: true, where test_stream_e2e pins it; "
+ "stripped from the prefill request so vLLM cannot reject the pair",
+ ),
+ ],
+ "kv_transfer_params": [
+ (
+ INERT,
+ {},
+ "set BY the router on the prefill request, never read from the " "client's body",
+ ),
+ ],
+}
+
+
+@pytest.fixture(scope="module")
+def router():
+ mp = pytest.MonkeyPatch()
+ for var in ("no_proxy", "NO_PROXY"):
+ mp.setenv(var, "127.0.0.1,localhost")
+ for var in ("http_proxy", "HTTP_PROXY", "https_proxy", "HTTPS_PROXY"):
+ mp.delenv(var, raising=False)
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode())
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ yield f"http://127.0.0.1:{_serve(build_app(ctx))}"
+ mp.undo()
+
+
+def _post(url: str, over: dict) -> httpx.Response:
+ body = {"model": "stub", "max_tokens": 64, "messages": [{"role": "user", "content": "hi"}]}
+ body.update(over)
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ return c.post(f"{url}/v1/chat/completions", json=body)
+
+
+def _cases():
+ for field, entries in ROUTER_FIELDS.items():
+ for i, entry in enumerate(entries):
+ yield pytest.param(field, entry, id=f"{field}-{i}")
+
+
+@pytest.mark.parametrize("field,entry", list(_cases()))
+def test_a_router_field_is_never_silently_dropped(router, field, entry):
+ """200 with the field not taking effect is the defect this pins.
+
+ A `served` case is not satisfied by the status alone: the predicate has to
+ find the field's effect in the response, because a silently ignored field
+ also answers 200.
+ """
+ kind = entry[0]
+ if kind is INERT:
+ _, over, why = entry
+ r = _post(router, over)
+ assert r.status_code == 200, f"{field}: {r.status_code} {r.text[:160]}"
+ assert why, "an inert field needs a stated reason"
+ return
+ if kind is REFUSED:
+ _, over, want = entry
+ r = _post(router, over)
+ assert (
+ r.status_code == want
+ ), f"{field}: expected {want}, got {r.status_code} {r.text[:160]}"
+ assert r.json().get("error_type"), (
+ f"{field}: refused without a typed error_type, so a client cannot "
+ f"tell what to change"
+ )
+ return
+ _, over, check = entry
+ r = _post(router, over)
+ assert r.status_code == 200, f"{field}: {r.status_code} {r.text[:200]}"
+ assert check(r.json(), _SEEN["decode_body"]), (
+ f"{field}: answered 200 but the field did not take effect -- "
+ f"{json.dumps(r.json())[:300]}"
+ )
+
+
+def test_every_router_field_is_accounted_for():
+ """A field the router reads but does not declare here fails this.
+
+ The point is that the next one is forgotten loudly. Seven review findings
+ were the same defect on this surface, each found one at a time.
+ """
+ src = pytest.importorskip("pathlib").Path(build_app.__globals__["__file__"]).read_text()
+ read = set(re.findall(r'body\.get\("([a-z_]+)"', src))
+ read |= set(re.findall(r'body\["([a-z_]+)"\]', src))
+ read |= set(re.findall(r'"([a-z_]+)" in body', src))
+ # Fields other modules read off the same body, resolved on the router's
+ # behalf, so they belong to this surface too.
+ read |= {"stop", "max_completion_tokens", "temperature"}
+ from tilert.pd_vllm.capabilities import STATIC_FIELD_NAMES
+
+ unaccounted = read - set(ROUTER_FIELDS) - set(STATIC_FIELD_NAMES)
+ assert not unaccounted, (
+ f"the router reads {sorted(unaccounted)} but ROUTER_FIELDS does not say "
+ f"whether each is executed, refused, or inert. Add an entry rather than "
+ f"leaving the next silent drop to be found in review."
+ )
+
+
+# --------------------------------------------------------------------------- #
+# The same rule on a deployment that cannot execute the field at all
+# --------------------------------------------------------------------------- #
+@pytest.fixture(scope="module")
+def tokenizerless_router(router):
+ """`--parser none` with no `--model-path`: a supported configuration.
+
+ Reuses the module fixture's proxy setup and backends. It has no tokenizer, so
+ it cannot match stop strings or name tokens -- and the rule is the same:
+ refuse, do not answer 200 with the field quietly inert.
+ """
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode())
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), None, "none")
+ return f"http://127.0.0.1:{_serve(build_app(ctx))}"
+
+
+@pytest.mark.parametrize(
+ "over,field",
+ [
+ ({"stop": ["STOP"]}, "stop"),
+ ({"logprobs": True, "top_logprobs": 1, "temperature": 0.6}, "logprobs"),
+ ],
+)
+def test_a_field_this_deployment_cannot_execute_is_refused(tokenizerless_router, over, field):
+ """501, not 200 with the field inert.
+
+ Both were silent once: `stop` was ignored and the reply ran past it, and
+ `logprobs` came back `null` under a 200 after both backends had computed the
+ values. Nothing in the response said the request had not been honoured.
+ """
+ r = _post(tokenizerless_router, over)
+ assert r.status_code == 501, f"{field}: got {r.status_code} {r.text[:200]}"
+ assert r.json()["error_type"] == "capability_unavailable"
+
+
+def test_the_same_deployment_still_serves_what_it_can(tokenizerless_router):
+ """The refusals are narrow: a request that asks for neither is served, with
+ `logprobs` declared null rather than absent (#43).
+ """
+ r = _post(tokenizerless_router, {})
+ assert r.status_code == 200, r.text[:200]
+ assert r.json()["choices"][0]["logprobs"] is None
diff --git a/tests/pd_vllm/test_openai_envelope.py b/tests/pd_vllm/test_openai_envelope.py
new file mode 100644
index 0000000..02cfab4
--- /dev/null
+++ b/tests/pd_vllm/test_openai_envelope.py
@@ -0,0 +1,279 @@
+"""The response envelope must match what an OpenAI client expects to read.
+
+Three gaps, all of them things a client reads unconditionally:
+
+* ``usage.total_tokens`` was sent on the streaming path and omitted on the
+ non-streaming one, so the same deployment answered a client's
+ ``usage.total_tokens`` with a number or a KeyError depending on ``stream``.
+* ``usage.prompt_tokens`` is copied from the prefill response, where it can be
+ absent — and ``null`` violates the contract as surely as a missing key.
+* ``created`` was read from the clock per chunk, so one streamed response carried
+ several timestamps. vLLM threads a single ``created_time`` through every chunk;
+ a client that groups or de-duplicates by ``(id, created)`` needs that.
+
+Shapes are taken from vLLM rather than from the spec alone, because that is the
+thing this endpoint stands in for: ``UsageInfo`` declares all three fields as
+integers, and non-streaming responses go out as
+``JSONResponse(content=result.model_dump())`` with no ``exclude_none``, so a
+declared-but-unset field appears as null. Streaming is the opposite — chunks use
+``model_dump_json(exclude_unset=True)``, so an unset field is omitted there.
+
+CPU only -- no GPU, no tilert, no real vLLM.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_openai_envelope.py -v
+"""
+
+import pytest
+from fastapi.testclient import TestClient
+
+from tilert.pd_vllm import pd_router, presentation
+from tilert.pd_vllm.pd_router import build_usage
+
+# --------------------------------------------------------------------------- #
+# build_usage: vLLM's UsageInfo shape
+# --------------------------------------------------------------------------- #
+
+
+class _Stream:
+ """The two facts `blocking_choice` reads off a reply stream."""
+
+ stop_reason = None
+ completion_tokens = 3
+
+ def finish_reason(self, from_node):
+ return from_node
+
+
+def test_usage_carries_all_three_fields():
+ """``UsageInfo`` declares prompt_tokens, completion_tokens and total_tokens.
+
+ Omitting the total is what made the two paths disagree.
+ """
+ assert set(build_usage(7, 3)) == {"prompt_tokens", "completion_tokens", "total_tokens"}
+
+
+def test_the_total_is_the_sum():
+ u = build_usage(7, 3)
+ assert (u["prompt_tokens"], u["completion_tokens"], u["total_tokens"]) == (7, 3, 10)
+
+
+@pytest.mark.parametrize("prompt", [None, 0, "5"])
+def test_prompt_tokens_is_always_an_integer(prompt):
+ """It is copied from the prefill response, which may not carry it. ``null``
+
+ breaks a client arithmetic-ing over usage just as a missing key does.
+ """
+ u = build_usage(prompt, 3)
+ assert isinstance(u["prompt_tokens"], int)
+ assert isinstance(u["total_tokens"], int)
+
+
+def test_an_absent_prompt_count_is_zero_not_null():
+ u = build_usage(None, 4)
+ assert u["prompt_tokens"] == 0
+ assert u["total_tokens"] == 4
+
+
+# --------------------------------------------------------------------------- #
+# The two paths agree, over real HTTP
+# --------------------------------------------------------------------------- #
+class _Resp:
+ def __init__(self, payload, status=200):
+ self._payload = payload
+ self.status_code = status
+
+ def json(self):
+ return self._payload
+
+ def raise_for_status(self):
+ if self.status_code >= 400:
+ raise RuntimeError(f"HTTP {self.status_code}")
+
+
+PREFILL = {
+ "id": "cmpl-env",
+ "choices": [{"logprobs": {"content": [{"token": "token_id:7"}]}}],
+ "usage": {"prompt_tokens": 5},
+ "model": "m",
+}
+DECODED = {"rid": "x", "token_ids": [7, 8], "seq_len": 8, "timing_ms": {"finish_reason": "stop"}}
+BODY = {"messages": [{"role": "user", "content": "hi"}]}
+
+
+class _StubTokenizer:
+ _VOCAB = {7: "hi", 8: " there"}
+
+ def decode(self, ids, skip_special_tokens=False):
+ return "".join(self._VOCAB.get(i, "") for i in ids)
+
+
+def _client(monkeypatch):
+ monkeypatch.setattr(
+ pd_router.requests,
+ "get",
+ lambda url, timeout=None, **kw: _Resp(
+ {"capabilities": {"penalties": True, "ignore_eos": True}}
+ ),
+ )
+
+ def fake_post(url, json=None, timeout=None, **kw):
+ return _Resp(DECODED if url.endswith("/pd/decode") else PREFILL)
+
+ monkeypatch.setattr(pd_router.requests, "post", fake_post)
+ pool = pd_router.Pool([pd_router.DecodeNode("127.0.0.1", 5556, 5557)])
+ ctx = pd_router.RouterCtx(
+ "http://vllm.invalid", pool, tokenizer=_StubTokenizer(), parser_name="none"
+ )
+ return TestClient(pd_router.build_app(ctx))
+
+
+def test_non_streaming_usage_has_a_total(monkeypatch):
+ r = _client(monkeypatch).post("/v1/chat/completions", json=BODY)
+ assert r.status_code == 200, r.text
+ usage = r.json()["usage"]
+ assert usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"]
+
+
+def test_non_streaming_usage_fields_are_integers(monkeypatch):
+ r = _client(monkeypatch).post("/v1/chat/completions", json=BODY)
+ usage = r.json()["usage"]
+ for field, value in usage.items():
+ assert isinstance(value, int), f"{field} is {type(value).__name__}"
+
+
+def test_usage_survives_a_prefill_that_reports_none(monkeypatch):
+ """The prefill response is another service's output; it may omit usage."""
+ monkeypatch.setattr(
+ pd_router.requests, "get", lambda url, timeout=None, **kw: _Resp({"capabilities": {}})
+ )
+
+ def fake_post(url, json=None, timeout=None, **kw):
+ if url.endswith("/pd/decode"):
+ return _Resp(DECODED)
+ return _Resp({**PREFILL, "usage": None})
+
+ monkeypatch.setattr(pd_router.requests, "post", fake_post)
+ pool = pd_router.Pool([pd_router.DecodeNode("127.0.0.1", 5556, 5557)])
+ ctx = pd_router.RouterCtx(
+ "http://vllm.invalid", pool, tokenizer=_StubTokenizer(), parser_name="none"
+ )
+ client = TestClient(pd_router.build_app(ctx))
+ usage = client.post("/v1/chat/completions", json=BODY).json()["usage"]
+ assert usage["prompt_tokens"] == 0
+ assert usage["total_tokens"] == usage["completion_tokens"]
+
+
+# --------------------------------------------------------------------------- #
+# Declared-but-null fields on the non-streaming choice
+# --------------------------------------------------------------------------- #
+def test_logprobs_is_declared_null_when_not_requested(monkeypatch):
+ """vLLM's non-streaming response is dumped without ``exclude_none``, so a
+ client can read ``choices[0].logprobs`` unconditionally.
+ """
+ r = _client(monkeypatch).post("/v1/chat/completions", json=BODY)
+ choice = r.json()["choices"][0]
+ assert "logprobs" in choice and choice["logprobs"] is None
+
+
+def test_stop_reason_is_declared(monkeypatch):
+ """vLLM carries it for legacy reasons, and it is how a client tells "stopped
+ on a stop string" from "stopped on EOS".
+
+ Always null here: the stop strings and stop_token_ids that would populate it
+ are refused by the capability gate, so there is nothing it could name. The
+ field is present rather than absent so the distinction is readable at all.
+ """
+ r = _client(monkeypatch).post("/v1/chat/completions", json=BODY)
+ choice = r.json()["choices"][0]
+ assert "stop_reason" in choice and choice["stop_reason"] is None
+
+
+def test_a_requested_logprobs_still_wins_over_the_null_default(monkeypatch):
+ """The null is a default, not an override."""
+
+ def fake_post(url, json=None, timeout=None, **kw):
+ if url.endswith("/pd/decode"):
+ return _Resp({**DECODED, "logprobs": {"lp": [-0.5, -0.6], "tp": [[], []]}})
+ return _Resp(PREFILL)
+
+ monkeypatch.setattr(
+ pd_router.requests,
+ "get",
+ lambda url, timeout=None, **kw: _Resp(
+ {"capabilities": {"penalties": True, "ignore_eos": True}}
+ ),
+ )
+ monkeypatch.setattr(pd_router.requests, "post", fake_post)
+ pool = pd_router.Pool([pd_router.DecodeNode("127.0.0.1", 5556, 5557)])
+ ctx = pd_router.RouterCtx(
+ "http://vllm.invalid", pool, tokenizer=_StubTokenizer(), parser_name="none"
+ )
+ r = TestClient(pd_router.build_app(ctx)).post(
+ "/v1/chat/completions",
+ json={**BODY, "logprobs": True, "top_logprobs": 0, "temperature": 0.6},
+ )
+ assert r.status_code == 200, r.text
+ assert r.json()["choices"][0]["logprobs"] is not None
+
+
+# --------------------------------------------------------------------------- #
+# created is stamped once per response
+# --------------------------------------------------------------------------- #
+def test_created_is_an_integer_timestamp(monkeypatch):
+ body = _client(monkeypatch).post("/v1/chat/completions", json=BODY).json()
+ assert isinstance(body["created"], int)
+
+
+def test_no_path_reads_the_clock_per_chunk():
+ """Source-level, because reproducing a multi-second stream to catch a
+ one-second drift would be a slow and flaky test for a property that is
+ plainly visible in the code: vLLM threads a single ``created_time`` through
+ every chunk, and each ``int(time.time())`` inside a chunk builder is a
+ response that can carry more than one timestamp.
+ """
+ import pathlib
+
+ src = pathlib.Path(pd_router.__file__).read_text()
+ # One assignment is expected (the per-response stamp); a call inside a chunk
+ # payload is not.
+ for marker in ('"created": int(time.time())', '"created": int(time.time()), "model"'):
+ assert marker not in src, f"a chunk builder still reads the clock: {marker}"
+ assert src.count("int(time.time())") <= 1, (
+ "more than one clock read: each is a chance for one response to carry " "two timestamps"
+ )
+
+
+@pytest.mark.parametrize(
+ "choice_of,kw,why",
+ [
+ (
+ presentation.textless_choice,
+ {"from_node": "length"},
+ "no tokenizer: ids are all the reply can carry",
+ ),
+ (
+ presentation.blocking_choice,
+ {"from_node": "length", "logprobs_asked": False},
+ "with a tokenizer: alongside the text",
+ ),
+ ],
+)
+def test_completions_keeps_token_ids(choice_of, kw, why):
+ """``token_ids`` is NOT a non-standard wart to remove.
+
+ vLLM declares it on its own choice model — "not part of the OpenAI spec but
+ is useful for tracing the tokens in agent scenarios" — so carrying it on
+ /v1/completions matches the thing this endpoint stands in for.
+
+ Asserted on the field, not on the source line that sets it: the grep this
+ replaces failed when the line moved to `presentation.py` unchanged.
+ """
+ if choice_of is presentation.blocking_choice:
+ kw["stream"] = _Stream()
+ kw["got"] = presentation.collect([])
+ choice, _ = choice_of(is_chat=False, token_ids=[5, 6, 7], **kw)
+ assert choice["token_ids"] == [5, 6, 7], why
+ chat, _ = choice_of(is_chat=True, token_ids=[5, 6, 7], **dict(kw))
+ assert "token_ids" not in chat, "chat carries the message instead"
diff --git a/tests/pd_vllm/test_pd_prompt_token_ids.py b/tests/pd_vllm/test_pd_prompt_token_ids.py
new file mode 100644
index 0000000..8c169fa
--- /dev/null
+++ b/tests/pd_vllm/test_pd_prompt_token_ids.py
@@ -0,0 +1,174 @@
+"""The decode node must receive the prompt ids so repetition_penalty keeps its scope.
+
+`repetition_penalty` is scoped over prompt UNION output in both HF and vLLM. On
+a single machine TileRT matches that -- `seed_prompt_tokens()` fills a second
+bitmap. Under PD the wire carried only `last_prompt_token` and the decode server
+has no tokeniser, so the prompt half stayed empty and the knob silently degraded
+to output-only scope, inconsistent with the same request served non-PD.
+
+vLLM does not need a wire field for this: its PD proxy sends the SAME request to
+both instances, so the decode instance runs `add_request` with the full
+`prompt_token_ids` and the connector only carries KV block locations
+(`nixl/pull_scheduler.py::get_num_new_matched_tokens` reads
+`request.prompt_token_ids` locally). This change is the explicit equivalent of
+what vLLM gets for free.
+
+No GPU, no tilert, no vllm: `prefill_connector` cannot be imported here (it
+pulls in vllm), so the send-side gate is tested through `wire`, where it lives.
+The engine-side seeding (``seed_prompt_tokens`` after ``reset``) belongs to a
+decode runtime with a penalty pre-pass; no public profile ships one today.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest tests/pd_vllm/test_pd_prompt_token_ids.py -v
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from tilert.pd_vllm import wire
+from tilert.pd_vllm.receive_server import ReceivedRequest
+
+# --------------------------------------------------------------------------- #
+# send-side gate: who pays for the payload
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+ "sampling",
+ [
+ None,
+ {},
+ {"temperature": 0.7},
+ {"repetition_penalty": 1.0}, # the kernel's no-op value
+ {"presence_penalty": 1.5}, # output-scoped, does not need the prompt
+ {"repetition_penalty": "not-a-number"},
+ ],
+)
+def test_no_prompt_ids_when_they_would_be_unused(sampling) -> None:
+ """A request that cannot use the prompt half must not ship it.
+
+ Mirrors vLLM's `needs_prompt_token_ids` gate. Shipping unconditionally would
+ put a per-request `4 * prompt_len` payload on every plain request.
+ """
+ assert wire.wants_prompt_token_ids(sampling) is False
+
+
+@pytest.mark.parametrize("rep", [1.05, 1.5, 2.0, 0.5])
+def test_prompt_ids_shipped_when_repetition_penalty_is_live(rep) -> None:
+ assert wire.wants_prompt_token_ids({"repetition_penalty": rep}) is True
+
+
+def test_gate_ignores_presence_but_honours_a_paired_repetition() -> None:
+ assert wire.wants_prompt_token_ids({"presence_penalty": 2.0, "repetition_penalty": 1.3}) is True
+
+
+# --------------------------------------------------------------------------- #
+# receive-side parse
+# --------------------------------------------------------------------------- #
+
+
+def test_received_request_defaults_to_no_prompt_ids() -> None:
+ """Ranks 1-7 and older prefill builds send no such field.
+
+ The field is optional on the wire (JSON, `req.get(...)`), so a prefill connector
+ without this change must still work -- it just gets output-only scope.
+ """
+ r = ReceivedRequest(rid="x", seq_len=8, last_prompt_token=5, first_token_id=None, sampling=None)
+ assert r.prompt_token_ids == []
+
+
+def test_received_request_carries_prompt_ids() -> None:
+ r = ReceivedRequest(
+ rid="x",
+ seq_len=8,
+ last_prompt_token=5,
+ first_token_id=None,
+ sampling={"repetition_penalty": 1.3},
+ prompt_token_ids=[1, 2, 3],
+ )
+ assert r.prompt_token_ids == [1, 2, 3]
+
+
+# --------------------------------------------------------------------------- #
+# Multi-rank arrival order.
+#
+# Found by a live PD run, not by the tests above: `ReceivedRequest` is built by
+# whichever rank connects FIRST for a given rid, and only rank 0 carries the
+# prompt ids. Ranks connect in arbitrary order (observed: 5, 1, 0, 7, 5, 1 across
+# consecutive requests), so keying the ids off the creation path dropped them
+# ~7/8 of the time -- a NON-DETERMINISTIC feature that unit tests over the
+# dataclass and the gate could not see.
+# --------------------------------------------------------------------------- #
+
+
+def _absorb(messages: list[dict]) -> list[int]:
+ """Replay `receive_server`'s per-connection absorb step for one rid.
+
+ Mirrors the body of the `with self._lock` block: the first message builds the
+ request, every message may contribute the prompt ids.
+ """
+ cur = None
+ for req in messages:
+ if cur is None or cur.rid != req["rid"]:
+ cur = ReceivedRequest(
+ rid=req["rid"],
+ seq_len=req["seq_len"],
+ last_prompt_token=req.get("last_prompt_token", 0),
+ first_token_id=req.get("first_token_id"),
+ sampling=req.get("sampling"),
+ prompt_token_ids=list(req.get("prompt_token_ids") or []),
+ )
+ if not cur.prompt_token_ids and req.get("prompt_token_ids"):
+ cur.prompt_token_ids = list(req["prompt_token_ids"])
+ return cur.prompt_token_ids if cur else []
+
+
+def _msgs(rank_order: list[int], ids: list[int]) -> list[dict]:
+ """One request message per rank, in the given connection order.
+
+ Only rank 0 carries the ids -- what `prefill_connector._send` does.
+ """
+ return [
+ {
+ "rid": "r1",
+ "rank": r,
+ "seq_len": 8,
+ "last_prompt_token": 5,
+ "sampling": {"repetition_penalty": 1.5},
+ **({"prompt_token_ids": ids} if r == 0 else {}),
+ }
+ for r in rank_order
+ ]
+
+
+IDS = [11, 22, 33]
+
+
+@pytest.mark.parametrize(
+ "first",
+ [0, 1, 2, 3, 4, 5, 6, 7],
+ ids=[f"rank{r}_first" for r in range(8)],
+)
+def test_prompt_ids_survive_whatever_rank_connects_first(first: int) -> None:
+ """The regression this file exists for, second edition.
+
+ Before the fix only `rank0_first` passed; the other seven silently produced
+ an empty prompt bitmap and repetition degraded to output-only scope.
+ """
+ order = [first] + [r for r in range(8) if r != first]
+ assert _absorb(_msgs(order, IDS)) == IDS, f"lost the ids when rank {first} led"
+
+
+def test_absorb_is_idempotent_and_does_not_grow() -> None:
+ """Rank 0 appearing once must not be double-counted, and no rank may clear it."""
+ order = [3, 0, 1, 2, 4, 5, 6, 7]
+ assert _absorb(_msgs(order, IDS)) == IDS
+ # ranks after rank 0 send no ids -- they must not reset the field
+ assert _absorb(_msgs(order + [4, 5], IDS)) == IDS
+
+
+def test_no_ids_anywhere_stays_empty() -> None:
+ """A penalty-free request (or an old prefill connector) leaves the bitmap clear."""
+ msgs = [{"rid": "r1", "rank": r, "seq_len": 8} for r in range(8)]
+ assert _absorb(msgs) == []
diff --git a/tests/pd_vllm/test_pool_queue.py b/tests/pd_vllm/test_pool_queue.py
new file mode 100644
index 0000000..b7f56ef
--- /dev/null
+++ b/tests/pd_vllm/test_pool_queue.py
@@ -0,0 +1,108 @@
+"""A decode node must be waited for, not refused, when ``queue_timeout`` > 0.
+
+A TileRT decode engine serves one sequence at a time, so the router's pool
+is as deep as the node count. An agentic session fans a single conversation
+out into concurrent sub-conversations, which puts more requests in flight
+than the pool has nodes -- with fail-fast reservation the surplus turns
+into ``429``s even though the pool can serve them a moment later.
+``queue_timeout`` makes the surplus wait instead; ``0`` keeps the original
+fail-fast behaviour, so fixed-sequence-length runs are unaffected.
+
+The waiting side matters as much as the timeout: ``release`` has to wake a
+waiter, otherwise a queued request sleeps out its full timeout even though a
+node freed up immediately.
+
+No GPU, no tilert, no vllm: ``Pool`` is pure ``threading``.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_pool_queue.py -v
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+
+from tilert.pd_vllm.decode_pool import DecodeNode, Pool
+
+TIMEOUT = 0.3
+
+
+def _pool(queue_timeout: float, nodes: int = 1) -> Pool:
+ return Pool(
+ [DecodeNode("h%d" % i, 5556 + i, 5557 + i) for i in range(nodes)],
+ queue_timeout=queue_timeout,
+ )
+
+
+def test_fail_fast_is_the_default() -> None:
+ """Constructed without the argument, the pool refuses as it always did."""
+ pool = Pool([DecodeNode("h0", 5556, 5557)])
+ assert pool.queue_timeout == 0.0
+ assert pool.acquire() is not None
+ t0 = time.monotonic()
+ assert pool.acquire() is None
+ assert time.monotonic() - t0 < TIMEOUT, "fail-fast must not block"
+
+
+def test_timeout_zero_still_hands_out_free_nodes() -> None:
+ pool = _pool(0.0, nodes=2)
+ first, second = pool.acquire(), pool.acquire()
+ assert first is not None and second is not None and first is not second
+ assert pool.acquire() is None
+ pool.release(first)
+ assert pool.acquire() is first
+
+
+def test_waits_until_a_node_is_released() -> None:
+ """The queued caller gets the node a releaser frees, not a 429."""
+ pool = _pool(30.0)
+ held = pool.acquire()
+ assert held is not None
+
+ def _release_soon() -> None:
+ time.sleep(0.05)
+ pool.release(held)
+
+ threading.Thread(target=_release_soon, daemon=True).start()
+ t0 = time.monotonic()
+ got = pool.acquire()
+ waited = time.monotonic() - t0
+ assert got is held, "the freed node must be handed to the waiter"
+ assert waited < 30.0, "release must wake the waiter, not let it time out"
+
+
+def test_gives_up_after_the_timeout() -> None:
+ """Nothing frees up, so the caller is refused -- but only after waiting."""
+ pool = _pool(TIMEOUT)
+ assert pool.acquire() is not None
+ t0 = time.monotonic()
+ assert pool.acquire() is None
+ assert time.monotonic() - t0 >= TIMEOUT
+
+
+def test_every_waiter_is_served_when_nodes_come_back() -> None:
+ """One release must not wake a waiter that then loses the node again."""
+ pool = _pool(30.0, nodes=2)
+ held = [pool.acquire(), pool.acquire()]
+ assert all(n is not None for n in held)
+ got: list[DecodeNode | None] = []
+ lock = threading.Lock()
+
+ def _waiter() -> None:
+ n = pool.acquire()
+ with lock:
+ got.append(n)
+
+ threads = [threading.Thread(target=_waiter, daemon=True) for _ in range(2)]
+ for t in threads:
+ t.start()
+ for n in held:
+ time.sleep(0.05)
+ assert n is not None
+ pool.release(n)
+ for t in threads:
+ t.join(timeout=10)
+ assert len(got) == 2 and all(n is not None for n in got)
+ assert len({id(n) for n in got}) == 2, "two waiters must not share one node"
diff --git a/tests/pd_vllm/test_prefill_body_rewrite.py b/tests/pd_vllm/test_prefill_body_rewrite.py
new file mode 100644
index 0000000..00b1c77
--- /dev/null
+++ b/tests/pd_vllm/test_prefill_body_rewrite.py
@@ -0,0 +1,148 @@
+"""The prefill request must not inherit client fields that fight its overrides.
+
+``build_prefill_body`` forwards the client body verbatim apart from the handful
+of fields the split needs (``max_tokens=1``, ``stream=False``, logprobs,
+``kv_transfer_params``). Two client fields survive that rewrite and break it:
+
+``stream_options``
+ vLLM refuses ``stream_options`` unless ``stream`` is true --
+ ``ChatCompletionRequest.validate_stream_options`` is a ``mode="before"``
+ model_validator, so the request is rejected with 400 before the model is
+ even looked at. We force ``stream=False``, so any client that sent
+ ``stream_options`` got its prefill 400'd and the router turned that into a
+ 502 for the caller.
+
+``max_completion_tokens``
+ vLLM prefers it over ``max_tokens`` (``ChatCompletionRequest`` resolves
+ ``max_completion_tokens`` first when both are present), so it overrides our
+ ``max_tokens=1`` and the prefill instance decodes the client's whole output
+ length -- the split still "works", it just stops being a split.
+
+Both are unconditional in ``vllm bench serve --backend openai-chat``, which is
+why the official benchmark failed every request against a PD deployment.
+
+No GPU, no tilert, no vllm: ``pd_router`` imports fastapi/requests/uvicorn but
+nothing that needs a device, and ``build_prefill_body`` is pure.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_prefill_body_rewrite.py -v
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from tilert.pd_vllm.decode_pool import DecodeNode
+from tilert.pd_vllm.pd_router import build_prefill_body
+
+CHAT = "/v1/chat/completions"
+COMPLETIONS = "/v1/completions"
+
+
+@pytest.fixture
+def node() -> DecodeNode:
+ return DecodeNode("10.0.0.2", 5556, 5557)
+
+
+def _bench_body() -> dict:
+ """What ``vllm bench serve --backend openai-chat`` puts on the wire.
+
+ Mirrors ``vllm/benchmarks/lib/endpoint_request_func.py``: streaming with
+ usage accounting, output length via ``max_completion_tokens``.
+ """
+ return {
+ "model": "glm5.1",
+ "messages": [{"role": "user", "content": "hi"}],
+ "temperature": 0.0,
+ "max_completion_tokens": 3000,
+ "stream": True,
+ "stream_options": {"include_usage": True},
+ }
+
+
+# --------------------------------------------------------------------------- #
+# the two fields that must not survive
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize("path", [CHAT, COMPLETIONS])
+@pytest.mark.parametrize("field", ["stream_options", "max_completion_tokens"])
+def test_contradicting_field_is_dropped(path, field, node) -> None:
+ out = build_prefill_body(path, _bench_body(), node)
+ assert field not in out
+
+
+def test_stream_options_cannot_coexist_with_our_stream_false(node) -> None:
+ """The exact pair vLLM rejects with 400."""
+ out = build_prefill_body(CHAT, _bench_body(), node)
+ assert out["stream"] is False
+ assert "stream_options" not in out
+
+
+def test_prefill_decodes_exactly_one_token(node) -> None:
+ """max_tokens=1 must be the only output-length field left standing."""
+ out = build_prefill_body(CHAT, _bench_body(), node)
+ assert out["max_tokens"] == 1
+ assert "max_completion_tokens" not in out
+
+
+def test_max_completion_tokens_does_not_leak_into_max_tokens(node) -> None:
+ """Dropping it must not be implemented by copying it over max_tokens."""
+ body = _bench_body()
+ body.pop("stream_options")
+ out = build_prefill_body(CHAT, body, node)
+ assert out["max_tokens"] == 1
+
+
+@pytest.mark.parametrize("path", [CHAT, COMPLETIONS])
+def test_absent_fields_need_no_special_case(path, node) -> None:
+ """A plain non-streaming client must be rewritten without raising."""
+ out = build_prefill_body(path, {"model": "m", "prompt": "hi", "max_tokens": 128}, node)
+ assert out["max_tokens"] == 1
+ assert out["stream"] is False
+
+
+# --------------------------------------------------------------------------- #
+# everything else about the rewrite is unchanged
+# --------------------------------------------------------------------------- #
+
+
+def test_chat_asks_for_top_logprobs(node) -> None:
+ """The router reads the first token id out of chat logprobs."""
+ out = build_prefill_body(CHAT, _bench_body(), node)
+ assert out["logprobs"] is True
+ assert out["top_logprobs"] == 1
+
+
+def test_completions_logprobs_is_a_count(node) -> None:
+ out = build_prefill_body(COMPLETIONS, {"model": "m", "prompt": "hi"}, node)
+ assert out["logprobs"] == 1
+ assert "top_logprobs" not in out
+
+
+def test_kv_transfer_params_point_at_the_chosen_node(node) -> None:
+ out = build_prefill_body(CHAT, _bench_body(), node)
+ assert out["kv_transfer_params"] == {
+ "tilert_host": "10.0.0.2",
+ "tilert_ctrl_port": 5556,
+ }
+
+
+def test_unrelated_client_fields_are_forwarded(node) -> None:
+ body = _bench_body()
+ body["top_p"] = 0.95
+ body["repetition_penalty"] = 1.1
+ out = build_prefill_body(CHAT, body, node)
+ assert out["top_p"] == 0.95
+ assert out["repetition_penalty"] == 1.1
+ assert out["messages"] == body["messages"]
+
+
+def test_client_body_is_not_mutated(node) -> None:
+ """The caller reuses `body` for the decode request; it must survive intact."""
+ body = _bench_body()
+ build_prefill_body(CHAT, body, node)
+ assert body["stream"] is True
+ assert body["stream_options"] == {"include_usage": True}
+ assert body["max_completion_tokens"] == 3000
diff --git a/tests/pd_vllm/test_presentation.py b/tests/pd_vllm/test_presentation.py
new file mode 100644
index 0000000..0cc38fc
--- /dev/null
+++ b/tests/pd_vllm/test_presentation.py
@@ -0,0 +1,203 @@
+"""The reply's shape, decided once and rendered twice, without HTTP.
+
+`presentation` is what the two response paths used to spell out for themselves.
+Everything here was previously only reachable by driving a real router over a
+socket, which is why the ordering rules below were each found by a client rather
+than by a test.
+
+ CUDA_VISIBLE_DEVICES= python -m pytest tests/pd_vllm/test_presentation.py -v
+"""
+
+from __future__ import annotations
+
+import json
+
+import pytest
+
+from tilert.pd_vllm.presentation import (
+ SseWriter,
+ collect,
+ finish_reason,
+ sse_chunk,
+ sse_delta,
+)
+from tilert.pd_vllm.reply import CONTENT, REASONING, TOOL_CALL, Emission
+
+CHUNK = {"chunk_id": "cmpl-1", "model": "m", "created": 7}
+
+
+def _chunk(delta, **kw):
+ return sse_chunk(delta, **CHUNK, **kw)
+
+
+def _data(frame):
+ assert frame.startswith("data: ") and frame.endswith("\n\n")
+ body = frame[6:-2]
+ return body if body == "[DONE]" else json.loads(body)
+
+
+class _Stream:
+ """The three facts presentation reads off a ReplyStream."""
+
+ def __init__(self, held=(), stop_reason=None, tokens=3):
+ self._held = list(held)
+ self.stop_reason = stop_reason
+ self.completion_tokens = tokens
+
+ def finish(self):
+ out, self._held = self._held, [] # idempotent, as the real one is
+ return out
+
+ def finish_reason(self, from_node):
+ return "stop" if self.stop_reason is not None else from_node
+
+
+# --------------------------------------------------------------------------- #
+# collect: channels, and where logprob entries go
+# --------------------------------------------------------------------------- #
+def test_entries_follow_content_and_nothing_else():
+ """#22: `logprobs` covers `message.content`.
+
+ Reasoning text is text, not a channel entries can be attributed to.
+ """
+ got = collect(
+ [
+ Emission(REASONING, "thinking"),
+ Emission(CONTENT, "hi", [{"token": "hi"}]),
+ Emission(CONTENT, " there", [{"token": " there"}]),
+ ]
+ )
+ assert got.content == "hi there"
+ assert got.reasoning == "thinking"
+ assert [e["token"] for e in got.entries] == ["hi", " there"]
+ assert got.tool_calls == []
+
+
+# --------------------------------------------------------------------------- #
+# finish_reason: one precedence for both paths
+# --------------------------------------------------------------------------- #
+@pytest.mark.parametrize(
+ "saw_tool,stop,from_node,want,why",
+ [
+ (
+ True,
+ "END",
+ "length",
+ "tool_calls",
+ "tool_calls outranks a stop, as it does in vLLM's serving path",
+ ),
+ (False, "END", "length", "stop", "a stop outranks the node: it cannot see text"),
+ (False, None, "length", "length", "otherwise the node's own reason"),
+ (True, None, "length", "tool_calls", "tool_calls outranks that too"),
+ ],
+)
+def test_the_finish_reason_precedence(saw_tool, stop, from_node, want, why):
+ assert (
+ finish_reason(saw_tool=saw_tool, from_node=from_node, stream=_Stream(stop_reason=stop))
+ == want
+ ), why
+
+
+# --------------------------------------------------------------------------- #
+# SseWriter: the order the frames have to come out in
+# --------------------------------------------------------------------------- #
+def test_the_role_chunk_goes_out_once_and_only_before_content():
+ """#34: a request that emits nothing sends no role either."""
+ out = SseWriter(_Stream(), _chunk)
+ frames = list(out.frames([Emission(CONTENT, "a"), Emission(CONTENT, "b")]))
+ assert [_data(f)["choices"][0]["delta"] for f in frames] == [
+ {"role": "assistant"},
+ {"content": "a"},
+ {"content": "b"},
+ ]
+
+
+def test_a_reply_that_emits_nothing_sends_no_role():
+ out = SseWriter(_Stream(), _chunk)
+ assert list(out.frames([])) == []
+ assert not out.role_sent
+
+
+def test_an_emission_with_neither_text_nor_entries_makes_no_frame():
+ """The parser consuming a tag produces one; a chunk for it is noise."""
+ out = SseWriter(_Stream(), _chunk)
+ assert list(out.frames([Emission(CONTENT, "")])) == []
+
+
+def test_an_empty_emission_carrying_entries_is_still_a_frame():
+ """A token whose text a stop removed still has a logprob to report."""
+ out = SseWriter(_Stream(), _chunk)
+ frames = list(out.frames([Emission(CONTENT, "", [{"token": "x"}])]))
+ assert len(frames) == 2, "the role chunk, then the entries"
+ assert _data(frames[1])["choices"][0]["logprobs"]["content"]
+
+
+def test_a_tool_call_is_remembered_for_the_finish_reason():
+ out = SseWriter(_Stream(), _chunk)
+ list(
+ out.frames(
+ [
+ Emission(
+ TOOL_CALL,
+ "",
+ tool_call={"index": 0, "id": "c1", "name": "f", "arguments": "{}"},
+ )
+ ]
+ )
+ )
+ assert out.saw_tool, "the finish reason needs it after the loop"
+
+
+def test_held_text_is_flushed_before_anything_else():
+ """With a stop the matcher may hold len(stop)-1 characters the reply earned;
+ a marker emitted first lands as `prefix[decode error]suffix`.
+ """
+ out = SseWriter(_Stream(held=[Emission(CONTENT, "tail")]), _chunk)
+ frames = [_data(f) for f in out.flush_held()]
+ assert [f["choices"][0]["delta"] for f in frames] == [
+ {"role": "assistant"},
+ {"content": "tail"},
+ ]
+
+
+def test_flushing_twice_is_safe():
+ """Three exits need the flush and two of them return early."""
+ out = SseWriter(_Stream(held=[Emission(CONTENT, "tail")]), _chunk)
+ first = list(out.flush_held())
+ assert list(out.flush_held()) == [], "nothing is repeated"
+ assert len(first) == 2
+
+
+def test_a_stream_that_cannot_be_completed_still_ends_properly():
+ """The 200 is spent, so the only honest ending is: the reply's own text, a
+ finish_reason, a typed error event, [DONE].
+ """
+ out = SseWriter(_Stream(held=[Emission(CONTENT, "earned")]), _chunk)
+ frames = [
+ _data(f) for f in out.fail_closed({"error": "boom", "error_type": "decode_truncated"})
+ ]
+ assert frames[0]["choices"][0]["delta"] == {"role": "assistant"}
+ assert frames[1]["choices"][0]["delta"] == {"content": "earned"}
+ assert frames[2]["choices"][0]["finish_reason"] == "stop"
+ assert frames[3]["error"]["error_type"] == "decode_truncated"
+ assert frames[4] == "[DONE]"
+
+
+# --------------------------------------------------------------------------- #
+# frames
+# --------------------------------------------------------------------------- #
+def test_stop_reason_rides_only_the_closing_chunk():
+ """Where a client reads it, and where vLLM puts it."""
+ assert "stop_reason" not in _data(_chunk({"content": "x"}))["choices"][0]
+ closing = _data(_chunk({}, finish="stop", stop_reason="END"))
+ assert closing["choices"][0]["stop_reason"] == "END"
+
+
+def test_a_channel_becomes_its_own_delta_field():
+ assert sse_delta(Emission(CONTENT, "a")) == {"content": "a"}
+ assert sse_delta(Emission(REASONING, "b")) == {"reasoning_content": "b"}
+ tool = sse_delta(
+ Emission(TOOL_CALL, "", tool_call={"index": 0, "id": "c1", "name": "f", "arguments": "{}"})
+ )
+ assert tool["tool_calls"][0]["function"]["name"] == "f"
+ assert tool["tool_calls"][0]["type"] == "function"
diff --git a/tests/pd_vllm/test_receive_admission.py b/tests/pd_vllm/test_receive_admission.py
new file mode 100644
index 0000000..a431832
--- /dev/null
+++ b/tests/pd_vllm/test_receive_admission.py
@@ -0,0 +1,881 @@
+"""No RDMA write without an admission, and no buffer reuse under a live writer.
+
+The decode node's receive buffer holds one request at a time. The sender used to
+send its request metadata and then write immediately, never reading the reply --
+so a receiver that answered "busy" was overruled, and this request's KV landed
+inside the request the node was already serving. Nothing downstream detects that:
+the victim decodes from a mix of two prompts' state and answers confidently.
+
+Two invariants are pinned here, and they are separate:
+
+1. **Admission gates the write.** The sender performs zero ``transport.write``
+ calls unless it received an accept whose rid, rank and generation match what
+ it asked for.
+2. **A tenancy is not replaced while its writers are live.** A request that timed
+ out or was abandoned holds the slot until every admitted rank has stopped, so
+ the buffer is never handed over underneath an in-flight RDMA.
+
+The receive server is exercised over a real TCP socket pair with a fake transport
+and a fake profile, so the framing and the multi-rank ordering are real. No GPU,
+no mooncake, no vLLM.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_receive_admission.py -v
+"""
+
+import socket
+import threading
+import time
+import types
+
+import pytest
+
+from tilert.pd_vllm import receive_server as rs
+from tilert.pd_vllm import wire
+
+SENDER_RANKS = frozenset({0, 1})
+MAX_SEQ = 128
+
+
+# --------------------------------------------------------------------------- #
+# A ReceiveServer with the GPU/RDMA parts replaced, but the real control plane
+# --------------------------------------------------------------------------- #
+class _FakeTransport:
+ name = "mooncake"
+
+ def init(self, hostname):
+ pass
+
+ def register(self, ptr, size, dev):
+ pass
+
+ def local_meta(self):
+ return {"mooncake_session_id": "fake-session"}
+
+
+@pytest.fixture
+def server(monkeypatch):
+ """Real ReceiveServer: real socket, real framing, no torch and no RDMA."""
+ monkeypatch.setattr(rs, "make_transport", lambda name: _FakeTransport(), raising=False)
+ monkeypatch.setattr("tilert.pd_vllm.transport.make_transport", lambda name: _FakeTransport())
+ monkeypatch.setattr(
+ rs.torch, "zeros", lambda *a, **k: types.SimpleNamespace(data_ptr=lambda: 0x1000)
+ )
+ monkeypatch.setattr(rs.wire, "local_ip", lambda *a, **k: "127.0.0.1")
+
+ profile = types.SimpleNamespace(
+ name="fake",
+ layout_version=7,
+ sender_ranks=SENDER_RANKS,
+ buffer_bytes=lambda max_seq_len: 4096,
+ hello_layout=lambda base_ptr, max_seq_len: {"kv_base": base_ptr},
+ )
+ srv = rs.ReceiveServer(
+ profile,
+ max_seq_len=MAX_SEQ,
+ ctrl_port=0,
+ hostname="127.0.0.1",
+ device="cpu",
+ request_timeout=5.0,
+ )
+ yield srv
+ srv.close()
+
+
+class _Rank:
+ """One prefill rank's side of the control channel."""
+
+ def __init__(self, server: rs.ReceiveServer, rank: int):
+ self.rank = rank
+ self.sock = socket.create_connection(("127.0.0.1", server._srv.getsockname()[1]), timeout=5)
+ self.hello = wire.recv_msg(self.sock)
+
+ def request(self, rid="rid-1", seq_len=8, **extra):
+ wire.send_msg(
+ self.sock,
+ {"rid": rid, "rank": self.rank, "seq_len": seq_len, "last_prompt_token": 5, **extra},
+ )
+ return wire.recv_msg(self.sock)
+
+ def done(self, rid, generation):
+ wire.send_msg(self.sock, wire.done_msg(rid, self.rank, generation))
+
+ def close(self):
+ self.sock.close()
+
+
+def _settle():
+ """Let the server's per-connection threads reach their next lock step."""
+ for _ in range(200):
+ if not any(t.name.startswith("Thread-") and t.is_alive() for t in threading.enumerate()):
+ break
+ threading.Event().wait(0.005)
+ threading.Event().wait(0.05)
+
+
+# --------------------------------------------------------------------------- #
+# The hello advertises the control-plane version
+# --------------------------------------------------------------------------- #
+def test_hello_carries_the_protocol_version(server):
+ r = _Rank(server, 0)
+ assert r.hello["protocol_version"] == wire.PROTOCOL_VERSION
+ assert wire.PROTOCOL_VERSION >= 2, "admission arrived in v2"
+ r.close()
+
+
+def test_the_protocol_version_is_separate_from_the_layout_version(server):
+ """They version different things: one the control flow, one the buffer geometry.
+
+ Folding them together would make every layout bump look like a control-plane bump to the
+ sender's assertion.
+ """
+ r = _Rank(server, 0)
+ assert r.hello["layout_version"] == 7
+ assert r.hello["protocol_version"] != 7
+ r.close()
+
+
+# --------------------------------------------------------------------------- #
+# Admission: accept carries rid/rank/generation
+# --------------------------------------------------------------------------- #
+def test_a_first_rank_is_admitted_with_a_generation(server):
+ r = _Rank(server, 0)
+ ack = r.request()
+ assert ack["accepted"] is True
+ assert ack["rid"] == "rid-1"
+ assert ack["rank"] == 0
+ assert isinstance(ack["generation"], int)
+ r.close()
+
+
+def test_same_rid_ranks_are_admitted_in_arbitrary_order(server):
+ """Ranks connect in no particular order and only one of them creates the
+ tenancy, so every other rank must still be admitted to the same generation.
+ """
+ r1 = _Rank(server, 1)
+ a1 = r1.request()
+ r0 = _Rank(server, 0)
+ a0 = r0.request()
+ assert a1["accepted"] and a0["accepted"]
+ assert a1["generation"] == a0["generation"]
+ r0.close()
+ r1.close()
+
+
+def test_a_duplicate_rank_is_refused(server):
+ """Counting one rank twice could complete the request before every shard has
+ actually landed.
+ """
+ r0 = _Rank(server, 0)
+ ack = r0.request()
+ r0.done("rid-1", ack["generation"])
+ _settle()
+ dup = _Rank(server, 0)
+ assert dup.request()["accepted"] is False
+ r0.close()
+ dup.close()
+
+
+def test_a_second_rid_is_refused_while_the_first_is_transferring(server):
+ """The reply the sender used to ignore."""
+ held = _Rank(server, 0)
+ held.request(rid="rid-1")
+ other = _Rank(server, 0)
+ ack = other.request(rid="rid-2")
+ assert ack["accepted"] is False
+ assert ack["error"] == "busy"
+ assert ack["busy_rid"] == "rid-1"
+ held.close()
+ other.close()
+
+
+def test_an_oversized_seq_len_is_refused(server):
+ r = _Rank(server, 0)
+ ack = r.request(seq_len=MAX_SEQ + 1)
+ assert ack["accepted"] is False
+ assert "seq_len" in ack["error"]
+ r.close()
+
+
+def test_an_oversized_request_claims_no_tenancy(server):
+ """A refusal before admission must leave the slot free for the next request."""
+ r = _Rank(server, 0)
+ r.request(seq_len=MAX_SEQ + 1)
+ r.close()
+ _settle()
+ assert server.state_snapshot()["state"] == rs.FREE
+
+
+# --------------------------------------------------------------------------- #
+# Completion and generations
+# --------------------------------------------------------------------------- #
+def test_the_request_completes_when_every_sender_rank_is_done(server):
+ ranks = [_Rank(server, i) for i in sorted(SENDER_RANKS)]
+ acks = [r.request() for r in ranks]
+ gen = acks[0]["generation"]
+ for r in ranks:
+ r.done("rid-1", gen)
+ got = server.completed.get(timeout=5)
+ assert got.rid == "rid-1"
+ assert got.done_ranks == set(SENDER_RANKS)
+ assert got.state == rs.COMPLETE
+ for r in ranks:
+ r.close()
+
+
+def test_a_partially_done_request_does_not_complete(server):
+ r0 = _Rank(server, 0)
+ ack = r0.request()
+ r0.done("rid-1", ack["generation"])
+ _settle()
+ assert server.completed.empty()
+ assert server.state_snapshot()["state"] == rs.TRANSFERRING
+ r0.close()
+
+
+def test_a_rejected_straggler_cannot_disturb_the_current_tenancy(server):
+ """A rank turned away at admission must not affect whoever holds the slot,
+ even if it goes on to send a done anyway.
+ """
+ ranks = [_Rank(server, i) for i in sorted(SENDER_RANKS)]
+ acks = [r.request() for r in ranks]
+ gen = acks[0]["generation"]
+ for r in ranks:
+ r.done("rid-1", gen)
+ first = server.completed.get(timeout=5)
+ server.release("rid-1")
+ _settle()
+
+ # A second tenancy takes the slot ...
+ nxt = [_Rank(server, i) for i in sorted(SENDER_RANKS)]
+ new_acks = [r.request(rid="rid-2") for r in nxt]
+ assert new_acks[0]["generation"] != first.generation
+
+ # ... and a straggler from the first one is refused and then ignored.
+ straggler = _Rank(server, 0)
+ assert straggler.request(rid="rid-1")["accepted"] is False
+ wire.send_msg(straggler.sock, wire.done_msg("rid-1", 0, first.generation))
+ _settle()
+ snap = server.state_snapshot()
+ assert snap["rid"] == "rid-2"
+ assert snap["state"] == rs.TRANSFERRING
+ for r in ranks + nxt + [straggler]:
+ r.close()
+
+
+def test_a_done_from_a_superseded_generation_of_the_same_rid_is_ignored(server):
+ """The rid alone is not enough to identify a tenancy.
+
+ ``rid`` is derived from the vLLM request id, so a retry of the same request
+ carries the SAME rid against a NEW receive-buffer tenancy. A done left over
+ from the previous attempt therefore matches on rid and must be rejected on
+ generation -- otherwise it credits a rank that has written nothing for the
+ current tenancy, and the request can be handed to the engine before that
+ rank's shard has actually landed.
+
+ Driven through the real socket path: the extra rank-0 connection is admitted
+ to generation 1 (nothing has marked rank 0 done yet) and deliberately outlives
+ it, which is how a stale done reaches a live tenancy.
+ """
+ a = _Rank(server, 0)
+ gen1 = a.request()["generation"]
+ stale = _Rank(server, 0) # second rank-0 channel
+ assert stale.request()["generation"] == gen1
+ c = _Rank(server, 1)
+ c.request()
+
+ a.done("rid-1", gen1)
+ c.done("rid-1", gen1)
+ assert server.completed.get(timeout=5).done_ranks == set(SENDER_RANKS)
+ a.close()
+ c.close()
+ server.release("rid-1")
+ _settle()
+
+ # The retry announces itself, as /pd/decode does in production: release()
+ # leaves a tombstone so a straggler from the finished attempt cannot open a
+ # tenancy nobody is waiting for, and only a consumer saying "I want this
+ # rid" distinguishes the retry from that straggler.
+ server.expect("rid-1")
+
+ # Same rid, new tenancy: rank 0 is admitted but has NOT written yet, while
+ # rank 1 has finished. The tenancy is one rank short of complete.
+ d = _Rank(server, 0)
+ gen2 = d.request()["generation"]
+ assert gen2 != gen1
+ e = _Rank(server, 1)
+ e.request()
+ e.done("rid-1", gen2)
+ _settle()
+ assert server.completed.empty(), "precondition: gen2 is not complete yet"
+
+ # The leftover channel reports done for the OLD generation. Counted, it
+ # supplies the missing rank and the request is handed to the engine while
+ # rank 0's shard for THIS tenancy has never been written.
+ stale.done("rid-1", gen1)
+ _settle()
+
+ assert server.completed.empty(), (
+ "a stale done completed the tenancy: the engine would decode from a "
+ "buffer whose rank-0 shard was never written for this request"
+ )
+ snap = server.state_snapshot()
+ assert snap["generation"] == gen2
+ assert snap["state"] == rs.TRANSFERRING
+ for r in (stale, d, e):
+ r.close()
+
+
+def test_generations_are_never_reused(server):
+ seen = set()
+ for i in range(3):
+ r = _Rank(server, 0)
+ ack = r.request(rid=f"rid-{i}")
+ assert ack["accepted"], ack
+ seen.add(ack["generation"])
+ r.close()
+ _settle()
+ server.release(f"rid-{i}") # scoped to the rid this round created
+ _settle()
+ assert len(seen) == 3
+
+
+# --------------------------------------------------------------------------- #
+# The tenancy is not replaced under a live writer
+# --------------------------------------------------------------------------- #
+def test_release_under_a_live_writer_cancels_instead_of_freeing(server):
+ """The abandon path (rejected request, drained stale entry) must not hand the
+ buffer over while a rank is still writing into it.
+ """
+ r0 = _Rank(server, 0)
+ r0.request()
+ _settle()
+ server.release("rid-1")
+ snap = server.state_snapshot()
+ assert snap["state"] == rs.CANCELLING
+ assert snap["rid"] == "rid-1"
+ assert snap["active_writers"] == 1
+ r0.close()
+
+
+def test_a_cancelling_tenancy_still_refuses_a_new_rid(server):
+ r0 = _Rank(server, 0)
+ r0.request()
+ _settle()
+ server.release("rid-1")
+ other = _Rank(server, 0)
+ assert other.request(rid="rid-2")["accepted"] is False
+ r0.close()
+ other.close()
+
+
+def test_the_slot_frees_once_the_writer_goes_away(server):
+ """Every sender connection carries request_timeout as its socket timeout, so
+ the drain is bounded -- a disconnect just gets there sooner.
+ """
+ r0 = _Rank(server, 0)
+ r0.request()
+ _settle()
+ server.release("rid-1")
+ assert server.state_snapshot()["state"] == rs.CANCELLING
+ r0.close()
+ _settle()
+ assert server.state_snapshot()["state"] == rs.FREE
+
+
+def test_a_new_rid_is_admitted_after_the_drain_completes(server):
+ r0 = _Rank(server, 0)
+ r0.request()
+ _settle()
+ server.release("rid-1")
+ r0.close()
+ _settle()
+ nxt = _Rank(server, 0)
+ assert nxt.request(rid="rid-2")["accepted"] is True
+ nxt.close()
+
+
+def test_a_dead_writer_releases_its_claim(server):
+ """A sender that dies mid-RDMA must not pin the slot forever."""
+ r0 = _Rank(server, 0)
+ r0.request()
+ _settle()
+ assert server.state_snapshot()["active_writers"] == 1
+ r0.close()
+ _settle()
+ assert server.state_snapshot()["active_writers"] == 0
+
+
+# --------------------------------------------------------------------------- #
+# Timeout: bounded, but never at the cost of the live-writer rule
+# --------------------------------------------------------------------------- #
+def test_a_tenancy_whose_senders_all_died_ages_out(server):
+ """Otherwise the node is out of service for good.
+
+ Every rank is admitted and then dies without reporting done, so the tenancy
+ stays TRANSFERRING with no writers. Nothing will call ``release()`` for it --
+ the router gave up on the prefill leg, so ``/pd/decode`` never arrives for
+ that rid. The age-out is the only way back.
+ """
+ server.request_timeout = 0.05
+ r0 = _Rank(server, 0)
+ r0.request()
+ r0.close()
+ _settle()
+ snap = server.state_snapshot()
+ assert snap["state"] == rs.TRANSFERRING and snap["active_writers"] == 0
+
+ threading.Event().wait(0.1)
+ nxt = _Rank(server, 0)
+ assert nxt.request(rid="rid-2")["accepted"] is True
+ nxt.close()
+
+
+def _aged(state, writers, server):
+ """A tenancy whose first connection is far older than request_timeout."""
+ return rs.ReceivedRequest(
+ rid="rid-1",
+ seq_len=8,
+ last_prompt_token=5,
+ first_token_id=None,
+ sampling=None,
+ state=state,
+ active_writers=writers,
+ t_first_conn=time.time() - server.request_timeout - 60,
+ )
+
+
+@pytest.mark.parametrize("state", [rs.RESERVED, rs.TRANSFERRING])
+def test_the_age_out_never_overrides_a_live_writer(server, state):
+ """The original code aged out on wall clock ALONE, which is exactly how a
+ request got replaced while its ranks were still writing. No amount of
+ elapsed time may make that reusable.
+
+ Tested on ``_reusable`` directly rather than over sockets: ``request_timeout``
+ is also each sender connection's socket timeout, so shrinking it to force an
+ age-out kills the very live writer the case is about.
+ """
+ assert server._reusable(_aged(state, writers=1, server=server)) is False
+
+
+@pytest.mark.parametrize("state", [rs.RESERVED, rs.TRANSFERRING])
+def test_an_aged_out_tenancy_with_no_writer_is_reusable(server, state):
+ assert server._reusable(_aged(state, writers=0, server=server)) is True
+
+
+@pytest.mark.parametrize(
+ "state,writers,expected",
+ [
+ (rs.COMPLETE, 0, True), # finished; the decode server owns it now
+ (rs.COMPLETE, 1, True), # all ranks reported done, sockets still closing
+ (rs.CANCELLING, 0, True), # abandoned and drained
+ (rs.CANCELLING, 1, False), # abandoned, still draining
+ ],
+)
+def test_the_reuse_rule_by_state(server, state, writers, expected):
+ cur = rs.ReceivedRequest(
+ rid="rid-1",
+ seq_len=8,
+ last_prompt_token=5,
+ first_token_id=None,
+ sampling=None,
+ state=state,
+ active_writers=writers,
+ t_first_conn=time.time(),
+ )
+ assert server._reusable(cur) is expected
+
+
+def test_a_fresh_tenancy_is_not_aged_out_immediately(server):
+ """The age-out must not turn into "first come, first served, briefly"."""
+ r0 = _Rank(server, 0)
+ r0.request(rid="rid-1")
+ r0.close()
+ _settle()
+ # request_timeout is 5.0s in the fixture, so this is still young.
+ other = _Rank(server, 0)
+ assert other.request(rid="rid-2")["accepted"] is False
+ other.close()
+
+
+# --------------------------------------------------------------------------- #
+# The sender side: zero RDMA writes on any rejection
+# --------------------------------------------------------------------------- #
+class _RecordingTransport:
+ name = "mooncake"
+
+ def __init__(self):
+ self.writes = []
+
+ def write(self, hello, srcs, dsts, lens):
+ self.writes.append((srcs, dsts, lens))
+
+
+def _import_connector(monkeypatch):
+ """Import ``prefill_connector`` without a real vLLM.
+
+ It subclasses vLLM's connector base at module level, so a serve-only
+ environment cannot import it -- and ``importorskip`` would SKIP the most
+ important tests in this file (zero RDMA writes on a rejection) exactly where
+ they run: CI has no vLLM. So the base classes are stubbed instead, the way
+ the engine-free tests in this directory stub ``tilert``.
+
+ The stubs are installed through ``monkeypatch`` rather than at import time so
+ the fake ``vllm`` entry disappears at teardown: left in ``sys.modules`` it
+ would make ``pytest.importorskip("vllm")`` succeed elsewhere in the session
+ (test_oai_parser.py) and those tests would fail on a module that has no
+ ``vllm.parser``.
+ """
+ import dataclasses
+ import sys
+
+ base_path = "vllm.distributed.kv_transfer.kv_connector.v1.base"
+ try: # a real vLLM (the router's own environment) is used as-is
+ __import__(base_path)
+ except Exception:
+ for name in (
+ "vllm",
+ "vllm.distributed",
+ "vllm.distributed.kv_transfer",
+ "vllm.distributed.kv_transfer.kv_connector",
+ "vllm.distributed.kv_transfer.kv_connector.v1",
+ ):
+ if name not in sys.modules:
+ monkeypatch.setitem(sys.modules, name, types.ModuleType(name))
+ base = types.ModuleType(base_path)
+
+ class KVConnectorBase_V1: # noqa: N801 (mirrors vLLM's own name)
+ def __init__(self, vllm_config=None, role=None, kv_cache_config=None):
+ pass
+
+ @dataclasses.dataclass
+ class KVConnectorMetadata:
+ pass
+
+ class SupportsHMA:
+ pass
+
+ base.KVConnectorBase_V1 = KVConnectorBase_V1
+ base.KVConnectorMetadata = KVConnectorMetadata
+ base.SupportsHMA = SupportsHMA
+ monkeypatch.setitem(sys.modules, base_path, base)
+
+ from tilert.pd_vllm import prefill_connector as pc
+
+ return pc
+
+
+def _sender(monkeypatch, replies, *, protocol_version=wire.PROTOCOL_VERSION):
+ """Drive ``TileRTConnector._send`` against a scripted receiver.
+
+ Returns (transport, sent_messages). ``replies`` are the messages the fake
+ receiver returns after the hello, in order.
+ """
+ pc = _import_connector(monkeypatch)
+
+ hello = {
+ "magic": wire.MAGIC,
+ "protocol_version": protocol_version,
+ "layout_version": 7,
+ "transport": "mooncake",
+ "max_seq_len": MAX_SEQ,
+ "busy": False,
+ "kv_base": 0x1000,
+ }
+ inbox = [hello, *replies]
+ sent = []
+
+ monkeypatch.setattr(pc.wire, "recv_msg", lambda conn: inbox.pop(0))
+ monkeypatch.setattr(pc.wire, "send_msg", lambda conn, obj: sent.append(obj))
+
+ class _Sock:
+ def __init__(self, *a, **k):
+ pass
+
+ def setsockopt(self, *a):
+ pass
+
+ def settimeout(self, *a):
+ pass
+
+ def connect(self, *a):
+ pass
+
+ def close(self):
+ pass
+
+ monkeypatch.setattr("socket.socket", _Sock)
+
+ transport = _RecordingTransport()
+ conn_obj = object.__new__(pc.TileRTConnector)
+ conn_obj._transport = transport
+ conn_obj._tp_rank = 0
+ # _send declares this sender's remaining retry budget so the receiver can
+ # size a tombstone to outlast it. The retry test overrides it below.
+ conn_obj._admission_attempts = pc._ADMISSION_ATTEMPTS
+ conn_obj._staging = types.SimpleNamespace(data_ptr=lambda: 0x2000)
+ conn_obj._profile = types.SimpleNamespace(
+ layout_version=7, rdma_plan=lambda hello, sections, rank, seq, base: ([1], [2], [3])
+ )
+
+ meta = pc._ReqMeta(
+ req_id="r",
+ rid="rid-1",
+ num_tokens=8,
+ last_prompt_token=5,
+ block_ids_per_group=[],
+ tilert_host="127.0.0.1",
+ tilert_ctrl_port=1,
+ )
+ conn_obj._send({"meta": meta, "sections": {"seq": 8}, "seq": 8})
+ return transport, sent
+
+
+def test_the_sender_writes_after_an_accept(monkeypatch):
+ transport, sent = _sender(monkeypatch, [wire.accept_msg("rid-1", 0, 12)])
+ assert len(transport.writes) == 1
+ # ... and the done echoes the generation it was admitted under.
+ assert sent[-1] == {"done": True, "rid": "rid-1", "rank": 0, "generation": 12}
+
+
+def test_the_sender_writes_nothing_on_a_busy_rejection(monkeypatch):
+ """The original bug: this reply existed and was never read."""
+ transport, sent = _sender(monkeypatch, [wire.reject_msg("busy", busy_rid="other")])
+ assert transport.writes == []
+ assert not any("done" in m for m in sent)
+
+
+def test_the_sender_writes_nothing_on_a_seq_len_rejection(monkeypatch):
+ transport, _ = _sender(monkeypatch, [wire.reject_msg("seq_len exceeds max_seq_len")])
+ assert transport.writes == []
+
+
+@pytest.mark.parametrize(
+ "ack",
+ [
+ {}, # empty
+ {"accepted": False}, # bare refusal
+ {"rid": "rid-1", "rank": 0, "generation": 1}, # no accepted flag
+ {"accepted": True, "rid": "other", "rank": 0, "generation": 1},
+ {"accepted": True, "rid": "rid-1", "rank": 3, "generation": 1},
+ {"accepted": True, "rid": "rid-1", "rank": 0}, # no generation
+ {"accepted": True, "rid": "rid-1", "rank": 0, "generation": "x"},
+ ],
+)
+def test_only_a_fully_matching_admission_permits_a_write(monkeypatch, ack):
+ """Checked field by field, not just for an ``error`` key: an admission for a
+ different rid or rank is an admission for a different tenancy, and writing on
+ it corrupts exactly the same way a busy rejection would.
+ """
+ transport, _ = _sender(monkeypatch, [ack])
+ assert transport.writes == []
+
+
+def test_a_protocol_version_mismatch_fails_before_writing(monkeypatch):
+ """A v1 receiver never sends an accept, so waiting for one would hang every
+ request; failing loud names the actual problem.
+ """
+ with pytest.raises(AssertionError, match="protocol mismatch"):
+ _sender(monkeypatch, [wire.accept_msg("rid-1", 0, 1)], protocol_version=1)
+
+
+def test_a_layout_version_mismatch_still_fails_before_writing(monkeypatch):
+ pc = _import_connector(monkeypatch)
+
+ hello = {
+ "magic": wire.MAGIC,
+ "protocol_version": wire.PROTOCOL_VERSION,
+ "layout_version": 99,
+ "transport": "mooncake",
+ "max_seq_len": MAX_SEQ,
+ "busy": False,
+ }
+ monkeypatch.setattr(pc.wire, "recv_msg", lambda conn: hello)
+ monkeypatch.setattr(pc.wire, "send_msg", lambda conn, obj: None)
+
+ class _Sock:
+ def __init__(self, *a, **k):
+ pass
+
+ def setsockopt(self, *a):
+ pass
+
+ def settimeout(self, *a):
+ pass
+
+ def connect(self, *a):
+ pass
+
+ def close(self):
+ pass
+
+ monkeypatch.setattr("socket.socket", _Sock)
+ transport = _RecordingTransport()
+ conn_obj = object.__new__(pc.TileRTConnector)
+ conn_obj._transport = transport
+ conn_obj._tp_rank = 0
+ # _send declares this sender's remaining retry budget so the receiver can
+ # size a tombstone to outlast it. The retry test overrides it below.
+ conn_obj._admission_attempts = pc._ADMISSION_ATTEMPTS
+ conn_obj._staging = types.SimpleNamespace(data_ptr=lambda: 0x2000)
+ conn_obj._profile = types.SimpleNamespace(layout_version=7)
+ meta = pc._ReqMeta(
+ req_id="r",
+ rid="rid-1",
+ num_tokens=8,
+ last_prompt_token=5,
+ block_ids_per_group=[],
+ tilert_host="127.0.0.1",
+ tilert_ctrl_port=1,
+ )
+ with pytest.raises(AssertionError, match="layout version"):
+ conn_obj._send({"meta": meta, "sections": {"seq": 8}, "seq": 8})
+ assert transport.writes == []
+
+
+# --------------------------------------------------------------------------- #
+# A rejected admission must not silently drop the shard (codex, PR #40)
+# --------------------------------------------------------------------------- #
+def _retrying_sender(monkeypatch, reply_sequence, attempts=3):
+ """Drive ``_send_with_retry`` against a receiver that answers in sequence.
+
+ Returns (transport, attempt_count). Sleeping is stubbed out so the backoff
+ does not slow the suite.
+ """
+ pc = _import_connector(monkeypatch)
+
+ hello = {
+ "magic": wire.MAGIC,
+ "protocol_version": wire.PROTOCOL_VERSION,
+ "layout_version": 7,
+ "transport": "mooncake",
+ "max_seq_len": MAX_SEQ,
+ "busy": False,
+ "kv_base": 0x1000,
+ }
+ replies = list(reply_sequence)
+ calls = {"n": 0}
+
+ def recv(conn):
+ # Each attempt reopens the channel: hello, then that attempt's verdict.
+ if calls["hello_pending"]:
+ calls["hello_pending"] = False
+ return hello
+ calls["hello_pending"] = True
+ return replies.pop(0)
+
+ calls["hello_pending"] = True
+ monkeypatch.setattr(pc.wire, "recv_msg", recv)
+ monkeypatch.setattr(pc.wire, "send_msg", lambda conn, obj: None)
+ monkeypatch.setattr(
+ pc._time if hasattr(pc, "_time") else pc, "sleep", lambda s: None, raising=False
+ )
+
+ class _Sock:
+ def __init__(self, *a, **k):
+ calls["n"] += 1
+
+ def setsockopt(self, *a):
+ pass
+
+ def settimeout(self, *a):
+ pass
+
+ def connect(self, *a):
+ pass
+
+ def close(self):
+ pass
+
+ monkeypatch.setattr("socket.socket", _Sock)
+ monkeypatch.setattr("time.sleep", lambda s: None)
+
+ transport = _RecordingTransport()
+ conn_obj = object.__new__(pc.TileRTConnector)
+ conn_obj._transport = transport
+ conn_obj._tp_rank = 0
+ # _send declares this sender's remaining retry budget so the receiver can
+ # size a tombstone to outlast it. The retry test overrides it below.
+ conn_obj._admission_attempts = pc._ADMISSION_ATTEMPTS
+ conn_obj._staging = types.SimpleNamespace(data_ptr=lambda: 0x2000)
+ conn_obj._admission_attempts = attempts
+ conn_obj._profile = types.SimpleNamespace(
+ layout_version=7, rdma_plan=lambda hello, sections, rank, seq, base: ([1], [2], [3])
+ )
+ meta = pc._ReqMeta(
+ req_id="r",
+ rid="rid-1",
+ num_tokens=8,
+ last_prompt_token=5,
+ block_ids_per_group=[],
+ tilert_host="127.0.0.1",
+ tilert_ctrl_port=1,
+ )
+ conn_obj._send_with_retry({"meta": meta, "sections": {"seq": 8}, "seq": 8})
+ return transport, calls["n"]
+
+
+def test_a_busy_slot_is_retried_and_then_written(monkeypatch):
+ """A rank turned away while a previous transfer drains would otherwise drop
+ its shard for good, and nothing tells the router: the prefill response still
+ succeeds and /pd/decode waits out its whole kv_transfer_timeout.
+ """
+ transport, attempts = _retrying_sender(
+ monkeypatch,
+ [
+ wire.reject_msg("busy", busy_rid="other"),
+ wire.accept_msg("rid-1", 0, 9),
+ ],
+ )
+ assert len(transport.writes) == 1
+ assert attempts == 2
+
+
+def test_a_draining_slot_is_retried(monkeypatch):
+ transport, attempts = _retrying_sender(
+ monkeypatch,
+ [
+ wire.reject_msg("cancelling", rid="rid-1"),
+ wire.accept_msg("rid-1", 0, 9),
+ ],
+ )
+ assert len(transport.writes) == 1
+
+
+def test_retries_are_bounded_and_never_write(monkeypatch):
+ """Still zero writes when every attempt is refused -- the no-write rule is
+ not traded away for promptness.
+ """
+ transport, attempts = _retrying_sender(monkeypatch, [wire.reject_msg("busy")] * 3, attempts=3)
+ assert transport.writes == []
+ assert attempts == 3
+
+
+@pytest.mark.parametrize(
+ "reject",
+ [
+ wire.reject_msg("seq_len exceeds max_seq_len"),
+ wire.reject_msg("duplicate_rank", rank=0),
+ ],
+)
+def test_a_permanent_rejection_is_not_retried(monkeypatch, reject):
+ """These are about THIS request and would be refused again; retrying only
+ delays the failure.
+ """
+ transport, attempts = _retrying_sender(monkeypatch, [reject], attempts=3)
+ assert transport.writes == []
+ assert attempts == 1
+
+
+def test_the_request_declares_the_senders_admission_budget(monkeypatch):
+ """The receiver sizes a tombstone from this, and cannot derive it itself.
+
+ Its own `request_timeout` bounds one connection; each retry opens a new
+ one, so only the sender knows how long this rid may keep coming back.
+ """
+ _, sent = _sender(monkeypatch, [wire.accept_msg("rid-1", 0, 12)])
+ req = sent[0]
+ assert "admission_window_s" in req, req
+ # The default 5 attempts back off 0.2 + 0.4 + 0.8 + 1.6.
+ assert abs(req["admission_window_s"] - 3.0) < 1e-9, req
diff --git a/tests/pd_vllm/test_reply.py b/tests/pd_vllm/test_reply.py
new file mode 100644
index 0000000..92257a2
--- /dev/null
+++ b/tests/pd_vllm/test_reply.py
@@ -0,0 +1,613 @@
+"""``ReplyStream``: the five transformations, driven directly.
+
+No HTTP, no asyncio, no vLLM — a list of token ids in, a list of emissions out.
+That is the point of the module: the invariants that used to have to hold at
+every point either response path emitted something are properties of one object
+here, so they can be stated once and checked once.
+
+The properties that matter, and why:
+
+* **A token's logprob entry travels with the emission carrying its text.** Every
+ earlier attempt attached entries at the emit site, and every emit site added
+ was another chance to attach the wrong ones or none.
+* **``completion_tokens`` counts what the reply contains.** A stop removes text,
+ and the tokens whose text it removed are not in the reply.
+* **Chunking does not matter.** The same tokens fed one at a time, in pairs, or
+ all at once produce the same emissions — which is what makes the streaming and
+ non-streaming replies to one request agree.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest tests/pd_vllm/test_reply.py -v
+"""
+
+import pytest
+
+from tilert.pd_vllm.logprobs import (
+ LOGPROB_UNAVAILABLE,
+ LogprobsRequest,
+)
+from tilert.pd_vllm.reply import (
+ CONTENT,
+ REASONING,
+ TOOL_CALL,
+ ReplyStream,
+ as_logprobs,
+)
+
+# "alpha " / "STOP" / " beta", plus the pieces the interesting cases need:
+# 13+14 spell the stop across two tokens, 17 carries text AND the stop, and 15
+# is a special that spells out only when specials are kept.
+_VOCAB = {
+ 10: "alpha ",
+ 11: "STOP",
+ 12: " beta",
+ 13: "ST",
+ 14: "OP",
+ 16: "abc",
+ 17: " abcSTOP",
+ 20: "|",
+ 21: "R",
+ 22: "|C",
+}
+_SPECIAL = {15: "<|s|>"}
+# 18 alone decodes to the replacement character; 18+19 decode to "。STOP".
+_PAIRS = {(18, 19): "。STOP"}
+
+
+class _Tok:
+ def decode(self, ids, skip_special_tokens=False):
+ out, k, ids = [], 0, list(ids)
+ while k < len(ids):
+ pair = tuple(ids[k : k + 2])
+ if pair in _PAIRS:
+ out.append(_PAIRS[pair])
+ k += 2
+ continue
+ i = ids[k]
+ if i in _SPECIAL:
+ if not skip_special_tokens:
+ out.append(_SPECIAL[i])
+ elif i in (18, 19):
+ out.append("�")
+ else:
+ out.append(_VOCAB.get(i, ""))
+ k += 1
+ return "".join(out)
+
+
+def _drive(ids, *, stop=(), include=False, logprobs=False, session=None, batch=None):
+ """Run the assembler and return (emissions, assembler)."""
+ req = LogprobsRequest(top_n=1) if logprobs else None
+ asm = ReplyStream(
+ _Tok(), stop=stop, include_stop_in_output=include, parser_session=session, logprobs_req=req
+ )
+ ems = []
+ groups = batch or [[i] for i in ids]
+ for g in groups:
+ ems += asm.push(
+ g, [-0.5] * len(g) if logprobs else None, [[(t, -0.5)] for t in g] if logprobs else None
+ )
+ return ems + asm.finish(), asm
+
+
+def _content(ems):
+ return "".join(e.text for e in ems if e.channel == CONTENT)
+
+
+def _entries(ems):
+ return [x for e in ems if e.channel == CONTENT for x in e.logprobs]
+
+
+# --------------------------------------------------------------------------- #
+# Text: where the cut lands
+# --------------------------------------------------------------------------- #
+def test_no_stop_passes_everything_through():
+ ems, asm = _drive([10, 12])
+ assert _content(ems) == "alpha beta"
+ assert asm.completion_tokens == 2
+ assert asm.stop_reason is None
+
+
+@pytest.mark.parametrize(
+ "ids,want_text,why",
+ [
+ ([10, 11, 12], "alpha ", "the stop is its own token"),
+ ([16, 13, 14], "abc", "the stop spans two tokens"),
+ ([17], " abc", "the stop begins inside a token, whose prefix survives"),
+ ([16, 15, 11], "abc", "a stripped special sits before the stop"),
+ ([16, 18, 19], "abc", "a byte fragment is absorbed by the stop"),
+ ],
+)
+def test_where_a_stop_cuts(ids, want_text, why):
+ stop = ["。STOP"] if ids == [16, 18, 19] else ["STOP"]
+ ems, asm = _drive(ids, stop=stop)
+ assert _content(ems) == want_text, why
+ assert asm.stop_reason == stop[0]
+
+
+@pytest.mark.parametrize(
+ "ids,stop,want,why",
+ [
+ (
+ [16, 13, 14],
+ ["STOP"],
+ 3,
+ "13 and 14 spell the stop; their text is gone " "and they still ran",
+ ),
+ ([17], ["STOP"], 1, "the stop begins inside the only token"),
+ ([16, 18, 19], ["。STOP"], 3, "a byte fragment absorbed by the stop"),
+ ([16, 15, 11], ["STOP"], 3, "a stripped special contributes no text"),
+ ([10, 12], ["ZZZZ"], 2, "no match at all"),
+ ],
+)
+def test_the_cut_does_not_reduce_the_token_count(ids, stop, want, why):
+ """`completion_tokens` counts what ran, not what came back.
+
+ vLLM's is `len()` of its detokeniser's UNtruncated id list -- measured on
+ 0.25.1, not assumed: a stop of `" abcSTOP"` leaves `output_text` ending at
+ `"xy"` while still reporting all twelve ids. Counting only the tokens whose
+ text survived would under-report a cost the client is billed for.
+ """
+ ems, asm = _drive(ids, stop=stop)
+ assert asm.completion_tokens == want, why
+
+
+def test_tokens_arriving_after_the_stop_are_not_counted():
+ """The reply ended at the stop.
+
+ The node kept generating only because it cannot see text, and the router cancels it -- those
+ tokens are not part of what was asked for, and in vLLM they would never have been generated.
+ """
+ ems, asm = _drive([10, 11, 12, 12, 12], stop=["STOP"])
+ assert _content(ems) == "alpha "
+ assert asm.completion_tokens == 2, "10 and the token that completed the stop"
+
+
+def test_include_stop_str_in_output_keeps_the_stop_and_its_token():
+ ems, asm = _drive([10, 11, 12], stop=["STOP"], include=True)
+ assert _content(ems) == "alpha STOP"
+
+
+def test_nothing_is_held_back_when_the_stop_stays_in_the_output():
+ """There is nothing to remove from released text, so nothing to hold it for.
+
+ vLLM computes its hold-back under the same condition. Without this, every
+ delta of an `include_stop_str_in_output` request is delayed by up to
+ len(stop)-1 characters for no reason.
+ """
+ from tilert.pd_vllm.stop_strings import StopWindow
+
+ kept = StopWindow(["STOP"], True)
+ kept.push("hi ST")
+ assert kept.take() == "hi ST"
+ cut = StopWindow(["STOP"], False)
+ cut.push("hi ST")
+ assert cut.take() == "hi"
+
+
+def test_an_unmatched_stop_releases_the_held_tail():
+ """The tracker holds back what could still become a stop; it must come out."""
+ ems, asm = _drive([10, 12], stop=["ZZZZ"])
+ assert _content(ems) == "alpha beta"
+ assert asm.stop_reason is None
+
+
+def test_a_byte_fragment_whose_character_survives_is_kept():
+ """Byte-level BPE splits "。" across both ids, so both contributed it."""
+ ems, asm = _drive([16, 18, 19], stop=["STOP"])
+ assert _content(ems) == "abc。", "the stop is cut; the character is not"
+
+
+def test_the_finish_reason_is_overridden_only_by_a_stop():
+ _, asm = _drive([10, 12], stop=["ZZZZ"])
+ assert asm.finish_reason("length") == "length"
+ _, asm = _drive([10, 11], stop=["STOP"])
+ assert asm.finish_reason("length") == "stop"
+
+
+# --------------------------------------------------------------------------- #
+# Logprobs: one entry per token the reply contains, no more and no fewer
+# --------------------------------------------------------------------------- #
+@pytest.mark.parametrize(
+ "ids,stop,why",
+ [
+ ([10, 12], (), "no stop"),
+ ([10, 11, 12], ("STOP",), "the stop is its own token"),
+ ([16, 13, 14], ("STOP",), "the stop spans two tokens"),
+ ([17], ("STOP",), "the stop begins inside a token"),
+ ([16, 15, 11], ("STOP",), "a stripped special is kept"),
+ ([15, 11], ("STOP",), "a kept token produced no visible text at all"),
+ ([10, 12], ("ZZZZ",), "the tail is released at the end"),
+ ],
+)
+def test_entries_match_the_token_count(ids, stop, why):
+ ems, asm = _drive(ids, stop=stop, logprobs=True)
+ assert len(_entries(ems)) == asm.completion_tokens, (
+ f"{why}: {len(_entries(ems))} entries for " f"{asm.completion_tokens} tokens"
+ )
+
+
+def test_no_entries_when_logprobs_were_not_requested():
+ ems, _ = _drive([10, 12], logprobs=False)
+ assert _entries(ems) == []
+
+
+def test_the_first_token_takes_its_logprob_from_the_prefill():
+ """The decode node echoed it rather than sampling, so it sends null."""
+ asm = ReplyStream(
+ _Tok(), logprobs_req=LogprobsRequest(top_n=1), first_token_logprob=(-0.125, [(10, -0.125)])
+ )
+ ems = asm.push([10, 12], [None, -0.5], [[], [(12, -0.5)]]) + asm.finish()
+ entries = _entries(ems)
+ assert len(entries) == 2
+ assert entries[0]["logprob"] == pytest.approx(-0.125)
+
+
+def test_the_first_tokens_candidates_come_from_the_prefill_too():
+ """The row describes the PROMPT's last distribution, not the decode token.
+
+ The decode node sends an empty row for that position along with the null
+ logprob. Filling the logprob from the prefill entry but leaving the row
+ empty would report a token with no alternatives; taking the row from the
+ decode position would report the alternatives of a distribution that never
+ produced this token.
+ """
+ asm = ReplyStream(
+ _Tok(), logprobs_req=LogprobsRequest(top_n=1), first_token_logprob=(-0.125, [(16, -0.125)])
+ )
+ ems = asm.push([10, 12], [None, -0.5], [[], [(12, -0.5)]]) + asm.finish()
+ row = _entries(ems)[0]["top_logprobs"]
+ assert [c["token"] for c in row] == ["abc"], "prefill's alternative id (16)"
+ assert row[0]["logprob"] == pytest.approx(-0.125)
+
+
+def test_the_first_token_without_a_prefill_entry_takes_the_sentinel():
+ """No prefill value available -> the documented -9999.0, never a fake one.
+
+ Every other position is still correct, so one sentinel entry beats failing
+ the whole completion.
+ """
+ ems, _ = _drive([10, 12], logprobs=True)
+ asm = ReplyStream(_Tok(), logprobs_req=LogprobsRequest(top_n=1))
+ ems = asm.push([10, 12], [None, -0.5], [[], [(12, -0.5)]]) + asm.finish()
+ assert _entries(ems)[0]["logprob"] == LOGPROB_UNAVAILABLE
+
+
+def test_the_prefill_value_does_not_leak_onto_a_later_content_token():
+ """It belongs to position 0 and to no other token.
+
+ When position 0 goes to `reasoning` the client never sees its entry, and it
+ is dropped. What must not happen is the value landing on whichever token
+ reached `content` first -- that would report the prompt's last distribution
+ as if it described the reply's first content token.
+ """
+ asm = ReplyStream(
+ _Tok(),
+ parser_session=_ReasoningThenContent(),
+ logprobs_req=LogprobsRequest(top_n=1),
+ first_token_logprob=(-0.125, [(16, -0.125)]),
+ )
+ # 10 -> "alpha " (reasoning), 20 -> "|", 12 -> " beta" (content).
+ ems = (
+ asm.push([10, 20, 12], [None, -0.25, -0.5], [[], [(20, -0.25)], [(12, -0.5)]])
+ + asm.finish()
+ )
+ entries = _entries(ems)
+ assert [e["logprob"] for e in entries] == [
+ pytest.approx(-0.5)
+ ], "only the content token has an entry, with its own value"
+
+
+def test_an_entry_can_ride_an_emission_with_no_text():
+ """A token that produced no visible text still owes an entry.
+
+ Here 15 is a special the caller strips and 11 is the stop itself; neither
+ contributes to `content`, and both were generated.
+ """
+ ems, asm = _drive([15, 11], stop=["STOP"], logprobs=True)
+ assert _content(ems) == ""
+ assert len(_entries(ems)) == 2
+ assert any(e.channel == CONTENT and not e.text and e.logprobs for e in ems)
+
+
+# --------------------------------------------------------------------------- #
+# The property both response paths rest on
+# --------------------------------------------------------------------------- #
+@pytest.mark.parametrize("size", [1, 2, 3, 4, 7])
+@pytest.mark.parametrize("stop", [(), ("STOP",), ("ZZZZ",)])
+def test_chunking_changes_nothing(size, stop):
+ ids = [10, 13, 14, 12]
+ whole, a1 = _drive(ids, stop=stop, logprobs=True, batch=[ids])
+ groups = [ids[i : i + size] for i in range(0, len(ids), size)]
+ part, a2 = _drive(ids, stop=stop, logprobs=True, batch=groups)
+ assert _content(whole) == _content(part)
+ assert len(_entries(whole)) == len(_entries(part))
+ assert a1.completion_tokens == a2.completion_tokens
+ assert a1.stop_reason == a2.stop_reason
+
+
+# --------------------------------------------------------------------------- #
+# Channels
+# --------------------------------------------------------------------------- #
+class _AllReasoning:
+ def feed(self, text):
+ return [{"kind": "reasoning", "text": text}]
+
+ def finish(self):
+ return []
+
+
+class _ReasoningThenContent:
+ """Everything before "|" is reasoning, everything after is content."""
+
+ def __init__(self):
+ self._switched = False
+
+ def feed(self, text):
+ out = []
+ for part in text.split("|"):
+ out.append({"kind": "content" if self._switched else "reasoning", "text": part})
+ self._switched = self._switched or "|" in text
+ return [e for e in out if e["text"]]
+
+ def finish(self):
+ return []
+
+
+class _AllContent:
+ def feed(self, text):
+ return [{"kind": "content", "text": text}]
+
+ def finish(self):
+ return []
+
+
+class _OneToolCall:
+ def feed(self, text):
+ return [{"kind": "tool", "index": 0, "id": "call_x", "name": "f", "arguments": '{"a":1}'}]
+
+ def finish(self):
+ return []
+
+
+def test_stop_with_a_parser_and_logprobs_is_refused_not_guessed():
+ """The combination this design refuses, asserted at the object that cannot
+ serve it.
+
+ With a stop configured the matcher releases text in chunks of its own, so one
+ chunk can span the end of a reasoning segment and the parser answers with a
+ reasoning event followed by a content event. There is then no signal saying
+ which of the chunk's tokens produced which: the parser reports no difference
+ between buffering an ambiguous marker prefix and consuming a complete marker.
+ Dropping the entry under-reports; keeping it puts a consumed marker's token
+ text into `logprobs.content`.
+
+ Refusing is the design decision. `refuse_unattributable_logprobs` answers 501
+ before any backend work; this guard catches a routing bug that got past it,
+ where 502 is the honest answer.
+ """
+ with pytest.raises(ValueError, match="attributable"):
+ ReplyStream(
+ _Tok(),
+ stop=["ZZZZ"],
+ parser_session=_ReasoningThenContent(),
+ logprobs_req=LogprobsRequest(top_n=1),
+ )
+
+
+@pytest.mark.parametrize(
+ "stop,session,lp,why",
+ [
+ ((), None, True, "no parser, no stop"),
+ (("ZZZZ",), None, True, "no parser: the reply IS content"),
+ ((), "session", True, "parser without stop: nothing is held back"),
+ (("ZZZZ",), "session", False, "parser with stop but nothing to attribute"),
+ ],
+)
+def test_every_other_combination_is_served(stop, session, lp, why):
+ """Only the one row is refused. The other three need no offset arithmetic --
+
+ see the table in the module docstring.
+ """
+ ReplyStream(
+ _Tok(),
+ stop=stop,
+ parser_session=_AllContent() if session else None,
+ logprobs_req=LogprobsRequest(top_n=1) if lp else None,
+ )
+
+
+def test_reasoning_logprobs_never_reach_the_content_channel():
+ """logprobs cover message.content alone."""
+ ems, asm = _drive([10, 12], logprobs=True, session=_AllReasoning())
+ assert _content(ems) == ""
+ assert "".join(e.text for e in ems if e.channel == REASONING) == "alpha beta"
+ assert _entries(ems) == []
+ assert asm.completion_tokens == 2, "the tokens were still generated"
+
+
+def test_a_tool_call_arrives_whole():
+ """The parser emits each index once, so a non-streaming caller can collect
+ them without reassembling arguments.
+ """
+ ems, _ = _drive([10], session=_OneToolCall())
+ calls = [e.tool_call for e in ems if e.channel == TOOL_CALL]
+ assert calls == [{"index": 0, "id": "call_x", "name": "f", "arguments": '{"a":1}'}]
+
+
+def test_emissions_are_immutable():
+ """A caller cannot rewrite an emission's logprobs after the fact."""
+ ems, _ = _drive([10], logprobs=True)
+ with pytest.raises(AttributeError): # dataclasses.FrozenInstanceError
+ ems[0].logprobs = []
+
+
+def test_as_logprobs_shapes_the_envelope():
+ assert as_logprobs([{"token": "x"}]) == {"content": [{"token": "x"}], "refusal": None}
+
+
+def test_finish_is_idempotent():
+ asm = ReplyStream(_Tok(), stop=["ZZZZ"])
+ asm.push([10, 12])
+ first = asm.finish()
+ assert first and asm.finish() == []
+
+
+def test_a_split_multibyte_tokens_entry_is_held_not_dropped():
+ """ "No text yet" is not "text the parser sent elsewhere".
+
+ 18 decodes to nothing on its own; 18+19 decode to "。STOP". So token 18
+ produces no visible text and its character arrives with 19. Dropping its
+ entry there loses an entry for a token whose text does reach `content`;
+ holding it attributes both to whichever channel 19's text goes to.
+ """
+ asm = ReplyStream(_Tok(), parser_session=_AllContent(), logprobs_req=LogprobsRequest(top_n=1))
+ ems = (
+ asm.push([16, 18, 19], [-0.1, -0.2, -0.3], [[(16, -0.1)], [(18, -0.2)], [(19, -0.3)]])
+ + asm.finish()
+ )
+ assert _content(ems) == "abc。STOP"
+ assert [pytest.approx(e["logprob"]) for e in _entries(ems)] == [
+ pytest.approx(-0.1),
+ pytest.approx(-0.2),
+ pytest.approx(-0.3),
+ ], "all three tokens contributed text that reached content"
+
+
+def test_a_held_entry_still_follows_its_text_to_reasoning():
+ """Held is not "kept": the token that completes the text decides.
+
+ Fixing the held case must not undo the reason entries are dropped at all --
+ a reasoning token's entry must never reach `content`.
+ """
+ asm = ReplyStream(_Tok(), parser_session=_AllReasoning(), logprobs_req=LogprobsRequest(top_n=1))
+ ems = asm.push([18, 19], [-0.2, -0.3], [[(18, -0.2)], [(19, -0.3)]]) + asm.finish()
+ assert _content(ems) == ""
+ assert _entries(ems) == [], "the text went to reasoning, so neither is owed"
+
+
+def test_a_terminal_byte_fragment_still_reaches_the_reply():
+ """18 alone: generation stopped mid-character, and that is not nothing.
+
+ The tokenizer's own decode of `[18]` is the replacement character, and the
+ blocking path used to produce it by decoding the whole id list. Holding it
+ inside the detokeniser and never flushing dropped a character the reply had,
+ while `token_ids`, `completion_tokens` and the entry all still counted the
+ token.
+ """
+ asm = ReplyStream(_Tok(), parser_session=_AllContent(), logprobs_req=LogprobsRequest(top_n=1))
+ ems = asm.push([18], [-0.2], [[(18, -0.2)]]) + asm.finish()
+ assert _content(ems) == "\ufffd", "what the tokenizer itself decodes"
+ assert len(_entries(ems)) == 1, "and the token that produced it is described"
+ assert asm.completion_tokens == 1
+
+
+def test_a_split_multibyte_tokens_entry_waits_for_its_character():
+ """`ends_at` is where a token's text is COMPLETE, not where it was queued.
+
+ 18 decodes to nothing on its own; 18+19 decode to "。". Marking 18 complete
+ at the current offset makes its entry due before the character is out, so it
+ rides an empty chunk while the chunk carrying "。" describes one token too
+ few. Both entries belong with the character.
+ """
+ # No stop, so nothing is held back and every token's text goes out at once.
+ # That is what exposes it: with a hold-back the entry is delayed anyway.
+ ems, _ = _drive([16, 18, 19], logprobs=True)
+ for e in ems:
+ if e.channel == CONTENT and e.logprobs and not e.text:
+ raise AssertionError("an entry rode an emission with no text")
+ assert len(_entries(ems)) == 3, "all three tokens are still described"
+ assert _content(ems) == "abc。STOP", "18+19 spell 。STOP in the stub vocab"
+ # Not asserted: that the entries' `token` strings rebuild the emission's
+ # text. A byte fragment decodes to U+FFFD on its own, which is why the
+ # OpenAI schema carries `bytes` at all -- so this cannot hold for a split
+ # character, and requiring it would pin a property the format disclaims.
+
+
+def test_the_terminal_flush_agrees_with_decoding_the_whole_id_list():
+ """The property the flush restores, stated directly.
+
+ Whatever the reply shows must be what the tokenizer makes of the ids the
+ reply reports -- for a generation that ends mid-character too. Without the
+ flush this was `"abc"` against `"abc\ufffd"`.
+ """
+ for ids in ([16, 18], [18], [16, 18, 19], [16]):
+ asm = ReplyStream(_Tok())
+ ems = asm.push(ids) + asm.finish()
+ assert _content(ems) == _Tok().decode(asm.token_ids), (
+ f"{ids}: reply {_content(ems)!r} vs " f"decode {_Tok().decode(asm.token_ids)!r}"
+ )
+
+
+def test_the_detokenisers_flush_is_idempotent_on_its_own():
+ """`ReplyStream.finish()` already guards against a second call, but the
+ detokeniser should not depend on that guard to avoid emitting text twice --
+ a second caller would otherwise duplicate the terminal character.
+ """
+ from tilert.pd_vllm.oai_parser import IncrementalDetok
+
+ d = IncrementalDetok(_Tok(), skip_special_tokens=True)
+ assert d.push([18]) == "", "held: the window ends mid-character"
+ assert d.finish() == "\ufffd"
+ assert d.finish() == "", "and not a second time"
+
+
+def test_a_stop_cut_tokens_entry_rides_its_surviving_prefix():
+ """17 decodes to " abcSTOP"; the stop begins inside it.
+
+ `ends_at` points past the untruncated token, so the offset comparison alone
+ deferred the entry to a later empty chunk while the visible `" abc"` went out
+ describing nothing. Once the window has stopped no more text can arrive, so
+ the token is as visible as it will ever be.
+ """
+ ems, asm = _drive([17], stop=["STOP"], logprobs=True)
+ assert _content(ems) == " abc"
+ carrying = [e for e in ems if e.channel == CONTENT and e.logprobs]
+ assert len(carrying) == 1, "exactly one emission owns the entry"
+ assert carrying[0].text == " abc", "and it is the one with the text"
+ assert asm.completion_tokens == 1
+
+
+@pytest.mark.parametrize(
+ "ids,stop,want_text,why",
+ [
+ ([17], ["STOP"], " abc", "the stop begins inside a token"),
+ ([16, 13, 14], ["STOP"], "abc", "the stop spans two tokens"),
+ ([10, 11, 12], ["STOP"], "alpha ", "the stop is its own token"),
+ ],
+)
+def test_no_entry_rides_an_empty_chunk_when_text_survived(ids, stop, want_text, why):
+ """An entry on an empty chunk is only right when NO token produced visible
+ text -- a stripped special, or a stop that consumed everything. Whenever some
+ text survived, every entry belongs with text.
+ """
+ ems, asm = _drive(ids, stop=stop, logprobs=True)
+ assert _content(ems) == want_text, why
+ orphans = [e for e in ems if e.channel == CONTENT and e.logprobs and not e.text]
+ assert not orphans, f"{why}: {len(orphans)} entries rode an empty chunk"
+ assert len(_entries(ems)) == asm.completion_tokens
+
+
+def test_an_emission_never_ends_inside_a_token_whose_entry_is_still_due():
+ """Measured on a live pair before it was a test.
+
+ The window's holdback is a CHARACTER count, so it lands mid-token: "abc" +
+ "alpha " with a 4-character stop makes 6 characters visible, three of them
+ "alpha "'s. Those went out describing nothing, and when the stop then matched
+ at exactly that cursor there was no text left for the entry to ride:
+
+ delta=' *' entries=[] <- hardware, chunk 9
+ delta='' entries=[' **', 'An', 'alyze'] <- hardware, chunk 10
+
+ Rounding the ceiling down to a token boundary makes the surviving prefix and
+ its entry the same emission. The two fully-cut tokens keep riding it: their
+ own text is gone, which is the one case an entry has nowhere else to go.
+ """
+ ems, asm = _drive([16, 10, 11], stop=["ha S"], logprobs=True)
+ assert _content(ems) == "abcalp"
+ orphans = [e for e in ems if e.channel == CONTENT and e.logprobs and not e.text]
+ assert not orphans, f"{len(orphans)} entries rode an empty chunk"
+ texted = [e for e in ems if e.channel == CONTENT and e.text]
+ assert texted[-1].logprobs, "the last chunk with text carries its entries"
+ assert asm.stop_reason == "ha S"
+ assert asm.completion_tokens == 3
diff --git a/tests/pd_vllm/test_request_capabilities.py b/tests/pd_vllm/test_request_capabilities.py
new file mode 100644
index 0000000..0e7a27c
--- /dev/null
+++ b/tests/pd_vllm/test_request_capabilities.py
@@ -0,0 +1,724 @@
+"""Every generation field a client can send must be executed or refused.
+
+The PD split samples token 1 on the vLLM prefill instance and tokens 2..N on the
+decode node. The client body reaches vLLM almost verbatim, while the decode node
+gets only the keys ``pd_router._sampling_of`` forwards. A field in the gap is
+applied to token 1 and dropped for the rest of the reply, and the response is a
+200 that quietly violates what was asked for.
+
+These tests pin the closed set: for each such field, either the request is
+refused before anything observable happens, or the field is forwarded to a node
+that declared it can execute it.
+
+CPU only -- no GPU, no tilert, no real vLLM.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_request_capabilities.py -v
+"""
+
+import queue
+import types
+
+import pytest
+from fastapi.testclient import TestClient
+
+from tilert.pd_vllm import pd_router
+from tilert.pd_vllm.capabilities import (
+ PROFILE_FIELD_NAMES,
+ STATIC_FIELD_NAMES,
+ CapabilityUnavailable,
+ InvalidParameter,
+ NodeCapabilities,
+ engine_capabilities,
+ validate_generation_request,
+)
+from tilert.pd_vllm.decode_server import build_app as build_decode_app
+from tilert.pd_vllm.engine_iface import StubEngine
+
+# One non-neutral value per statically unsupported field, i.e. a value that asks
+# for behaviour the decode node cannot produce.
+LIVE_VALUES = {
+ "stop_token_ids": [151643],
+ "min_tokens": 16,
+ "frequency_penalty": 0.5,
+ "min_p": 0.05,
+ "seed": 42,
+ "logit_bias": {"151643": -100.0},
+ "bad_words": ["foo"],
+ "allowed_token_ids": [1, 2, 3],
+ "structured_outputs": {"json": {"type": "object"}},
+ "n": 2,
+ "best_of": 4,
+ "use_beam_search": True,
+ "prompt_logprobs": 1,
+ "logprob_token_ids": [7],
+ "skip_special_tokens": False,
+}
+
+# The value vLLM itself would have used, so honouring and ignoring it agree.
+NEUTRAL_VALUES = {
+ "stop_token_ids": [],
+ "min_tokens": 0,
+ "frequency_penalty": 0.0,
+ "min_p": 0.0,
+ "seed": None,
+ "logit_bias": {},
+ "bad_words": [],
+ "allowed_token_ids": None,
+ "structured_outputs": None,
+ "n": 1,
+ "best_of": 1,
+ "use_beam_search": False,
+ "prompt_logprobs": None,
+ "logprob_token_ids": [],
+ "skip_special_tokens": True,
+}
+
+FULL_CAPS = NodeCapabilities(penalties=True, ignore_eos=True)
+NO_CAPS = NodeCapabilities()
+
+
+def _body(**extra):
+ return {"messages": [{"role": "user", "content": "hi"}], **extra}
+
+
+# --------------------------------------------------------------------------- #
+# The field tables are complete and self-consistent
+# --------------------------------------------------------------------------- #
+def test_every_static_field_has_a_live_and_a_neutral_fixture():
+ """A field added to the gate without a fixture here would be untested."""
+ assert set(STATIC_FIELD_NAMES) == set(LIVE_VALUES) == set(NEUTRAL_VALUES)
+
+
+def test_no_forwarded_field_is_left_ungated():
+ """``_sampling_of``'s whitelist and the gate must partition the fields.
+
+ A field that is forwarded AND statically refused is a contradiction; a field
+ that is neither forwarded nor gated is the original bug. The only fields
+ allowed to be forwarded ungated are the three the decode sampler implements
+ unconditionally.
+ """
+ forwarded = {
+ "temperature",
+ "top_p",
+ "top_k",
+ "repetition_penalty",
+ "presence_penalty",
+ "ignore_eos",
+ }
+ unconditional = {"temperature", "top_p", "top_k"}
+ assert forwarded.isdisjoint(STATIC_FIELD_NAMES)
+ assert forwarded - unconditional == set(PROFILE_FIELD_NAMES)
+
+
+# --------------------------------------------------------------------------- #
+# validate_generation_request: the unit contract
+# --------------------------------------------------------------------------- #
+@pytest.mark.parametrize("field", sorted(STATIC_FIELD_NAMES))
+def test_a_live_static_field_is_refused(field):
+ with pytest.raises(CapabilityUnavailable) as e:
+ validate_generation_request(_body(**{field: LIVE_VALUES[field]}), FULL_CAPS)
+ # The message must name the field, so an operator reading a client's error
+ # knows which one to drop.
+ assert field in str(e.value)
+
+
+@pytest.mark.parametrize("field", sorted(STATIC_FIELD_NAMES))
+def test_a_neutral_static_field_is_accepted(field):
+ """Mentioning a field is not asking for it.
+
+ Refusing the neutral value would reject clients and SDKs that send the whole schema with
+ defaults filled in -- which is what made a blanket "unknown field" rejection unusable.
+ """
+ validate_generation_request(_body(**{field: NEUTRAL_VALUES[field]}), FULL_CAPS)
+
+
+@pytest.mark.parametrize("field", sorted(STATIC_FIELD_NAMES))
+def test_an_absent_static_field_is_accepted(field):
+ validate_generation_request(_body(), FULL_CAPS)
+
+
+def test_n_greater_than_one_is_rejected():
+ with pytest.raises(CapabilityUnavailable):
+ validate_generation_request(_body(n=2), FULL_CAPS)
+
+
+def test_n_of_one_is_accepted():
+ validate_generation_request(_body(n=1), FULL_CAPS)
+
+
+def test_n_below_one_is_a_client_error_not_a_capability_gap():
+ """``n=0`` is not "a feature we lack"; no backend would serve it."""
+ with pytest.raises(InvalidParameter):
+ validate_generation_request(_body(n=0), FULL_CAPS)
+
+
+def test_a_string_where_a_number_belongs_is_a_client_error():
+ with pytest.raises(InvalidParameter):
+ validate_generation_request(_body(frequency_penalty="high"), FULL_CAPS)
+
+
+def test_a_bool_is_not_a_number():
+ """bool is an int subclass in Python; ``min_p: true`` is a type error."""
+ with pytest.raises(InvalidParameter):
+ validate_generation_request(_body(min_p=True), FULL_CAPS)
+
+
+def test_a_number_where_a_flag_belongs_is_a_client_error():
+ with pytest.raises(InvalidParameter):
+ validate_generation_request(_body(use_beam_search=1), FULL_CAPS)
+
+
+def test_explicit_null_is_treated_as_absent():
+ """An explicit null is how SDKs spell "unset", and vLLM resolves it to the
+ same default as an absent key.
+ """
+ for field in STATIC_FIELD_NAMES:
+ validate_generation_request(_body(**{field: None}), FULL_CAPS)
+
+
+def test_error_payload_shape_matches_the_other_gates():
+ for exc in (CapabilityUnavailable("x"), InvalidParameter("y")):
+ payload = exc.to_payload()
+ assert set(payload) == {"error", "error_type"}
+ assert CapabilityUnavailable.http_status == 501
+ assert InvalidParameter.http_status == 400
+
+
+# --------------------------------------------------------------------------- #
+# Profile-dependent fields follow the node's declaration
+# --------------------------------------------------------------------------- #
+@pytest.mark.parametrize(
+ "field,value",
+ [
+ ("repetition_penalty", 1.2),
+ ("presence_penalty", 0.4),
+ ],
+)
+def test_penalties_need_a_node_that_declares_them(field, value):
+ validate_generation_request(_body(**{field: value}), FULL_CAPS)
+ with pytest.raises(CapabilityUnavailable):
+ validate_generation_request(_body(**{field: value}), NO_CAPS)
+
+
+@pytest.mark.parametrize(
+ "field,neutral",
+ [
+ ("repetition_penalty", 1.0),
+ ("presence_penalty", 0.0),
+ ],
+)
+def test_neutral_penalties_pass_on_a_node_without_them(field, neutral):
+ validate_generation_request(_body(**{field: neutral}), NO_CAPS)
+
+
+def test_ignore_eos_is_allowed_only_where_it_is_honoured():
+ validate_generation_request(_body(ignore_eos=True), FULL_CAPS)
+ with pytest.raises(CapabilityUnavailable):
+ validate_generation_request(_body(ignore_eos=True), NO_CAPS)
+ # False asks for nothing, so it needs no capability.
+ validate_generation_request(_body(ignore_eos=False), NO_CAPS)
+
+
+def test_unknown_capabilities_fail_closed():
+ """``None`` means the router could not establish what the node can do.
+
+ Guessing "supported" would restore the silent-wrong-answer; guessing
+ "unsupported" costs an honest 501.
+ """
+ with pytest.raises(CapabilityUnavailable):
+ validate_generation_request(_body(ignore_eos=True), None)
+
+
+# --------------------------------------------------------------------------- #
+# NodeCapabilities: parsing and pool intersection
+# --------------------------------------------------------------------------- #
+def test_intersection_keeps_only_what_every_node_can_do():
+ mixed = FULL_CAPS.intersect(NodeCapabilities(penalties=True))
+ assert mixed == NodeCapabilities(penalties=True, ignore_eos=False)
+
+
+@pytest.mark.parametrize(
+ "payload",
+ [
+ None,
+ "nope",
+ 42,
+ [],
+ {},
+ {"capabilities": "nope"},
+ {"capabilities": {}},
+ {"capabilities": {"penalties": "yes"}},
+ ],
+)
+def test_an_unusable_payload_declares_nothing(payload):
+ assert NodeCapabilities.from_payload(payload) == NO_CAPS
+
+
+def test_a_well_formed_payload_is_parsed():
+ caps = NodeCapabilities.from_payload({"capabilities": {"penalties": True, "ignore_eos": False}})
+ assert caps == NodeCapabilities(penalties=True, ignore_eos=False)
+
+
+def test_a_bare_capability_dict_is_also_accepted():
+ assert NodeCapabilities.from_payload({"penalties": True}) == NodeCapabilities(penalties=True)
+
+
+# --------------------------------------------------------------------------- #
+# engine_capabilities: read from the live engine, absent predicate = no support
+# --------------------------------------------------------------------------- #
+def test_the_stub_engine_declares_everything():
+ assert engine_capabilities(StubEngine()) == FULL_CAPS
+
+
+def test_an_engine_without_the_predicates_declares_nothing():
+ assert engine_capabilities(object()) == NO_CAPS
+
+
+def test_a_predicate_that_raises_is_treated_as_unsupported():
+ class _Broken:
+ def supports_penalties(self):
+ raise RuntimeError("engine went away")
+
+ def supports_ignore_eos(self):
+ return True
+
+ assert engine_capabilities(_Broken()) == NodeCapabilities(ignore_eos=True)
+
+
+def test_a_demoted_adapter_reports_the_demotion_not_the_claim():
+ """An adapter that probes for the penalty pre-pass demotes its claim when
+ the installed engine lacks it. ``/capabilities`` must show the demotion, or
+ the router would pre-approve a request the engine then refuses.
+ """
+
+ class _Demoted:
+ _supports_penalties = False
+
+ def supports_penalties(self):
+ return self._supports_penalties
+
+ def supports_ignore_eos(self):
+ return True
+
+ assert engine_capabilities(_Demoted()).penalties is False
+
+
+# --------------------------------------------------------------------------- #
+# decode_server: /capabilities, and the pre-wire-wait refusal
+# --------------------------------------------------------------------------- #
+class _FakeReq:
+ rid = "rid-1"
+ seq_len = 8
+ last_prompt_token = 5
+
+
+class _FakeServer:
+ def __init__(self):
+ self.completed: queue.Queue = queue.Queue()
+ self.completed.put(_FakeReq())
+ self.profile = types.SimpleNamespace(
+ convert=lambda *a, **k: "converted", num_ranks=8, name="stub"
+ )
+ self.buffer = None
+ self.base_ptr = 0
+ self.max_seq_len = 4096
+ self.released = 0
+
+ def expect(self, rid=None):
+ # /pd/decode announces its rid so the real ReceiveServer can drop a
+ # tombstone left by a previous attempt at the same request. Recorded,
+ # so a test can assert the announcement happened.
+ self.expected_rids = getattr(self, "expected_rids", [])
+ self.expected_rids.append(rid)
+
+ def release(self, rid=None):
+ # Scoped like the real ReceiveServer.release: the decode server
+ # names the rid it owns, because the slot may since have been
+ # handed to a later request.
+ self.released_rids = getattr(self, "released_rids", [])
+ self.released_rids.append(rid)
+ self.released += 1
+
+
+class _NoPenaltyEngine(StubEngine):
+ def supports_penalties(self) -> bool:
+ return False
+
+
+def _decode_client(engine=None):
+ return TestClient(build_decode_app(_FakeServer(), engine or StubEngine()))
+
+
+def test_capabilities_endpoint_reports_the_engine():
+ r = _decode_client().get("/capabilities")
+ assert r.status_code == 200
+ body = r.json()
+ assert body["engine"] == "StubEngine"
+ assert body["capabilities"]["penalties"] is True
+ assert body["capabilities"]["ignore_eos"] is True
+
+
+def test_capabilities_endpoint_follows_the_engine_that_cannot():
+ r = _decode_client(_NoPenaltyEngine()).get("/capabilities")
+ assert r.json()["capabilities"]["penalties"] is False
+
+
+def test_decode_refuses_a_penalty_the_engine_cannot_apply():
+ """501 rather than a 200 whose tokens were never penalised."""
+ r = _decode_client(_NoPenaltyEngine()).post(
+ "/pd/decode",
+ json={
+ "rid": "rid-1",
+ "first_token_id": 7,
+ "max_tokens": 8,
+ "sampling": {"repetition_penalty": 1.3},
+ },
+ )
+ assert r.status_code == 501
+ assert r.json()["error_type"] == "capability_unavailable"
+
+
+def test_decode_serves_a_neutral_penalty_on_the_same_engine():
+ r = _decode_client(_NoPenaltyEngine()).post(
+ "/pd/decode",
+ json={
+ "rid": "rid-1",
+ "first_token_id": 7,
+ "max_tokens": 8,
+ "sampling": {"repetition_penalty": 1.0, "temperature": 0.6},
+ },
+ )
+ assert r.status_code == 200
+
+
+def test_a_refused_request_hands_back_the_receive_slot():
+ """Otherwise one refused request stalls the NEXT one for the full
+ kv_transfer_timeout -- the failure ``_abandon_pending_kv`` exists for.
+ """
+ server = _FakeServer()
+ client = TestClient(build_decode_app(server, _NoPenaltyEngine()))
+ r = client.post(
+ "/pd/decode",
+ json={
+ "rid": "rid-1",
+ "first_token_id": 7,
+ "max_tokens": 8,
+ "sampling": {"presence_penalty": 0.5},
+ },
+ )
+ assert r.status_code == 501
+ assert server.released >= 1
+ # ... and the node is not left busy.
+ assert client.get("/decode_status").json()["status"] == "idle"
+
+
+# --------------------------------------------------------------------------- #
+# pd_router: refused before prefill, before a node is acquired
+# --------------------------------------------------------------------------- #
+class _Resp:
+ def __init__(self, payload, status=200):
+ self._payload = payload
+ self.status_code = status
+
+ def json(self):
+ return self._payload
+
+ def raise_for_status(self):
+ if self.status_code >= 400:
+ raise RuntimeError(f"HTTP {self.status_code}")
+
+
+def _router(monkeypatch, caps=FULL_CAPS, on_post=None):
+ """Router whose single decode node declares ``caps``."""
+
+ def fake_get(url, timeout=None, **kw):
+ assert url.endswith("/capabilities"), url
+ return _Resp({"capabilities": caps.to_payload()})
+
+ def default_post(url, json=None, timeout=None, **kw):
+ raise AssertionError(f"network must not be touched: {url}")
+
+ monkeypatch.setattr(pd_router.requests, "get", fake_get)
+ monkeypatch.setattr(pd_router.requests, "post", on_post or default_post)
+ pool = pd_router.Pool([pd_router.DecodeNode("127.0.0.1", 5556, 5557)])
+ ctx = pd_router.RouterCtx("http://vllm.invalid", pool, tokenizer=None, parser_name="none")
+ return TestClient(pd_router.build_app(ctx)), pool
+
+
+@pytest.mark.parametrize("field", sorted(STATIC_FIELD_NAMES))
+def test_a_live_static_field_is_refused_before_prefill(monkeypatch, field):
+ """The refusal must cost nothing: no vLLM call, so no prompt is prefilled
+ and no KV is pushed to a node that will never consume it.
+ """
+ client, pool = _router(monkeypatch)
+ r = client.post("/v1/chat/completions", json=_body(**{field: LIVE_VALUES[field]}))
+ assert r.status_code == 501
+ assert r.json()["error_type"] == "capability_unavailable"
+ # ... and no decode node was reserved.
+ assert all(not n.busy for n in pool.nodes)
+
+
+def test_a_refused_stream_request_also_reserves_nothing(monkeypatch):
+ client, pool = _router(monkeypatch)
+ r = client.post("/v1/chat/completions", json=_body(seed=7, stream=True))
+ assert r.status_code == 501
+ assert all(not n.busy for n in pool.nodes)
+
+
+def test_a_client_error_is_400_not_501(monkeypatch):
+ client, _ = _router(monkeypatch)
+ r = client.post("/v1/chat/completions", json=_body(n=0))
+ assert r.status_code == 400
+ assert r.json()["error_type"] == "invalid_parameter"
+
+
+def test_completions_endpoint_is_gated_too(monkeypatch):
+ client, _ = _router(monkeypatch)
+ r = client.post("/v1/completions", json={"prompt": "hi", "seed": 1})
+ assert r.status_code == 501
+
+
+def test_a_node_that_declares_penalties_gets_the_request(monkeypatch):
+ captured = {}
+
+ def on_post(url, json=None, timeout=None, **kw):
+ if url.endswith("/pd/decode"):
+ captured["decode"] = json
+ return _Resp(
+ {"rid": "x", "token_ids": [7], "seq_len": 8, "timing_ms": {"finish_reason": "stop"}}
+ )
+ captured["prefill"] = json
+ return _Resp(
+ {
+ "id": "cmpl-abc",
+ "choices": [{"logprobs": {"content": [{"token": "token_id:7"}]}}],
+ "usage": {"prompt_tokens": 3},
+ "model": "m",
+ }
+ )
+
+ client, _ = _router(monkeypatch, on_post=on_post)
+ r = client.post("/v1/chat/completions", json=_body(repetition_penalty=1.2, ignore_eos=True))
+ assert r.status_code == 200
+ assert captured["decode"]["sampling"]["repetition_penalty"] == 1.2
+ assert captured["decode"]["sampling"]["ignore_eos"] is True
+
+
+def test_a_node_that_declares_nothing_refuses_the_same_request(monkeypatch):
+ client, _ = _router(monkeypatch, caps=NO_CAPS)
+ r = client.post("/v1/chat/completions", json=_body(repetition_penalty=1.2))
+ assert r.status_code == 501
+
+
+def test_a_failed_capability_probe_refuses_rather_than_assumes(monkeypatch):
+ def boom(url, timeout=None, **kw):
+ raise OSError("connection refused")
+
+ monkeypatch.setattr(pd_router.requests, "get", boom)
+ monkeypatch.setattr(pd_router.requests, "post", lambda *a, **k: pytest.fail("must not prefill"))
+ pool = pd_router.Pool([pd_router.DecodeNode("127.0.0.1", 5556, 5557)])
+ ctx = pd_router.RouterCtx("http://vllm.invalid", pool, tokenizer=None, parser_name="none")
+ r = TestClient(pd_router.build_app(ctx)).post(
+ "/v1/chat/completions", json=_body(ignore_eos=True)
+ )
+ assert r.status_code == 501
+
+
+def test_a_plain_request_survives_an_unreachable_probe(monkeypatch):
+ """Fail-closed must cost only the optional fields: a request that asks for
+ nothing profile-dependent still goes through when the probe fails.
+ """
+
+ def boom(url, timeout=None, **kw):
+ raise OSError("connection refused")
+
+ def on_post(url, json=None, timeout=None, **kw):
+ if url.endswith("/pd/decode"):
+ return _Resp(
+ {"rid": "x", "token_ids": [7], "seq_len": 8, "timing_ms": {"finish_reason": "stop"}}
+ )
+ return _Resp(
+ {
+ "id": "cmpl-abc",
+ "choices": [{"logprobs": {"content": [{"token": "token_id:7"}]}}],
+ "usage": {"prompt_tokens": 3},
+ "model": "m",
+ }
+ )
+
+ monkeypatch.setattr(pd_router.requests, "get", boom)
+ monkeypatch.setattr(pd_router.requests, "post", on_post)
+ pool = pd_router.Pool([pd_router.DecodeNode("127.0.0.1", 5556, 5557)])
+ ctx = pd_router.RouterCtx("http://vllm.invalid", pool, tokenizer=None, parser_name="none")
+ r = TestClient(pd_router.build_app(ctx)).post(
+ "/v1/chat/completions", json=_body(temperature=0.6)
+ )
+ assert r.status_code == 200
+
+
+def test_the_pool_intersects_every_node(monkeypatch):
+ """Validation runs before a node is chosen, so a field is only safe if EVERY
+ node could have executed it.
+ """
+
+ def fake_get(url, timeout=None, **kw):
+ # :5557 supports penalties, :5559 does not.
+ supports = ":5557/" in url
+ return _Resp({"capabilities": {"penalties": supports, "ignore_eos": True}})
+
+ monkeypatch.setattr(pd_router.requests, "get", fake_get)
+ pool = pd_router.Pool(
+ [
+ pd_router.DecodeNode("127.0.0.1", 5556, 5557),
+ pd_router.DecodeNode("127.0.0.1", 5558, 5559),
+ ]
+ )
+ assert pool.capabilities() == NodeCapabilities(penalties=False, ignore_eos=True)
+
+
+def test_a_probe_result_is_cached(monkeypatch):
+ calls = {"n": 0}
+
+ def fake_get(url, timeout=None, **kw):
+ calls["n"] += 1
+ return _Resp({"capabilities": FULL_CAPS.to_payload()})
+
+ monkeypatch.setattr(pd_router.requests, "get", fake_get)
+ pool = pd_router.Pool([pd_router.DecodeNode("127.0.0.1", 5556, 5557)])
+ for _ in range(5):
+ pool.capabilities()
+ assert calls["n"] == 1
+
+
+def test_a_failed_probe_is_not_cached(monkeypatch):
+ """A node that comes back must be picked up on the next request, not after
+ a whole TTL -- the router is meant to survive an independent restart.
+ """
+ calls = {"n": 0}
+
+ def fake_get(url, timeout=None, **kw):
+ calls["n"] += 1
+ if calls["n"] == 1:
+ raise OSError("down")
+ return _Resp({"capabilities": FULL_CAPS.to_payload()})
+
+ monkeypatch.setattr(pd_router.requests, "get", fake_get)
+ pool = pd_router.Pool([pd_router.DecodeNode("127.0.0.1", 5556, 5557)])
+ assert pool.capabilities() == NO_CAPS
+ assert pool.capabilities() == FULL_CAPS
+
+
+def test_an_empty_pool_declares_nothing():
+ assert pd_router.Pool([]).capabilities() == NO_CAPS
+
+
+# --------------------------------------------------------------------------- #
+# Review findings on PR #40 (codex)
+# --------------------------------------------------------------------------- #
+@pytest.mark.parametrize("field", ["seed", "prompt_logprobs"])
+def test_zero_is_a_request_not_an_absence(field):
+ """ "Falsy" and "asks for nothing" are different questions.
+
+ ``seed: 0`` is a valid seed and ``prompt_logprobs: 0`` asks for the prompt
+ tokens' own log probabilities. Treating either as neutral lets the request
+ through, and the prefill instance then applies it to the first token while
+ the decode node cannot -- the exact split this gate exists to stop.
+ """
+ with pytest.raises(CapabilityUnavailable):
+ validate_generation_request(_body(**{field: 0}), FULL_CAPS)
+
+
+@pytest.mark.parametrize("field", ["seed", "prompt_logprobs"])
+def test_only_an_explicit_null_is_neutral_for_those_fields(field):
+ validate_generation_request(_body(**{field: None}), FULL_CAPS)
+
+
+@pytest.mark.parametrize(
+ "field", ["logit_bias", "bad_words", "stop_token_ids", "logprob_token_ids"]
+)
+def test_an_empty_collection_is_still_neutral(field):
+ """Unlike the two above, an empty collection genuinely asks for nothing --
+
+ the distinction is whether the field has a meaningful zero, not whether the
+ value is falsy.
+ """
+ validate_generation_request(_body(**{field: NEUTRAL_VALUES[field]}), FULL_CAPS)
+
+
+@pytest.mark.parametrize(
+ "body",
+ [
+ {"stop": ["\n\n"]},
+ {"stop": "END"},
+ {"stop": ["A"], "include_stop_str_in_output": True},
+ ],
+)
+def test_the_gate_does_not_own_stop_strings(body):
+ """The router matches them over the reply text, so the gate lets them by.
+
+ Refusing them here would make the router's implementation unreachable.
+ Their validation lives in ``request_gate.resolve_stop_request``, which is where
+ the tokenizer that executes them is; see test_stop_semantics.py.
+ """
+ validate_generation_request(_body(**body), FULL_CAPS)
+
+
+# ── the three always-supported fields are resolved, so they are type-checked ──
+@pytest.mark.parametrize("value", [20, "20", 20.0])
+def test_an_integral_top_k_is_accepted_whatever_its_spelling(value):
+ """As permissive as vLLM: rejecting ``"20"`` would turn a request vLLM
+ serves into a 400.
+ """
+ validate_generation_request(_body(top_k=value), FULL_CAPS)
+
+
+@pytest.mark.parametrize("value", [1.9, "1.9", True])
+def test_a_non_integral_top_k_is_refused_not_truncated(value):
+ """The router resolves top_k and writes the result into the prefill request,
+ which overwrites what the client sent -- so vLLM no longer gets to reject a
+ bad value. Truncating 1.9 to 1 would have both legs serve a materially
+ different request and report success.
+ """
+ with pytest.raises(InvalidParameter):
+ validate_generation_request(_body(top_k=value), FULL_CAPS)
+
+
+@pytest.mark.parametrize(
+ "field,value",
+ [
+ ("temperature", "0.7"),
+ ("top_p", "0.95"),
+ ("temperature", 0),
+ ("top_p", 1),
+ ],
+)
+def test_a_usable_temperature_or_top_p_is_accepted(field, value):
+ validate_generation_request(_body(**{field: value}), FULL_CAPS)
+
+
+@pytest.mark.parametrize(
+ "field,value",
+ [
+ ("temperature", "hot"),
+ ("top_p", "wide"),
+ ("temperature", True),
+ ("top_p", True),
+ ("temperature", -1),
+ ("top_p", -0.5),
+ ],
+)
+def test_an_unusable_temperature_or_top_p_is_a_client_error(field, value):
+ with pytest.raises(InvalidParameter):
+ validate_generation_request(_body(**{field: value}), FULL_CAPS)
+
+
+@pytest.mark.parametrize("field", ["temperature", "top_p", "top_k"])
+def test_the_always_supported_fields_are_never_refused_for_capability(field):
+ """They are type-checked, not gated: every profile's decode path applies all
+ three on every request.
+ """
+ assert field not in STATIC_FIELD_NAMES
+ assert field not in PROFILE_FIELD_NAMES
+ validate_generation_request(_body(**{field: 1}), NO_CAPS)
diff --git a/tests/pd_vllm/test_request_gate.py b/tests/pd_vllm/test_request_gate.py
new file mode 100644
index 0000000..a4fb121
--- /dev/null
+++ b/tests/pd_vllm/test_request_gate.py
@@ -0,0 +1,142 @@
+"""What the gate accepts, refuses, and refuses to look at.
+
+`gate_request` is a pure function, so these run without HTTP, a tokenizer or a
+backend. What they pin is not that the checks exist -- other suites cover each
+one -- but the two things a single gate can silently get wrong: WHICH fields it
+validates, and IN WHICH ORDER, since only one error reaches the client.
+
+Both were found by a differential against the pre-refactor router rather than by
+reasoning: the same request matrix through both, byte-compared.
+
+ CUDA_VISIBLE_DEVICES= python -m pytest tests/pd_vllm/test_request_gate.py -v
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from tilert.pd_vllm.capabilities import (
+ CapabilityUnavailable,
+ InvalidParameter,
+)
+from tilert.pd_vllm.request_gate import gate_request
+
+CHAT = "/v1/chat/completions"
+COMP = "/v1/completions"
+
+
+class _Tok:
+ def decode(self, ids, skip_special_tokens=False):
+ return ""
+
+
+_TOK = _Tok()
+
+
+def gate(path, body, tokenizer=_TOK, parser=False):
+ return gate_request(path, body, tokenizer=tokenizer, parser_active=lambda thinking: parser)
+
+
+# --------------------------------------------------------------------------- #
+# chat_template_kwargs: a chat field, validated only there
+# --------------------------------------------------------------------------- #
+def test_a_non_dict_chat_template_kwargs_is_refused_on_chat():
+ """vLLM types the field `dict[str, Any] | None` on `ChatCompletionRequest`
+
+ and answers 422 for a string. Reading `.get` off one raised AttributeError
+ here once, which surfaced as a 500 for a client mistake.
+ """
+ with pytest.raises(InvalidParameter):
+ gate(CHAT, {"chat_template_kwargs": "oops"})
+
+
+@pytest.mark.parametrize(
+ "extra,why",
+ [
+ ({}, "no grammar: the field reaches nothing at all"),
+ (
+ {"response_format": {"type": "json_object"}},
+ "with a grammar the node gets `enable_thinking`, and True is what the "
+ "absent field would have meant",
+ ),
+ ],
+)
+def test_the_same_value_is_served_on_completions(extra, why):
+ """Measured on vLLM 0.25.1: `CompletionRequest` has no
+ `chat_template_kwargs` field and `extra="allow"`, so it accepts and ignores
+ one. Refusing it here would 400 a request vLLM serves.
+
+ The pre-refactor router served it without a grammar and answered 502 WITH
+ one -- it raised inside the handler, after a node was acquired and prefill
+ had run, and a generic handler flattened it. Both now serve.
+ """
+ req = gate(COMP, {"chat_template_kwargs": "oops", **extra})
+ assert req.thinking is True, why
+ assert req.is_chat is False
+
+
+def test_thinking_is_read_once_and_carried():
+ """It decided the parser session and the decode body separately before, from
+ two reads of the same body.
+ """
+ assert gate(CHAT, {"chat_template_kwargs": {"enable_thinking": False}}).thinking is False
+ assert gate(CHAT, {}).thinking is True
+
+
+# --------------------------------------------------------------------------- #
+# order: only one error reaches the client
+# --------------------------------------------------------------------------- #
+def test_an_unusable_stop_is_reported_before_an_unusable_template_kwarg():
+ """Two bad fields, one answer.
+
+ The order is the pre-refactor order, kept because a client reading the message would
+ otherwise see a different one after a refactor that was supposed to change nothing.
+ """
+ with pytest.raises(InvalidParameter) as e:
+ gate(CHAT, {"stop": [""], "chat_template_kwargs": "oops"})
+ assert "stop" in str(e.value), str(e.value)
+
+
+def test_a_grammar_is_reported_before_anything_else():
+ from tilert.pd_vllm.grammar_spec import GrammarError
+
+ with pytest.raises(GrammarError):
+ gate(CHAT, {"response_format": "not an object", "stop": [""]})
+
+
+# --------------------------------------------------------------------------- #
+# the refusals themselves, at the gate rather than after backend work
+# --------------------------------------------------------------------------- #
+def test_stop_without_a_tokenizer_is_refused():
+ with pytest.raises(CapabilityUnavailable):
+ gate(CHAT, {"stop": ["END"]}, tokenizer=None)
+
+
+def test_logprobs_without_a_tokenizer_is_refused():
+ with pytest.raises(CapabilityUnavailable):
+ gate(CHAT, {"logprobs": True, "temperature": 0.6}, tokenizer=None)
+
+
+def test_stop_with_logprobs_and_a_parser_is_refused():
+ with pytest.raises(CapabilityUnavailable):
+ gate(CHAT, {"stop": ["END"], "logprobs": True, "temperature": 0.6}, parser=True)
+
+
+def test_the_same_three_without_the_parser_are_served():
+ req = gate(CHAT, {"stop": ["END"], "logprobs": True, "temperature": 0.6}, parser=False)
+ assert req.stop == ["END"]
+ assert req.logprobs_req is not None
+
+
+def test_the_gate_does_not_probe_capabilities():
+ """`validate_generation_request` stays at the call sites on purpose: its
+ probe blocks per unreachable node, and nothing here depends on it, so a 400
+ for a malformed request must not wait on one.
+
+ Asserted as an interface fact -- the gate takes no capabilities argument --
+ because the cost of getting it wrong is latency nothing would measure.
+ """
+ import inspect
+
+ params = inspect.signature(gate_request).parameters
+ assert set(params) == {"path", "body", "tokenizer", "parser_active"}
diff --git a/tests/pd_vllm/test_sampling_alignment.py b/tests/pd_vllm/test_sampling_alignment.py
new file mode 100644
index 0000000..9b19996
--- /dev/null
+++ b/tests/pd_vllm/test_sampling_alignment.py
@@ -0,0 +1,255 @@
+"""The two PD legs must sample under the same rules.
+
+The vLLM prefill instance samples token 1; the decode node samples tokens 2..N.
+Anything they resolve independently can disagree, and the response carries no
+sign of it. Two such disagreements are pinned here:
+
+* ``top_p`` -- the adapters each defaulted it to 0.95 while vLLM resolved its own
+ default through ``client value > generation_config.json > 1.0``. A client that
+ sent ``temperature`` but no ``top_p`` had its first token drawn from one
+ nucleus and the rest from another.
+* ``ignore_eos`` -- forwarded to the decode node and honoured by the adapter; an
+ adapter that dropped it made a fixed-length benchmark stop at the first EOS and
+ report throughput for a shorter reply than it asked for.
+
+CPU only -- no GPU, no tilert, no real vLLM.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_sampling_alignment.py -v
+"""
+
+import pathlib
+import types
+
+import pytest
+from fastapi.testclient import TestClient
+
+from tilert.pd_vllm import pd_router
+from tilert.pd_vllm.generation_defaults import GenerationDefaults
+from tilert.pd_vllm.profiles.mla_nsa import MlaNsaEngineAdapter
+from tilert.pd_vllm.sampling import VLLM_DEFAULT_TOP_P, resolve_top_p
+
+ADAPTER_MODULES = (
+ "tilert.pd_vllm.profiles.mla_nsa",
+ "tilert.pd_vllm.profiles.glm5_rocm_engine",
+)
+
+
+# --------------------------------------------------------------------------- #
+# resolve_top_p: the single resolution point
+# --------------------------------------------------------------------------- #
+def test_the_baseline_is_vllms_framework_default():
+ """vLLM's ``_DEFAULT_SAMPLING_PARAMS["top_p"]`` is 1.0, not 0.95."""
+ assert VLLM_DEFAULT_TOP_P == 1.0
+ assert resolve_top_p({}) == 1.0
+
+
+def test_an_explicit_value_wins_over_the_default():
+ assert resolve_top_p({"top_p": 0.8}, 0.95) == 0.8
+
+
+def test_an_explicit_null_falls_through_to_the_default():
+ """SDKs spell "unset" as an explicit null; vLLM resolves it as absent."""
+ assert resolve_top_p({"top_p": None}, 0.7) == 0.7
+
+
+def test_the_deployment_default_stands_in_for_generation_config():
+ """A deployment whose checkpoint recommends 0.95 sets it once, and BOTH legs
+ are handed that number -- which is the property that matters, not the value.
+ """
+ assert resolve_top_p({}, 0.95) == 0.95
+
+
+def test_a_string_is_coerced_like_the_other_resolvers():
+ assert resolve_top_p({"top_p": "0.5"}) == 0.5
+
+
+def test_a_bool_is_not_a_top_p():
+ with pytest.raises(ValueError):
+ resolve_top_p({"top_p": True})
+
+
+@pytest.mark.parametrize("module", ADAPTER_MODULES)
+def test_no_adapter_carries_its_own_top_p_default(module):
+ """Source-level, in the shape test_top_k_resolution.py uses: a re-introduced
+ literal default is the exact regression this file exists for, and it would
+ otherwise only show up as a quality drift nobody can attribute.
+ """
+ src = pathlib.Path(pytest.importorskip(module).__file__).read_text()
+ assert (
+ 'sampling.get("top_p"' not in src
+ ), f"{module} must consume the resolved value, not default it again"
+ assert "0.95" not in src
+
+
+# --------------------------------------------------------------------------- #
+# Both legs receive the same resolved value
+# --------------------------------------------------------------------------- #
+class _Resp:
+ def __init__(self, payload, status=200):
+ self._payload = payload
+ self.status_code = status
+
+ def json(self):
+ return self._payload
+
+ def raise_for_status(self):
+ if self.status_code >= 400:
+ raise RuntimeError(f"HTTP {self.status_code}")
+
+
+def _both_legs(monkeypatch, body, defaults=None):
+ """Drive one request through the router; return (prefill_body, decode_body)."""
+ seen = {}
+
+ def fake_get(url, timeout=None, **kw):
+ return _Resp({"capabilities": {"penalties": True, "ignore_eos": True}})
+
+ def fake_post(url, json=None, timeout=None, **kw):
+ if url.endswith("/pd/decode"):
+ seen["decode"] = json
+ return _Resp(
+ {"rid": "x", "token_ids": [7], "seq_len": 8, "timing_ms": {"finish_reason": "stop"}}
+ )
+ seen["prefill"] = json
+ return _Resp(
+ {
+ "id": "cmpl-abc",
+ "choices": [{"logprobs": {"content": [{"token": "token_id:7"}]}}],
+ "usage": {"prompt_tokens": 3},
+ "model": "m",
+ }
+ )
+
+ monkeypatch.setattr(pd_router.requests, "get", fake_get)
+ monkeypatch.setattr(pd_router.requests, "post", fake_post)
+ pool = pd_router.Pool([pd_router.DecodeNode("127.0.0.1", 5556, 5557)])
+ ctx = pd_router.RouterCtx(
+ "http://vllm.invalid", pool, tokenizer=None, parser_name="none", gen_defaults=defaults
+ )
+ r = TestClient(pd_router.build_app(ctx)).post(
+ "/v1/chat/completions", json={"messages": [{"role": "user", "content": "hi"}], **body}
+ )
+ assert r.status_code == 200, r.text
+ return seen["prefill"], seen["decode"]
+
+
+def test_an_absent_top_p_is_pinned_identically_on_both_legs(monkeypatch):
+ prefill, decode = _both_legs(monkeypatch, {"temperature": 0.7})
+ assert prefill["top_p"] == decode["sampling"]["top_p"] == 1.0
+
+
+def test_an_explicit_top_p_reaches_both_legs_unchanged(monkeypatch):
+ prefill, decode = _both_legs(monkeypatch, {"top_p": 0.8})
+ assert prefill["top_p"] == decode["sampling"]["top_p"] == 0.8
+
+
+def test_the_deployment_default_reaches_both_legs(monkeypatch):
+ """A deployment's resolved defaults must land on both legs identically.
+
+ This is what keeps a checkpoint's recommended sampling (say temperature
+ 0.6 / top_p 0.95 / top_k 20) from applying to the first token only.
+ """
+ prefill, decode = _both_legs(
+ monkeypatch, {}, defaults=GenerationDefaults(temperature=0.6, top_p=0.95, top_k=20)
+ )
+ for field, want in (("temperature", 0.6), ("top_p", 0.95), ("top_k", 20)):
+ assert prefill[field] == decode["sampling"][field] == want, field
+
+
+def test_the_prefill_leg_is_pinned_even_for_a_greedy_request(monkeypatch):
+ """Greedy ignores top_p, but leaving it unset on one leg only would let a
+ generation_config default reappear the moment temperature rises.
+ """
+ prefill, decode = _both_legs(monkeypatch, {"temperature": 0.0})
+ assert prefill["top_p"] == decode["sampling"]["top_p"] == 1.0
+
+
+@pytest.mark.parametrize(
+ "flag",
+ [
+ "--model",
+ "--generation-config",
+ "--default-temperature",
+ "--default-top-p",
+ "--default-top-k",
+ "--default-repetition-penalty",
+ ],
+)
+def test_the_cli_exposes_the_default_resolution_knobs(flag):
+ """Mirrors vLLM's own surface (--generation-config plus per-key overrides).
+
+ Without them, matching a checkpoint's recommended sampling would mean
+ editing code -- and the last time that was true, three adapters ended up
+ each carrying their own literal.
+ """
+ src = pathlib.Path(pd_router.__file__).read_text()
+ assert f'"{flag}"' in src
+
+
+# --------------------------------------------------------------------------- #
+# ignore_eos is honoured by the adapter
+# --------------------------------------------------------------------------- #
+def _mla_adapter():
+ ad = object.__new__(MlaNsaEngineAdapter)
+ ad.gen = types.SimpleNamespace(update_sampling_params=lambda **kw: None)
+ ad.with_mtp = False
+ ad.mtp_seq_len = 4
+ ad.max_seq_len = 4096
+ ad._seq_len = 8
+ ad.stop_ids = {7, 8}
+ ad._ignore_eos = False
+ ad.last_stats = {}
+ return ad
+
+
+def test_the_mla_nsa_adapter_still_honours_the_flag():
+ ad = _mla_adapter()
+ ad.decode(5, 0, {"ignore_eos": True}, cancel_event=None, grammar_session=None)
+ assert ad._ignore_eos is True
+
+
+@pytest.mark.parametrize("adapter_factory", [_mla_adapter])
+def test_every_adapter_declares_whether_it_honours_ignore_eos(adapter_factory):
+ """The router refuses the field on any node that does not declare it, so an
+ adapter that honours it and stays silent loses the feature.
+ """
+ ad = adapter_factory()
+ assert ad.supports_ignore_eos() is True
+
+
+# --------------------------------------------------------------------------- #
+# Penalties: declared per adapter, and refused rather than silently dropped
+# --------------------------------------------------------------------------- #
+def test_the_mla_nsa_adapter_declares_no_penalty_support():
+ assert _mla_adapter().supports_penalties() is False
+
+
+@pytest.mark.parametrize(
+ "sampling",
+ [
+ {"repetition_penalty": 1.2},
+ {"presence_penalty": 0.4},
+ ],
+)
+def test_the_mla_nsa_adapter_refuses_a_penalty_it_cannot_apply(sampling):
+ """It used to accept the parameter and decode unpenalised -- the one thing
+ an adapter must never do.
+ """
+ with pytest.raises(NotImplementedError):
+ _mla_adapter().decode(5, 0, sampling, cancel_event=None, grammar_session=None)
+
+
+@pytest.mark.parametrize(
+ "sampling",
+ [
+ {},
+ {"repetition_penalty": 1.0},
+ {"presence_penalty": 0.0},
+ {"repetition_penalty": None, "presence_penalty": None},
+ ],
+)
+def test_the_mla_nsa_adapter_serves_neutral_penalties(sampling):
+ ad = _mla_adapter()
+ assert ad.decode(5, 0, sampling, cancel_event=None, grammar_session=None) == [5]
diff --git a/tests/pd_vllm/test_slot_availability.py b/tests/pd_vllm/test_slot_availability.py
new file mode 100644
index 0000000..23148b6
--- /dev/null
+++ b/tests/pd_vllm/test_slot_availability.py
@@ -0,0 +1,1013 @@
+"""A decode node's single slot must free promptly, and "busy" must reach the
+client as 429.
+
+Two failures that share a cause -- the slot outliving the request that needed it:
+
+* **Cancel during the KV transfer was a no-op.** `cancel_event` was created only
+ once decoding began, so a cancel arriving during the wire-wait found it `None`
+ and answered 404 while the wait ran to `timeout_s` (120 s by default). With
+ bs=1 that took the whole node out of service for two minutes every time a
+ client hung up early.
+* **A busy decode node reached the client as 502.** `raise_for_status` turned
+ "retry shortly" into "a component is broken", which is also what the release
+ notes then told the operator to go and do.
+
+CPU only -- no GPU, no tilert, no vLLM.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_slot_availability.py -v
+"""
+
+import json
+import queue
+import threading
+import time
+import types
+from unittest import mock
+
+import pytest
+from fastapi.testclient import TestClient
+
+from tilert.pd_vllm import decode_pool, decode_server, pd_router
+from tilert.pd_vllm.decode_pool import DecodeNode, NodeLease, Pool
+from tilert.pd_vllm.decode_response import (
+ TYPED_ERROR,
+ DecodeReader,
+ terminal_verdict,
+)
+from tilert.pd_vllm.decode_server import build_app as build_decode_app
+from tilert.pd_vllm.engine_iface import StubEngine
+
+
+# --------------------------------------------------------------------------- #
+# decode server: the wire-wait observes the cancel
+# --------------------------------------------------------------------------- #
+class _FakeReq:
+ rid = "rid-1"
+ seq_len = 8
+ last_prompt_token = 5
+
+
+class _SlowServer:
+ """A receive server whose KV never arrives, so /pd/decode sits in the wait."""
+
+ def __init__(self, deliver_after=None):
+ self.completed: queue.Queue = queue.Queue()
+ self.profile = types.SimpleNamespace(
+ convert=lambda *a, **k: "converted", num_ranks=8, name="stub"
+ )
+ self.buffer = None
+ self.base_ptr = 0
+ self.max_seq_len = 4096
+ self.released = 0
+ if deliver_after is not None:
+ threading.Timer(deliver_after, lambda: self.completed.put(_FakeReq())).start()
+
+ def expect(self, rid=None):
+ # /pd/decode announces its rid so the real ReceiveServer can drop a
+ # tombstone left by a previous attempt at the same request. Recorded,
+ # so a test can assert the announcement happened.
+ self.expected_rids = getattr(self, "expected_rids", [])
+ self.expected_rids.append(rid)
+
+ def release(self, rid=None):
+ # Scoped like the real ReceiveServer.release: the decode server
+ # names the rid it owns, because the slot may since have been
+ # handed to a later request.
+ self.released_rids = getattr(self, "released_rids", [])
+ self.released_rids.append(rid)
+ self.released += 1
+
+
+def _decode_client(server=None, engine=None):
+ return TestClient(build_decode_app(server or _SlowServer(), engine or StubEngine()))
+
+
+def test_a_cancel_during_the_wire_wait_returns_promptly():
+ """The whole point: this used to block for timeout_s."""
+ server = _SlowServer()
+ client = _decode_client(server)
+ result = {}
+
+ def _decode():
+ t0 = time.time()
+ r = client.post(
+ "/pd/decode",
+ json={"rid": "rid-1", "first_token_id": 7, "max_tokens": 8, "timeout_s": 30.0},
+ )
+ result["status"] = r.status_code
+ result["body"] = r.json()
+ result["elapsed"] = time.time() - t0
+
+ t = threading.Thread(target=_decode)
+ t.start()
+ # Let it reach the wire-wait, then hang up.
+ time.sleep(0.3)
+ assert client.post("/pd/cancel", json={"rid": "rid-1"}).status_code == 200
+ t.join(timeout=10)
+ assert not t.is_alive(), "the wire-wait ignored the cancel"
+ assert (
+ result["elapsed"] < 5
+ ), f"took {result['elapsed']:.1f}s; the cancel should land within one poll"
+ assert result["status"] == 499
+ assert result["body"]["error_type"] == "request_cancelled"
+
+
+def test_a_cancel_during_the_wire_wait_frees_the_slot():
+ """Otherwise the next request is turned away with 429 for two minutes."""
+ server = _SlowServer()
+ client = _decode_client(server)
+ t = threading.Thread(
+ target=lambda: client.post(
+ "/pd/decode",
+ json={"rid": "rid-1", "first_token_id": 7, "max_tokens": 8, "timeout_s": 30.0},
+ )
+ )
+ t.start()
+ time.sleep(0.3)
+ client.post("/pd/cancel", json={"rid": "rid-1"})
+ t.join(timeout=10)
+ assert client.get("/decode_status").json()["status"] == "idle"
+ assert server.released >= 1
+
+
+def test_cancel_is_accepted_from_the_moment_the_request_is_admitted():
+ """It used to 404 until decoding began, which is after the wire-wait."""
+ server = _SlowServer()
+ client = _decode_client(server)
+ t = threading.Thread(
+ target=lambda: client.post(
+ "/pd/decode",
+ json={"rid": "rid-1", "first_token_id": 7, "max_tokens": 8, "timeout_s": 30.0},
+ )
+ )
+ t.start()
+ time.sleep(0.3)
+ r = client.post("/pd/cancel", json={"rid": "rid-1"})
+ assert r.status_code == 200, r.json()
+ assert r.json()["cancelled"] == "rid-1"
+ t.join(timeout=10)
+
+
+def test_a_cancel_for_another_rid_is_still_a_404():
+ server = _SlowServer()
+ client = _decode_client(server)
+ t = threading.Thread(
+ target=lambda: client.post(
+ "/pd/decode",
+ json={"rid": "rid-1", "first_token_id": 7, "max_tokens": 8, "timeout_s": 30.0},
+ )
+ )
+ t.start()
+ time.sleep(0.3)
+ assert client.post("/pd/cancel", json={"rid": "other"}).status_code == 404
+ client.post("/pd/cancel", json={"rid": "rid-1"})
+ t.join(timeout=10)
+
+
+def test_an_uncancelled_wait_still_times_out_as_504():
+ """The cancel path must not swallow the timeout it shares a loop with: a
+ 504 points at the RDMA path, a cancel means the client left.
+ """
+ r = _decode_client(_SlowServer()).post(
+ "/pd/decode", json={"rid": "rid-1", "first_token_id": 7, "max_tokens": 8, "timeout_s": 0.6}
+ )
+ assert r.status_code == 504
+ assert r.json()["error"] == "kv_transfer_timeout"
+
+
+def test_kv_that_arrives_before_any_cancel_is_served():
+ server = _SlowServer(deliver_after=0.2)
+ r = _decode_client(server).post(
+ "/pd/decode", json={"rid": "rid-1", "first_token_id": 7, "max_tokens": 4, "timeout_s": 30.0}
+ )
+ assert r.status_code == 200
+ assert r.json()["token_ids"][0] == 7
+
+
+def test_the_cancel_sentinel_is_distinct_from_a_timeout():
+ """Two different answers from one loop; conflating them would report a
+ client hang-up as an RDMA fault.
+ """
+ assert decode_server._CANCELLED is not None
+ assert decode_server._CANCELLED is not False
+
+
+# --------------------------------------------------------------------------- #
+# router: a busy decode node is retried once, then surfaced as 429
+# --------------------------------------------------------------------------- #
+class _Resp:
+ def __init__(self, payload, status=200):
+ self._payload = payload
+ self.status_code = status
+
+ def json(self):
+ return self._payload
+
+ def raise_for_status(self):
+ if self.status_code >= 400:
+ raise RuntimeError(f"HTTP {self.status_code}")
+
+
+PREFILL = {
+ "id": "cmpl-abc",
+ "choices": [{"logprobs": {"content": [{"token": "token_id:7"}]}}],
+ "usage": {"prompt_tokens": 3},
+ "model": "m",
+}
+DECODED = {"rid": "x", "token_ids": [7], "seq_len": 8, "timing_ms": {"finish_reason": "stop"}}
+BODY = {"messages": [{"role": "user", "content": "hi"}]}
+
+
+def _router(monkeypatch, decode_statuses):
+ """Router whose decode node answers with each status in turn.
+
+ Returns (client, calls) where calls counts /pd/decode attempts.
+ """
+ statuses = list(decode_statuses)
+ calls = {"decode": 0}
+
+ monkeypatch.setattr(
+ pd_router.requests,
+ "get",
+ lambda url, timeout=None, **kw: _Resp(
+ {"capabilities": {"penalties": True, "ignore_eos": True}}
+ ),
+ )
+ monkeypatch.setattr(pd_router.time, "sleep", lambda s: None)
+
+ def fake_post(url, json=None, timeout=None, **kw):
+ if url.endswith("/pd/decode"):
+ calls["decode"] += 1
+ status = statuses.pop(0) if statuses else 200
+ if status == 200:
+ return _Resp(DECODED)
+ return _Resp({"error": "busy", "current_rid": "other"}, status)
+ return _Resp(PREFILL)
+
+ monkeypatch.setattr(pd_router.requests, "post", fake_post)
+ pool = pd_router.Pool([pd_router.DecodeNode("127.0.0.1", 5556, 5557)])
+ ctx = pd_router.RouterCtx("http://vllm.invalid", pool, tokenizer=None, parser_name="none")
+ return TestClient(pd_router.build_app(ctx)), calls, pool
+
+
+def test_a_transient_busy_is_retried_and_succeeds(monkeypatch):
+ """The router frees its own reservation as soon as it stops reading, while
+ the node's slot unwinds a little later -- a request dispatched into that
+ window meets the node's admission.
+ """
+ client, calls, _ = _router(monkeypatch, [429, 200])
+ r = client.post("/v1/chat/completions", json=BODY)
+ assert r.status_code == 200
+ assert calls["decode"] == 2
+
+
+def test_a_node_still_busy_after_the_retry_is_429_not_502(monkeypatch):
+ """429 is retryable and truthful; 502 sends the operator to restart a
+ healthy component.
+ """
+ client, calls, _ = _router(monkeypatch, [429, 429])
+ r = client.post("/v1/chat/completions", json=BODY)
+ assert r.status_code == 429
+ assert r.json()["error_type"] == "decode_busy"
+ assert calls["decode"] == pd_router._DECODE_BUSY_ATTEMPTS
+
+
+def test_the_retry_is_bounded(monkeypatch):
+ """One short retry absorbs the handover window; more would mask a node that
+ is genuinely stuck behind a status the client can act on.
+ """
+ client, calls, _ = _router(monkeypatch, [429] * 10)
+ client.post("/v1/chat/completions", json=BODY)
+ assert calls["decode"] == 2
+
+
+def test_a_busy_node_is_returned_to_the_pool(monkeypatch):
+ """A 429 must not leak the reservation, or the pool drains one node per
+ busy reply.
+ """
+ client, _, pool = _router(monkeypatch, [429, 429])
+ client.post("/v1/chat/completions", json=BODY)
+ assert all(not n.busy for n in pool.nodes)
+
+
+def test_a_healthy_node_is_not_retried(monkeypatch):
+ client, calls, _ = _router(monkeypatch, [200])
+ assert client.post("/v1/chat/completions", json=BODY).status_code == 200
+ assert calls["decode"] == 1
+
+
+def test_a_non_busy_decode_failure_is_still_502(monkeypatch):
+ """The retry is for 429 only; an untyped 500 is a component call that did
+ not work, and saying so is correct.
+ """
+ client, calls, _ = _router(monkeypatch, [500])
+ r = client.post("/v1/chat/completions", json=BODY)
+ assert r.status_code == 502
+ assert calls["decode"] == 1
+
+
+def test_a_typed_decode_error_still_propagates(monkeypatch):
+ """The busy branch must not swallow the classified statuses."""
+
+ def fake_post(url, json=None, timeout=None, **kw):
+ if url.endswith("/pd/decode"):
+ return _Resp({"error": "no", "error_type": "invalid_grammar"}, 400)
+ return _Resp(PREFILL)
+
+ monkeypatch.setattr(
+ pd_router.requests, "get", lambda url, timeout=None, **kw: _Resp({"capabilities": {}})
+ )
+ monkeypatch.setattr(pd_router.requests, "post", fake_post)
+ pool = pd_router.Pool([pd_router.DecodeNode("127.0.0.1", 5556, 5557)])
+ ctx = pd_router.RouterCtx("http://vllm.invalid", pool, tokenizer=None, parser_name="none")
+ client = TestClient(pd_router.build_app(ctx))
+ r = client.post("/v1/chat/completions", json=BODY)
+ assert r.status_code == 400
+ assert r.json()["error_type"] == "invalid_grammar"
+
+
+def test_request_cancelled_is_propagated_not_masked():
+ """The decode node's 499 says the caller asked it to stop; flattening that
+ into a 502 would read as a component fault.
+
+ Asserted through the verdict both response paths take, rather than by
+ membership in a table: the table is where the answer comes from, but the
+ verdict is what the handlers act on.
+ """
+ reader = DecodeReader(stream=None, logprobs_req=None, rid="rid-1")
+ reader.feed(json.dumps({"error": "client left", "error_type": "request_cancelled"}))
+ verdict, payload, status = terminal_verdict(reader)
+ assert verdict == TYPED_ERROR, verdict
+ assert status == 499, status
+ assert payload["error_type"] == "request_cancelled"
+
+
+# --------------------------------------------------------------------------- #
+# Streaming: the status is decided before the response begins
+# --------------------------------------------------------------------------- #
+#
+# The decode request is now sent BEFORE StreamingResponse is returned. Inside the
+# generator the response has already begun and the status is spent, so a busy node
+# could only have been reported as an SSE error inside a 200 -- or, as it was, a
+# stream truncated with no terminator at all. These run against real servers so
+# the httpx streaming path is the real one.
+import json as _json # noqa: E402
+
+import httpx as _httpx # noqa: E402
+import uvicorn as _uvicorn # noqa: E402
+from fastapi import FastAPI as _FastAPI # noqa: E402
+
+
+def _serve(app):
+ cfg = _uvicorn.Config(app, host="127.0.0.1", port=0, log_level="error")
+ server = _uvicorn.Server(cfg)
+ threading.Thread(target=server.run, daemon=True).start()
+ for _ in range(200):
+ if server.started:
+ return server.servers[0].sockets[0].getsockname()[1]
+ time.sleep(0.05)
+ raise RuntimeError("stub server did not start")
+
+
+@pytest.fixture
+def no_proxy():
+ mp = pytest.MonkeyPatch()
+ for var in ("no_proxy", "NO_PROXY"):
+ mp.setenv(var, "127.0.0.1,localhost")
+ for var in ("http_proxy", "HTTP_PROXY", "https_proxy", "HTTPS_PROXY", "all_proxy", "ALL_PROXY"):
+ mp.delenv(var, raising=False)
+ yield
+ mp.undo()
+
+
+class _StubTokenizer:
+ """The streaming path detokenises every token, so it needs one."""
+
+ _VOCAB = {7: "hi", 11: " there"}
+
+ def decode(self, ids, skip_special_tokens=False):
+ return "".join(self._VOCAB.get(i, "") for i in ids)
+
+
+def _stream_stack(decode_statuses):
+ """vLLM stub + a decode stub answering `decode_statuses` in turn + router."""
+ statuses = list(decode_statuses)
+ calls = {"decode": 0}
+
+ vllm = _FastAPI()
+
+ @vllm.post("/v1/chat/completions")
+ async def _prefill():
+ return PREFILL
+
+ decode = _FastAPI()
+
+ @decode.get("/capabilities")
+ def _caps():
+ return {"capabilities": {"penalties": True, "ignore_eos": True, "logprobs": True}}
+
+ @decode.post("/pd/decode")
+ def _dec():
+ from fastapi.responses import JSONResponse, StreamingResponse
+
+ calls["decode"] += 1
+ status = statuses.pop(0) if statuses else 200
+ if status != 200:
+ return JSONResponse({"error": "busy", "current_rid": "other"}, status_code=status)
+
+ def _body():
+ yield _json.dumps({"t": [11]}) + "\n"
+ yield _json.dumps(
+ {"done": True, "n": 1, "seq_len": 8, "finish_reason": "stop", "timing_ms": {}}
+ ) + "\n"
+
+ return StreamingResponse(_body(), media_type="application/x-ndjson")
+
+ @decode.post("/pd/cancel")
+ def _cancel(b: dict):
+ return {"cancelled": b.get("rid")}
+
+ vllm_port = _serve(vllm)
+ decode_port = _serve(decode)
+ node = pd_router.DecodeNode("127.0.0.1", 5556, decode_port)
+ pool = pd_router.Pool([node])
+ ctx = pd_router.RouterCtx(
+ f"http://127.0.0.1:{vllm_port}", pool, tokenizer=_StubTokenizer(), parser_name="none"
+ )
+ router_port = _serve(pd_router.build_app(ctx))
+ return f"http://127.0.0.1:{router_port}", calls, pool
+
+
+STREAM_BODY = {"messages": [{"role": "user", "content": "hi"}], "stream": True}
+
+
+def test_a_streamed_request_gets_a_real_429_not_a_truncated_stream(no_proxy):
+ """It used to raise inside the generator, so the client got a 200 whose body
+ simply stopped -- no finish_reason, no [DONE], nothing to distinguish it from
+ a network cut.
+ """
+ url, calls, _ = _stream_stack([429, 429])
+ with _httpx.Client(timeout=30, trust_env=False) as c:
+ r = c.post(f"{url}/v1/chat/completions", json=STREAM_BODY)
+ assert r.status_code == 429
+ assert r.json()["error_type"] == "decode_busy"
+ assert calls["decode"] == 2
+
+
+def test_a_streamed_request_retries_a_transient_busy(no_proxy):
+ url, calls, _ = _stream_stack([429, 200])
+ chunks = []
+ with _httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream("POST", f"{url}/v1/chat/completions", json=STREAM_BODY) as r:
+ assert r.status_code == 200
+ for line in r.iter_lines():
+ if line.startswith("data: "):
+ chunks.append(line[6:])
+ assert chunks[-1] == "[DONE]"
+ assert calls["decode"] == 2
+
+
+def test_a_streamed_request_still_streams_normally(no_proxy):
+ """The restructure moved the send out of the generator; the happy path must
+ be untouched.
+ """
+ url, calls, _ = _stream_stack([200])
+ seen = []
+ with _httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream("POST", f"{url}/v1/chat/completions", json=STREAM_BODY) as r:
+ assert r.status_code == 200
+ for line in r.iter_lines():
+ if line.startswith("data: ") and line[6:] != "[DONE]":
+ seen.append(_json.loads(line[6:]))
+ assert seen, "no chunks at all"
+ assert any(
+ c["choices"] and c["choices"][0].get("finish_reason") for c in seen
+ ), "no terminating chunk"
+ assert calls["decode"] == 1
+
+
+def test_a_busy_streamed_request_returns_the_node_to_the_pool(no_proxy):
+ url, _, pool = _stream_stack([429, 429])
+ with _httpx.Client(timeout=30, trust_env=False) as c:
+ c.post(f"{url}/v1/chat/completions", json=STREAM_BODY)
+ assert all(not n.busy for n in pool.nodes)
+
+
+def test_a_streamed_typed_decode_error_keeps_its_status(no_proxy):
+ """A classified 400 must survive the new pre-flight, not become a 502."""
+ statuses = [400]
+ vllm = _FastAPI()
+
+ @vllm.post("/v1/chat/completions")
+ async def _prefill():
+ return PREFILL
+
+ decode = _FastAPI()
+
+ @decode.get("/capabilities")
+ def _caps():
+ return {"capabilities": {}}
+
+ @decode.post("/pd/decode")
+ def _dec():
+ from fastapi.responses import JSONResponse
+
+ statuses.pop(0)
+ return JSONResponse({"error": "bad", "error_type": "invalid_grammar"}, status_code=400)
+
+ vllm_port = _serve(vllm)
+ decode_port = _serve(decode)
+ ctx = pd_router.RouterCtx(
+ f"http://127.0.0.1:{vllm_port}",
+ pd_router.Pool([pd_router.DecodeNode("127.0.0.1", 5556, decode_port)]),
+ tokenizer=_StubTokenizer(),
+ parser_name="none",
+ )
+ port = _serve(pd_router.build_app(ctx))
+ with _httpx.Client(timeout=30, trust_env=False) as c:
+ r = c.post(f"http://127.0.0.1:{port}/v1/chat/completions", json=STREAM_BODY)
+ assert r.status_code == 400
+ assert r.json()["error_type"] == "invalid_grammar"
+
+
+# --------------------------------------------------------------------------- #
+# Review findings on PR #41 (codex)
+# --------------------------------------------------------------------------- #
+def test_release_is_scoped_to_the_rid_that_owns_the_slot():
+ """A caller only knows about its own request, and the slot may have moved on.
+
+ Reachable whenever a transfer arrives after its consumer gave up: the entry
+ is enqueued with nobody waiting, the NEXT request drains it as unmatched, and
+ an unscoped release there would free that next request's own tenancy —
+ which then never completes, because a cancelled tenancy drops its ranks'
+ `done` messages.
+ """
+ import types as _t
+
+ from tilert.pd_vllm import receive_server as rs
+
+ srv = object.__new__(rs.ReceiveServer)
+ srv._lock = threading.Lock()
+ srv._cancelled = {}
+ srv.request_timeout = 120.0
+ srv._current = rs.ReceivedRequest(
+ rid="current",
+ seq_len=8,
+ last_prompt_token=5,
+ first_token_id=None,
+ sampling=None,
+ state=rs.TRANSFERRING,
+ active_writers=1,
+ )
+
+ srv.release("someone-else")
+ assert srv._current is not None, "released a tenancy it does not own"
+ assert srv._current.state == rs.TRANSFERRING, "cancelled the wrong tenancy"
+
+ srv.release("current")
+ assert srv._current.state == rs.CANCELLING, "its own release had no effect"
+ del _t
+
+
+def test_a_late_transfer_does_not_disturb_the_next_request():
+ """End-to-end shape of the same bug, through /pd/decode.
+
+ rid-A is cancelled before any sender arrives, so nothing is released. Its KV
+ then lands with no consumer. rid-B is in flight when the next request drains
+ that stale entry — and must survive it.
+ """
+ server = _SlowServer()
+ client = _decode_client(server)
+
+ t = threading.Thread(
+ target=lambda: client.post(
+ "/pd/decode",
+ json={"rid": "rid-A", "first_token_id": 7, "max_tokens": 4, "timeout_s": 30.0},
+ )
+ )
+ t.start()
+ time.sleep(0.3)
+ client.post("/pd/cancel", json={"rid": "rid-A"})
+ t.join(timeout=10)
+
+ # rid-A's transfer arrives late, with nobody waiting for it.
+ server.completed.put(_FakeReq()) # _FakeReq.rid == "rid-1"
+ # The next request drains it as unmatched; the release must name rid-1, not
+ # whatever the slot holds now.
+ r = client.post(
+ "/pd/decode", json={"rid": "rid-B", "first_token_id": 7, "max_tokens": 4, "timeout_s": 1.0}
+ )
+ assert r.status_code == 504, r.text
+ assert server.released_rids, "nothing was released at all"
+ assert "rid-1" in server.released_rids, (
+ f"the stale entry was not released by its own rid: " f"{server.released_rids}"
+ )
+
+
+def test_the_streaming_preflight_watches_for_a_disconnect():
+ """The decode node holds its headers through the whole KV wire-wait, so the
+ preflight `send` can sit for `timeout_s`.
+
+ At that point StreamingResponse does not exist, so neither the generator's
+ `finally` nor its `is_disconnected` poll is running — a client that hangs up
+ there would hold the router's reservation and the decode slot for the full
+ wait, which is the failure this endpoint is supposed to have stopped having.
+ """
+ import pathlib as _p
+
+ src = _p.Path(pd_router.__file__).read_text()
+ # Defined AND used: a helper that exists but is not called on the preflight
+ # path leaves the disconnect unwatched, which is the whole bug.
+ assert (
+ src.count("_send_watching_client") >= 2
+ ), "the disconnect-aware helper is defined but never called"
+ assert (
+ "await _send_watching_client(" in src
+ ), "the preflight does not go through the disconnect-aware helper"
+ assert (
+ "client.send(" not in src.split("for attempt in range")[1][:600]
+ ), "the preflight still sends directly, bypassing the disconnect watch"
+ assert (
+ "request.is_disconnected()" in src.split("async def _send_watching_client")[1][:1200]
+ ), "the helper does not poll for a disconnect"
+ # CancelledError is a BaseException; suppress(Exception) would let it escape
+ # and turn a clean 499 into a 500.
+ assert "asyncio.CancelledError" in src
+
+
+@pytest.mark.parametrize(
+ "dispatched,terminated,want_cancel,why",
+ [
+ (True, False, True, "a request went out and the node never said done"),
+ (
+ True,
+ True,
+ False,
+ "the node reported done; cancelling could land on the " "NEXT request for that slot",
+ ),
+ (False, False, False, "nothing went out, so there is nothing to cancel"),
+ ],
+)
+def test_a_lease_cancels_exactly_when_the_node_may_still_be_working(
+ dispatched, terminated, want_cancel, why
+):
+ """Releasing the router's own reservation is not enough -- the decode node is
+ a separate process still holding its slot.
+
+ This was spelled out at five call sites across the two handlers, and the
+ fixes for it landed at some of them; `NodeLease` is the rule once. It used to
+ be checked by grepping pd_router.py for a cancel call, which passed for any
+ spelling and failed for a correct rename.
+ """
+ fired = []
+ node = DecodeNode("127.0.0.1", 5556, 5557)
+ pool = Pool([node])
+ assert pool.acquire() is node
+ lease = NodeLease(pool, node)
+ lease.rid = "rid-1"
+ lease.dispatched = dispatched
+ with mock.patch.object(decode_pool, "cancel_decode", lambda n, rid: fired.append(rid)):
+ lease.release(terminated=terminated)
+ lease.release(terminated=terminated) # idempotent
+ for _ in range(50):
+ if fired:
+ break
+ time.sleep(0.02)
+ assert fired == (["rid-1"] if want_cancel else []), why
+ assert not node.busy, "the slot goes back exactly once, on every exit"
+ assert pool.acquire() is node, "and the node is reusable"
+
+
+# --------------------------------------------------------------------------- #
+# the abandon drain observes the cancel too
+# --------------------------------------------------------------------------- #
+def test_a_cancel_during_the_abandon_drain_returns_promptly():
+ """A rejected request's drain used to ignore the cancel it reported taking.
+
+ When decode-side validation refuses a request its KV may never arrive (the
+ prefill leg died, say), so _abandon_pending_kv drains before releasing. That
+ drain ran to _ABANDON_DRAIN_S while /pd/cancel answered 200 -- the event is
+ armed and this rid is still current -- so the client was told the request was
+ cancelled and the next one was refused 429 for another 30 s.
+ """
+ server = _SlowServer()
+ client = _decode_client(server)
+ result = {}
+
+ def _decode():
+ t0 = time.time()
+ # An unknown grammar type is refused post-admission, which is what
+ # sends us down the abandon path.
+ r = client.post(
+ "/pd/decode",
+ json={
+ "rid": "rid-1",
+ "first_token_id": 7,
+ "max_tokens": 8,
+ "grammar_spec": {"type": "__nope__"},
+ "timeout_s": 30.0,
+ },
+ )
+ result["status"] = r.status_code
+ result["elapsed"] = time.time() - t0
+
+ t = threading.Thread(target=_decode)
+ t.start()
+ time.sleep(0.3) # let it reach the abandon drain
+ assert (
+ client.post("/pd/cancel", json={"rid": "rid-1"}).status_code == 200
+ ), "the cancel was refused, so this test is not exercising the drain"
+ t.join(timeout=15)
+ assert not t.is_alive(), "the abandon drain ignored the cancel"
+ assert result["elapsed"] < 5, (
+ f"took {result['elapsed']:.1f}s; a cancelled abandon drain must land "
+ f"within one poll, not at _ABANDON_DRAIN_S "
+ f"({decode_server._ABANDON_DRAIN_S:.0f}s)"
+ )
+ # The request is still refused for its own reason -- the cancel only ends
+ # the wait, it does not change the answer.
+ assert result["status"] == 400, result
+
+
+def test_the_slot_is_free_after_a_cancelled_abandon_drain():
+ server = _SlowServer()
+ client = _decode_client(server)
+ t = threading.Thread(
+ target=lambda: client.post(
+ "/pd/decode",
+ json={
+ "rid": "rid-1",
+ "first_token_id": 7,
+ "max_tokens": 8,
+ "grammar_spec": {"type": "__nope__"},
+ "timeout_s": 30.0,
+ },
+ )
+ )
+ t.start()
+ time.sleep(0.3)
+ client.post("/pd/cancel", json={"rid": "rid-1"})
+ t.join(timeout=15)
+ assert client.get("/decode_status").json()["status"] == "idle"
+
+
+# --------------------------------------------------------------------------- #
+# a rid nobody is waiting for must not open a tenancy
+# --------------------------------------------------------------------------- #
+def _bare_receive_server(request_timeout=120.0):
+ """A ReceiveServer with just the admission state, no sockets or buffer."""
+ from tilert.pd_vllm import receive_server as rs
+
+ srv = object.__new__(rs.ReceiveServer)
+ srv._lock = threading.Lock()
+ srv._current = None
+ srv._cancelled = {}
+ srv._generation = 0
+ srv.request_timeout = request_timeout
+ return srv
+
+
+def _admit(srv, rid, rank=0, seq_len=8):
+ return srv._admit({"seq_len": seq_len, "last_prompt_token": 5}, rid, rank)
+
+
+def test_a_rank_arriving_after_its_request_was_released_is_refused():
+ """The hole rid-scoping alone left open.
+
+ release() with no writer yet was a plain return: nothing recorded that the
+ rid was dead. The straggler then found a FREE buffer, was admitted, and held
+ it while the next request's ranks were turned away "busy" until they
+ exhausted their retries -- that request then waited out its whole
+ kv_transfer_timeout.
+ """
+ srv = _bare_receive_server()
+ srv.release("rid-dead") # cancelled before any sender arrived
+ assert srv._current is None, "precondition: nothing holds the slot"
+
+ reply = _admit(srv, "rid-dead")
+ assert reply["accepted"] is False, "the straggler opened a tenancy"
+ # `cancelling`, which every connector already retries; a new reason would
+ # be permanent to an older prefill and drop the shard during a rolling
+ # upgrade. The distinction lives in `detail`.
+ assert reply["error"] == "cancelling"
+ assert reply["detail"] == "no_consumer"
+ assert srv._current is None, "a refused straggler must claim nothing"
+
+
+def test_a_released_rid_does_not_block_the_next_request():
+ """The refusal has to be scoped to the dead rid, or it is a worse bug."""
+ srv = _bare_receive_server()
+ srv.release("rid-dead")
+
+ reply = _admit(srv, "rid-next")
+ assert reply["accepted"] is True, "a different rid was caught by the tombstone"
+ assert srv._current.rid == "rid-next"
+
+
+def test_a_re_announced_rid_is_admitted_again():
+ """vLLM reuses the request id when it reschedules a preempted request.
+
+ So a tombstone must not outlive the request coming back. /pd/decode calling
+ expect() is the announcement, and it is the only thing that separates the
+ retry from a straggler of the attempt before it.
+ """
+ srv = _bare_receive_server()
+ srv.release("rid-1")
+ assert _admit(srv, "rid-1")["accepted"] is False
+
+ srv.expect("rid-1") # what /pd/decode does on admission
+ reply = _admit(srv, "rid-1")
+ assert reply["accepted"] is True, "the retry lost its shard to a tombstone"
+ assert srv._current.rid == "rid-1"
+
+
+def test_a_tombstone_ages_out():
+ """The backstop, in case no consumer ever re-announces the rid.
+
+ Bounded by request_timeout because that is the senders' socket timeout: past
+ it no rank can still be trying to join, so keeping the entry only leaks it.
+ """
+ srv = _bare_receive_server(request_timeout=0.05)
+ srv.release("rid-1")
+ assert _admit(srv, "rid-1")["accepted"] is False
+ time.sleep(0.1)
+ assert _admit(srv, "rid-1")["accepted"] is True
+
+
+def test_a_straggler_is_refused_transiently_so_a_retry_can_still_land():
+ """The tombstone refusal must be a reason senders already retry.
+
+ A permanent one drops the shard of a rescheduled request whose /pd/decode
+ has not arrived yet -- the same 120 s stall this fix is for, caused by the
+ fix. A NEW reason is permanent to every connector built before it, so during
+ a rolling upgrade a new decode node would do exactly that to an old prefill;
+ the tombstone therefore reuses one the sender already knows.
+
+ The reason set is read from the connector's source: it pulls in vLLM, and
+ this file is CPU-only. Parsed, not pattern-matched.
+ """
+ import ast as _ast
+ import pathlib as _p
+
+ src = _p.Path(pd_router.__file__).with_name("prefill_connector.py").read_text()
+ for node in _ast.walk(_ast.parse(src)):
+ if isinstance(node, _ast.Assign) and any(
+ getattr(t, "id", None) == "_TRANSIENT_REJECTS" for t in node.targets
+ ):
+ value = node.value
+ if isinstance(value, _ast.Call): # frozenset({...})
+ value = value.args[0]
+ reasons = _ast.literal_eval(value)
+ break
+ else:
+ raise AssertionError("_TRANSIENT_REJECTS not found in the connector")
+
+ srv = _bare_receive_server()
+ srv.release("rid-1")
+ assert (
+ _admit(srv, "rid-1")["error"] in reasons
+ ), "the tombstone refusal is not one the sender comes back from"
+
+
+def test_pd_decode_announces_its_rid():
+ """Otherwise a rescheduled request's senders meet the old tombstone."""
+ server = _SlowServer()
+ client = _decode_client(server)
+ client.post(
+ "/pd/decode", json={"rid": "rid-1", "first_token_id": 7, "max_tokens": 8, "timeout_s": 0.2}
+ )
+ assert getattr(server, "expected_rids", []) == ["rid-1"]
+
+
+# --------------------------------------------------------------------------- #
+# The two ways the retry budgets can disagree
+# --------------------------------------------------------------------------- #
+def test_the_send_never_runs_inside_the_forward_window():
+ """Sending from the forward window cannot work with admission.
+
+ The admission retries would all run before the prefill response returns,
+ and the router cannot call /pd/decode -- the only thing that clears a
+ tombstone for a rescheduled rid -- until it has. A sender there is
+ therefore guaranteed to spend its whole budget against a tombstone it
+ cannot outlast, and then drop the shard.
+
+ So `wait_for_save` only ever queues, and the selectable synchronous path
+ that used to bypass the queue is gone.
+ """
+ import ast as _ast
+ import pathlib as _p
+
+ src = _p.Path(pd_router.__file__).with_name("prefill_connector.py").read_text()
+ assert "self._sync_send" not in src, (
+ "the synchronous send path is back; its retries cannot outlast a "
+ "tombstone, because /pd/decode comes after the prefill response"
+ )
+
+ tree = _ast.parse(src)
+ fn = next(
+ n for n in _ast.walk(tree) if isinstance(n, _ast.FunctionDef) and n.name == "wait_for_save"
+ )
+ called = {
+ n.func.attr
+ for n in _ast.walk(fn)
+ if isinstance(n, _ast.Call) and isinstance(n.func, _ast.Attribute)
+ }
+ assert (
+ "_send" not in called and "_send_with_retry" not in called
+ ), f"wait_for_save sends inline: {sorted(called)}"
+
+
+def test_wait_for_save_queues_a_complete_job():
+ """Run it, do not read it.
+
+ Reading the source could tell that `put` is called but not what it is
+ called with -- and a `job` that is never built raises NameError on every
+ single transfer, which no assertion about the source would have noticed.
+ """
+ import queue as _q
+ import types as _t
+
+ from tilert.pd_vllm import prefill_connector as pc
+
+ meta = pc._ReqMeta(
+ req_id="r",
+ rid="rid-1",
+ num_tokens=8,
+ last_prompt_token=5,
+ block_ids_per_group=[],
+ tilert_host="127.0.0.1",
+ tilert_ctrl_port=1,
+ )
+
+ conn = object.__new__(pc.TileRTConnector)
+ conn._send_q = _q.Queue()
+ conn._tp_rank = 0
+ conn._reg = object()
+ conn._staging = _t.SimpleNamespace(data_ptr=lambda: 0x1000)
+ conn._max_seq = 4096
+ conn._profile = _t.SimpleNamespace(
+ sender_ranks=(0,), extract=lambda reg, m, rank, staging, max_seq: {"seq": 8, "x": 1}
+ )
+ conn._ensure_worker_ready = lambda: None
+ conn._get_connector_metadata = lambda: _t.SimpleNamespace(requests=[meta])
+ # The isinstance check in wait_for_save is against the real metadata type.
+ md = pc.TileRTMetadata()
+ md.requests = [meta]
+ conn._get_connector_metadata = lambda: md
+
+ conn.wait_for_save()
+
+ job = conn._send_q.get_nowait()
+ assert job["meta"] is meta
+ assert job["seq"] == 8
+ assert job["sections"] == {"seq": 8, "x": 1}
+
+
+def test_a_tombstone_outlasts_the_senders_own_retry_budget():
+ """request_timeout bounds ONE connection, not the sequence of them.
+
+ Each retry opens a new socket, so a sender configured with enough attempts
+ is still trying after a tombstone sized to one socket timeout has expired
+ -- and is then admitted for a request nobody wants. The sender declares its
+ budget; the receiver sizes the tombstone to outlast it.
+ """
+ srv = _bare_receive_server(request_timeout=0.05)
+ srv.release("rid-1")
+ # Without the declaration, this tombstone expires almost immediately.
+ time.sleep(0.1)
+ assert _admit(srv, "rid-1")["accepted"] is True, "precondition"
+
+ srv = _bare_receive_server(request_timeout=0.05)
+ srv.release("rid-2")
+ reply = srv._admit(
+ {"seq_len": 8, "last_prompt_token": 5, "admission_window_s": 30.0}, "rid-2", 0
+ )
+ assert reply["accepted"] is False
+ time.sleep(0.1) # past request_timeout, inside the window
+ assert (
+ srv._admit({"seq_len": 8, "last_prompt_token": 5, "admission_window_s": 30.0}, "rid-2", 1)[
+ "accepted"
+ ]
+ is False
+ ), "the tombstone expired while the sender was still retrying"
+
+
+def test_a_second_rank_cannot_shorten_the_tombstone():
+ """Never take the smaller of two declared budgets."""
+ srv = _bare_receive_server(request_timeout=0.05)
+ srv.release("rid-1")
+ srv._admit({"seq_len": 8, "last_prompt_token": 5, "admission_window_s": 30.0}, "rid-1", 0)
+ long_deadline = srv._cancelled["rid-1"]
+ srv._admit({"seq_len": 8, "last_prompt_token": 5, "admission_window_s": 0.01}, "rid-1", 1)
+ assert srv._cancelled["rid-1"] == long_deadline
+
+
+def test_a_sender_that_declares_nothing_keeps_the_default():
+ """An older connector sends no budget; the default must still apply."""
+ srv = _bare_receive_server(request_timeout=60.0)
+ srv.release("rid-1")
+ before = srv._cancelled["rid-1"]
+ srv._admit({"seq_len": 8, "last_prompt_token": 5}, "rid-1", 0)
+ assert srv._cancelled["rid-1"] == before
diff --git a/tests/pd_vllm/test_stop_matches_vllm.py b/tests/pd_vllm/test_stop_matches_vllm.py
new file mode 100644
index 0000000..cd8caac
--- /dev/null
+++ b/tests/pd_vllm/test_stop_matches_vllm.py
@@ -0,0 +1,142 @@
+"""``stop`` validation, compared value by value against vLLM's own.
+
+This PR's whole justification for matching vLLM is that the same request should
+behave the same on both stacks -- so the comparison is made by running vLLM's
+validation next to ours, not by reading its source and reimplementing what it
+seems to say. Reading it is how three of these cases were got wrong:
+
+* `include_stop_str_in_output` -- I refused `1` and `"true"`, which vLLM accepts
+ (pydantic coerces them), and silently accepted an explicit `null`, which it
+ refuses. Exactly backwards.
+* `include_stop_str_in_output` without `stop` -- I refused it as meaningless.
+ vLLM serves it; with no stop strings the flag is a no-op.
+
+Both directions matter, and they are not symmetric. Being STRICTER than vLLM
+turns a request a native endpoint serves into a 400 -- the mistake the
+capability gate's `top_k: "20"` handling already warns about. Being MORE LENIENT
+serves something vLLM would have refused, which for `stop` means returning an
+unrestricted completion to a client who asked for a restricted one.
+
+vLLM validates in two layers and both count:
+
+* ``ChatCompletionRequest`` (pydantic) checks types;
+* ``to_sampling_params()`` builds a ``SamplingParams``, whose ``_verify_args``
+ checks values -- and this is where an empty stop string dies, with
+ ``ValueError("stop cannot contain an empty string.")``.
+
+Comparing against only the first layer says vLLM accepts ``stop: [""]``. It does
+not; the serving layer turns that ValueError into a 400.
+
+Needs vllm, so it is skipped in the default dev env:
+
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_stop_matches_vllm.py -v
+
+No GPU and no weights -- request validation only.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+pytest.importorskip("vllm")
+
+from vllm.entrypoints.openai.chat_completion.protocol import ( # noqa: E402
+ ChatCompletionRequest,
+)
+
+from tilert.pd_vllm.capabilities import CapabilityError # noqa: E402
+from tilert.pd_vllm.request_gate import resolve_stop_request # noqa: E402
+
+
+class _Tok:
+ """A tokenizer only has to exist here; nothing decodes anything."""
+
+ def decode(self, ids, skip_special_tokens=False):
+ return ""
+
+
+def _vllm(body: dict):
+ """``(accepted, (stop, include))`` from vLLM's own two layers."""
+ try:
+ req = ChatCompletionRequest(model="m", messages=[{"role": "user", "content": "hi"}], **body)
+ except Exception:
+ return False, None
+ try:
+ params = req.to_sampling_params(max_tokens=8, default_sampling_params={})
+ except Exception:
+ return False, None
+ stop = params.stop
+ stop = [stop] if isinstance(stop, str) else list(stop or [])
+ return True, (stop, params.include_stop_str_in_output)
+
+
+def _ours(body: dict):
+ try:
+ return True, resolve_stop_request(dict(body), _Tok())
+ except CapabilityError:
+ return False, None
+
+
+# Every case is a request a client can actually send. The three that were wrong
+# before this test existed are marked.
+CASES = [
+ {},
+ {"stop": None},
+ {"stop": "END"},
+ {"stop": ["A", "B"]},
+ {"stop": []},
+ {"stop": ["A", "A"]},
+ {"stop": ["\n\n"]},
+ {"stop": 5},
+ {"stop": [1, 2]},
+ {"stop": {"a": 1}},
+ {"stop": [""]}, # refused, layer 2
+ {"stop": ["", "END"]}, # refused, layer 2
+ {"stop": ""}, # refused, layer 2
+ {"stop": ["A"], "include_stop_str_in_output": True},
+ {"stop": ["A"], "include_stop_str_in_output": False},
+ {"stop": ["A"], "include_stop_str_in_output": None}, # was wrong: accepted
+ {"stop": ["A"], "include_stop_str_in_output": 1}, # was wrong: refused
+ {"stop": ["A"], "include_stop_str_in_output": 0}, # was wrong: refused
+ {"stop": ["A"], "include_stop_str_in_output": 2},
+ {"stop": ["A"], "include_stop_str_in_output": "true"}, # was wrong: refused
+ {"stop": ["A"], "include_stop_str_in_output": "False"},
+ {"stop": ["A"], "include_stop_str_in_output": "yes"},
+ {"stop": ["A"], "include_stop_str_in_output": "maybe"},
+ {"stop": ["A"], "include_stop_str_in_output": []},
+ {"include_stop_str_in_output": True}, # was wrong: refused
+ {"include_stop_str_in_output": False},
+]
+
+
+@pytest.mark.parametrize("body", CASES, ids=lambda b: repr(b))
+def test_the_router_accepts_exactly_what_vllm_accepts(body):
+ """Accept/refuse must agree, and so must the normalised value.
+
+ The error TYPES differ by construction -- pydantic raises ValidationError,
+ the router raises InvalidParameter -- and that is not what is being
+ compared. What matters is whether the request is served at all, and with
+ which stop strings and which flag if it is.
+ """
+ v_ok, v_val = _vllm(body)
+ o_ok, o_val = _ours(body)
+ assert o_ok == v_ok, (
+ f"vLLM {'accepts' if v_ok else 'refuses'} this and the router "
+ f"{'accepts' if o_ok else 'refuses'} it"
+ )
+ if v_ok:
+ assert o_val == v_val, "accepted by both, normalised differently"
+
+
+def test_an_empty_stop_string_is_refused_by_vllms_second_layer():
+ """Pins where it happens, because comparing the wrong layer misleads.
+
+ `ChatCompletionRequest` accepts `stop: [""]`; `SamplingParams` refuses it.
+ A test written against the model alone would conclude vLLM serves it and
+ would push the router into serving an unrestricted completion.
+ """
+ req = ChatCompletionRequest(model="m", messages=[{"role": "user", "content": "hi"}], stop=[""])
+ assert req.stop == [""], "the type check passes"
+ with pytest.raises(ValueError, match="empty string"):
+ req.to_sampling_params(max_tokens=8, default_sampling_params={})
diff --git a/tests/pd_vllm/test_stop_real_tokenizer.py b/tests/pd_vllm/test_stop_real_tokenizer.py
new file mode 100644
index 0000000..76bdd51
--- /dev/null
+++ b/tests/pd_vllm/test_stop_real_tokenizer.py
@@ -0,0 +1,171 @@
+"""``stop`` against a real byte-level BPE, rather than a stub vocabulary.
+
+test_stop_strings.py and test_reply.py drive a stub whose every id
+decodes to one fixed string. A real tokenizer does not behave that way, and the
+differences are exactly what the hold-back and the entry-to-text alignment exist
+for:
+
+* a stop string is almost never token-aligned, so the cut usually lands inside a
+ token -- and that token's prefix is still part of the reply;
+* one character can span two tokens, so the first of them decodes to nothing
+ emittable and the second carries the whole character;
+* a special decodes to a tag or to nothing depending on the policy, and the
+ matcher and the emitted text have to agree about which.
+
+Skipped automatically when no tokenizer is on disk, so a dev box stays green:
+
+ TOKENIZER_PATH=/path/to/tokenizer \
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_stop_real_tokenizer.py -v
+
+No GPU and no weights -- the tokenizer alone.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+
+transformers = pytest.importorskip("transformers")
+
+# Any byte-level BPE tokenizer will do (GLM-5 / DeepSeek-V3.2 checkpoints
+# both qualify); the rules under test must not depend on which one.
+_PATHS = [p for p in (os.environ.get("TOKENIZER_PATH", ""),) if p and os.path.isdir(p)]
+if not _PATHS:
+ pytest.skip("no tokenizer on disk", allow_module_level=True)
+
+from tilert.pd_vllm.logprobs import LogprobsRequest # noqa: E402
+from tilert.pd_vllm.reply import ( # noqa: E402
+ CONTENT,
+ ReplyStream,
+)
+
+_ASCII = "The answer is 42.\nObservation: done\nMore text follows here."
+# Chosen so the stops below cut inside a token and land on a multi-byte
+# character -- neither is possible with a one-id-one-string stub.
+_CJK = "今天天气很好。所以我们出去散步了。"
+
+
+@pytest.fixture(scope="module", params=_PATHS, ids=lambda p: p.split("/")[-1])
+def tok(request):
+ return transformers.AutoTokenizer.from_pretrained(request.param, trust_remote_code=True)
+
+
+def _drive(tok, text, stop, *, include=False, chunk=1, logprobs=False, ids=None):
+ """Push `text`'s real token ids through an assembler; return the reply."""
+ ids = tok.encode(text, add_special_tokens=False) if ids is None else ids
+ asm = ReplyStream(
+ tok,
+ stop=stop,
+ include_stop_in_output=include,
+ logprobs_req=LogprobsRequest(top_n=1) if logprobs else None,
+ )
+ ems = []
+ for i in range(0, len(ids), chunk):
+ g = ids[i : i + chunk]
+ ems += asm.push(
+ g, [-0.5] * len(g) if logprobs else None, [[(t, -0.5)] for t in g] if logprobs else None
+ )
+ ems += asm.finish()
+ return {
+ "text": "".join(e.text for e in ems if e.channel == CONTENT),
+ "entries": [x for e in ems if e.channel == CONTENT for x in e.logprobs],
+ "asm": asm,
+ "ids": ids,
+ }
+
+
+def test_a_multi_token_stop_cuts_at_the_right_character(tok):
+ got = _drive(tok, _ASCII, ["Observation:"])
+ assert got["text"] == "The answer is 42.\n"
+ assert got["asm"].stop_reason == "Observation:"
+
+
+def test_include_keeps_exactly_the_stop_string(tok):
+ got = _drive(tok, _ASCII, ["Observation:"], include=True)
+ assert got["text"] == "The answer is 42.\nObservation:"
+
+
+@pytest.mark.parametrize("chunk", [1, 2, 3, 5, 8, 64])
+def test_chunking_real_tokens_changes_nothing(tok, chunk):
+ """MTP delivers several tokens per message, and the batch size varies with acceptance.
+
+ A reply that depended on it would be non-deterministic.
+ """
+ one = _drive(tok, _ASCII, ["Observation:"], logprobs=True, chunk=1)
+ got = _drive(tok, _ASCII, ["Observation:"], logprobs=True, chunk=chunk)
+ assert got["text"] == one["text"]
+ assert len(got["entries"]) == len(one["entries"])
+ assert got["asm"].completion_tokens == one["asm"].completion_tokens
+
+
+def test_one_entry_per_token_counted(tok):
+ got = _drive(tok, _ASCII, ["Observation:"], logprobs=True)
+ assert len(got["entries"]) == got["asm"].completion_tokens
+ assert got["asm"].completion_tokens < len(
+ got["ids"]
+ ), "the stream did not run to the end -- the reply stopped earlier"
+
+
+def test_the_count_covers_the_tokens_the_stop_consumed(tok):
+ """A stop string is several tokens long on a real vocabulary, and those
+ tokens ran even though their text is gone.
+
+ "Observation:" is three tokens here. Counting only the visible ones would
+ bill for less than the model did -- and the visible text is not on a token
+ boundary in general, so no count can describe it. vLLM reports the
+ untruncated id list for exactly this reason.
+ """
+ got = _drive(tok, _ASCII, ["Observation:"], logprobs=True)
+ n = got["asm"].completion_tokens
+ assert tok.decode(got["ids"][:n], skip_special_tokens=True).startswith(
+ got["text"]
+ ), "the visible text is a prefix of what was counted"
+ assert len(tok.decode(got["ids"][:n], skip_special_tokens=True)) > len(
+ got["text"]
+ ), "and the stop's own tokens are inside the count"
+
+
+def test_the_reply_is_a_prefix_of_the_full_decode(tok):
+ got = _drive(tok, _ASCII, ["Observation:"])
+ full = tok.decode(got["ids"], skip_special_tokens=True)
+ assert full.startswith(got["text"])
+
+
+def test_a_stop_that_cuts_inside_a_token(tok):
+ """ "。所以" starts mid-token, so the token that begins the stop also holds text the reply keeps.
+
+ Dropping the whole token would lose "好".
+ """
+ got = _drive(tok, _CJK, ["。所以"])
+ assert got["text"] == "今天天气很好"
+ assert got["asm"].stop_reason == "。所以"
+
+
+def test_a_stop_landing_on_a_multi_byte_character(tok):
+ got = _drive(tok, _CJK, ["散步"])
+ assert got["text"] == "今天天气很好。所以我们出去"
+
+
+def test_without_a_stop_nothing_is_held_or_cut(tok):
+ got = _drive(tok, _ASCII, [])
+ assert got["text"] == tok.decode(got["ids"], skip_special_tokens=True)
+ assert got["asm"].completion_tokens == len(got["ids"])
+
+
+def test_an_unmatched_stop_releases_the_tail(tok):
+ got = _drive(tok, _ASCII, ["ZZZZZZ"])
+ assert got["text"] == tok.decode(got["ids"], skip_special_tokens=True)
+ assert got["asm"].stop_reason is None
+
+
+def test_a_special_does_not_surface_as_content(tok):
+ """With no parser nothing downstream would consume it, so the assembler
+ strips it -- and the stop matcher has to see the same text, or a stop
+ spelled like a tag would end one channel's reply and not the other's.
+ """
+ ids = tok.encode("hi", add_special_tokens=False) + [tok.eos_token_id]
+ got = _drive(tok, "", ["ZZZZZZ"], ids=ids)
+ assert got["text"].strip() == "hi"
+ assert "<" not in got["text"]
diff --git a/tests/pd_vllm/test_stop_strings.py b/tests/pd_vllm/test_stop_strings.py
new file mode 100644
index 0000000..c1840eb
--- /dev/null
+++ b/tests/pd_vllm/test_stop_strings.py
@@ -0,0 +1,1679 @@
+"""``stop`` served on the router's side of the detokeniser.
+
+`stop` is a property of the decoded TEXT, so it cannot be served by handing
+extra ids to the decode loop: a stop string routinely spans two tokens, and
+byte-level BPE can split one character across tokens. The decode node emits ids;
+the router is where text exists.
+
+Three layers, tested at the level each one actually decides something:
+
+* :func:`check_stop_strings` / :class:`StopWindow` -- the matcher, ported from
+ vLLM, and the hold-back that keeps text off the wire while it could still turn
+ out to be the start of a stop.
+* :func:`resolve_stop_request` -- what the router accepts.
+* the router -- that both response paths answer the same request the same way.
+
+Per-token attribution (logprob entries, `completion_tokens`) is a property of
+`ReplyStream` and is pinned in test_reply.py, which needs no
+HTTP. This file does not repeat it.
+
+No GPU, no tilert, no vllm.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest tests/pd_vllm/test_stop_strings.py -v
+"""
+
+from __future__ import annotations
+
+import json
+import threading
+import time
+
+import httpx
+import pytest
+import uvicorn
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+
+from tilert.pd_vllm import pd_router
+from tilert.pd_vllm.capabilities import (
+ CapabilityUnavailable,
+ InvalidParameter,
+)
+from tilert.pd_vllm.decode_pool import DecodeNode, Pool
+from tilert.pd_vllm.pd_router import (
+ RouterCtx,
+ build_app,
+ build_prefill_body,
+)
+from tilert.pd_vllm.request_gate import resolve_stop_request
+from tilert.pd_vllm.stop_strings import (
+ StopWindow,
+ check_stop_strings,
+ resolve_stop,
+)
+
+
+def drive(w: StopWindow, delta: str) -> str:
+ """Absorb a delta and take whatever the client may now see.
+
+ `push` and `take` are separate on purpose -- absorbing text and deciding how
+ much of it is safe to send are different questions, and that split is why the
+ window needs no released/unreleased queue. Tests that only care about the
+ combination go through here.
+ """
+ w.push(delta)
+ return w.take()
+
+
+# --------------------------------------------------------------------------- #
+# The matcher: same request, same character, as vLLM
+# --------------------------------------------------------------------------- #
+@pytest.mark.parametrize(
+ "text,new,stop,want",
+ [
+ ("hello STOP", 5, ["STOP"], ("STOP", 6)),
+ ("hello STOP", 5, ["ZZ"], None),
+ ("hello", 5, [], None),
+ ("hello STOP", 0, ["STOP"], None),
+ ],
+)
+def test_the_matcher_reports_the_cut(text, new, stop, want):
+ assert check_stop_strings(text, new, stop, False) == want
+
+
+def test_a_stop_straddling_the_delta_boundary_is_found():
+ """The search starts before the new text, not at it.
+
+ "STO" was already emitted and "P" just arrived, so a matcher looking only at
+ the new characters would miss the stop entirely -- and the reply would run
+ past it. This is the rule the vLLM port exists for.
+ """
+ assert check_stop_strings("aSTOP", 1, ["STOP"], False) == ("STOP", 1)
+
+
+def test_the_stop_that_completes_earliest_wins():
+ """MTP appends a batch of tokens at once, so several stops can match in one step.
+
+ The earliest-completing one is what a one-token-at-a-time stream would have hit, so the
+ reply does not depend on the batch size.
+ """
+ got = check_stop_strings("aXbYc", 5, ["Y", "X"], False)
+ assert got == ("X", 1)
+
+
+def test_ties_go_to_stop_list_order():
+ assert check_stop_strings("abAB", 4, ["AB", "B"], False) == ("AB", 2)
+
+
+def test_include_in_output_cuts_after_the_stop_instead_of_before():
+ assert check_stop_strings("hi STOP!", 8, ["STOP"], True) == ("STOP", 7)
+
+
+def test_include_in_output_needs_no_cut_when_the_stop_ends_the_text():
+ """-1 means "the text is already right", which is the common case: the stop
+ completes on the token that just arrived.
+ """
+ assert check_stop_strings("hi STOP", 7, ["STOP"], True) == ("STOP", -1)
+
+
+# --------------------------------------------------------------------------- #
+# The hold-back: text on the wire cannot be recalled
+# --------------------------------------------------------------------------- #
+def test_the_tail_that_could_start_a_stop_is_held():
+ t = StopWindow(["STOP"])
+ assert drive(t, "hello STO") == "hello ", "STO could still become STOP"
+ assert drive(t, "P") == ""
+ assert t.stopped == "STOP"
+
+
+def test_the_held_tail_is_released_once_it_cannot_be_a_stop():
+ """The hold-back is a fixed len(stop)-1 characters, not the matched prefix.
+
+ "STORY" proves "STO" was not a stop, but "RY" could itself begin one, so
+ three characters stay held. Releasing eagerly would need the matcher to
+ report how much of the tail is still live, which buys nothing: the delay is
+ bounded by the longest stop string.
+ """
+ t = StopWindow(["STOP"])
+ assert drive(t, "hello STO") == "hello "
+ assert drive(t, "RY") == "ST"
+ assert t.take(final=True) == "ORY"
+
+
+def test_the_held_tail_is_released_at_end_of_stream():
+ t = StopWindow(["STOP"])
+ assert drive(t, "hi ST") == "hi"
+ assert t.take(final=True) == " ST"
+
+
+def test_nothing_is_held_when_no_stop_was_asked_for():
+ """The pass-through matters: holding back would delay every token of every
+ request that does not use stop, which is nearly all of them. It also does
+ not accumulate the text, which cost a copy of the whole reply per token.
+ """
+ t = StopWindow([])
+ assert drive(t, "hello") == "hello"
+ assert t.take(final=True) == ""
+ assert len(t._text) == 0, "nothing retained"
+
+
+def test_pushes_after_a_stop_return_nothing():
+ """The decode node cannot see text, so it keeps sending. The reply ended."""
+ t = StopWindow(["STOP"])
+ drive(t, "aSTOP")
+ assert drive(t, " more") == "", "the reply ended; later text is not it"
+ assert t.take(final=True) == ""
+
+
+def test_every_character_comes_out_exactly_once():
+ """Held text must be released later, not dropped, and not sent twice."""
+ t = StopWindow(["STOP"])
+ assert (drive(t, "hello STO") + drive(t, "RY") + t.take(final=True)) == "hello STORY"
+
+
+@pytest.mark.parametrize(
+ "raw,want",
+ [
+ (None, []),
+ ("END", ["END"]),
+ (["A", "B"], ["A", "B"]),
+ ([], []),
+ ],
+)
+def test_stop_is_normalised_to_a_list(raw, want):
+ assert resolve_stop({"stop": raw}) == want
+
+
+@pytest.mark.parametrize("raw", [[""], ["", "END"], ""])
+def test_an_empty_stop_string_is_rejected_not_dropped(raw):
+ """It matches at position 0 of everything, so it is not a neutral value the
+ way `[]` is -- dropping it serves an unrestricted completion to a client who
+ asked for a restricted one.
+
+ vLLM raises `ValueError("stop cannot contain an empty string.")`, verified
+ against 0.25.1. The router strips `stop` from the prefill request, so vLLM
+ no longer gets the chance to say so and the router has to.
+ """
+ with pytest.raises(ValueError):
+ resolve_stop({"stop": raw})
+
+
+# --------------------------------------------------------------------------- #
+# What the router accepts
+# --------------------------------------------------------------------------- #
+class _Tok:
+ def decode(self, ids, skip_special_tokens=False):
+ return "".join(_VOCAB.get(i, "") for i in ids)
+
+
+def test_stop_needs_no_tokenizer_when_it_is_not_asked_for():
+ assert resolve_stop_request({}, None) == ([], False)
+
+
+def test_stop_strings_without_a_tokenizer_are_refused_not_ignored():
+ """`--parser none` with no `--model-path`: there is no text to match
+ against, and silently ignoring the field would serve a reply that runs past
+ the client's stop.
+ """
+ with pytest.raises(CapabilityUnavailable):
+ resolve_stop_request({"stop": ["END"]}, None)
+
+
+@pytest.mark.parametrize(
+ "body",
+ [
+ {"stop": 5},
+ {"stop": [1, 2]},
+ {"stop": [""]},
+ ],
+)
+def test_an_unusable_stop_request_is_a_400(body):
+ with pytest.raises(InvalidParameter):
+ resolve_stop_request(body, _Tok())
+
+
+def test_the_flag_without_a_stop_is_a_no_op_not_an_error():
+ """vLLM serves it, so refusing it here would turn a request that works
+ against a native endpoint into a 400. With no stop strings the flag has
+ nothing to include.
+ """
+ assert resolve_stop_request({"include_stop_str_in_output": True}, _Tok()) == ([], True)
+
+
+# `include_stop_str_in_output` is coerced by the same validator vLLM's request
+# model uses, so the two stacks accept and refuse the same values. Verified
+# against `ChatCompletionRequest` on vLLM 0.25.1 / pydantic 2.13.4 -- being
+# STRICTER here would turn a request vLLM serves into a 400, which is the
+# mistake the capability gate's top_k handling already warns about.
+@pytest.mark.parametrize(
+ "value,want",
+ [
+ (True, True),
+ (False, False),
+ (1, True),
+ (0, False),
+ ("true", True),
+ ("False", False),
+ ],
+)
+def test_the_flag_is_coerced_exactly_as_vllm_coerces_it(value, want):
+ assert resolve_stop_request({"stop": ["A"], "include_stop_str_in_output": value}, _Tok()) == (
+ ["A"],
+ want,
+ )
+
+
+@pytest.mark.parametrize("value", [None, 2, "maybe", [], {}])
+def test_a_value_vllm_would_refuse_is_a_400(value):
+ """An explicit `null` included: vLLM's field is `bool = False`, not
+ `bool | None`, so pydantic answers 422 for it. The router strips the field
+ before prefill, so vLLM never gets the chance to say so.
+ """
+ with pytest.raises(InvalidParameter):
+ resolve_stop_request({"stop": ["A"], "include_stop_str_in_output": value}, _Tok())
+
+
+def test_the_prefill_request_does_not_carry_the_stop_fields():
+ """The prefill instance generates one token.
+
+ Matching there could truncate it or report finish_reason="stop" for a prefill that in fact
+ succeeded -- and the router is going to match over the whole reply anyway.
+ """
+ out = build_prefill_body(
+ "/v1/chat/completions",
+ {"messages": [], "stop": ["END"], "include_stop_str_in_output": True},
+ DecodeNode(host="h", ctrl_port=1, http_port=2),
+ )
+ assert "stop" not in out
+ assert "include_stop_str_in_output" not in out
+
+
+# --------------------------------------------------------------------------- #
+# The router: one request, two paths, one answer
+# --------------------------------------------------------------------------- #
+# "Hello, world. STOP tail" over eight tokens, with the stop split across two of
+# them so the hold-back and the cross-token match are both exercised.
+_VOCAB = {
+ 201: "Hello",
+ 202: ", world",
+ 203: ". ",
+ 204: "ST",
+ 205: "OP",
+ 206: " tail",
+ 207: " more",
+ 208: "!",
+}
+# The decode node's reply INCLUDES the prefill's token: `PDEngine.decode`
+# returns "completion ids (incl. first_token_id)" and fires on_token for it too,
+# so both response paths see it. Its logprob is null -- the node echoed the
+# token rather than sampling it.
+_FIRST = 201
+_IDS = [201, 202, 203, 204, 205, 206, 207, 208]
+_FULL = "".join(_VOCAB[i] for i in _IDS)
+
+
+def _serve(app):
+ cfg = uvicorn.Config(app, host="127.0.0.1", port=0, log_level="error")
+ server = uvicorn.Server(cfg)
+ threading.Thread(target=server.run, daemon=True).start()
+ for _ in range(200):
+ if server.started:
+ return server.servers[0].sockets[0].getsockname()[1]
+ time.sleep(0.05)
+ raise RuntimeError("stub server did not start")
+
+
+def _make_vllm():
+ app = FastAPI()
+
+ @app.post("/v1/chat/completions")
+ async def chat(request: Request):
+ body = await request.json()
+ assert "stop" not in body, "the prefill leg must not match stop strings"
+ return {
+ "id": "cmpl-stop",
+ "model": "stub",
+ "choices": [
+ {
+ "logprobs": {
+ "content": [
+ {
+ "token": f"token_id:{_FIRST}",
+ "logprob": -0.125,
+ "top_logprobs": [
+ {"token": f"token_id:{_FIRST}", "logprob": -0.125}
+ ],
+ }
+ ]
+ }
+ }
+ ],
+ "usage": {"prompt_tokens": 4},
+ "kv_transfer_params": {},
+ }
+
+ @app.post("/v1/completions")
+ async def completions(request: Request):
+ body = await request.json()
+ assert "stop" not in body, "the prefill leg must not match stop strings"
+ # /v1/completions reports the first token under `tokens`, not `content`.
+ return {
+ "id": "cmpl-stop",
+ "model": "stub",
+ "choices": [{"logprobs": {"tokens": [f"token_id:{_FIRST}"]}}],
+ "usage": {"prompt_tokens": 4},
+ "kv_transfer_params": {},
+ }
+
+ return app
+
+
+# What the decode node saw, so a test can tell early cancellation from mere
+# truncation of a reply that was generated in full.
+_SEEN: dict = {"emitted": 0, "cancelled": [], "stream": None, "decode_body": {}}
+
+
+def _make_decode(
+ filler: int = 0,
+ batched: bool = False,
+ error_after: int | None = None,
+ error_type: str | None = None,
+ truncate_after: int | None = None,
+ drop_tp: bool = False,
+ garbage_after: int | None = None,
+ lose_lines_after: int | None = None,
+ null_lp: bool = False,
+ status: int | None = None,
+):
+ app = FastAPI()
+
+ @app.get("/capabilities")
+ def capabilities():
+ return {
+ "profile": "stub",
+ "engine": "StubEngine",
+ "capabilities": {"penalties": True, "ignore_eos": True, "logprobs": True},
+ }
+
+ def _lp(n, ids):
+ out = {"lp": [None] * len(ids) if null_lp else [-0.5] * len(ids)}
+ if not drop_tp:
+ out["tp"] = [[[t, -0.5 - k] for k in range(n)] for t in ids]
+ return out
+
+ @app.post("/pd/decode")
+ async def decode(request: Request):
+ body = await request.json()
+ if status is not None:
+ # A node answering with a status rather than a stream: the shape
+ # both handlers have to classify before any body exists.
+ return JSONResponse({"error": "engine down"}, status_code=status)
+ n = body.get("top_logprobs")
+ _SEEN["stream"] = bool(body.get("stream"))
+ _SEEN["emitted"] = 0
+ _SEEN["decode_body"] = body
+ if not body.get("stream"):
+ out = {
+ "rid": body["rid"],
+ "token_ids": _IDS,
+ "seq_len": 8,
+ "timing_ms": {"finish_reason": "length"},
+ }
+ if n is not None:
+ out["logprobs"] = _lp(n, _IDS)
+ out["logprobs"]["lp"][0] = None # echoed, not sampled
+ return out
+
+ if batched:
+ # What a real node does: it drains its queue with `get_nowait()`
+ # before writing, so one line carries everything the engine has
+ # produced -- measured at 30 tokens on a live pair. A stop then
+ # lands part-way into a batch.
+ def gen_batched():
+ ids = _IDS + [207] * filler
+ line = {"t": ids}
+ if n is not None:
+ line["lp"] = [None] * len(ids) if null_lp else [None] + [-0.5] * (len(ids) - 1)
+ line["tp"] = [[[t, -0.5]] for t in ids]
+ yield json.dumps(line) + "\n"
+ yield json.dumps({"done": True, "finish_reason": "length"}) + "\n"
+
+ return StreamingResponse(gen_batched(), media_type="application/x-ndjson")
+
+ def gen():
+ # `filler` stands in for a client's large max_tokens: the node keeps
+ # going because it cannot see text.
+ for k, tid in enumerate(_IDS + [207] * filler):
+ if lose_lines_after is not None and k >= lose_lines_after:
+ # Silently drop the rest, then still send a normal `done`.
+ continue
+ if garbage_after is not None and k == garbage_after:
+ # Not JSON: `DecodeReader.feed` raises on it.
+ yield "{this is not json\n"
+ return
+ if truncate_after is not None and k == truncate_after:
+ # A clean EOF with no done/error: a proxy cutting the body,
+ # or a node dying mid-reply.
+ return
+ if error_after is not None and k == error_after:
+ err = {"error": "engine exploded"}
+ if error_type is not None:
+ err["error_type"] = error_type
+ yield json.dumps(err) + "\n"
+ return
+ line = {"t": [tid]}
+ if n is not None:
+ line["lp"] = [None if (k == 0 or null_lp) else -0.5]
+ if not drop_tp:
+ line["tp"] = [[[tid, -0.5 - j] for j in range(n)]]
+ _SEEN["emitted"] = k + 1
+ yield json.dumps(line) + "\n"
+ # `n` is already the requested top_logprobs in this closure, so the
+ # token count needs its own name.
+ declared = _SEEN["emitted"]
+ if lose_lines_after is not None:
+ # The node's own count, as a real one reports it -- unchanged by
+ # the lines that went missing.
+ declared = len(_IDS) + filler
+ yield json.dumps({"done": True, "n": declared, "finish_reason": "length"}) + "\n"
+
+ return StreamingResponse(gen(), media_type="application/x-ndjson")
+
+ @app.post("/pd/cancel")
+ async def cancel(request: Request):
+ _SEEN["cancelled"].append((await request.json()).get("rid"))
+ return {"ok": True}
+
+ return app
+
+
+@pytest.fixture(scope="module")
+def _proxy_off():
+ """Keep a local proxy from hijacking the stub traffic.
+
+ The router reaches both stubs with `requests`, which honours `http_proxy`
+ from the environment. On a box that has one set -- which is also why a real
+ deployment must set `no_proxy` for its internal addresses -- the call is
+ proxied, never arrives, and the test hangs to its read timeout.
+ """
+ mp = pytest.MonkeyPatch()
+ for var in ("no_proxy", "NO_PROXY"):
+ mp.setenv(var, "127.0.0.1,localhost")
+ for var in ("http_proxy", "HTTP_PROXY", "https_proxy", "HTTPS_PROXY", "all_proxy", "ALL_PROXY"):
+ mp.delenv(var, raising=False)
+ yield
+ mp.undo()
+
+
+@pytest.fixture(scope="module")
+def router(_proxy_off):
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode())
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ return f"http://127.0.0.1:{_serve(build_app(ctx))}"
+
+
+def _body(**over):
+ b = {
+ "model": "stub",
+ "messages": [{"role": "user", "content": "hi"}],
+ "max_tokens": 64,
+ "temperature": 0.0,
+ }
+ b.update(over)
+ return b
+
+
+def _post(router, **over):
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ r = c.post(f"{router}/v1/chat/completions", json=_body(**over))
+ return r
+
+
+def _stream(router, **over):
+ """Drive the streaming path; return (choice-shaped dict, usage)."""
+ text, finish, stop_reason, usage, entries = "", None, "sentinel", None, []
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream(
+ "POST", f"{router}/v1/chat/completions", json=_body(stream=True, **over)
+ ) as r:
+ r.raise_for_status()
+ for line in r.iter_lines():
+ if not line.startswith("data: ") or line[6:] == "[DONE]":
+ continue
+ ch = json.loads(line[6:])
+ if ch.get("usage") is not None:
+ usage = ch["usage"]
+ for choice in ch["choices"]:
+ text += choice["delta"].get("content", "") or ""
+ entries += (choice.get("logprobs") or {}).get("content") or []
+ if choice.get("finish_reason"):
+ finish = choice["finish_reason"]
+ stop_reason = choice.get("stop_reason", "absent")
+ return {
+ "text": text,
+ "finish_reason": finish,
+ "stop_reason": stop_reason,
+ "entries": entries,
+ }, usage
+
+
+def test_a_stop_string_truncates_the_reply(router):
+ got = _post(router, stop=["STOP"]).json()
+ assert got["choices"][0]["message"]["content"] == "Hello, world. "
+
+
+def test_the_stop_spans_two_tokens(router):
+ """ "ST" and "OP" arrive as separate tokens, so a matcher working per token
+ would never see the string.
+ """
+ assert "STOP" in _FULL
+ assert _VOCAB[204] + _VOCAB[205] == "STOP"
+ assert (
+ _post(router, stop=["STOP"]).json()["choices"][0]["message"]["content"] == "Hello, world. "
+ )
+
+
+def test_include_stop_str_in_output_keeps_it(router):
+ got = _post(router, stop=["STOP"], include_stop_str_in_output=True).json()
+ assert got["choices"][0]["message"]["content"] == "Hello, world. STOP"
+
+
+def test_stop_reason_names_the_string_that_ended_the_reply(router):
+ """A client cannot recover it from the text: the string was cut out."""
+ choice = _post(router, stop=["STOP"]).json()["choices"][0]
+ assert choice["stop_reason"] == "STOP"
+ assert choice["finish_reason"] == "stop"
+
+
+def test_stop_reason_is_null_when_no_stop_string_matched(router):
+ choice = _post(router, stop=["ZZZZ"]).json()["choices"][0]
+ assert choice["stop_reason"] is None
+ assert choice["finish_reason"] == "length", "what the decode node said"
+ assert choice["message"]["content"] == _FULL
+
+
+def test_usage_counts_what_ran_up_to_the_stop(router):
+ """The stop truncates the text, not the count.
+
+ 204 ("ST") and 205 ("OP") contributed no visible text and are still counted:
+ they ran. vLLM does the same -- its `completion_tokens` is `len()` of the
+ detokeniser's untruncated id list. What is NOT counted is the tokens after
+ the stop, which the node only generated because it cannot see text.
+ """
+ got = _post(router, stop=["STOP"]).json()
+ assert got["usage"]["completion_tokens"] == 5, "201..205"
+ assert _post(router).json()["usage"]["completion_tokens"] == 8
+
+
+def test_completions_reports_every_id_the_node_sent(router):
+ """`token_ids` says what ran, `text` says what came back, and a stop makes
+ them differ.
+
+ vLLM's `CompletionOutput.token_ids` is the detokeniser's untruncated list
+ for the same reason. Slicing it to the visible text would report fewer
+ tokens than `usage` bills for, and could not represent a stop that cut
+ inside a token anyway.
+ """
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ r = c.post(
+ f"{router}/v1/completions",
+ json={"model": "stub", "prompt": "hi", "max_tokens": 64, "stop": ["STOP"]},
+ )
+ assert r.status_code == 200, r.text[:300]
+ choice = r.json()["choices"][0]
+ assert choice["text"] == "Hello, world. "
+ assert choice["token_ids"] == [201, 202, 203, 204, 205]
+ assert len(choice["token_ids"]) == r.json()["usage"]["completion_tokens"]
+
+
+@pytest.mark.parametrize(
+ "over",
+ [
+ {},
+ {"stop": ["STOP"]},
+ {"stop": ["STOP"], "include_stop_str_in_output": True},
+ {"stop": ["ZZZZ"]},
+ {"stop": [". ", "STOP"]},
+ ],
+)
+def test_streaming_and_non_streaming_answer_the_same(router, over):
+ """The property the assembler exists for, checked over real HTTP.
+
+ Text, finish_reason, stop_reason and completion_tokens all have to match:
+ a client switching `stream` must not get a different reply.
+ """
+ blocking = _post(router, **over).json()
+ streamed, usage = _stream(router, **over)
+ choice = blocking["choices"][0]
+ assert streamed["text"] == choice["message"]["content"]
+ assert streamed["finish_reason"] == choice["finish_reason"]
+ assert streamed["stop_reason"] == choice["stop_reason"]
+ assert usage is None or usage == blocking["usage"]
+
+
+@pytest.mark.parametrize("over", [{}, {"stop": ["STOP"]}])
+def test_the_two_paths_agree_on_logprobs_too(router, over):
+ """One entry per token the reply contains, the same entries either way."""
+ extra = {"logprobs": True, "top_logprobs": 1, "temperature": 0.6}
+ blocking = _post(router, **over, **extra).json()
+ streamed, _ = _stream(router, **over, **extra)
+ want = blocking["choices"][0]["logprobs"]["content"]
+ assert len(want) == blocking["usage"]["completion_tokens"]
+ assert [e["token"] for e in streamed["entries"]] == [e["token"] for e in want]
+ assert [e["logprob"] for e in streamed["entries"]] == [e["logprob"] for e in want]
+
+
+def test_the_streamed_usage_opt_in_agrees_with_the_blocking_one(router):
+ _, usage = _stream(router, stop=["STOP"], stream_options={"include_usage": True})
+ assert usage["completion_tokens"] == 5
+
+
+def test_a_stop_string_is_no_longer_refused(router):
+ """It was a 501 until the router grew a matcher."""
+ assert _post(router, stop=["STOP"]).status_code == 200
+
+
+@pytest.mark.parametrize(
+ "over",
+ [
+ {"stop": 5},
+ {"stop": [""]},
+ {"stop": ["A"], "include_stop_str_in_output": None},
+ ],
+)
+def test_an_unusable_stop_request_is_rejected_before_any_backend(router, over):
+ assert _post(router, **over).status_code == 400
+
+
+# --------------------------------------------------------------------------- #
+# The point of serving stop at all: it has to stop something
+# --------------------------------------------------------------------------- #
+@pytest.fixture(scope="module")
+def long_router(_proxy_off):
+ """A node that keeps emitting long past the stop, as a real one does."""
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode(filler=200))
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ return f"http://127.0.0.1:{_serve(build_app(ctx))}"
+
+
+def _wait_for_cancel(timeout=5.0):
+ """The cancel goes out on a daemon thread so a dead node cannot delay the
+ reply, so it lands after the response the test just read.
+ """
+ deadline = time.time() + timeout
+ while time.time() < deadline and not _SEEN["cancelled"]:
+ time.sleep(0.02)
+ return list(_SEEN["cancelled"])
+
+
+def test_a_non_streaming_stop_stops_the_decode_node(long_router):
+ """Not just the reply: the node too.
+
+ The node cannot see text, so a router that asks for the whole sequence up
+ front can only trim the answer after the fact -- the node still runs to
+ max_tokens and holds its slot for all of it, and the client waits for it.
+ So a request with a stop reads the node's streaming protocol even though its
+ own reply is not streamed, stops reading at the match, and cancels.
+
+ The stub emits 208 tokens; the stop lands on the fifth.
+ """
+ _SEEN["cancelled"].clear()
+ got = _post(long_router, stop=["STOP"]).json()
+ assert got["choices"][0]["message"]["content"] == "Hello, world. "
+ assert _SEEN["stream"] is True, "asked the node to stream"
+ assert _SEEN["emitted"] < 20, f"read {_SEEN['emitted']} of 208 tokens -- did not stop early"
+ assert _wait_for_cancel(), "and told the node to stop"
+
+
+def test_a_request_without_a_stop_does_not_pay_for_the_stream(long_router):
+ """Nothing to match, so nothing to react to mid-generation.
+
+ The blocking protocol stays the default, and a reply that ran to completion is not
+ cancelled.
+ """
+ _SEEN["cancelled"].clear()
+ assert _post(long_router).status_code == 200
+ assert _SEEN["stream"] is False
+ time.sleep(0.3)
+ assert not _SEEN["cancelled"], "it finished; there is nothing to cancel"
+
+
+@pytest.fixture(scope="module")
+def batched_router(_proxy_off):
+ """A node that delivers the whole sequence in one line, as a fast one does."""
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode(filler=25, batched=True))
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ return f"http://127.0.0.1:{_serve(build_app(ctx))}"
+
+
+def test_a_batch_that_straddles_the_stop_is_not_counted_whole(batched_router):
+ """The tokens behind the stop, inside the same line, are not the reply.
+
+ Found on a live pair: the node's first line carried 30 tokens, the stop
+ landed on the sixth, and `token_ids` reported all 30 against a
+ `completion_tokens` of 6. One source for the id list, the count and the
+ entries is what keeps them from disagreeing.
+ """
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ r = c.post(
+ f"{batched_router}/v1/completions",
+ json={
+ "model": "stub",
+ "prompt": "hi",
+ "max_tokens": 512,
+ "stop": ["STOP"],
+ "temperature": 0.0,
+ },
+ )
+ assert r.status_code == 200, r.text[:300]
+ j = r.json()
+ choice = j["choices"][0]
+ assert choice["text"] == "Hello, world. "
+ assert choice["token_ids"] == [201, 202, 203, 204, 205], "up to the stop"
+ assert len(choice["token_ids"]) == j["usage"]["completion_tokens"]
+
+
+def test_the_entries_of_a_straddled_batch_match_the_count(batched_router):
+ """Same for logprobs: one entry per token counted, no more."""
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ r = c.post(
+ f"{batched_router}/v1/chat/completions",
+ json={
+ "model": "stub",
+ "messages": [{"role": "user", "content": "hi"}],
+ "max_tokens": 512,
+ "stop": ["STOP"],
+ "temperature": 0.6,
+ "logprobs": True,
+ "top_logprobs": 1,
+ },
+ )
+ assert r.status_code == 200, r.text[:300]
+ j = r.json()
+ ents = j["choices"][0]["logprobs"]["content"]
+ assert len(ents) == j["usage"]["completion_tokens"] == 5
+
+
+# --------------------------------------------------------------------------- #
+# Ordering and fail-fast
+# --------------------------------------------------------------------------- #
+@pytest.fixture(scope="module")
+def erroring_router(_proxy_off):
+ """A node that fails part-way through, while a tail is still held back."""
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode(error_after=4))
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ return f"http://127.0.0.1:{_serve(build_app(ctx))}"
+
+
+def test_a_decode_error_does_not_jump_ahead_of_held_text(erroring_router):
+ """The error marker must come after the reply's own text, not through it.
+
+ With a stop that never matches, the matcher is holding up to len(stop)-1
+ characters when the node fails. Emitting the marker before releasing them
+ gives the client `prefix[decode error]suffix` -- text out of order, which is
+ worse than the error itself.
+ """
+ text = ""
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream(
+ "POST", f"{erroring_router}/v1/chat/completions", json=_body(stream=True, stop=["ZZZZ"])
+ ) as r:
+ r.raise_for_status()
+ for line in r.iter_lines():
+ if not line.startswith("data: ") or line[6:] == "[DONE]":
+ continue
+ for ch in json.loads(line[6:])["choices"]:
+ text += ch["delta"].get("content", "") or ""
+ assert "[decode error:" in text, "the failure is reported"
+ body, _, tail = text.partition("[decode error:")
+ assert "]" in tail
+ after = tail.split("]", 1)[1]
+ assert after == "", f"reply text arrived after the error marker: {after!r}"
+ # 201..204 -> "Hello" ", world" ". " "ST"; the trailing newline is the
+ # marker's own prefix. The held "ST" is the point: it is inside `body`.
+ assert body == "Hello, world. ST\n", "all of the reply came first"
+
+
+def test_a_malformed_stop_does_not_wait_on_a_hung_node():
+ """`stop: 5` is answerable from the request alone.
+
+ The capability probe talks HTTP to every uncached node and blocks for
+ `Pool.CAPS_TIMEOUT_S` (2 s) per node that accepts the connection and then
+ says nothing -- which is what a wedged node looks like, and is different
+ from a refused connection, which fails instantly. Validating after the
+ probe makes a deterministic 400 pay for a backend its answer never
+ depended on.
+ """
+ import socket
+
+ from fastapi.testclient import TestClient
+
+ # Accept the connection, then never answer: `requests` blocks on read until
+ # its timeout. A closed port would be refused immediately and prove nothing.
+ listener = socket.socket()
+ listener.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+ listener.bind(("127.0.0.1", 0))
+ listener.listen(8)
+ hung_port = listener.getsockname()[1]
+ held = []
+
+ def _swallow():
+ while True:
+ try:
+ held.append(listener.accept()[0])
+ except OSError:
+ return
+
+ threading.Thread(target=_swallow, daemon=True).start()
+ try:
+ node = DecodeNode("127.0.0.1", 5556, hung_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{hung_port}", Pool([node]), _Tok(), "none")
+ client = TestClient(build_app(ctx))
+ t0 = time.time()
+ r = client.post(
+ "/v1/chat/completions",
+ json={"model": "stub", "stop": 5, "messages": [{"role": "user", "content": "hi"}]},
+ )
+ dt = time.time() - t0
+ assert r.status_code == 400, r.text[:200]
+ assert dt < Pool.CAPS_TIMEOUT_S, (
+ f"took {dt:.1f}s against a {Pool.CAPS_TIMEOUT_S}s probe timeout "
+ f"-- it waited on the capability probe"
+ )
+ finally:
+ listener.close()
+ for sock in held:
+ sock.close()
+
+
+def test_logprobs_without_a_tokenizer_are_refused_not_nulled(_proxy_off):
+ """A 200 with `logprobs: null` reports success for a field the client asked
+ for and did not get -- and both backends compute the values first.
+
+ `--parser none` with no `--model-path` is a supported configuration; it just
+ cannot name tokens. Refused at the door, before any backend work.
+ """
+ from fastapi.testclient import TestClient
+
+ node = DecodeNode("127.0.0.1", 5556, _serve(_make_decode()))
+ ctx = RouterCtx(
+ f"http://127.0.0.1:{_serve(_make_vllm())}", Pool([node]), None, "none"
+ ) # no tokenizer
+ client = TestClient(build_app(ctx))
+ r = client.post(
+ "/v1/chat/completions",
+ json={
+ "model": "stub",
+ "logprobs": True,
+ "top_logprobs": 1,
+ "messages": [{"role": "user", "content": "hi"}],
+ },
+ )
+ assert r.status_code == 501, r.text[:200]
+ assert r.json()["error_type"] == "capability_unavailable"
+ # Without logprobs the same router still serves the request.
+ ok = client.post(
+ "/v1/chat/completions",
+ json={"model": "stub", "messages": [{"role": "user", "content": "hi"}]},
+ )
+ assert ok.status_code == 200, ok.text[:200]
+ assert ok.json()["choices"][0]["logprobs"] is None
+
+
+@pytest.fixture(scope="module")
+def grammar_violation_router(_proxy_off):
+ """A node that reports a grammar violation while a tail is still held."""
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode(error_after=4, error_type="grammar_violation"))
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ return f"http://127.0.0.1:{_serve(build_app(ctx))}"
+
+
+def test_a_grammar_violation_does_not_drop_held_text(grammar_violation_router):
+ """The fail-closed branch returns early, so it has to flush too.
+
+ It emits an SSE error event and `[DONE]` and returns without reaching the
+ flush after the loop -- so with a non-matching stop the characters the
+ matcher was still holding vanish from a reply the client otherwise keeps.
+ The generic decode-error branch already flushed; this one did not.
+ """
+ text, saw_error = "", False
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream(
+ "POST",
+ f"{grammar_violation_router}/v1/chat/completions",
+ json=_body(stream=True, stop=["ZZZZ"]),
+ ) as r:
+ r.raise_for_status()
+ for line in r.iter_lines():
+ if not line.startswith("data: ") or line[6:] == "[DONE]":
+ continue
+ payload = json.loads(line[6:])
+ if payload.get("error"):
+ saw_error = True
+ continue
+ for ch in payload.get("choices", []):
+ text += ch["delta"].get("content", "") or ""
+ assert saw_error, "the violation is still reported"
+ # 201..204 -> "Hello" ", world" ". " "ST"; "ST" is the held tail.
+ assert text == "Hello, world. ST", f"held text was dropped: {text!r}"
+
+
+@pytest.fixture(scope="module")
+def truncating_router(_proxy_off):
+ """A node whose body ends with no `done` and no `error`."""
+ vllm_port = _serve(_make_vllm())
+ # Cut AFTER the point where "STOP" completes (201..205), so the two cases
+ # below differ by the stop matching rather than by how much was emitted.
+ decode_port = _serve(_make_decode(truncate_after=7))
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ return f"http://127.0.0.1:{_serve(build_app(ctx))}"
+
+
+def test_a_stream_ending_without_a_terminal_message_is_not_a_success(truncating_router):
+ """A partial completion must not be reported as a finished one.
+
+ A stop-bearing non-streaming request reads the node's NDJSON protocol. If
+ that body reaches a clean EOF before `done`, `error` or a stop match -- a
+ proxy cutting it short, a node dying mid-reply -- assembling what arrived
+ would answer 200 with `finish_reason: "stop"` and a null `stop_reason` for a
+ reply that was truncated.
+ """
+ r = _post(truncating_router, stop=["ZZZZ"])
+ assert r.status_code == 502, r.text[:200]
+ assert r.json()["error_type"] == "decode_truncated"
+
+
+def test_a_stop_match_is_still_a_legitimate_early_exit(truncating_router):
+ """The stop case leaves the loop early on purpose and must stay a 200."""
+ r = _post(truncating_router, stop=["STOP"])
+ assert r.status_code == 200, r.text[:200]
+ assert r.json()["choices"][0]["stop_reason"] == "STOP"
+
+
+def test_the_role_chunk_still_precedes_every_delta(router):
+ """Guards a merge resolution rather than a feature of this PR.
+
+ #34 made the opening `delta.role` chunk lazy -- sent when there is
+ something to send, not unconditionally -- and this PR rewrote the loop it
+ sits in. Resolving that conflict wrongly would let content deltas reach the
+ client before the role, which nothing else here would notice.
+ """
+ order = []
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream("POST", f"{router}/v1/chat/completions", json=_body(stream=True)) as r:
+ r.raise_for_status()
+ for line in r.iter_lines():
+ if not line.startswith("data: ") or line[6:] == "[DONE]":
+ continue
+ for ch in json.loads(line[6:]).get("choices", []):
+ d = ch["delta"]
+ if d.get("role"):
+ order.append("role")
+ elif d.get("content"):
+ order.append("content")
+ assert order.count("role") == 1, f"role sent {order.count('role')} times"
+ assert order[0] == "role", f"a delta preceded the role: {order[:3]}"
+
+
+# --------------------------------------------------------------------------- #
+# The retained window
+# --------------------------------------------------------------------------- #
+class _WholeTextTracker:
+ """The obvious implementation: keep everything, never trim.
+
+ The reference the windowed tracker has to agree with. Ported from what
+ `StopWindow` was before the window, so a divergence shows up as a
+ difference in released text rather than as a performance number.
+ """
+
+ def __init__(self, stop, include=False):
+ self.stop, self.include = list(stop), include
+ self._hold = 0 if include else (max((len(x) for x in self.stop), default=1) - 1)
+ self.text, self._released, self.stopped = "", 0, None
+
+ def push(self, delta):
+ """The old shape on purpose: absorb and release in one call.
+
+ The reference keeps the WHOLE text and never trims, which is what the
+ window has to stay equivalent to.
+ """
+ if self.stopped is not None or not delta:
+ return ""
+ if not self.stop:
+ return delta
+ self.text += delta
+ hit = check_stop_strings(self.text, len(delta), self.stop, self.include)
+ if hit is not None:
+ self.stopped, cut = hit
+ if cut != -1:
+ self.text = self.text[:cut]
+ out = self.text[self._released :]
+ self._released = len(self.text)
+ return out
+ end = max(self._released, len(self.text) - self._hold)
+ out = self.text[self._released : end]
+ self._released = end
+ return out
+
+ def finish(self):
+ if self.stopped is not None:
+ return ""
+ out = self.text[self._released :]
+ self._released = len(self.text)
+ return out
+
+
+@pytest.mark.parametrize("slack", [0, 4096])
+@pytest.mark.parametrize(
+ "stop,include",
+ [
+ (["STOP"], False),
+ (["STOP"], True),
+ (["Observation:", "\n\n"], False),
+ (["NEVERMATCHES"], False),
+ ],
+)
+def test_the_window_releases_exactly_what_keeping_everything_would(stop, include, slack):
+ """The window is an allocation change and must not be a behaviour change.
+
+ Driven with deltas that straddle the boundary in every way that matters:
+ one character at a time, in chunks, and with the stop split across two.
+
+ `slack=0` makes the trim fire on every push, which is the only way a short
+ run reaches it at all -- at the shipped 4096 the window never fills here, so
+ the trim would go untested and a look-back cut short by it would not show.
+ That case matters most with `include_stop_str_in_output`, where nothing is
+ held back and a straddling stop therefore sits in ALREADY RELEASED text.
+ """
+ import random
+
+ rng = random.Random(20260828)
+ alphabet = ["a", "ST", "OP", "\n", "Observation", ":", "STOP", "x", "\n\n"]
+ for _ in range(40):
+ deltas = [rng.choice(alphabet) for _ in range(60)]
+ a, b = StopWindow(stop, include), _WholeTextTracker(stop, include)
+ a._slack = slack # per instance now, not a class attribute
+ out_a = "".join(drive(a, d) for d in deltas) + a.take(final=True)
+ out_b = "".join(b.push(d) for d in deltas) + b.finish()
+ assert out_a == out_b, f"released text differs on {deltas[:6]}"
+ assert a.stopped == b.stopped
+
+
+def test_the_retained_text_does_not_grow_with_the_reply():
+ """A stop that never matches must not make the router hold the completion.
+
+ Keeping all of it cost a copy of the response per token -- quadratic in the
+ length, 0.75 s and 200 KB over a 200k-token generation -- for text nothing
+ reads once it has gone out.
+ """
+ t = StopWindow(["NEVERMATCHES"])
+ for _ in range(200_000):
+ drive(t, "x")
+ assert len(t._text) < 8_192, f"retained {len(t._text):,} characters"
+
+
+def test_a_stop_still_matches_after_the_window_has_trimmed():
+ """The look-back a straddling match needs survives trimming."""
+ t = StopWindow(["STOP"])
+ released = "".join(drive(t, "x") for _ in range(20_000))
+ released += drive(t, "ST") # could still become the stop, so held back
+ released += drive(t, "OP") # completes it, long after the first trim
+ assert t.stopped == "STOP"
+ assert released == "x" * 20_000, "the stop and its prefix are cut"
+
+
+@pytest.fixture(scope="module")
+def failing_status_router(_proxy_off):
+ """A node whose /pd/decode answers 500 with an unclassified body."""
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode(status=500))
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ return f"http://127.0.0.1:{_serve(build_app(ctx))}"
+
+
+def test_both_paths_refuse_an_unclassified_status_the_same_way(failing_status_router):
+ """Same node behaviour, same answer, whether or not a stream was asked for.
+
+ They differed: the blocking path let `raise_for_status` reach a generic
+ handler and reported its message, the streaming one reported `decode call
+ failed` with the status. A client could not write one error handler for a
+ node that was down.
+ """
+ blocking = _post(failing_status_router, stop=["STOP"])
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ streaming = c.post(
+ f"{failing_status_router}/v1/chat/completions", json=_body(stream=True, stop=["STOP"])
+ )
+ assert blocking.status_code == 502, blocking.text[:200]
+ assert streaming.status_code == 502, streaming.text[:200]
+ assert blocking.json() == streaming.json(), "one shape for one fault"
+ body = blocking.json()
+ assert body["error"] == "decode call failed", body
+ assert body["status"] == 500, body
+ assert body["rid"], body
+
+
+@pytest.fixture(scope="module")
+def null_logprobs_router(_proxy_off):
+ """A node whose `lp` rows are correctly sized and entirely null."""
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode(null_lp=True))
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ return f"http://127.0.0.1:{_serve(build_app(ctx))}"
+
+
+@pytest.mark.parametrize(
+ "stop,why",
+ [
+ (["STOP"], "streaming protocol -- a stop puts a blocking request on it"),
+ (None, "blocking protocol -- one object, same validator"),
+ ],
+)
+def test_a_null_logprob_past_the_echoed_token_is_refused(null_logprobs_router, stop, why):
+ """Only the first token may report null: it was sampled by prefill, so no
+ decode-side value exists and the router fills it from the prefill reply.
+
+ Anywhere else a null is a node that cannot produce what was asked for, and
+ the length check alone passed it through: `build_logprobs` turned it into
+ -9999.0, the value OpenAI documents for "very unlikely", so the client read
+ an engine fault as a measurement. This is the invariant the decode server
+ already states by raising `LogprobsUnavailable` past position 0; the router
+ is where an older or faulty node has to be caught.
+ """
+ kw = {"stop": stop} if stop else {}
+ r = _post(null_logprobs_router, logprobs=True, top_logprobs=1, temperature=0.6, **kw)
+ assert r.status_code == 501, (why, r.text[:200])
+ assert r.json()["error_type"] == "logprobs_unavailable", why
+
+
+@pytest.fixture(scope="module")
+def no_candidates_router(_proxy_off):
+ """A node that answers with `lp` but no `tp`."""
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode(drop_tp=True))
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ return f"http://127.0.0.1:{_serve(build_app(ctx))}"
+
+
+def test_missing_candidate_rows_are_refused_not_padded(no_candidates_router):
+ """A client that asked for `top_logprobs: 1` and got empty rows was told the request succeeded.
+
+ An absent positional row also loses the alignment, so it is not the same thing as a
+ genuinely empty row.
+ """
+ r = _post(no_candidates_router, stop=["STOP"], logprobs=True, top_logprobs=1, temperature=0.6)
+ assert r.status_code == 501, r.text[:200]
+ assert r.json()["error_type"] == "logprobs_unavailable"
+
+
+def test_no_candidates_asked_for_means_empty_rows_are_correct(no_candidates_router):
+ """`logprobs: true` without `top_logprobs` resolves to `top_n == 0`, and a
+ node that sends no `tp` is then answering exactly what was asked.
+ """
+ r = _post(no_candidates_router, stop=["STOP"], logprobs=True, temperature=0.6)
+ assert r.status_code == 200, r.text[:200]
+ ents = r.json()["choices"][0]["logprobs"]["content"]
+ assert ents and all(e["top_logprobs"] == [] for e in ents)
+
+
+# --------------------------------------------------------------------------- #
+# The constraint: the one shape that is refused rather than guessed
+# --------------------------------------------------------------------------- #
+class _Session:
+ """The minimum an output parser session has to be: everything is content."""
+
+ def feed(self, text):
+ return [{"kind": "content", "text": text}]
+
+ def finish(self):
+ return []
+
+
+class _Parsing:
+ """A router context that has an output parser, like `--parser glm47`."""
+
+ def stream(self):
+ return _Session()
+
+
+@pytest.fixture(scope="module")
+def parsing_router(_proxy_off):
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode())
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ ctx._parsers = {True: _Parsing(), False: _Parsing()} # pretend --parser glm47
+ return f"http://127.0.0.1:{_serve(build_app(ctx))}"
+
+
+@pytest.mark.parametrize(
+ "over,want,why",
+ [
+ (
+ {"stop": ["STOP"], "logprobs": True, "top_logprobs": 1, "temperature": 0.6},
+ 501,
+ "all three: undecidable, so refused",
+ ),
+ ({"stop": ["STOP"]}, 200, "stop alone"),
+ (
+ {"logprobs": True, "top_logprobs": 1, "temperature": 0.6},
+ 200,
+ "parser with logprobs but no stop: nothing is held back",
+ ),
+ ],
+)
+def test_only_stop_with_a_parser_and_logprobs_is_refused(parsing_router, over, want, why):
+ """`logprobs` covers `message.content`, so an entry has to be attributed to a
+ channel, and that is exact only while the parser is fed one token's text at a
+ time. A stop makes the router hold back, so what reaches the parser spans
+ token boundaries -- and for a token whose text the parser did not emit there
+ is no signal to decide the channel.
+
+ Refusing is the design decision, and the other rows are why it is narrow: the
+ combinations that need no arithmetic are all served.
+ """
+ r = _post(parsing_router, **over)
+ assert r.status_code == want, f"{why}: got {r.status_code} {r.text[:160]}"
+ if want == 501:
+ assert r.json()["error_type"] == "capability_unavailable"
+
+
+def test_the_refusal_happens_before_any_backend_work(parsing_router):
+ """501 costs nothing: no prefill, no KV transfer, no decode slot."""
+ _SEEN["stream"] = None
+ r = _post(parsing_router, stop=["STOP"], logprobs=True, top_logprobs=1, temperature=0.6)
+ assert r.status_code == 501
+ assert _SEEN["stream"] is None, "the decode node was contacted"
+
+
+def test_a_chunks_entries_never_run_ahead_of_its_own_text(router):
+ """An entry belongs to the chunk carrying its token's text, not an earlier
+ one.
+
+ With a stop configured the router holds back `len(stop)-1` characters, so a
+ chunk's text can end part-way into a token -- and that token's entry has to
+ wait for the chunk that finishes it. Releasing entries as soon as they exist
+ keeps the totals right and still tells the client that a token's probability
+ describes text it has not been sent yet.
+
+ Checked as a prefix relation over every chunk boundary, which is the strongest
+ statement that survives a chunk ending mid-token.
+ """
+ text = entries_text = ""
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream(
+ "POST",
+ f"{router}/v1/chat/completions",
+ json=_body(stream=True, stop=["ZZZZ"], logprobs=True, top_logprobs=1, temperature=0.6),
+ ) as r:
+ r.raise_for_status()
+ for line in r.iter_lines():
+ if not line.startswith("data: ") or line[6:] == "[DONE]":
+ continue
+ for ch in json.loads(line[6:]).get("choices", []):
+ text += ch["delta"].get("content", "") or ""
+ for e in (ch.get("logprobs") or {}).get("content") or []:
+ entries_text += e["token"]
+ assert text.startswith(entries_text) or entries_text.startswith(text), (
+ f"entries ran ahead: entries={entries_text!r} " f"text={text!r}"
+ )
+ assert len(entries_text) <= len(text), (
+ f"entries describe {len(entries_text)} characters but "
+ f"only {len(text)} have been sent"
+ )
+ assert entries_text == text, "and they agree once the stream ends"
+
+
+def test_a_streamed_logprobs_line_the_router_cannot_use_fails_closed(no_candidates_router):
+ """The non-streaming path answers 501; the streaming path cannot.
+
+ The 200 headers are already out by the time the node's first token line
+ arrives, so the status is spent. Padding the missing candidate rows would
+ stream invented sentinels to a client that asked for real ones, and
+ truncating the stream would leave it without a terminator. The only honest
+ ending is the reply's own text, a finish_reason, a typed error event, and
+ `[DONE]` -- the same shape the grammar-violation branch uses.
+ """
+ finish, err, saw_done, entries = None, None, False, 0
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream(
+ "POST",
+ f"{no_candidates_router}/v1/chat/completions",
+ json=_body(stream=True, stop=["STOP"], logprobs=True, top_logprobs=1, temperature=0.6),
+ ) as r:
+ assert r.status_code == 200, "the status was spent before this"
+ for line in r.iter_lines():
+ if not line.startswith("data: "):
+ continue
+ if line[6:] == "[DONE]":
+ saw_done = True
+ continue
+ payload = json.loads(line[6:])
+ if payload.get("error"):
+ err = payload["error"]
+ continue
+ for ch in payload.get("choices", []):
+ entries += len((ch.get("logprobs") or {}).get("content") or [])
+ if ch.get("finish_reason"):
+ finish = ch["finish_reason"]
+ assert err is not None, "the failure is reported"
+ assert err["error_type"] == "logprobs_unavailable", err
+ assert finish == "stop", "the choice is closed before the error event"
+ assert saw_done, "and the stream is terminated"
+ assert entries == 0, "no invented sentinels reached the client"
+
+
+def test_a_streamed_reply_that_ends_without_a_terminal_message_fails_closed(truncating_router):
+ """The same refusal the non-streaming path makes, in SSE form.
+
+ A body that reaches a clean EOF before `done`, `error` or a stop match is a
+ truncated generation. Emitting a normal finish chunk and `[DONE]` would
+ report it as complete; the non-streaming path answers 502 `decode_truncated`
+ and this one says the same thing with the only means left after the 200.
+ """
+ finish, err, saw_done = None, None, False
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream(
+ "POST",
+ f"{truncating_router}/v1/chat/completions",
+ json=_body(stream=True, stop=["ZZZZ"]),
+ ) as r:
+ assert r.status_code == 200
+ for line in r.iter_lines():
+ if not line.startswith("data: "):
+ continue
+ if line[6:] == "[DONE]":
+ saw_done = True
+ continue
+ payload = json.loads(line[6:])
+ if payload.get("error"):
+ err = payload["error"]
+ continue
+ for ch in payload.get("choices", []):
+ if ch.get("finish_reason"):
+ finish = ch["finish_reason"]
+ assert err is not None and err["error_type"] == "decode_truncated", err
+ assert finish == "stop", "the choice is closed before the error event"
+ assert saw_done
+
+
+def test_a_streamed_stop_match_is_still_a_success(truncating_router):
+ """The stop leaves the loop early on purpose, so it stays a normal stream."""
+ err, text = None, ""
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream(
+ "POST",
+ f"{truncating_router}/v1/chat/completions",
+ json=_body(stream=True, stop=["STOP"]),
+ ) as r:
+ for line in r.iter_lines():
+ if not line.startswith("data: ") or line[6:] == "[DONE]":
+ continue
+ payload = json.loads(line[6:])
+ if payload.get("error"):
+ err = payload["error"]
+ for ch in payload.get("choices", []):
+ text += ch["delta"].get("content", "") or ""
+ assert err is None, f"a matched stop is not a failure: {err}"
+ assert text == "Hello, world. "
+
+
+def test_a_cancelled_finish_reason_is_normalised_on_both_protocol_forms():
+ """`cancelled` is the router's own doing, not a client-visible outcome.
+
+ The streaming form normalised it and the blocking form did not, so a cancel
+ landing while a non-streaming request was in flight surfaced the decode
+ protocol's internal reason as `choices[0].finish_reason`.
+ """
+ from tilert.pd_vllm.decode_response import DecodeReader
+
+ blocking = DecodeReader()
+ blocking.feed_blocking({"token_ids": [], "timing_ms": {"finish_reason": "cancelled"}})
+ assert blocking.finish_reason == "stop"
+
+ streamed = DecodeReader()
+ streamed.feed(json.dumps({"done": True, "finish_reason": "cancelled"}))
+ assert streamed.finish_reason == "stop"
+
+ kept = DecodeReader()
+ kept.feed_blocking({"token_ids": [], "timing_ms": {"finish_reason": "length"}})
+ assert kept.finish_reason == "length", "only `cancelled` is translated"
+
+
+# --------------------------------------------------------------------------- #
+# Failing closed: every way out after the 200 is spent
+# --------------------------------------------------------------------------- #
+@pytest.fixture(scope="module")
+def garbage_router(_proxy_off):
+ """A node that sends a line `DecodeReader.feed` cannot parse."""
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode(garbage_after=4))
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ return f"http://127.0.0.1:{_serve(build_app(ctx))}"
+
+
+def test_a_stream_that_raises_mid_flight_still_terminates(garbage_router):
+ """A malformed line reaches the generator's `except`, not the EOF check.
+
+ The 200 is spent by then, so exiting on the exception leaves the client with
+ a partial response and no terminator -- the same failure the clean-EOF check
+ refuses, reached by a different route.
+ """
+ finish, err, saw_done = None, None, False
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream(
+ "POST", f"{garbage_router}/v1/chat/completions", json=_body(stream=True, stop=["ZZZZ"])
+ ) as r:
+ assert r.status_code == 200
+ for line in r.iter_lines():
+ if not line.startswith("data: "):
+ continue
+ if line[6:] == "[DONE]":
+ saw_done = True
+ continue
+ payload = json.loads(line[6:])
+ if payload.get("error"):
+ err = payload["error"]
+ continue
+ for ch in payload.get("choices", []):
+ if ch.get("finish_reason"):
+ finish = ch["finish_reason"]
+ assert err is not None, "the failure is reported"
+ assert err["error_type"] == "decode_stream_failed", err
+ assert finish == "stop", "the choice is closed first"
+ assert saw_done, "and the stream is terminated"
+
+
+def test_a_decode_post_that_raises_still_cancels_the_node(monkeypatch, _proxy_off):
+ """The node may have admitted the request before the call failed.
+
+ A timeout or a reset while waiting for headers never reaches the reader's
+ `finally`, so without covering the POST phase the node holds its slot until
+ its own timeout -- and #41 exists because a node holding a slot surfaces as a
+ 429 for whoever comes next.
+ """
+ from fastapi.testclient import TestClient
+
+ cancelled = []
+ real_post = pd_router.requests.post
+
+ def flaky_post(url, **kw):
+ if url.endswith("/pd/decode"):
+ raise ConnectionResetError("reset while waiting for headers")
+ if url.endswith("/pd/cancel"):
+ cancelled.append(kw.get("json", {}).get("rid"))
+ return real_post(url, **kw)
+ return real_post(url, **kw)
+
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode())
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ client = TestClient(build_app(ctx))
+ monkeypatch.setattr(pd_router.requests, "post", flaky_post)
+ r = client.post("/v1/chat/completions", json=_body(stop=["STOP"]))
+ assert r.status_code == 502, r.text[:200]
+ for _ in range(200):
+ if cancelled:
+ break
+ time.sleep(0.02)
+ assert cancelled, "the node was left holding its slot"
+
+
+@pytest.fixture(scope="module")
+def lossy_router(_proxy_off):
+ """A node whose `done` count exceeds the token lines that arrived."""
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode(filler=6, lose_lines_after=3))
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ return f"http://127.0.0.1:{_serve(build_app(ctx))}"
+
+
+def test_a_terminal_count_that_disagrees_is_a_truncated_decode(lossy_router):
+ """The node states its own count on the terminal line; a mismatch means
+ token lines were lost.
+
+ It still sends a normal `done`, so nothing else in the exchange looks wrong:
+ accepting it reports a shortened generation with a successful finish reason
+ and an understated `usage`.
+ """
+ r = _post(lossy_router, stop=["ZZZZ"])
+ assert r.status_code == 502, r.text[:200]
+ assert r.json()["error_type"] == "decode_truncated"
+
+
+def test_the_two_refusals_do_not_share_a_status(lossy_router, no_candidates_router):
+ """501 is "the router cannot serve this"; 502 is "the node sent something incomplete".
+
+ Collapsing them would tell a caller to change its request when the backend is at fault.
+ """
+ incomplete = _post(lossy_router, stop=["ZZZZ"])
+ unusable = _post(
+ no_candidates_router, stop=["STOP"], logprobs=True, top_logprobs=1, temperature=0.6
+ )
+ assert incomplete.status_code == 502, incomplete.text[:120]
+ assert unusable.status_code == 501, unusable.text[:120]
+
+
+@pytest.mark.parametrize("value", ["bad", 5, [1, 2]])
+def test_a_malformed_chat_template_kwargs_is_a_400(router, value):
+ """vLLM declares it `dict[str, Any] | None` and answers 422 for anything else.
+
+ Reading `.get` off a string raised AttributeError instead, which the typed handlers do not
+ catch -- so a client mistake surfaced as a 500.
+ """
+ r = _post(router, chat_template_kwargs=value)
+ assert r.status_code == 400, f"got {r.status_code} {r.text[:160]}"
+
+
+def test_a_well_formed_chat_template_kwargs_still_passes(router):
+ for value in ({}, {"enable_thinking": False}, None):
+ r = _post(router, chat_template_kwargs=value)
+ assert r.status_code == 200, f"{value!r}: {r.status_code} {r.text[:120]}"
+
+
+# --------------------------------------------------------------------------- #
+# A typed node error keeps its status on both protocols
+# --------------------------------------------------------------------------- #
+@pytest.mark.parametrize(
+ "error_type,want",
+ [
+ ("logprobs_unavailable", 501),
+ ("capability_unavailable", 501),
+ ("invalid_grammar", 400),
+ ("invalid_parameter", 400),
+ ("request_cancelled", 499),
+ ("grammar_backend_unavailable", 500),
+ (None, 502),
+ ("something_new", 502),
+ ],
+)
+def test_a_typed_node_error_keeps_the_status_it_deserves(_proxy_off, error_type, want):
+ """The two protocols carry the type differently and must not disagree.
+
+ Over the blocking protocol the node answers an HTTP status and the router
+ forwards it. Over the streaming one the error arrives inside a 200 body --
+ the status is already spent -- so the router reconstructs it from the type.
+ Mapping every propagated type to one status is how `logprobs_unavailable`
+ came back 400 here and 501 there for the same inability, and adding a `stop`
+ string is what moves a non-streaming request onto this protocol.
+
+ An unclassified error stays a 502: the node did not say what went wrong, so
+ a component fault is the honest answer.
+ """
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode(error_after=3, error_type=error_type))
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ url = f"http://127.0.0.1:{_serve(build_app(ctx))}"
+ r = _post(url, stop=["ZZZZ"]) # stop -> the streaming protocol
+ assert r.status_code == want, f"{error_type}: got {r.status_code}"
+ if error_type is not None:
+ assert r.json().get("error_type") == error_type
+
+
+@pytest.mark.parametrize(
+ "error_type,inline",
+ [
+ ("logprobs_unavailable", False),
+ ("capability_unavailable", False),
+ ("grammar_violation", False),
+ (None, True),
+ ],
+)
+def test_a_streamed_node_error_fails_closed_when_it_was_classified(_proxy_off, error_type, inline):
+ """A typed error is the node's answer about the contract, not more content.
+
+ Emitting it as a `[decode error: ...]` chunk and then a normal finish reports
+ a successful completion that broke what the request asked for -- unconstrained
+ output for a grammar, or a reply without the logprobs it requested. An
+ UNCLASSIFIED error keeps the inline marker: the node did not say what went
+ wrong, and a visible marker beats a bare error event for a client that is
+ already rendering text.
+ """
+ vllm_port = _serve(_make_vllm())
+ decode_port = _serve(_make_decode(error_after=3, error_type=error_type))
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), _Tok(), "none")
+ url = f"http://127.0.0.1:{_serve(build_app(ctx))}"
+
+ text, err, saw_done = "", None, False
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream("POST", f"{url}/v1/chat/completions", json=_body(stream=True)) as r:
+ assert r.status_code == 200
+ for line in r.iter_lines():
+ if not line.startswith("data: "):
+ continue
+ if line[6:] == "[DONE]":
+ saw_done = True
+ continue
+ payload = json.loads(line[6:])
+ if payload.get("error"):
+ err = payload["error"]
+ continue
+ for ch in payload.get("choices", []):
+ text += ch["delta"].get("content", "") or ""
+ assert saw_done, "the stream terminates either way"
+ if inline:
+ assert "[decode error:" in text, "an unclassified error stays inline"
+ assert err is None
+ else:
+ assert err is not None and err.get("error_type") == error_type, err
+ assert "[decode error:" not in text, "a classified error must not be dressed up as content"
diff --git a/tests/pd_vllm/test_stream_e2e.py b/tests/pd_vllm/test_stream_e2e.py
new file mode 100644
index 0000000..345e904
--- /dev/null
+++ b/tests/pd_vllm/test_stream_e2e.py
@@ -0,0 +1,708 @@
+r"""A client `stream: true` must still stream without `stream_options` on prefill.
+
+`build_prefill_body` strips two fields a streaming client always sends, which
+invites the obvious worry: did we just turn streaming off? No -- the fields were
+going to the wrong backend. Three HTTP conversations are in play and only the
+middle one is forced non-streaming:
+
+ client --stream:true--> router --stream:false--> vLLM prefill (1 token + KV)
+ \--stream:true--> decode node (tokens 2..N)
+ <---- SSE text/event-stream
+
+`stream_options` is meaningful only between client and router (the client wants
+usage accounting) but was being copied onto the prefill request, the one hop
+that *must* be non-streaming. The client's intent is not lost: the router
+satisfies `include_usage` itself in a trailing usage-only chunk, and the
+client's output length reaches the decode node where it belongs.
+
+The other tests in this directory pin `build_prefill_body` in isolation. This
+one runs the real `build_app` against stub backends over real HTTP, so it also
+covers the seams those unit tests cannot see: SSE framing, incremental
+detokenisation, and the NDJSON->SSE conversion from the decode node back to the
+client.
+
+The fake vLLM reproduces `ChatCompletionRequest.validate_stream_options`
+verbatim, so deleting the drop-list makes this test fail the same way a real
+deployment does -- with a 4xx and zero chunks, not a subtle assertion.
+
+No GPU, no tilert, no vllm.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_stream_e2e.py -v
+"""
+
+from __future__ import annotations
+
+import json
+import threading
+import time
+
+import httpx
+import pytest
+import uvicorn
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse, StreamingResponse
+
+from tilert.pd_vllm.decode_pool import DecodeNode, Pool
+from tilert.pd_vllm.pd_router import RouterCtx, build_app, build_prefill_body, should_include_usage
+
+# What `vllm bench serve --backend openai-chat` puts on the wire
+# (vllm/benchmarks/lib/endpoint_request_func.py): streaming, usage accounting,
+# and the output length under the *new* OpenAI field name.
+BENCH_BODY = {
+ "model": "stub-model",
+ "messages": [{"role": "user", "content": "hi"}],
+ "temperature": 0.0,
+ "max_completion_tokens": 3000,
+ "stream": True,
+ "stream_options": {"include_usage": True},
+}
+
+# BENCH_BODY carries temperature 0.0, which is what the benchmark sends and
+# what the non-logprobs cases must keep. Logprobs cannot be served there: the
+# engine's export is only valid down to MIN_LOGPROBS_TEMPERATURE, and at
+# temperature ~= 0 it takes a greedy graph that never runs the top-p sampler at
+# all. Requests asking for logprobs therefore override it.
+LOGPROBS_TEMPERATURE = 0.6
+
+FIRST_TOKEN_ID = 100 # sampled by prefill, handed to decode
+DECODE_TOKEN_IDS = [101, 102, 103]
+EXPECTED_TEXT = "ello world"
+
+
+def STUB_LOGPROB(token_id: int) -> float:
+ return -0.5 - 0.1 * (token_id % 5) # what StubTokenizer makes of the decode tokens
+
+
+class Captured:
+ """Bodies the stub backends received, for asserting on both hops."""
+
+ def __init__(self):
+ self.prefill: dict = {}
+ self.decode: dict = {}
+
+
+class StubTokenizer:
+ _VOCAB = {101: "ello", 102: " wor", 103: "ld"}
+
+ def decode(self, ids, skip_special_tokens=False):
+ return "".join(self._VOCAB.get(i, "") for i in ids)
+
+
+def _make_vllm(cap: Captured) -> FastAPI:
+ """Prefill instance: validates like vLLM, replies like vLLM."""
+ app = FastAPI()
+
+ @app.post("/v1/chat/completions")
+ async def chat(request: Request):
+ cap.prefill = await request.json()
+
+ # vllm/entrypoints/openai/protocol.py, ChatCompletionRequest:
+ # a mode="before" model_validator, so this fires during body parsing --
+ # before the model is looked at. That is why the bug is model-agnostic.
+ if cap.prefill.get("stream_options") and not cap.prefill.get("stream"):
+ return JSONResponse(
+ {
+ "error": {
+ "message": "Stream options can only be defined " "when `stream=True`.",
+ "param": "stream_options",
+ "code": 400,
+ }
+ },
+ status_code=400,
+ )
+
+ # vLLM runs with --return-tokens-as-token-ids, so the router reads the
+ # first token id out of the logprobs as "token_id:N".
+ return JSONResponse(
+ {
+ "id": "cmpl-stub",
+ "model": "stub-model",
+ "choices": [
+ {
+ "index": 0,
+ "finish_reason": "length",
+ "logprobs": {"content": [{"token": f"token_id:{FIRST_TOKEN_ID}"}]},
+ "message": {"role": "assistant", "content": "H"},
+ }
+ ],
+ "usage": {"prompt_tokens": 7, "completion_tokens": 1},
+ }
+ )
+
+ return app
+
+
+def _make_decode(cap: Captured) -> FastAPI:
+ """Decode node: NDJSON token batches then a done line, like /pd/decode."""
+ app = FastAPI()
+
+ @app.get("/capabilities")
+ def capabilities():
+ """A real decode node declares what it can execute, and the router
+ refuses the optional fields no node claims. A stub without this endpoint
+ would make every such request 501 here for a reason that has nothing to
+ do with what these tests are about, so it declares full support -- the
+ gate itself is tested in test_request_capabilities.py.
+ """
+ return {
+ "profile": "stub",
+ "engine": "StubEngine",
+ "capabilities": {"penalties": True, "ignore_eos": True, "logprobs": True},
+ }
+
+ @app.post("/pd/decode")
+ async def decode(request: Request):
+ cap.decode = await request.json()
+
+ # The real decode server has two branches; mirror both so the
+ # non-streaming handler can be exercised too.
+ if not cap.decode.get("stream"):
+ out = {
+ "rid": cap.decode["rid"],
+ "token_ids": DECODE_TOKEN_IDS,
+ "seq_len": 8,
+ "timing_ms": {"finish_reason": "stop"},
+ }
+ n = cap.decode.get("top_logprobs")
+ if n is not None:
+ out["logprobs"] = {
+ "lp": [STUB_LOGPROB(t) for t in DECODE_TOKEN_IDS],
+ "tp": [
+ [[t + k, STUB_LOGPROB(t) - k] for k in range(n)] for t in DECODE_TOKEN_IDS
+ ],
+ }
+ return out
+
+ n = cap.decode.get("top_logprobs")
+
+ def gen():
+ for tid in DECODE_TOKEN_IDS:
+ line = {"t": [tid]}
+ if n is not None:
+ line["lp"] = [STUB_LOGPROB(tid)]
+ line["tp"] = [[[tid + k, STUB_LOGPROB(tid) - k] for k in range(n)]]
+ yield json.dumps(line) + "\n"
+ yield json.dumps({"done": True, "finish_reason": "stop"}) + "\n"
+
+ return StreamingResponse(gen(), media_type="application/x-ndjson")
+
+ @app.post("/pd/cancel")
+ async def cancel():
+ # The router fires this on any incomplete stream; accept and ignore.
+ return {"ok": True}
+
+ return app
+
+
+def _serve(app) -> tuple[uvicorn.Server, int]:
+ """Run `app` on an ephemeral port; return the server and the port."""
+ cfg = uvicorn.Config(app, host="127.0.0.1", port=0, log_level="error")
+ server = uvicorn.Server(cfg)
+ threading.Thread(target=server.run, daemon=True).start()
+ for _ in range(200):
+ if server.started:
+ return server, server.servers[0].sockets[0].getsockname()[1]
+ time.sleep(0.05)
+ raise RuntimeError("stub server did not start")
+
+
+@pytest.fixture(scope="module")
+def _no_proxy_for_loopback():
+ """Keep a local proxy from hijacking the stub traffic.
+
+ The router reaches both stubs with `requests` / `httpx`, which honour
+ `http_proxy` from the environment. On a box that has one set (common — this
+ is also why a real router deployment must set `no_proxy` for its internal
+ addresses), the outbound call is proxied, never arrives, and the test hangs
+ to its read timeout instead of failing usefully. Pin it here so the test
+ does not depend on how the caller's shell is configured.
+ """
+ mp = pytest.MonkeyPatch()
+ for var in ("no_proxy", "NO_PROXY"):
+ mp.setenv(var, "127.0.0.1,localhost")
+ for var in ("http_proxy", "HTTP_PROXY", "https_proxy", "HTTPS_PROXY", "all_proxy", "ALL_PROXY"):
+ mp.delenv(var, raising=False)
+ yield
+ mp.undo()
+
+
+@pytest.fixture(scope="module")
+def stack(_no_proxy_for_loopback):
+ """vLLM stub + decode stub + the real router, wired together."""
+ cap = Captured()
+ _, vllm_port = _serve(_make_vllm(cap))
+ _, decode_port = _serve(_make_decode(cap))
+
+ node = DecodeNode("127.0.0.1", 5556, decode_port)
+ ctx = RouterCtx(f"http://127.0.0.1:{vllm_port}", Pool([node]), StubTokenizer(), "none")
+ _, router_port = _serve(build_app(ctx))
+
+ yield f"http://127.0.0.1:{router_port}", cap
+
+
+@pytest.fixture(scope="module")
+def streamed(stack):
+ """Drive one streaming chat request end to end; return chunks + capture.
+
+ `trust_env=False` keeps a proxy in the environment from hijacking the
+ loopback call -- the router's own outbound `requests` calls need `no_proxy`
+ for the same reason in a real deployment.
+ """
+ url, cap = stack
+ chunks: list[str] = []
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream("POST", f"{url}/v1/chat/completions", json=BENCH_BODY) as r:
+ status, ctype = r.status_code, r.headers.get("content-type", "")
+ for line in r.iter_lines():
+ if line.startswith("data: "):
+ chunks.append(line[6:])
+ return status, ctype, chunks, cap
+
+
+# --------------------------------------------------------------------------- #
+# client <-> router -- the client asked for SSE and gets SSE
+# --------------------------------------------------------------------------- #
+
+
+def test_streaming_request_is_not_rejected(streamed) -> None:
+ """The regression itself: this was 400 (or 502 before #17), with 0 chunks."""
+ status, _, _, _ = streamed
+ assert status == 200
+
+
+def test_response_is_server_sent_events(streamed) -> None:
+ _, ctype, _, _ = streamed
+ assert ctype.startswith("text/event-stream")
+
+
+def test_tokens_arrive_as_separate_chunks(streamed) -> None:
+ """More chunks than tokens would be wrong; one chunk would mean buffering."""
+ _, _, chunks, _ = streamed
+ assert len(chunks) > len(DECODE_TOKEN_IDS)
+
+
+def test_stream_terminates_with_done_sentinel(streamed) -> None:
+ _, _, chunks, _ = streamed
+ assert chunks[-1] == "[DONE]"
+
+
+def test_deltas_reassemble_into_the_decoded_text(streamed) -> None:
+ """Incremental detokenisation must concatenate back to the whole string."""
+ _, _, chunks, _ = streamed
+ # The trailing usage chunk carries `choices: []`, so index defensively.
+ text = "".join(
+ ch["delta"].get("content", "")
+ for c in chunks[:-1]
+ for ch in json.loads(c).get("choices", [])
+ )
+ assert text == EXPECTED_TEXT
+
+
+def test_usage_is_reported_despite_dropping_stream_options(streamed) -> None:
+ """The client asked for include_usage; the router answers it itself."""
+ _, _, chunks, _ = streamed
+ usage = next((u for c in chunks[:-1] if (u := json.loads(c).get("usage"))), None)
+ assert usage == {"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10}
+
+
+def test_usage_rides_a_chunk_of_its_own(streamed) -> None:
+ """Usage sits after the finish_reason chunk, on `choices: []`.
+
+ A client that stops reading at the first chunk bearing a finish_reason --
+ the shape vLLM and the OpenAI API emit, and what `vllm bench serve` does --
+ would never see usage carried on that same chunk.
+ """
+ _, _, chunks, _ = streamed
+ payloads = [json.loads(c) for c in chunks[:-1]]
+ usage_idx = next(i for i, p in enumerate(payloads) if "usage" in p)
+ finish_idx = next(
+ i for i, p in enumerate(payloads) if any(ch.get("finish_reason") for ch in p["choices"])
+ )
+ assert finish_idx < usage_idx
+ assert payloads[usage_idx]["choices"] == []
+ assert "usage" not in payloads[finish_idx]
+
+
+@pytest.fixture(scope="module")
+def streamed_without_usage_opt_in(stack):
+ """One streaming request that does NOT ask for usage accounting."""
+ url, _ = stack
+ body = {k: v for k, v in BENCH_BODY.items() if k != "stream_options"}
+ chunks = []
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream("POST", f"{url}/v1/chat/completions", json=body) as r:
+ assert r.status_code == 200
+ for line in r.iter_lines():
+ if line.startswith("data: ") and line[6:] != "[DONE]":
+ chunks.append(json.loads(line[6:]))
+ return chunks
+
+
+def test_no_empty_choices_chunk_without_the_opt_in(streamed_without_usage_opt_in) -> None:
+ """`choices: []` must stay opt-in.
+
+ A client that never sent `stream_options` has not declared it can handle
+ that shape, and the canonical streaming loop -- `chunk.choices[0].delta`,
+ straight out of the OpenAI docs -- raises IndexError on it. OpenAI gates
+ the usage chunk behind include_usage for exactly this reason.
+ """
+ assert all(p["choices"] for p in streamed_without_usage_opt_in)
+
+
+def test_the_prefill_drop_does_not_consume_the_clients_opt_in() -> None:
+ """Dropping stream_options for vLLM must not disarm the usage gate.
+
+ Two features read the same field with opposite intent: vLLM must NOT see it
+ (it 400s the pair against the forced stream=False, #14) while the gate MUST,
+ and the gate runs after the prefill call. That holds only because
+ ``build_prefill_body`` drops from a copy -- stripping in place would switch
+ every streaming client to "no usage" with nothing failing.
+ """
+ body = {"model": "m", "messages": [], "stream": True, "stream_options": {"include_usage": True}}
+ prefill = build_prefill_body("/v1/chat/completions", body, DecodeNode("h", 5556, 8000))
+ assert "stream_options" not in prefill # vLLM must not see it
+ assert should_include_usage(body) is True # the gate still must
+
+
+def test_no_usage_at_all_without_the_opt_in(streamed_without_usage_opt_in) -> None:
+ """Absent include_usage means no usage anywhere -- not usage relocated.
+
+ The router used to answer usage unconditionally, riding the chunk that
+ closes the choice. That is worse than non-conformant: a client written as
+ ``if chunk.choices: ... elif chunk.usage:`` takes the choices branch and
+ never reads it, which is how the InferenceX bench reported `Total
+ generated tokens: 0`. vLLM and SGLang both gate outright; so do we.
+ """
+ assert not any("usage" in p for p in streamed_without_usage_opt_in)
+
+
+@pytest.mark.parametrize(
+ "body,wanted",
+ [
+ ({"stream_options": {"include_usage": True}}, True),
+ ({"stream_options": {"include_usage": False}}, False),
+ ({"stream_options": {}}, False),
+ ({}, False),
+ # A client serialising an unset option as null, and one sending the wrong
+ # kind: neither may crash the stream on an attribute the router assumed.
+ ({"stream_options": None}, False),
+ ({"stream_options": "true"}, False),
+ ],
+)
+def test_usage_opt_in_predicate(body, wanted) -> None:
+ assert should_include_usage(body) is wanted
+
+
+@pytest.mark.parametrize(
+ "body",
+ [
+ {},
+ {"stream_options": None},
+ {"stream_options": {"include_usage": False}},
+ ],
+)
+def test_force_flag_overrides_an_absent_or_false_opt_in(body) -> None:
+ """The deployment-level escape hatch, mirroring vLLM's
+ enable_force_include_usage and SGLang's
+ stream_response_default_include_usage.
+ """
+ assert should_include_usage(body, force=True) is True
+
+
+def test_force_flag_defaults_off() -> None:
+ """Off unless an operator says otherwise -- the whole point of gating."""
+ ctx = RouterCtx("http://unused", Pool([]), None, "none")
+ assert ctx.force_include_usage is False
+ assert should_include_usage({}, ctx.force_include_usage) is False
+
+
+# --------------------------------------------------------------------------- #
+# router <-> vLLM prefill -- the one hop that must not stream
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize("field", ["stream_options", "max_completion_tokens"])
+def test_contradicting_field_never_reaches_vllm(streamed, field) -> None:
+ _, _, _, cap = streamed
+ assert field not in cap.prefill
+
+
+def test_prefill_leg_is_non_streaming_and_one_token(streamed) -> None:
+ _, _, _, cap = streamed
+ assert cap.prefill["stream"] is False
+ assert cap.prefill["max_tokens"] == 1
+
+
+# --------------------------------------------------------------------------- #
+# router <-> decode node -- where the client's intent actually lands
+# --------------------------------------------------------------------------- #
+
+
+def test_decode_leg_streams(streamed) -> None:
+ _, _, _, cap = streamed
+ assert cap.decode["stream"] is True
+
+
+def test_client_output_length_reaches_the_decode_leg(streamed) -> None:
+ """Dropped from prefill, honoured here -- the split is still a split."""
+ _, _, _, cap = streamed
+ assert cap.decode["max_tokens"] == BENCH_BODY["max_completion_tokens"]
+
+
+def test_prefill_token_is_handed_to_decode(streamed) -> None:
+ """Token 1 comes from prefill's logprobs; decode continues from it."""
+ _, _, _, cap = streamed
+ assert cap.decode["first_token_id"] == FIRST_TOKEN_ID
+
+
+def test_both_length_fields_resolve_the_way_vllm_resolves_them(stack) -> None:
+ """A client sending both must get vLLM's answer, not ours.
+
+ ``resolve_max_tokens`` is unit-tested on its own; this pins the wiring --
+ that the router reads the client body through it on the way to decode.
+ """
+ url, cap = stack
+ body = dict(BENCH_BODY, max_tokens=128, max_completion_tokens=3000)
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream("POST", f"{url}/v1/chat/completions", json=body) as r:
+ for _ in r.iter_lines():
+ pass
+ assert cap.decode["max_tokens"] == 3000
+
+
+def test_ignore_eos_reaches_the_decode_node(stack) -> None:
+ """Only the decode node can act on the flag.
+
+ The prefill request is pinned to max_tokens=1 and never reaches a stop
+ token, so the whole effect lives on the engine the decode node drives.
+ Dropping it is silent: the request succeeds and just stops at the first
+ EOS, which is what a fixed-length benchmark asked it not to do.
+ """
+ url, cap = stack
+ body = dict(BENCH_BODY, ignore_eos=True)
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream("POST", f"{url}/v1/chat/completions", json=body) as r:
+ for _ in r.iter_lines():
+ pass
+ assert cap.decode["sampling"]["ignore_eos"] is True
+
+
+# --------------------------------------------------------------------------- #
+# logprobs validation happens before either backend is contacted
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize(
+ "body,frag",
+ [
+ ({"top_logprobs": 3}, "logprobs must be set to true"),
+ ({"logprobs": True, "top_logprobs": 6}, "[0, 5]"),
+ ({"logprobs": True, "top_logprobs": -1}, "[0, 5]"),
+ ],
+)
+def test_bad_logprobs_request_is_rejected_without_touching_a_backend(stack, body, frag) -> None:
+ """A 400 must come from the router itself: no prefill, no decode, and the
+ stub captures must stay empty for this request.
+ """
+ url, cap = stack
+ cap.prefill.clear()
+ cap.decode.clear()
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ r = c.post(f"{url}/v1/chat/completions", json=dict(BENCH_BODY, stream=False, **body))
+ assert r.status_code == 400
+ payload = r.json()
+ assert payload["error_type"] == "invalid_logprobs"
+ assert frag in payload["error"]
+ assert cap.prefill == {} and cap.decode == {}
+
+
+def test_logprobs_on_completions_is_rejected(stack) -> None:
+ """The count-typed `logprobs` of /v1/completions is not served here."""
+ url, _ = stack
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ r = c.post(
+ f"{url}/v1/completions", json={"model": "stub-model", "prompt": "hi", "logprobs": 2}
+ )
+ assert r.status_code == 400
+ assert "chat/completions" in r.json()["error"]
+
+
+@pytest.fixture(scope="module")
+def logprobs_reply(stack):
+ """One non-streaming chat request asking for logprobs, end to end."""
+ url, cap = stack
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ r = c.post(
+ f"{url}/v1/chat/completions",
+ json=dict(
+ BENCH_BODY,
+ stream=False,
+ temperature=LOGPROBS_TEMPERATURE,
+ logprobs=True,
+ top_logprobs=3,
+ ),
+ )
+ return r, cap
+
+
+def test_logprobs_request_is_forwarded_to_the_decode_node(logprobs_reply):
+ _, cap = logprobs_reply
+ assert cap.decode["top_logprobs"] == 3
+
+
+def test_logprobs_reach_the_client_in_openai_shape(logprobs_reply):
+ r, _ = logprobs_reply
+ assert r.status_code == 200
+ lp = r.json()["choices"][0]["logprobs"]
+ assert set(lp) == {"content", "refusal"}
+ assert lp["refusal"] is None
+ assert len(lp["content"]) == len(DECODE_TOKEN_IDS)
+
+
+def test_each_entry_describes_its_token(logprobs_reply):
+ r, _ = logprobs_reply
+ content = r.json()["choices"][0]["logprobs"]["content"]
+ for tid, item in zip(DECODE_TOKEN_IDS, content):
+ assert item["logprob"] == STUB_LOGPROB(tid)
+ assert item["bytes"] == list(item["token"].encode("utf-8"))
+ assert len(item["top_logprobs"]) == 3
+
+
+def test_reassembled_tokens_match_the_message_content(logprobs_reply):
+ """The array must line up with what the client actually received."""
+ r, _ = logprobs_reply
+ body = r.json()
+ content = body["choices"][0]["logprobs"]["content"]
+ assert "".join(c["token"] for c in content) == body["choices"][0]["message"]["content"]
+
+
+def test_logprobs_is_null_not_absent_when_not_requested(stack):
+ """The field is declared, its value is null -- which is what vLLM sends.
+
+ Non-streaming responses go out as `JSONResponse(content=result.model_dump())`
+ there, with no `exclude_none`, so `ChatCompletionResponseChoice.logprobs`
+ appears as null. (Streaming is the other way round: chunks use
+ `exclude_unset=True`, so an unset logprobs is omitted -- which is what the
+ streaming assertions below check.)
+
+ The intent this replaces is unchanged: nothing is fabricated when logprobs
+ were not asked for. Only the spelling of "nothing" moved, from an absent key
+ to an explicit null a client can read unconditionally.
+ """
+ url, _ = stack
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ r = c.post(f"{url}/v1/chat/completions", json=dict(BENCH_BODY, stream=False))
+ choice = r.json()["choices"][0]
+ assert "logprobs" in choice
+ assert choice["logprobs"] is None
+
+
+# --------------------------------------------------------------------------- #
+# streaming logprobs: each chunk owns the tokens its text came from
+# --------------------------------------------------------------------------- #
+
+
+@pytest.fixture(scope="module")
+def streamed_logprobs(stack):
+ url, cap = stack
+ chunks = []
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream(
+ "POST",
+ f"{url}/v1/chat/completions",
+ json=dict(BENCH_BODY, temperature=LOGPROBS_TEMPERATURE, logprobs=True, top_logprobs=2),
+ ) as r:
+ assert r.status_code == 200
+ for line in r.iter_lines():
+ if line.startswith("data: ") and line[6:] != "[DONE]":
+ chunks.append(json.loads(line[6:]))
+ return chunks
+
+
+def _choices(chunks):
+ """Every choice across `chunks`.
+
+ Flattened rather than indexed at [0] because the trailing usage chunk
+ carries `choices: []`.
+ """
+ return [ch for c in chunks for ch in c["choices"]]
+
+
+def test_content_chunks_carry_logprobs(streamed_logprobs) -> None:
+ lp_choices = [ch for ch in _choices(streamed_logprobs) if ch.get("logprobs")]
+ assert lp_choices, "no chunk carried logprobs"
+ for ch in lp_choices:
+ lp = ch["logprobs"]
+ assert set(lp) == {"content", "refusal"}
+ assert lp["content"]
+
+
+def test_streamed_entries_cover_every_decode_token_once(streamed_logprobs) -> None:
+ """No token may be dropped or double-counted by the pending buffer."""
+ got = [
+ e["logprob"]
+ for ch in _choices(streamed_logprobs)
+ if ch.get("logprobs")
+ for e in ch["logprobs"]["content"]
+ ]
+ assert got == [STUB_LOGPROB(t) for t in DECODE_TOKEN_IDS]
+
+
+def test_streamed_logprob_text_matches_the_delta(streamed_logprobs) -> None:
+ """A chunk's entries must reassemble that chunk's own content delta."""
+ for ch in _choices(streamed_logprobs):
+ if not ch.get("logprobs"):
+ continue
+ joined = "".join(e["token"] for e in ch["logprobs"]["content"])
+ assert joined == ch["delta"].get("content", "")
+
+
+def test_streaming_without_logprobs_has_no_logprobs_key(stack) -> None:
+ url, _ = stack
+ chunks = []
+ with httpx.Client(timeout=30, trust_env=False) as c:
+ with c.stream("POST", f"{url}/v1/chat/completions", json=BENCH_BODY) as r:
+ for line in r.iter_lines():
+ if line.startswith("data: ") and line[6:] != "[DONE]":
+ chunks.append(json.loads(line[6:]))
+ assert not any("logprobs" in ch for ch in _choices(chunks))
+
+
+# --------------------------------------------------------------------------- #
+# envelope: one response, one timestamp
+# --------------------------------------------------------------------------- #
+def test_every_chunk_of_one_response_shares_a_created(streamed):
+ """vLLM threads a single ``created_time`` through every chunk it builds.
+
+ Reading the clock per chunk gave one response several timestamps, which
+ breaks a client that groups or de-duplicates by ``(id, created)``. Asserted
+ behaviourally here — the source-level guard in test_openai_envelope.py
+ catches a reintroduction that this stream is too short to expose.
+ """
+ _, _, chunks, _ = streamed
+ stamps = {json.loads(c)["created"] for c in chunks if c != "[DONE]"}
+ assert len(stamps) == 1, f"one response carried {len(stamps)} timestamps"
+
+
+def test_the_trailing_usage_chunk_shares_it_too(streamed):
+ """The usage chunk is built by a different function; it must agree."""
+ _, _, chunks, _ = streamed
+ parsed = [json.loads(c) for c in chunks if c != "[DONE]"]
+ usage_chunks = [p for p in parsed if p.get("usage") is not None]
+ assert usage_chunks, "no usage chunk to compare"
+ assert {p["created"] for p in parsed} == {usage_chunks[0]["created"]}
+
+
+def test_streamed_usage_carries_a_consistent_total(streamed):
+ _, _, chunks, _ = streamed
+ for c in chunks:
+ if c == "[DONE]":
+ continue
+ u = json.loads(c).get("usage")
+ if u is None:
+ continue
+ assert u["total_tokens"] == u["prompt_tokens"] + u["completion_tokens"]
+ for k, v in u.items():
+ assert isinstance(v, int), f"{k} is {type(v).__name__}"
diff --git a/tests/pd_vllm/test_top_k_resolution.py b/tests/pd_vllm/test_top_k_resolution.py
new file mode 100644
index 0000000..2d8050b
--- /dev/null
+++ b/tests/pd_vllm/test_top_k_resolution.py
@@ -0,0 +1,132 @@
+"""top_k must reach the engine the way vLLM would resolve it.
+
+The router sends ``top_k`` to *both* backends -- ``build_prefill_body`` copies
+the client body to the vLLM prefill instance, and ``_sampling_of`` forwards it
+to the decode node. vLLM applies its own rules on the prefill side, so if the
+decode side disagreed, a single request would sample its first token under
+vLLM's rules and tokens 2..N under ours. That is the same class of bug as the
+``max_completion_tokens`` precedence fix.
+
+vLLM 0.25.1's rule (``SamplingParams`` + ``gpu_input_batch.py``)::
+
+ top_k: int = 0 # "Set to 0 (or -1) to consider all tokens."
+ if 0 < top_k < vocab_size: applied
+ else: top_k = vocab_size # disabled
+
+``resolve_top_k`` mirrors that shape with the kernel's 256-candidate pool
+standing in for vocab_size, so the applied range is [1, 255].
+
+No GPU, no tilert, no vllm: ``sampling`` is pure dict handling.
+
+Run:
+ CUDA_VISIBLE_DEVICES= python -m pytest \
+ tests/pd_vllm/test_top_k_resolution.py -v
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from tilert.pd_vllm.sampling import TOP_K_DISABLED, resolve_top_k
+
+# --------------------------------------------------------------------------- #
+# the applied range
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize("k", [1, 2, 20, 50, 100, 101, 200, 254, 255])
+def test_values_the_kernel_can_apply_pass_through(k) -> None:
+ """[1, 255] is what the sampler's 256-candidate pool can actually cut to."""
+ assert resolve_top_k({"top_k": k}) == k
+
+
+def test_a_recommended_default_in_the_low_tens_is_applied() -> None:
+ """A checkpoint shipping ``top_k: 20`` in generation_config.json must have it
+ take effect, not be swallowed as 'disabled'.
+ """
+ assert resolve_top_k({"top_k": 20}) == 20
+
+
+# --------------------------------------------------------------------------- #
+# vLLM's disable sentinels
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize("k", [0, -1])
+def test_vllm_disable_sentinels(k) -> None:
+ """vLLM: "Set to 0 (or -1) to consider all tokens." Both must disable."""
+ assert resolve_top_k({"top_k": k}) == TOP_K_DISABLED
+
+
+def test_absent_is_disabled() -> None:
+ assert resolve_top_k({}) == TOP_K_DISABLED
+
+
+def test_explicit_null_is_disabled() -> None:
+ """A client serialising an unset option as null must not crash on int()."""
+ assert resolve_top_k({"top_k": None}) == TOP_K_DISABLED
+
+
+# --------------------------------------------------------------------------- #
+# out of range -- must degrade to disabled, never reach the engine unclamped
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize("k", [256, 257, 1000, 2048, 100000])
+def test_at_or_above_the_pool_bound_is_disabled(k) -> None:
+ """256 is the kernel's "no rank cut" value; anything above cannot be
+ honoured by a 256-candidate pool either. vLLM would apply these, so this is
+ a documented divergence -- but it widens the sampled set, never narrows it.
+ """
+ assert resolve_top_k({"top_k": k}) == TOP_K_DISABLED
+
+
+@pytest.mark.parametrize("k", [-2, -100])
+def test_below_vllms_valid_range_is_disabled_not_forwarded(k) -> None:
+ """The vLLM prefill instance 400s these before we are reached; if one ever
+ arrives, disable rather than hand a negative to the engine.
+ """
+ assert resolve_top_k({"top_k": k}) == TOP_K_DISABLED
+
+
+def test_no_value_ever_escapes_the_engines_accepted_range() -> None:
+ """The property that matters: whatever a client sends, the engine sees a value it can accept.
+
+ Guards every call site at once.
+ """
+ for raw in [None, -100, -2, -1, 0, 1, 20, 100, 255, 256, 999, 10**9]:
+ out = resolve_top_k({"top_k": raw})
+ assert 1 <= out <= TOP_K_DISABLED, (raw, out)
+
+
+def test_string_value_is_coerced() -> None:
+ assert resolve_top_k({"top_k": "20"}) == 20
+
+
+def test_other_sampling_keys_are_ignored() -> None:
+ """resolve_top_k reads only top_k; top_p et al. are the caller's business."""
+ assert resolve_top_k({"top_p": 0.95, "temperature": 0.6}) == TOP_K_DISABLED
+
+
+# --------------------------------------------------------------------------- #
+# every engine adapter goes through it
+# --------------------------------------------------------------------------- #
+
+
+@pytest.mark.parametrize("module", ["tilert.pd_vllm.profiles.mla_nsa"])
+def test_adapter_uses_the_shared_resolver(module) -> None:
+ """The sampling call sites shared one unclamped expression
+ (``int(sampling.get("top_k", 256))``). Pin that none of them reintroduces
+ it: the source must reference resolve_top_k and not the old default.
+
+ Source inspection rather than a call, because constructing an adapter needs
+ tilert and a GPU.
+ """
+ import importlib.util
+ import pathlib
+
+ spec = importlib.util.find_spec(module)
+ assert spec is not None and spec.origin is not None, module
+ src = pathlib.Path(spec.origin).read_text()
+ assert "resolve_top_k(sampling)" in src, f"{module} bypasses the resolver"
+ assert 'sampling.get("top_k"' not in src, f"{module} reads top_k directly"
diff --git a/tilert/pd_vllm/capabilities.py b/tilert/pd_vllm/capabilities.py
new file mode 100644
index 0000000..0d751f7
--- /dev/null
+++ b/tilert/pd_vllm/capabilities.py
@@ -0,0 +1,461 @@
+"""The generation parameters the PD stack can execute, and the gate refusing the rest.
+
+The PD split computes one request in two places: the vLLM prefill instance
+samples token 1, the TileRT decode node samples tokens 2..N. The client body is
+forwarded to vLLM almost verbatim (``pd_router.build_prefill_body`` only
+overrides ``max_tokens`` / ``stream`` / ``logprobs`` / ``kv_transfer_params``),
+while the decode node receives only the keys ``pd_router._sampling_of`` selects.
+
+Any field in the gap between those two sets is applied to token 1 and silently
+dropped for the rest of the reply. That is the failure mode this module exists to
+prevent: a client that asked for ``stop`` or ``seed`` gets a 200 whose content
+violates what it asked for, and has no way to detect it. Refusing the request is
+strictly better -- the caller can drop the field, lower its expectations, or
+route to a native vLLM endpoint.
+
+Two tiers, because two different things are being decided:
+
+``_STATIC_FIELDS``
+ Semantics no TileRT decode runtime implements at all. Refused unconditionally
+ when non-neutral. No capability probe can change the answer, so this tier is
+ always enforced and cannot be wrong.
+
+``_PROFILE_FIELDS``
+ Semantics some profiles honour and others do not (penalties need a pre-pass
+ the MLA/NSA runtimes do not carry). Refused only when the serving
+ node's declared capabilities say it cannot execute them -- see
+ :class:`NodeCapabilities` and ``decode_server``'s ``/capabilities``.
+
+"Neutral" means "the value vLLM would have used had the field been absent", read
+off ``ChatCompletionRequest`` and ``_DEFAULT_SAMPLING_PARAMS`` in
+``vllm/entrypoints/openai/chat_completion/protocol.py``. A neutral value is
+accepted, because honouring it and ignoring it are the same computation -- a
+client sending ``frequency_penalty: 0`` is not asking for anything. This is what
+makes the gate deployable: it refuses requests whose *behaviour* would differ,
+not requests that merely mention a field.
+
+Fields whose semantics belong entirely to the prefill stage
+(``truncate_prompt_tokens``, ``echo``, the chat-template knobs, ``tools``) are
+deliberately absent: vLLM applies them where they are meant to apply, so there
+is no gap to close.
+
+So are fields the ROUTER executes itself, over the text it detokenises rather
+than by sampling -- ``stop`` and ``include_stop_str_in_output``. The decode
+loop's lack of them is real and no longer decides anything: an entry here means
+"refuse when non-neutral", so keeping one would refuse a feature the stack
+serves. They are validated where the tokenizer that makes them executable
+lives. Absence is not silence: ``tests/pd_vllm/test_no_silent_degradation.py``
+requires every field the router reads to be declared in ``STATIC_FIELD_NAMES``
+or in its own table, so a field can move between the two but cannot fall out of
+both.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+__all__ = [
+ "CapabilityError",
+ "CapabilityUnavailable",
+ "InvalidParameter",
+ "NodeCapabilities",
+ "PROFILE_FIELD_NAMES",
+ "STATIC_FIELD_NAMES",
+ "TYPED_FIELD_NAMES",
+ "engine_capabilities",
+ "validate_generation_request",
+]
+
+
+# --------------------------------------------------------------------------- #
+# Error hierarchy (same envelope shape as GrammarError / LogprobsUnsupported)
+# --------------------------------------------------------------------------- #
+class CapabilityError(Exception):
+ """Base for request-gate failures. Never degrade silently."""
+
+ error_type = "capability_error"
+ http_status = 500
+
+ def to_payload(self) -> dict[str, str]:
+ return {"error": str(self), "error_type": self.error_type}
+
+
+class InvalidParameter(CapabilityError):
+ """The field's type or value is not usable -> HTTP 400.
+
+ The client's mistake, not a missing feature: no version of this stack would
+ accept it. Mirrors ``InvalidGrammarError`` / ``LogprobsUnsupported``.
+ """
+
+ error_type = "invalid_parameter"
+ http_status = 400
+
+
+class CapabilityUnavailable(CapabilityError):
+ """A valid OpenAI/vLLM field this decode stage cannot execute -> HTTP 501.
+
+ Not a client error: the request is well-formed and a native vLLM endpoint
+ would serve it. 501 is the same status the stack already uses for
+ "the engine cannot produce what was asked for" (``logprobs_unavailable``).
+ """
+
+ error_type = "capability_unavailable"
+ http_status = 501
+
+
+# --------------------------------------------------------------------------- #
+# Per-node capability declaration
+# --------------------------------------------------------------------------- #
+@dataclass(frozen=True)
+class NodeCapabilities:
+ """What one decode node's profile + live engine can execute.
+
+ Defaults are all-False so that every path which cannot obtain a real answer
+ -- an unreachable node, a malformed payload, a node predating
+ ``/capabilities`` -- fails closed rather than assuming support.
+ """
+
+ penalties: bool = False
+ ignore_eos: bool = False
+
+ def intersect(self, other: NodeCapabilities) -> NodeCapabilities:
+ """The capabilities a request can rely on across a whole pool.
+
+ The router picks a node only after validation, so a field may be
+ accepted only if EVERY node could have executed it.
+ """
+ return NodeCapabilities(
+ penalties=self.penalties and other.penalties,
+ ignore_eos=self.ignore_eos and other.ignore_eos,
+ )
+
+ def to_payload(self) -> dict[str, bool]:
+ return {"penalties": self.penalties, "ignore_eos": self.ignore_eos}
+
+ @classmethod
+ def from_payload(cls, payload: object) -> NodeCapabilities:
+ """Parse a ``/capabilities`` response, treating anything unrecognised as unsupported.
+
+ A node that omits a key does not declare it.
+ """
+ if not isinstance(payload, dict):
+ return cls()
+ caps = payload.get("capabilities", payload)
+ if not isinstance(caps, dict):
+ return cls()
+ return cls(
+ penalties=caps.get("penalties") is True,
+ ignore_eos=caps.get("ignore_eos") is True,
+ )
+
+
+def engine_capabilities(engine: object) -> NodeCapabilities:
+ """What the engine actually running on this node can execute.
+
+ Read from the live engine rather than the profile because the adapters
+ DEMOTE their own claims after probing the installed ``tilert`` build: a
+ from-source serve paired with an older engine wheel exposes the same
+ sampling entry points while ignoring penalties, and an adapter that probes
+ for the pre-pass turns that into ``supports_penalties() == False``.
+ Reporting the profile's static claim here
+ would hand the router a promise the engine has already withdrawn.
+
+ Each capability is an optional predicate, matching the ``supports_logprobs``
+ convention: absent means unsupported, and a predicate that raises is treated
+ the same way rather than failing the endpoint.
+ """
+
+ def _ask(name: str) -> bool:
+ probe = getattr(engine, name, None)
+ if not callable(probe):
+ return False
+ try:
+ return bool(probe())
+ except Exception:
+ return False
+
+ return NodeCapabilities(
+ penalties=_ask("supports_penalties"),
+ ignore_eos=_ask("supports_ignore_eos"),
+ )
+
+
+# --------------------------------------------------------------------------- #
+# Field tables
+# --------------------------------------------------------------------------- #
+# Kinds describe how the neutral value is recognised, not the JSON type:
+# "empty" -> neutral when falsy (absent, null, "", [], {})
+# "unset" -> neutral ONLY when absent or null; every other value asks for
+# something. For a field whose zero is a real request -- `seed: 0`
+# is a valid seed, `prompt_logprobs: 0` asks for the prompt
+# tokens' own log probabilities -- "falsy" and "asks for nothing"
+# are different questions, and answering the first would let the
+# request through to be applied on the prefill leg only.
+# "number" -> neutral when numerically equal to `neutral`
+# "flag" -> neutral when the boolean equals `neutral`
+# "count" -> neutral when the integer equals `neutral`; below `minimum` is 400
+_EMPTY, _UNSET, _NUMBER, _FLAG, _COUNT = ("empty", "unset", "number", "flag", "count")
+
+# (kind, neutral, minimum, why-it-cannot-be-honoured)
+_STATIC_FIELDS: dict[str, tuple] = {
+ # ── stop conditions ────────────────────────────────────────────────────
+ # `stop` and `include_stop_str_in_output` are absent from this table on
+ # purpose: text-level stop matching happens in the router, over the text
+ # it detokenises, so the decode loop's lack of it does not matter. They
+ # are validated by the router instead, which is where the tokenizer that
+ # makes them executable lives.
+ "stop_token_ids": (
+ _EMPTY,
+ None,
+ None,
+ "the decode loop uses the model's own stop set and " "accepts no per-request ids",
+ ),
+ "min_tokens": (
+ _COUNT,
+ 0,
+ 0,
+ "the decode loop cannot suppress its stop set for a " "minimum length",
+ ),
+ # ── sampling knobs with no decode-side implementation ──────────────────
+ "frequency_penalty": (
+ _NUMBER,
+ 0.0,
+ None,
+ "the decode sampler implements repetition and " "presence penalties only",
+ ),
+ "min_p": (_NUMBER, 0.0, None, "the decode sampler implements top-p and top-k only"),
+ "seed": (
+ _UNSET,
+ None,
+ None,
+ "the decode sampler's seed is per-process, so a per-request seed "
+ "cannot make the reply reproducible",
+ ),
+ "logit_bias": (_EMPTY, None, None, "the decode sampler has no per-request logit bias"),
+ "bad_words": (_EMPTY, None, None, "the decode loop has no bad-words matcher"),
+ "allowed_token_ids": (_EMPTY, None, None, "the decode sampler has no per-request allow list"),
+ # ── constraints: response_format / regex / ebnf ARE translated (see
+ # grammar_spec); vLLM 0.24's structured_outputs entry point is not ────
+ "structured_outputs": (
+ _EMPTY,
+ None,
+ None,
+ "use response_format, which this stack translates " "into a decode-side grammar",
+ ),
+ # ── multiplicity: one KV state is transferred, so one sequence ──────────
+ "n": (
+ _COUNT,
+ 1,
+ 1,
+ "one prefill KV state is transferred per request, so the decode "
+ "node produces exactly one sequence",
+ ),
+ "best_of": (
+ _COUNT,
+ 1,
+ 1,
+ "the decode node produces exactly one sequence, so there is " "nothing to select from",
+ ),
+ "use_beam_search": (_FLAG, False, None, "the decode node runs single-sequence AR/MTP decode"),
+ # ── response shaping the router does not perform ───────────────────────
+ "prompt_logprobs": (
+ _UNSET,
+ None,
+ None,
+ "the prefill instance is asked for one token, so no " "prompt distribution is collected",
+ ),
+ "logprob_token_ids": (
+ _EMPTY,
+ None,
+ None,
+ "the decode export returns the top candidates, not " "a caller-chosen vocab subset",
+ ),
+ "skip_special_tokens": (
+ _FLAG,
+ True,
+ None,
+ "the reply is detokenised for the output parser, " "which always consumes special tokens",
+ ),
+}
+
+# (kind, neutral, capability attribute, why-it-may-be-unavailable)
+_PROFILE_FIELDS: dict[str, tuple] = {
+ "repetition_penalty": (
+ _NUMBER,
+ 1.0,
+ "penalties",
+ "this model's decode runtime has no penalty " "pre-pass",
+ ),
+ "presence_penalty": (
+ _NUMBER,
+ 0.0,
+ "penalties",
+ "this model's decode runtime has no penalty " "pre-pass",
+ ),
+ "ignore_eos": (
+ _FLAG,
+ False,
+ "ignore_eos",
+ "this model's decode loop does not clear its stop set",
+ ),
+}
+
+# Always executable, so never refused -- but still checked, because the router
+# now RESOLVES these and writes the result into the prefill request. That
+# overwrites whatever the client sent, which takes away vLLM's own chance to
+# reject a bad value: `top_k: 1.9` would be truncated to 1 and served by both
+# legs as a materially different request. (kind, minimum)
+_TYPED_FIELDS: dict[str, tuple] = {
+ "temperature": (_NUMBER, 0.0),
+ "top_p": (_NUMBER, 0.0),
+ "top_k": (_COUNT, None),
+}
+
+STATIC_FIELD_NAMES = tuple(_STATIC_FIELDS)
+PROFILE_FIELD_NAMES = tuple(_PROFILE_FIELDS)
+TYPED_FIELD_NAMES = tuple(_TYPED_FIELDS)
+
+
+# --------------------------------------------------------------------------- #
+# Neutrality tests
+# --------------------------------------------------------------------------- #
+def _as_number(field: str, value: object) -> float:
+ # bool is an int subclass in Python; `top_p: true` is a type error, not 1.0.
+ if isinstance(value, bool) or not isinstance(value, (int, float)):
+ raise InvalidParameter(f"{field} must be a number, got {type(value).__name__}")
+ return float(value)
+
+
+def _as_flag(field: str, value: object) -> bool:
+ if not isinstance(value, bool):
+ raise InvalidParameter(f"{field} must be a boolean, got {type(value).__name__}")
+ return value
+
+
+def _as_count(field: str, value: object, minimum: int) -> int:
+ if isinstance(value, bool) or not isinstance(value, int):
+ raise InvalidParameter(f"{field} must be an integer, got {type(value).__name__}")
+ if value < minimum:
+ raise InvalidParameter(f"{field} must be >= {minimum}, got {value}")
+ return value
+
+
+def _lenient_float(field: str, raw: object) -> float:
+ """A number, accepting the string form vLLM's own validation coerces.
+
+ Deliberately as permissive as vLLM: rejecting ``temperature: "0.7"`` here
+ would turn a request vLLM serves into a 400. Only bools and non-numerics are
+ refused.
+ """
+ if isinstance(raw, bool):
+ raise InvalidParameter(f"{field} must be a number, got bool")
+ try:
+ return float(raw) # type: ignore[arg-type]
+ except (TypeError, ValueError):
+ raise InvalidParameter(f"{field} must be a number, got {raw!r}") from None
+
+
+def _lenient_int(field: str, raw: object) -> int:
+ """An integer, permissive about representation and strict about integrality.
+
+ ``"20"`` and ``20.0`` coerce losslessly and vLLM accepts both. ``1.9`` does
+ not: truncating it to 1 would have the router write a materially different
+ request into BOTH legs and report success, where before this resolution the
+ value reached vLLM and was rejected there.
+ """
+ value = _lenient_float(field, raw)
+ if value != int(value):
+ raise InvalidParameter(f"{field} must be an integer, got {raw!r}")
+ return int(value)
+
+
+def _is_neutral(field: str, value: object, kind: str, neutral, minimum: int | None) -> bool:
+ """Whether ``value`` asks for anything beyond vLLM's own default.
+
+ ``None`` is always neutral: an explicit null is how SDKs spell "unset", and
+ vLLM resolves it to the same default as an absent key.
+ """
+ if value is None:
+ return True
+ if kind == _EMPTY:
+ return not value
+ if kind == _UNSET:
+ return False # None already returned True above
+ if kind == _NUMBER:
+ return _as_number(field, value) == neutral
+ if kind == _FLAG:
+ return _as_flag(field, value) is neutral
+ # _COUNT
+ return _as_count(field, value, minimum if minimum is not None else 0) == neutral
+
+
+# --------------------------------------------------------------------------- #
+# The gate
+# --------------------------------------------------------------------------- #
+def validate_generation_request(
+ body: dict,
+ capabilities: NodeCapabilities | None = None,
+ adopted: dict | None = None,
+) -> None:
+ """Refuse a request whose generation parameters the decode stage would drop.
+
+ Call BEFORE anything observable happens -- before the vLLM prefill request,
+ before a decode node is acquired, before the connector claims anything --
+ so a refused request costs nothing and holds nothing.
+
+ ``capabilities`` is what every node in the pool can execute (the
+ intersection; see :meth:`NodeCapabilities.intersect`). ``None`` means the
+ router could not establish it, which is treated as no support: a field whose
+ execution cannot be confirmed must not be accepted.
+
+ ``adopted`` carries the deployment's own defaults for the profile-dependent
+ fields (``generation_defaults``). They are checked exactly like a
+ client-supplied value, because they end up on the wire the same way: a
+ ``repetition_penalty`` taken from the model's ``generation_config.json`` is
+ still a penalty the decode node has to apply, and a node whose engine demoted
+ its penalty claim must refuse it rather than decode unpenalised. Without this
+ the gate would look only at the request and never see it.
+
+ Raises:
+ InvalidParameter: unusable type or value (400).
+ CapabilityUnavailable: valid field, no decode-side implementation (501).
+ """
+ for field, (kind, minimum) in _TYPED_FIELDS.items():
+ if body.get(field) is None:
+ continue
+ if kind == _NUMBER:
+ value = _lenient_float(field, body[field])
+ if minimum is not None and value < minimum:
+ raise InvalidParameter(f"{field} must be >= {minimum}, got {value}")
+ else:
+ _lenient_int(field, body[field])
+
+ for field, (kind, neutral, minimum, why) in _STATIC_FIELDS.items():
+ if field not in body:
+ continue
+ if _is_neutral(field, body[field], kind, neutral, minimum):
+ continue
+ raise CapabilityUnavailable(
+ f"{field} is not supported by the TileRT decode stage: {why}. "
+ f"The vLLM prefill instance would apply it to the first token and "
+ f"the decode node would ignore it for the rest of the reply, so "
+ f"the request is refused instead of served incorrectly."
+ )
+
+ caps = capabilities or NodeCapabilities()
+ for field, (kind, neutral, attr, why) in _PROFILE_FIELDS.items():
+ if field in body:
+ value, origin = body[field], "the request"
+ elif adopted is not None and field in adopted:
+ value, origin = adopted[field], "this deployment's defaults"
+ else:
+ continue
+ if _is_neutral(field, value, kind, neutral, None):
+ continue
+ if getattr(caps, attr):
+ continue
+ raise CapabilityUnavailable(
+ f"{field} (from {origin}) is not supported by the decode node "
+ f"serving this pool: {why}. It would apply to the first token only, "
+ f"so the request is refused instead of served incorrectly."
+ )
diff --git a/tilert/pd_vllm/decode_pool.py b/tilert/pd_vllm/decode_pool.py
new file mode 100644
index 0000000..cc4ca00
--- /dev/null
+++ b/tilert/pd_vllm/decode_pool.py
@@ -0,0 +1,205 @@
+"""Decode nodes: who is free, who holds one, and who tells one to stop.
+
+A TileRT decode engine serves ONE sequence at a time, so a node is a reservation
+rather than a connection: the router hands it out, the request holds it for its
+whole life, and it goes back exactly once. Two rules make that safe, and both
+were learned the hard way:
+
+* release EXACTLY once, on every exit -- including the ones that are not
+ exceptions. A cancelled task (a client hanging up mid-stream) does not raise
+ ``Exception``, so a handler that only caught that leaked the node until
+ restart.
+* cancel the node if, and only if, a request went out to it and it has not
+ reported ``done``. It admits the request before answering and then holds its
+ slot until its OWN timeout, so a POST that failed while waiting for headers
+ still needs cancelling -- while cancelling one the node never saw is harmless
+ but cancelling one that finished is not, it can land on the NEXT request.
+
+:class:`NodeLease` is those two rules in one object, because they were spelled
+out at five call sites across two handlers and each fix landed at some of them.
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+import time
+
+import requests
+
+from tilert.pd_vllm.capabilities import NodeCapabilities
+
+logger = logging.getLogger("pd_vllm.pool")
+
+__all__ = ["DecodeNode", "NodeLease", "Pool", "acquire_lease", "cancel_decode"]
+
+QUEUE_LOG_SECONDS = 0.1
+
+
+class DecodeNode:
+ def __init__(self, host: str, ctrl_port: int, http_port: int):
+ self.host = host
+ self.ctrl_port = ctrl_port
+ self.http_port = http_port
+ self.busy = False
+ # Declared capabilities, probed lazily. None = not established yet.
+ self.caps: NodeCapabilities | None = None
+ self.caps_at: float = 0.0
+
+ @property
+ def http_base(self) -> str:
+ return f"http://{self.host}:{self.http_port}"
+
+
+class Pool:
+ """Decode-node reservation.
+
+ ``queue_timeout`` > 0 makes ``acquire`` wait for a node instead of failing
+ fast. A decode engine serves one sequence at a time, so a client that puts
+ more than one request in flight per node — a multi-turn agentic session
+ fanning out into concurrent sub-conversations, for instance — otherwise
+ gets 429s for load the pool can serve a moment later. 0 keeps the
+ fail-fast behaviour.
+ """
+
+ # How long a probed capability set is trusted. Bounded so a node that is
+ # restarted onto a newer engine is picked up without restarting the router,
+ # which the deployment contract promises ("三个组件可独立重启").
+ CAPS_TTL_S = 60.0
+ CAPS_TIMEOUT_S = 2.0
+
+ def __init__(self, nodes: list[DecodeNode], queue_timeout: float = 0.0):
+ self.nodes = nodes
+ self.queue_timeout = queue_timeout
+ self._cv = threading.Condition()
+ self._caps_lock = threading.Lock()
+
+ def acquire(self) -> DecodeNode | None:
+ """Reserve a node, or None once ``queue_timeout`` elapses.
+
+ Blocks while waiting; both call sites already hop off the event loop
+ via ``run_in_threadpool``, so other streams keep being served.
+ """
+ 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._cv:
+ node.busy = False
+ self._cv.notify()
+
+ # ── capability probing ───────────────────────────────────────────────
+ def _node_caps(self, node: DecodeNode) -> NodeCapabilities:
+ """One node's declared capabilities, cached for ``CAPS_TTL_S``.
+
+ A probe that fails returns "nothing supported" WITHOUT caching, so the
+ answer is conservative right now and re-probed on the next request
+ rather than pinned for a whole TTL. Probing lazily (not at startup)
+ keeps the router startable before any decode node exists.
+ """
+ with self._caps_lock:
+ if node.caps is not None and time.time() - node.caps_at < self.CAPS_TTL_S:
+ return node.caps
+ try:
+ r = requests.get(f"{node.http_base}/capabilities", timeout=self.CAPS_TIMEOUT_S)
+ r.raise_for_status()
+ caps = NodeCapabilities.from_payload(r.json())
+ except Exception as e:
+ # Includes a node predating /capabilities (404): it cannot declare
+ # support, so it does not get credit for any.
+ logger.warning(
+ "capability probe failed for %s (%s); treating "
+ "every optional field as unsupported",
+ node.http_base,
+ e,
+ )
+ return NodeCapabilities()
+ with self._caps_lock:
+ node.caps, node.caps_at = caps, time.time()
+ logger.info("capabilities for %s: %s", node.http_base, caps.to_payload())
+ return caps
+
+ def capabilities(self) -> NodeCapabilities:
+ """What a request may rely on whichever node serves it.
+
+ The intersection across the pool, because validation runs before a node
+ is chosen. An empty pool yields no support, which is also correct: there
+ is nothing that could execute the field.
+ """
+ result: NodeCapabilities | None = None
+ for n in self.nodes:
+ caps = self._node_caps(n)
+ result = caps if result is None else result.intersect(caps)
+ return result or NodeCapabilities()
+
+
+def cancel_decode(node, rid: str) -> None:
+ """Tell a decode node to stop working on ``rid``.
+
+ Best effort: the node may already have finished, and a failed POST has no
+ remedy. Callers run it off the response path so a wedged node cannot delay
+ the reply.
+ """
+ try:
+ requests.post(f"{node.http_base}/pd/cancel", json={"rid": rid}, timeout=5)
+ except Exception:
+ logger.warning("cancel POST failed for %s", rid)
+
+
+class NodeLease:
+ """One node, held for one request.
+
+ ``rid`` and ``dispatched`` are set as the request learns them: the rid comes
+ from the prefill reply, and ``dispatched`` goes true the moment a POST leaves
+ for the node -- not when it succeeds, since a timeout while waiting for
+ headers may still have been admitted.
+
+ ``release`` is idempotent and takes the one fact the lease cannot know: did
+ the NODE say it was done. Cancelling runs on a plain thread, off the response
+ path, so a wedged node cannot delay the reply, and off the event loop, where
+ an ``await`` could be cancelled before it fires.
+ """
+
+ def __init__(self, pool: Pool, node: DecodeNode):
+ self.pool = pool
+ self.node = node
+ self.rid: str | None = None
+ self.dispatched = False
+ self._released = False
+
+ def release(self, *, terminated: bool = False) -> None:
+ if self._released:
+ return
+ self._released = True
+ self.pool.release(self.node)
+ if self.dispatched and not terminated and self.rid is not None:
+ threading.Thread(target=cancel_decode, args=(self.node, self.rid), daemon=True).start()
+
+ def __enter__(self) -> NodeLease:
+ return self
+
+ def __exit__(self, *exc) -> None:
+ # A path that knows the node terminated releases explicitly before
+ # leaving; this is the backstop for every other exit.
+ self.release()
+
+
+def acquire_lease(pool: Pool) -> tuple[NodeLease | 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 None:
+ return None, waited
+ if waited >= QUEUE_LOG_SECONDS:
+ logger.info("queued %.1fs for decode node %s", waited, node.host)
+ return NodeLease(pool, node), waited
diff --git a/tilert/pd_vllm/decode_response.py b/tilert/pd_vllm/decode_response.py
new file mode 100644
index 0000000..679efcc
--- /dev/null
+++ b/tilert/pd_vllm/decode_response.py
@@ -0,0 +1,309 @@
+"""Reading a ``/pd/decode`` HTTP response, independent of how the bytes arrive.
+
+Not the PD wire protocol -- that is prefill <-> decode over TCP and RDMA and
+lives in ``wire.py``. This is the HTTP answer the node gives the ROUTER.
+
+``/pd/decode`` returns one JSON object, or NDJSON lines (``{"t": [...]}``
+repeatedly, then ``{"done": ...}`` or ``{"error": ...}``). Both router paths read
+it and both used to carry their own copy: four message kinds, a logprobs line to
+validate, a terminal message to notice, a node to cancel when one never arrived.
+
+Only the TRANSPORT and the PRESENTATION genuinely differ -- ``requests`` on a
+worker thread vs ``httpx`` on the event loop, and an HTTP status vs an SSE event
+once the streaming response has sent its 200. So the caller keeps the loop and
+the presentation, and :class:`DecodeReader` takes the rest::
+
+ for line in :
+ for emission in reader.feed(line):
+
+ if reader.finished:
+ break
+ # then, once: reader.refusal / reader.node_error / reader.timing
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Container
+from dataclasses import dataclass, field
+from typing import Any
+
+__all__ = [
+ "BUSY",
+ "DecodeReader",
+ "OK",
+ "PROPAGATE",
+ "REFUSED",
+ "RETRY",
+ "SERVER_ERROR",
+ "TRUNCATED",
+ "TYPED_ERROR",
+ "UNTYPED_ERROR",
+ "PROPAGATED_ERROR_STATUS",
+ "PROPAGATED_ERROR_TYPES",
+ "terminal_verdict",
+ "decode_refusal",
+ "classify_decode_status",
+]
+
+# What to do about a `/pd/decode` response status.
+RETRY = "retry" # 429 and attempts remain
+BUSY = "busy" # 429 and they do not: an honest, retryable 429
+PROPAGATE = "propagate" # a typed error the client should see verbatim
+SERVER_ERROR = "server_error" # anything else: the node is broken
+
+
+def classify_decode_status(
+ status: int, payload: Any, *, attempts_left: bool, propagated_types: Container[str]
+) -> str:
+ """What a ``/pd/decode`` status means, without deciding how to say it.
+
+ A 429 is the node's admission control answering, so it is retryable rather
+ than a fault. A typed error is the node's considered answer about THIS
+ request and reaches the client verbatim; flattening it into 502 would claim
+ a component is broken when none is.
+ """
+ if status == 429:
+ return RETRY if attempts_left else BUSY
+ if status == 200:
+ raise ValueError("200 is not an error status")
+ if isinstance(payload, dict) and payload.get("error_type") in propagated_types:
+ return PROPAGATE
+ return SERVER_ERROR
+
+
+# What a node's own error types mean to the client. Taken from the classes that
+# raise them, because the blocking protocol carries an HTTP status while the
+# streaming one carries the error inside a spent 200.
+PROPAGATED_ERROR_STATUS = {
+ # grammar_spec.py
+ "invalid_grammar": 400,
+ "grammar_violation": 400,
+ "grammar_backend_unavailable": 500,
+ # decode_server.py: the engine cannot produce what was asked for.
+ "logprobs_unavailable": 501,
+ # capabilities.py. Reachable even though the router pre-validates: its
+ # capability cache can be up to Pool.CAPS_TTL_S stale, and the node is the
+ # authority.
+ "capability_unavailable": 501,
+ "invalid_parameter": 400,
+ # The caller asked the node to stop mid-request; not a component fault.
+ "request_cancelled": 499,
+}
+PROPAGATED_ERROR_TYPES = frozenset(PROPAGATED_ERROR_STATUS)
+
+
+# What the reader's FINAL state means.
+OK = "ok" # nothing to refuse
+REFUSED = "refused" # the router cannot use what arrived
+TRUNCATED = "truncated" # a clean EOF with no terminal message
+TYPED_ERROR = "typed_error" # the node classified it; the client can act
+UNTYPED_ERROR = "untyped_error" # the node broke
+
+
+def terminal_verdict(reader, *, client_gone: bool = False) -> tuple[str, dict, int]:
+ """What the reader's final state means, and what says it.
+
+ The order is the point. Both response paths had this chain, in DIFFERENT
+ orders -- one checked the node's error before the truncation test and the
+ other after. They happened to agree, because an error line also marks the
+ node terminated, so truncation cannot fire alongside one; nothing said so.
+
+ Rendering stays with the caller: only the blocking path can answer a status,
+ and only the streaming path has to say it inside a spent 200.
+ """
+ if reader.refusal is not None:
+ return REFUSED, reader.refusal, reader.refusal_status
+ if reader.node_error is not None:
+ error_type = reader.node_error.get("error_type")
+ if error_type in PROPAGATED_ERROR_TYPES:
+ return (TYPED_ERROR, reader.node_error, PROPAGATED_ERROR_STATUS[error_type])
+ return UNTYPED_ERROR, reader.node_error, 502
+ if not reader.node_terminated and not reader.stop_hit and not client_gone:
+ # Assembling it would report finish_reason "stop" for a reply whose body
+ # was cut short, or whose node died.
+ return (
+ TRUNCATED,
+ {
+ "error": "decode stream ended without a terminal message",
+ "error_type": "decode_truncated",
+ "rid": reader.rid,
+ },
+ 502,
+ )
+ return OK, {}, 200
+
+
+def decode_refusal(verdict: str, status: int, payload: Any, rid: str) -> tuple[dict, int]:
+ """How to SAY what :func:`classify_decode_status` decided.
+
+ One body and one status per verdict, for both response paths. They answered
+ a non-200 separately before, and the 502 differed between them: the blocking
+ path let `raise_for_status` reach a generic handler and reported its message,
+ the streaming one reported `decode call failed` with the status. Same node
+ behaviour, two shapes, depending on whether the client asked for a stream.
+
+ ``RETRY`` has no answer -- the caller retries -- so asking for one is a bug.
+ """
+ if verdict == BUSY:
+ return ({"error": "decode node busy", "error_type": "decode_busy", "rid": rid}, 429)
+ if verdict == PROPAGATE:
+ return (payload, status)
+ if verdict == SERVER_ERROR:
+ return ({"error": "decode call failed", "status": status, "rid": rid}, 502)
+ raise ValueError(f"{verdict} is not a refusal")
+
+
+@dataclass
+class DecodeReader:
+ """Read one ``/pd/decode`` response, line by line, into emissions.
+
+ ``feed`` returns whatever became emittable; the caller reads the rest after
+ the loop.
+
+ ``finished`` stop reading: the node terminated, or a stop ended the
+ reply while the node kept generating.
+ ``node_terminated`` whether the NODE said so. It owns its slot until then, so
+ any other exit has to cancel.
+ ``refusal`` / a payload that cannot be served, and the status it
+ ``refusal_status`` deserves: 501 when the router cannot use what arrived,
+ 502 when what arrived is incomplete.
+ ``node_error`` a typed error the node reported.
+ ``timing`` / from the terminal message; ``timing`` is empty when a stop
+ ``finish_reason`` ended the reply, since it rides that line.
+ """
+
+ stream: Any = None # a ReplyStream, or None with no tokenizer
+ logprobs_req: Any = None
+ rid: str = ""
+
+ token_ids: list[int] = field(default_factory=list)
+ timing: dict = field(default_factory=dict)
+ finish_reason: str = "stop"
+ node_terminated: bool = False
+ refusal: dict | None = None
+ # The status that refusal deserves. 501 when the router cannot serve what the
+ # node sent, 502 when what the node sent is incomplete -- different faults,
+ # and the caller should not have to infer which from the payload.
+ refusal_status: int = 501
+ node_error: dict | None = None
+ _stopped: bool = False
+
+ @property
+ def finished(self) -> bool:
+ """Whether the caller should stop reading."""
+ return self.node_terminated or self._stopped or self.refusal is not None
+
+ @property
+ def stop_hit(self) -> bool:
+ """Whether a stop string ended the reply rather than the node."""
+ return self._stopped
+
+ def feed(self, line: str) -> list:
+ """One NDJSON line -> whatever became emittable."""
+ if not line:
+ return []
+ return self._message(json.loads(line))
+
+ def feed_blocking(self, body: dict) -> list:
+ """The one-object form of the same protocol.
+
+ Ids and logprobs arrive together, so it is one token message plus the
+ terminal one.
+ """
+ self.timing = body.get("timing_ms", {})
+ self._set_finish(self.timing.get("finish_reason", "stop"))
+ self.node_terminated = True
+ lp = body.get("logprobs") or {}
+ return self._tokens({"t": body["token_ids"], **lp})
+
+ # ── internals ───────────────────────────────────────────────────────────
+
+ def _message(self, msg: dict) -> list:
+ if "t" in msg:
+ return self._tokens(msg)
+ if "done" in msg:
+ self.node_terminated = True
+ self.timing = msg.get("timing_ms", {})
+ self._set_finish(msg.get("finish_reason", "stop"))
+ declared = msg.get("n")
+ got = len(self.stream.token_ids) if self.stream is not None else len(self.token_ids)
+ if declared is not None and declared != got:
+ # The node states its own count on the terminal line, so a
+ # mismatch means token lines were lost on the way -- a proxy, or
+ # a node whose stream did not survive. Accepting it would report
+ # a shortened generation with a successful finish reason and an
+ # understated usage. Only reachable when nothing stopped the read
+ # early: a matched stop or an earlier refusal never gets here.
+ self.refusal = {
+ "error": f"decode node declared {declared} tokens and sent " f"{got}",
+ "error_type": "decode_truncated",
+ "rid": self.rid,
+ }
+ self.refusal_status = 502
+ return []
+ if "error" in msg:
+ self.node_terminated = True
+ self.node_error = msg
+ return []
+ return []
+
+ def _set_finish(self, reason: str) -> None:
+ """Both forms of the protocol normalise the same way.
+
+ `cancelled` is the router's own doing, not a client-visible outcome: the
+ reply it already has is complete. Setting it directly in one form and
+ not the other is how it leaked once.
+ """
+ self.finish_reason = "stop" if reason == "cancelled" else reason
+
+ def _tokens(self, msg: dict) -> list:
+ ids = msg["t"]
+ seen = len(self.stream.token_ids) if self.stream is not None else len(self.token_ids)
+ if self.logprobs_req is not None and not _logprobs_line_ok(
+ msg, len(ids), self.logprobs_req, seen
+ ):
+ # Refused, not padded: the documented sentinel would report a model
+ # that had nothing to say about its own tokens.
+ self.refusal = {
+ "error": "decode node returned no logprobs",
+ "error_type": "logprobs_unavailable",
+ "rid": self.rid,
+ }
+ return []
+ if self.stream is None:
+ # No tokenizer: ids are all the reply can carry.
+ self.token_ids += ids
+ return []
+ out = self.stream.push(ids, msg.get("lp"), msg.get("tp"))
+ if self.stream.stop_reason is not None:
+ self._stopped = True
+ return out # noqa: R504 (stop_reason is read after push)
+
+
+def _logprobs_line_ok(payload: dict, n_tokens: int, req, seen: int = 0) -> bool:
+ """Whether a node's logprobs cover every token in the same message.
+
+ ``lp`` always. ``tp`` only when candidates were asked for: ``logprobs: true``
+ alone resolves to ``top_n == 0``, where empty rows are the right answer. A
+ short or absent ``tp`` is not an empty row -- it loses the alignment too.
+
+ Null is a value, not a length: only the reply's FIRST token may report it,
+ because prefill sampled it and no decode-side value exists. `seen` is how
+ many tokens already arrived, so the exemption cannot follow a batch. Past
+ there, a null is a node that cannot produce what was asked for -- padding it
+ reaches the client as -9999.0, which OpenAI documents for "very unlikely"
+ and is indistinguishable from a measurement. The decode server states the
+ same invariant by raising ``LogprobsUnavailable``; this is where an older or
+ faulty node is caught.
+ """
+ lp = payload.get("lp")
+ if lp is None or len(lp) != n_tokens:
+ return False
+ if any(value is None for i, value in enumerate(lp) if not (seen == 0 and i == 0)):
+ return False
+ if req.top_n > 0:
+ tp = payload.get("tp")
+ if tp is None or len(tp) != n_tokens:
+ return False
+ return True
diff --git a/tilert/pd_vllm/decode_server.py b/tilert/pd_vllm/decode_server.py
index 3372694..cc90786 100644
--- a/tilert/pd_vllm/decode_server.py
+++ b/tilert/pd_vllm/decode_server.py
@@ -1,4 +1,4 @@
-"""PD decode server (W6): HTTP orchestration around receive -> convert -> inject -> decode.
+"""PD decode server: HTTP orchestration around receive -> convert -> inject -> decode.
Internal token-level API (the client-facing OpenAI layer lives in pd_router /
a later serving layer):
@@ -6,17 +6,27 @@
POST /pd/decode {rid, first_token_id, max_tokens, sampling?, timeout_s?}
Waits for the wire transfer of `rid` to complete, converts, injects
into the engine, decodes, returns {"rid", "token_ids", "timing_ms"}.
+ Refuses (501) before the wire-wait if `sampling` asks for something this
+ engine cannot execute -- see /capabilities.
GET /health {"status": "ok"}
+ GET /capabilities which optional sampling params this engine honours; the
+ router reads it to refuse such a request before the
+ prefill instance has run the prompt
GET /decode_status {"status": "idle"|"busy", "current_rid": ...}
bs=1: a busy server answers 429 immediately (the router's gated dispatch
should make that unreachable).
+
+Run (stub engine, plumbing test):
+ python -m tilert.pd_vllm.decode_server \
+ --engine stub --max-seq-len 4096 --ctrl-port 5556 --http-port 5557
"""
import argparse
import contextlib
import json
import logging
+import os
import queue as queue_mod
import socket
import threading
@@ -28,10 +38,38 @@
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import BaseModel
+from tilert.pd_vllm.capabilities import (
+ CapabilityError,
+ engine_capabilities,
+ validate_generation_request,
+)
+from tilert.pd_vllm.grammar_spec import (
+ GrammarError,
+ GrammarViolationError,
+)
from tilert.pd_vllm.receive_server import ReceiveServer
logger = logging.getLogger("pd_vllm.decode_server")
+# How long to wait for the KV of a request we are rejecting after its prefill has
+# already run, before giving up and releasing the receive slot anyway. The
+# transfer is already in flight when we reject, so this is normally milliseconds;
+# it only bites if the prefill instance died mid-push.
+_ABANDON_DRAIN_S = 30.0
+
+# Returned by _drain_own_kv when the wire-wait was cancelled rather than timing
+# out. A distinct sentinel because the two mean opposite things to an operator:
+# a timeout points at the RDMA path, a cancel means the client left.
+_CANCELLED = object()
+
+
+class LogprobsUnavailable(Exception):
+ """The active engine cannot produce logprobs -> HTTP 501.
+
+ Refused rather than answered without the field: a caller cannot tell that
+ apart from the model having had nothing to report.
+ """
+
class DecodeBody(BaseModel):
rid: str
@@ -40,6 +78,16 @@ class DecodeBody(BaseModel):
sampling: dict | None = None
timeout_s: float = 120.0
stream: bool = False
+ # Constrained decoding: engine grammar spec ({"type","value"}) + thinking
+ # gate. None -> unconstrained (the router omits it for plain requests).
+ grammar_spec: dict | None = None
+ enable_thinking: bool = True
+ # Per-token log probabilities. None -> not requested (the router omits it),
+ # 0 -> the chosen token's logprob only, N -> N candidates per position.
+ top_logprobs: int | None = None
+
+
+DECODE_POLL_S = max(0.0, float(os.environ.get("TILERT_DECODE_POLL_MS") or "200")) / 1000.0
def build_app(server: ReceiveServer, engine) -> FastAPI:
@@ -51,6 +99,28 @@ def build_app(server: ReceiveServer, engine) -> FastAPI:
def health():
return {"status": "ok"}
+ @app.get("/capabilities")
+ def capabilities():
+ """Which optional generation parameters this node can execute.
+
+ The router polls this so it can refuse an unsupported request before
+ prefilling it. Without it the same request still fails -- but only after
+ vLLM has run the prompt and pushed its KV over RDMA, which is the whole
+ cost of the request for none of the answer.
+
+ Sourced from the live engine (see ``engine_capabilities``), so a probe
+ demotion inside the adapter shows up here too. ``logprobs`` is reported
+ alongside for symmetry with the 501 ``/pd/decode`` already returns.
+ """
+ caps = engine_capabilities(engine)
+ payload = caps.to_payload()
+ payload["logprobs"] = bool(getattr(engine, "supports_logprobs", lambda: False)())
+ return {
+ "profile": getattr(server.profile, "name", None),
+ "engine": type(engine).__name__,
+ "capabilities": payload,
+ }
+
@app.get("/decode_status")
def decode_status():
busy = lock.locked()
@@ -80,11 +150,82 @@ def _cleanup():
engine.reset()
except Exception:
logger.exception("engine reset failed")
- server.release()
+ # Scoped to OUR rid: by now the slot may hold a later request whose
+ # transfer started while this one was giving up.
+ if state["current_rid"]:
+ server.release(state["current_rid"])
state["current_rid"] = None
state["cancel_event"] = None
lock.release()
+ def _drain_own_kv(rid: str, timeout_s: float, cancel=None):
+ """Pop from server.completed until OUR rid surfaces.
+
+ Returns the request, ``None`` on timeout, or ``_CANCELLED`` if ``cancel``
+ was set while waiting.
+
+ Observing ``cancel`` is the point: a client that hangs up during the KV
+ transfer used to leave this loop running to the full timeout_s (120 s by
+ default), and since bs=1 the whole node was unavailable for that long.
+ The cancel is checked every poll, so it takes effect within one 0.5 s
+ tick rather than two minutes.
+
+ Entries for other rids are stale transfers whose consumer never called
+ /pd/decode (or was rejected before it could) — drop and release them so
+ they stop holding the single receive slot.
+ """
+ deadline = time.time() + timeout_s
+ while True:
+ if cancel is not None and cancel.is_set():
+ return _CANCELLED
+ remaining = deadline - time.time()
+ if remaining <= 0:
+ return None
+ try:
+ cand = server.completed.get(timeout=min(remaining, 0.5))
+ except queue_mod.Empty:
+ continue
+ if cand.rid == rid:
+ return cand
+ logger.warning("dropping unmatched request %s (waiting for %s)", cand.rid, rid)
+ # The dropped entry's rid, not ours and not "whatever is current":
+ # this transfer arrived after its own consumer gave up, and the
+ # tenancy may already belong to a later request.
+ server.release(cand.rid)
+
+ def _abandon_pending_kv(rid: str, cancel=None) -> None:
+ """Release the receive slot for a request we are rejecting post-prefill.
+
+ The prefill request always runs before /pd/decode, so by the time we reject a
+ request (bad grammar spec, missing backend, unexpected prep failure) vLLM
+ has already pushed its KV, or is pushing it. Dropping only the lock
+ leaves that transfer owning the receive server's single slot: the NEXT
+ request's ranks are turned away with "busy", its KV never lands, and its
+ /pd/decode blocks until the 120 s kv_transfer_timeout — one bad grammar
+ stalls the following request for two minutes.
+
+ So drain our own entry before releasing. The transfer is already under
+ way, so this normally costs milliseconds; the bound stops a prefill that
+ died mid-push from holding the slot indefinitely.
+
+ ``cancel`` is the SAME event /pd/cancel sets, and this drain has to watch
+ it like the phase-1 one does. Without it a cancel arriving here is
+ answered 200 -- the event is armed and this rid is still current -- while
+ the drain runs on to _ABANDON_DRAIN_S, so the client is told the request
+ was cancelled and the next one is refused 429 for another 30 s. The
+ bound alone is not enough: it is sized for a transfer already in flight,
+ not for one whose prefill leg is never coming.
+ """
+ try:
+ if _drain_own_kv(rid, _ABANDON_DRAIN_S, cancel) is None:
+ logger.warning(
+ "abandoned %s: its KV did not arrive within " "%.0fs", rid, _ABANDON_DRAIN_S
+ )
+ except Exception:
+ logger.exception("draining KV for abandoned %s failed", rid)
+ finally:
+ _cleanup()
+
def _log_reqstat(body, req, n_tokens, timing):
logger.info(
"REQSTAT rid=%s seq=%d completion=%d %s",
@@ -101,25 +242,82 @@ def pd_decode(body: DecodeBody):
{"error": "busy", "current_rid": state["current_rid"]}, status_code=429
)
state["current_rid"] = body.rid
+ # We are the consumer this rid was waiting for, so drop any tombstone a
+ # previous attempt at the same request left behind -- vLLM reuses the
+ # request id when it reschedules, and its senders would otherwise be
+ # refused until the tombstone aged out.
+ server.expect(body.rid)
+ # Armed HERE, not once decoding starts. /pd/cancel needs something to
+ # set from the moment the request is admitted: the wire-wait below can
+ # be the longest phase of all, and a cancel arriving during it used to
+ # find cancel_event still None and answer 404 while the slot stayed
+ # held.
+ cancel = threading.Event()
+ state["cancel_event"] = cancel
t0 = time.time()
+
+ # phase 0: compile the grammar BEFORE convert / inject, so a bad spec (or
+ # a missing backend) fails without touching the GPU. Fail-closed
+ # classification survives the HTTP hop via error_type (the router
+ # propagates the status verbatim). Both error paths must still hand back
+ # the receive slot -- see _abandon_pending_kv.
+ # Same fail-fast point as the grammar spec: refuse before the wire-wait
+ # and any GPU work if this engine cannot produce what was asked for.
+ if (
+ body.top_logprobs is not None
+ and not getattr(engine, "supports_logprobs", lambda: False)()
+ ):
+ logger.info(
+ "logprobs requested but unsupported by %s (rid=%s)", type(engine).__name__, body.rid
+ )
+ _abandon_pending_kv(body.rid, cancel)
+ return JSONResponse(
+ {
+ "error": f"{type(engine).__name__} does not produce logprobs",
+ "error_type": "logprobs_unavailable",
+ },
+ status_code=501,
+ )
+ # Same fail-fast point for the sampling params this engine cannot apply.
+ # The router normally refuses these before prefilling (it reads
+ # /capabilities), but /pd/decode is directly reachable and the streaming
+ # branch cannot report a status once its headers are out — so the check
+ # belongs here, ahead of the wire-wait, rather than inside decode().
+ try:
+ validate_generation_request(body.sampling or {}, engine_capabilities(engine))
+ except CapabilityError as e:
+ logger.info("sampling rejected for %s: %s (%s)", body.rid, e, e.error_type)
+ _abandon_pending_kv(body.rid, cancel)
+ return JSONResponse(e.to_payload(), status_code=e.http_status)
+ try:
+ grammar_session = engine.prepare_grammar(body.grammar_spec, body.enable_thinking)
+ except GrammarError as e:
+ logger.info("grammar rejected for %s: %s (%s)", body.rid, e, e.error_type)
+ _abandon_pending_kv(body.rid, cancel)
+ return JSONResponse(e.to_payload(), status_code=e.http_status)
+ except Exception as e:
+ logger.exception("grammar prepare failed for %s", body.rid)
+ _abandon_pending_kv(body.rid, cancel)
+ return JSONResponse({"error": str(e)}, status_code=500)
+
# phase 1: wire wait + convert + inject (common to both modes)
try:
- # Drain until OUR rid arrives; drop stale completed entries
- # (e.g. a transfer whose consumer never called /pd/decode).
- req = None
- deadline = time.time() + body.timeout_s
- while time.time() < deadline:
- try:
- cand = server.completed.get(timeout=max(0.1, deadline - time.time()))
- except queue_mod.Empty:
- break
- if cand.rid == body.rid:
- req = cand
- break
- logger.warning(
- "dropping unmatched request %s " "(waiting for %s)", cand.rid, body.rid
+ req = _drain_own_kv(body.rid, body.timeout_s, cancel)
+ if req is _CANCELLED:
+ logger.info("cancelled during KV transfer for %s", body.rid)
+ _cleanup()
+ # 499, nginx's "client closed request": the caller asked us to
+ # stop, so this is neither our failure (5xx) nor a bad request
+ # (4xx). The router has stopped reading by now; the status is
+ # for a direct caller and for the log.
+ return JSONResponse(
+ {
+ "error": "cancelled during KV transfer",
+ "error_type": "request_cancelled",
+ "rid": body.rid,
+ },
+ status_code=499,
)
- server.release()
if req is None:
_cleanup()
return JSONResponse(
@@ -143,9 +341,38 @@ def pd_decode(body: DecodeBody):
"inject": round(1000 * (t_inj - t_conv), 1),
}
- # phase 2: decode
- cancel = threading.Event()
- state["cancel_event"] = cancel
+ # phase 2: decode (cancel was armed at admission)
+
+ # Logprobs sink. The callback appends BEFORE the token goes on the
+ # queue, so anything the generator dequeues already has its entry --
+ # indexing by emitted count needs no second lock.
+ want_lp = body.top_logprobs is not None
+ lp_sink: list[tuple[float | None, list]] = []
+
+ def _emit(tok, logprob=None, candidates=None):
+ if want_lp:
+ if logprob is None:
+ if lp_sink:
+ # Past position 0 a bare token is an engine fault:
+ # recording None would reach the client as the
+ # "very unlikely" sentinel, indistinguishable from a
+ # measurement.
+ raise LogprobsUnavailable(f"engine emitted token {tok} without a logprob")
+ # Position 0 is first_token_id, sampled by prefill, so no
+ # decode-side value exists. Hold the slot to keep one entry
+ # per token and send null; the router fills it.
+ lp_sink.append((None, []))
+ return tok
+ lp_sink.append((float(logprob), list(candidates or ())))
+ return tok
+
+ def _lp_slice(start: int, count: int) -> dict:
+ """The `lp`/`tp` fields for tokens [start, start+count)."""
+ rows = lp_sink[start : start + count]
+ return {
+ "lp": [r[0] for r in rows],
+ "tp": [[list(c) for c in r[1]] for r in rows],
+ }
if not body.stream:
try:
@@ -153,7 +380,10 @@ def pd_decode(body: DecodeBody):
first_token_id=body.first_token_id,
max_tokens=body.max_tokens,
sampling=body.sampling,
+ on_token=_emit if want_lp else None,
cancel_event=cancel,
+ grammar_session=grammar_session,
+ **({"top_logprobs": body.top_logprobs} if want_lp else {}),
)
timing = {
**pre_timing,
@@ -161,12 +391,28 @@ def pd_decode(body: DecodeBody):
**getattr(engine, "last_stats", {}),
}
_log_reqstat(body, req, len(tokens), timing)
- return {
+ out = {
"rid": body.rid,
"token_ids": tokens,
"seq_len": req.seq_len,
"timing_ms": timing,
}
+ if want_lp:
+ if len(lp_sink) != len(tokens):
+ raise LogprobsUnavailable(
+ f"engine returned {len(lp_sink)} logprob entries "
+ f"for {len(tokens)} tokens"
+ )
+ out["logprobs"] = _lp_slice(0, len(tokens))
+ return out
+ except LogprobsUnavailable as e:
+ logger.info("logprobs unavailable for %s: %s", body.rid, e)
+ return JSONResponse(
+ {"error": str(e), "error_type": "logprobs_unavailable"}, status_code=501
+ )
+ except GrammarViolationError as e:
+ logger.info("grammar violation for %s: %s", body.rid, e)
+ return JSONResponse(e.to_payload(), status_code=e.http_status)
except Exception as e:
logger.exception("decode failed for %s", body.rid)
return JSONResponse({"error": str(e), "rid": body.rid}, status_code=500)
@@ -176,6 +422,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:
@@ -183,13 +435,33 @@ def _run():
first_token_id=body.first_token_id,
max_tokens=body.max_tokens,
sampling=body.sampling,
- on_token=q.put,
+ on_token=(lambda *a: q.put(_emit(*a))) if want_lp else q.put,
cancel_event=cancel,
+ grammar_session=grammar_session,
+ **({"top_logprobs": body.top_logprobs} if want_lp else {}),
)
q.put(("done", tokens))
+ _signal_done()
+ except GrammarViolationError as e:
+ # 200 headers may already be sent; signal a typed error so the
+ # router can emit an SSE error event + [DONE] (fail-closed).
+ logger.info("stream grammar violation for %s: %s", body.rid, e)
+ q.put(("error", e.to_payload()))
+ _signal_done()
+ except LogprobsUnavailable as e:
+ # Typed, like the blocking branch: the engine cannot produce what
+ # was asked for, which is a capability answer (501), not a broken
+ # component (502). Falling into the generic handler below dropped
+ # the type, so adding a stop string -- which is what puts a
+ # non-streaming request on this protocol -- silently changed the
+ # status for the same inability.
+ logger.info("stream logprobs unavailable for %s: %s", body.rid, e)
+ q.put(("error", {"error": str(e), "error_type": "logprobs_unavailable"}))
+ _signal_done()
except Exception as e: # pragma: no cover
logger.exception("stream decode failed for %s", body.rid)
- q.put(("error", str(e)))
+ q.put(("error", {"error": str(e)}))
+ _signal_done()
worker = threading.Thread(target=_run, name="pd-decode", daemon=True)
@@ -204,11 +476,35 @@ 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] = []
+ n_emitted = 0
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):
+ line = {"t": [first]}
+ if want_lp:
+ line.update(_lp_slice(0, 1))
+ n_emitted = 1
+ yield json.dumps(line) + "\n"
+ else:
+ done_msg = first
+ last_activity = time.time()
+ break
+
while done_msg is None:
drained = False
while True:
@@ -223,7 +519,13 @@ async def _gen():
done_msg = item
break
if batch:
- yield json.dumps({"t": batch}) + "\n"
+ line = {"t": batch}
+ if want_lp:
+ # lp[i] / tp[i] line up with t[i]; written before
+ # the queue put, so visible for everything dequeued.
+ line.update(_lp_slice(n_emitted, len(batch)))
+ n_emitted += len(batch)
+ yield json.dumps(line) + "\n"
batch = []
if done_msg is None:
if drained:
@@ -232,7 +534,8 @@ async def _gen():
yield json.dumps({"error": "decode stalled"}) + "\n"
return
else:
- await asyncio.sleep(0.005)
+ with contextlib.suppress(asyncio.TimeoutError, TimeoutError):
+ await asyncio.wait_for(fin["ev"].wait(), timeout=DECODE_POLL_S)
kind, payload = done_msg
if kind == "done":
timing = {
@@ -251,7 +554,9 @@ async def _gen():
}
) + "\n"
else:
- yield json.dumps({"error": payload}) + "\n"
+ # payload is a typed dict {"error", ["error_type"]}; emit
+ # verbatim so the router can classify (grammar_violation).
+ yield json.dumps(payload) + "\n"
finally:
cancel.set()
# shield: cleanup must complete even inside a cancelled scope,
@@ -268,11 +573,15 @@ async def _gen():
return app # noqa: R504 (assembled across the function)
-def main() -> None:
- logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
+def build_parser() -> argparse.ArgumentParser:
ap = argparse.ArgumentParser()
ap.add_argument("--engine", choices=["stub", "tilert"], default="stub")
- ap.add_argument("--model", default="glm5", help="model profile")
+ ap.add_argument(
+ "--model",
+ default="glm5",
+ help="model profile (glm5 / glm5_2 / glm5_3 / dsv32); "
+ "must match the prefill side's tilert_model",
+ )
ap.add_argument("--max-seq-len", type=int, default=4096)
ap.add_argument("--ctrl-port", type=int, default=5556)
ap.add_argument("--http-port", type=int, default=5557)
@@ -289,13 +598,18 @@ def main() -> None:
default="fp8_ds_mla",
help="MLA cache dtype (must match vLLM prefill); " "MLA-family profiles only",
)
- args = ap.parse_args()
+ return ap
+
+
+def main() -> None:
+ logging.basicConfig(level=logging.INFO, format="%(asctime)s %(name)s %(message)s")
+ args = build_parser().parse_args()
from tilert.pd_vllm.profiles import base as profiles
profile = profiles.get_profile(args.model)
- # MLA-family profiles (glm5/dsv32) need the cache dtype to size the receive
- # buffer.
+ # MLA-family profiles (glm5/glm5_2/dsv32) need the cache dtype to size the
+ # receive buffer; profiles without the knob skip it.
if hasattr(profile, "configure"):
profile.configure(args.kv_cache_dtype)
logger.info(
@@ -304,6 +618,10 @@ def main() -> None:
args.kv_cache_dtype,
profile.layout_version,
)
+ # Depth-from-checkpoint members must resolve their layer count before the
+ # receive buffer is sized off it.
+ if hasattr(profile, "configure_weights") and args.model_weights_dir:
+ profile.configure_weights(args.model_weights_dir)
if args.engine == "stub":
from tilert.pd_vllm.engine_iface import StubEngine
diff --git a/tilert/pd_vllm/engine_iface.py b/tilert/pd_vllm/engine_iface.py
index 6ed9133..765170b 100644
--- a/tilert/pd_vllm/engine_iface.py
+++ b/tilert/pd_vllm/engine_iface.py
@@ -8,24 +8,81 @@
from collections.abc import Callable
from typing import Any, Protocol
+from tilert.pd_vllm.grammar_spec import (
+ GrammarBackendUnavailable,
+ GrammarViolationError,
+ InvalidGrammarError,
+)
+
class PDEngine(Protocol):
def inject(self, req: Any) -> None:
"""Restore engine state to 'prefilled seq_len tokens' from req."""
+ def prepare_grammar(self, grammar_spec: dict | None, enable_thinking: bool = True) -> Any:
+ """Compile ``grammar_spec`` into a per-request grammar session.
+
+ Runs BEFORE any KV wire-wait / inject / GPU work, so a bad spec fails
+ fast.
+
+ Returns an opaque session object (passed back to :meth:`decode`), or
+ None when ``grammar_spec`` is None (unconstrained). Raises
+ ``InvalidGrammarError`` (client 400) for a malformed/unsupported spec
+ and ``GrammarBackendUnavailable`` (server 500) when xgrammar is
+ absent — never silently degrades to unconstrained decoding.
+ """
+
def decode(
self,
first_token_id: int,
max_tokens: int,
sampling: dict | None,
- on_token: Callable[[int], None] | None = None,
+ on_token: Callable[..., None] | None = None,
cancel_event=None,
+ grammar_session: Any = None,
+ top_logprobs: int | None = None,
) -> list[int]:
"""AR/MTP decode from first_token_id; returns completion ids.
Includes first_token_id, excludes the stop token. on_token never fires
for stop tokens; cancel_event stops early; last_stats['finish_reason']
- is 'stop' | 'length' | 'cancelled'.
+ is 'stop' | 'length' | 'cancelled'. When ``grammar_session`` is set,
+ every emitted token is grammar-masked/validated; a rejected token
+ raises ``GrammarViolationError``.
+
+ ``top_logprobs`` is the number of candidates the caller wants per
+ position, or None for no logprobs. An engine that supports it calls
+ ``on_token(token_id, logprob, candidates)`` -- ``candidates`` being a
+ list of ``(token_id, logprob)`` longest-first -- instead of
+ ``on_token(token_id)``. The extra arguments are optional at the call
+ site, so an engine that ignores ``top_logprobs`` keeps working; the
+ decode server detects the absence and reports it rather than returning a
+ response with the field silently missing. Support is declared by
+ :meth:`supports_logprobs`.
+ """
+
+ def supports_logprobs(self) -> bool:
+ """Whether :meth:`decode` honours ``top_logprobs``.
+
+ Optional: an engine that does not define it is treated as unsupported.
+ """
+
+ def supports_penalties(self) -> bool:
+ """Whether :meth:`decode` honours ``repetition_penalty`` / ``presence_penalty``.
+
+ Optional, and read by ``decode_server``'s ``/capabilities`` so the
+ router can refuse a penalty request BEFORE prefilling it -- the
+ alternative is a 501 after the KV has already crossed the wire. An
+ engine that does not define it is treated as unsupported, which is the
+ safe direction: the request is refused rather than decoded unpenalised.
+ """
+
+ def supports_ignore_eos(self) -> bool:
+ """Whether :meth:`decode` honours ``ignore_eos``.
+
+ Optional, same contract as :meth:`supports_penalties`. An engine that
+ claims this must clear its stop set for the request, not merely accept
+ the key -- accepting and ignoring it is the failure this exists to stop.
"""
def reset(self) -> None:
@@ -33,7 +90,17 @@ def reset(self) -> None:
class StubEngine:
- """Echo engine for plumbing tests: no GPU, no tilert."""
+ """Echo engine for plumbing tests: no GPU, no tilert.
+
+ ``prepare_grammar`` simulates the real classification deterministically via
+ spec sentinels so the fail-closed HTTP mapping can be tested without a GPU:
+ - ``{"type": "__backend_missing__"}`` -> GrammarBackendUnavailable (500)
+ - an unknown/malformed spec -> InvalidGrammarError (400)
+ - ``{"type": "regex", "value": "__violate__"}`` -> decode raises
+ GrammarViolationError (400) on the first token
+ """
+
+ _KNOWN_SPEC_TYPES = ("json_schema", "json_object", "ebnf", "regex", "structural_tag")
def __init__(self, fixed_tokens: tuple[int, ...] = (11, 22, 33)):
self._fixed = fixed_tokens
@@ -43,11 +110,57 @@ def __init__(self, fixed_tokens: tuple[int, ...] = (11, 22, 33)):
def inject(self, req: Any) -> None:
self.injected = req
- def decode(self, first_token_id, max_tokens, sampling, on_token=None, cancel_event=None):
+ def prepare_grammar(self, grammar_spec, enable_thinking=True):
+ if grammar_spec is None:
+ return None
+ if not isinstance(grammar_spec, dict) or "type" not in grammar_spec:
+ raise InvalidGrammarError("grammar spec must be a dict with a 'type'")
+ kind = grammar_spec["type"]
+ if kind == "__backend_missing__":
+ raise GrammarBackendUnavailable("xgrammar backend not installed")
+ if kind not in self._KNOWN_SPEC_TYPES:
+ raise InvalidGrammarError(f"unsupported grammar spec type: {kind!r}")
+ return {"spec": grammar_spec, "enable_thinking": enable_thinking}
+
+ def supports_logprobs(self) -> bool:
+ return True
+
+ def supports_penalties(self) -> bool:
+ # The echo engine applies no sampling at all, but it must not be the
+ # reason a plumbing test cannot reach the penalty path.
+ return True
+
+ def supports_ignore_eos(self) -> bool:
+ return True
+
+ @staticmethod
+ def fake_logprob(token_id: int) -> float:
+ """Deterministic stand-in so tests can assert exact values."""
+ return -0.5 - 0.25 * (token_id % 4)
+
+ def decode(
+ self,
+ first_token_id,
+ max_tokens,
+ sampling,
+ on_token=None,
+ cancel_event=None,
+ grammar_session=None,
+ top_logprobs=None,
+ ):
+ spec = (grammar_session or {}).get("spec", {})
+ if spec.get("value") == "__violate__":
+ raise GrammarViolationError(f"first token {first_token_id} violates the grammar")
out = ([int(first_token_id)] + list(self._fixed))[:max_tokens]
if on_token:
for t in out:
- on_token(t)
+ if top_logprobs is None:
+ on_token(t)
+ else:
+ # Candidates are the token itself plus neighbours, so the
+ # ordering (chosen first, then descending) is checkable.
+ cands = [(t + k, self.fake_logprob(t) - 0.5 * k) for k in range(top_logprobs)]
+ on_token(t, self.fake_logprob(t), cands)
self.last_stats = {"finish_reason": "stop"}
return out
diff --git a/tilert/pd_vllm/generation_defaults.py b/tilert/pd_vllm/generation_defaults.py
new file mode 100644
index 0000000..e56100d
--- /dev/null
+++ b/tilert/pd_vllm/generation_defaults.py
@@ -0,0 +1,333 @@
+"""Where a request's sampling defaults come from, resolved once for both PD legs.
+
+A PD request is sampled in two places -- the vLLM prefill instance takes token 1,
+the decode node takes tokens 2..N -- and each used to work out its own defaults
+for a field the client left out. vLLM resolves
+
+ client explicit value > the model's generation_config.json > 1.0 / 1.0 / 0
+
+(``ModelConfig.get_diff_sampling_param``, resolved once at startup and applied per
+request in ``to_sampling_params``), while the decode adapters carried literals of
+their own. A checkpoint shipping ``temperature: 0.6`` / ``top_k: 20`` then had
+token 1 sampled at 0.6 / 20 and tokens 2..N at 1.0 / uncapped.
+
+This module is the single resolution point. It mirrors vLLM's chain, and the
+router writes the result into BOTH requests explicitly -- which is what makes the
+agreement hold: an explicit value overrides vLLM's own resolution, so the prefill
+instance's ``--generation-config`` flag can no longer move one leg without the
+other.
+
+Adopted fields
+--------------
+``temperature``, ``top_p``, ``top_k`` only. Every profile's decode path applies
+all three on every request, so a value taken from the model's config is
+guaranteed to reach both legs.
+
+``repetition_penalty`` is adopted only where the served model's decode runtime
+implements it, which the profile states statically (``declares_penalties``; the
+MLA/NSA members -- GLM-5, GLM-5.2, DSV3.2 -- say no). Told which model it serves
+(``--model``), the router adopts the config's value on a profile that declares
+penalties and refuses to start on the others -- because there the prefill
+instance would apply it to the first token and the decode node could not apply
+it to the rest. Without ``--model`` the family is unknown and the conservative
+answer is taken.
+
+``min_p`` is never adopted: no decode runtime implements it on any member.
+
+``max_new_tokens`` is not adopted either, and needs no guard: the prefill request
+is pinned to ``max_tokens=1`` and the decode length comes from the client or the
+decode node's own default, so the two legs cannot disagree about it.
+
+Difference from vLLM worth knowing: this reads ``generation_config.json``
+directly, while vLLM loads it through HF's ``GenerationConfig.to_diff_dict()``,
+which drops any value equal to HF's own default (``top_k=50`` among them). So a
+config stating a field at HF's default is honoured here and ignored there. It
+cannot split the two legs -- both receive whatever this resolves -- but it can
+make a PD deployment sample differently from a stock vLLM one.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from dataclasses import dataclass
+from typing import Any
+
+from tilert.pd_vllm.sampling import VLLM_DEFAULT_TOP_P
+
+__all__ = [
+ "GUARDED_FIELDS",
+ "NEUTRAL_MIN_P",
+ "NEUTRAL_REPETITION_PENALTY",
+ "GenerationDefaults",
+ "UnsupportedGenerationDefault",
+ "load",
+ "penalties_supported_by",
+]
+
+logger = logging.getLogger("pd_vllm.generation_defaults")
+
+# vLLM's neutral defaults, the last link of its chain
+# (``ChatCompletionRequest._DEFAULT_SAMPLING_PARAMS``). ``top_k`` is 0, vLLM's
+# documented "no rank cut" sentinel -- NOT the engine-side ``TOP_K_DISABLED``,
+# which is the kernel's candidate-pool bound. What travels on the wire is the
+# request-domain value; ``sampling.resolve_top_k`` maps it for the engine.
+VLLM_NEUTRAL_TEMPERATURE = 1.0
+VLLM_NEUTRAL_TOP_K = 0
+
+# Fields whose adoption depends on the model, with the value that means "asks for
+# nothing". A config carrying one at a non-neutral value that this deployment
+# cannot execute is a STARTUP error: the prefill instance would apply it to token
+# 1 while the decode node could not apply it to the rest, and no per-request check
+# can see it -- the client never sent it, so the capability gate never looks at
+# it. Better to refuse to start than to serve every request half-penalised.
+GUARDED_FIELDS = {"repetition_penalty": 1.0, "min_p": 0.0}
+
+# The value each guarded field takes when it is not adopted. Sent explicitly to
+# the prefill leg regardless, so vLLM cannot resolve one of its own from the
+# model config behind the router's back.
+NEUTRAL_REPETITION_PENALTY = 1.0
+NEUTRAL_MIN_P = 0.0
+
+
+class UnsupportedGenerationDefault(Exception):
+ """The model's generation_config asks for something PD cannot guarantee."""
+
+
+@dataclass(frozen=True)
+class GenerationDefaults:
+ """The sampling defaults this deployment applies to both legs."""
+
+ temperature: float = VLLM_NEUTRAL_TEMPERATURE
+ top_p: float = VLLM_DEFAULT_TOP_P
+ top_k: int = VLLM_NEUTRAL_TOP_K
+ # Adopted only on a family whose decode runtime implements it; otherwise the
+ # runtime no-op, which is what the prefill leg is then pinned to as well.
+ repetition_penalty: float = NEUTRAL_REPETITION_PENALTY
+ # Where the values came from, for the startup log line.
+ source: str = "vllm"
+
+ def resolve(self, body: dict) -> dict:
+ """The three fields for this request: client value, else the default.
+
+ ``None`` counts as absent, which is how SDKs spell "unset" and how vLLM
+ treats it. Values are returned in the REQUEST domain, so they can be
+ written into the prefill body and the decode payload unchanged.
+ """
+ return {
+ "temperature": _as_float("temperature", body.get("temperature"), self.temperature),
+ "top_p": _as_float("top_p", body.get("top_p"), self.top_p),
+ "top_k": _as_int("top_k", body.get("top_k"), self.top_k),
+ # Pinned on both legs even at its no-op value: left unset, the
+ # prefill instance would take one from generation_config while the
+ # decode node took the router's, which is the split this module
+ # exists to close.
+ "repetition_penalty": _as_float(
+ "repetition_penalty", body.get("repetition_penalty"), self.repetition_penalty
+ ),
+ # Always the no-op, and always sent. The router's own
+ # --generation-config only governs what the ROUTER reads; the vLLM
+ # server is launched separately and still defaults to loading the
+ # checkpoint's generation_config.json. So a checkpoint carrying
+ # min_p would have it applied to token 1 and ignored for the rest --
+ # including on the very path documented as the way out of the startup
+ # refusal. Pinning it closes vLLM's six-key allowlist: temperature,
+ # top_p, top_k and repetition_penalty are resolved above,
+ # max_new_tokens is overridden by max_tokens=1, and this is the last.
+ "min_p": NEUTRAL_MIN_P,
+ }
+
+ def describe(self) -> str:
+ return (
+ f"temperature={self.temperature}, top_p={self.top_p}, "
+ f"top_k={self.top_k}, "
+ f"repetition_penalty={self.repetition_penalty} "
+ f"(from {self.source})"
+ )
+
+
+def _as_float(field: str, raw, default: float) -> float:
+ if raw is None:
+ return float(default)
+ if isinstance(raw, bool):
+ raise ValueError(f"{field} must be a number, got bool")
+ return float(raw)
+
+
+def _as_int(field: str, raw, default: int) -> int:
+ if raw is None:
+ return int(default)
+ if isinstance(raw, bool):
+ raise ValueError(f"{field} must be an integer, got bool")
+ return int(raw)
+
+
+def _read_config(model_path: str) -> dict:
+ path = os.path.join(model_path, "generation_config.json")
+ if not os.path.isfile(path):
+ logger.info(
+ "no generation_config.json under %s; using vLLM's neutral " "sampling defaults",
+ model_path,
+ )
+ return {}
+ with open(path, encoding="utf-8") as fh:
+ config = json.load(fh)
+ if not isinstance(config, dict):
+ raise UnsupportedGenerationDefault(f"{path} does not contain a JSON object")
+ return config
+
+
+def penalties_supported_by(model: str) -> bool:
+ """Whether the named model's FAMILY implements repetition/presence penalties.
+
+ Read from the profile's static ``declares_penalties``, which is the only
+ answer available at startup: the per-node answer comes from
+ ``/capabilities`` and needs a running decode node. An unknown or unnamed
+ model is treated as unsupported -- the conservative direction, since the
+ cost is refusing to adopt a default rather than serving half-penalised.
+ """
+ if not model:
+ return False
+ try:
+ from tilert.pd_vllm.profiles import base as profiles
+
+ profile = profiles.get_profile(model)
+ except Exception as e: # unknown name, or a profile that cannot import here
+ logger.warning(
+ "cannot resolve model %r to a profile (%s); treating "
+ "penalties as unsupported for default resolution",
+ model,
+ e,
+ )
+ return False
+ return bool(getattr(profile, "declares_penalties", False))
+
+
+def _check_guarded(config: dict, source: str, *, model: str, penalties_ok: bool) -> None:
+ """Refuse a config asking for a field this deployment cannot execute.
+
+ ``min_p`` is refused on every member. ``repetition_penalty`` is refused only
+ where the family lacks the pre-pass -- a profile declaring penalties adopts it.
+ """
+ for field, neutral in sorted(GUARDED_FIELDS.items()):
+ raw = config.get(field)
+ if raw is None or float(raw) == neutral:
+ continue
+ if field == "repetition_penalty" and penalties_ok:
+ continue
+ if field == "min_p":
+ reason = "no decode runtime implements it on any model"
+ elif model:
+ reason = (
+ f"the decode runtime for {model!r} has no penalty "
+ f"pre-pass (the GLM-5 / GLM-5.2 / DSV3.2 members do not "
+ f"declare one)"
+ )
+ else:
+ reason = (
+ "the router was not told which model it serves, so it "
+ "cannot confirm the decode runtime applies it -- pass "
+ "--model"
+ )
+ raise UnsupportedGenerationDefault(
+ f"{source} sets {field}={raw}, but {reason}. The vLLM prefill "
+ f"instance would apply it to the first token while the decode node "
+ f"would not apply it to the rest, and no per-request check can "
+ f"catch it because the client never sent it.\n"
+ f"Resolve it explicitly, whichever is true:\n"
+ f" - the value is not wanted: launch with --generation-config "
+ f"vllm, or remove {field} from generation_config.json;\n"
+ f" - the value is wanted: have clients send {field} per request, "
+ f"so the capability gate accepts it on a node that supports it and "
+ f"refuses it on one that does not."
+ )
+
+
+def load(
+ model_path: str = "",
+ source: str = "auto",
+ *,
+ model: str = "",
+ temperature: float | None = None,
+ top_p: float | None = None,
+ top_k: int | None = None,
+ repetition_penalty: float | None = None,
+) -> GenerationDefaults:
+ """Resolve this deployment's sampling defaults, once, at startup.
+
+ ``source`` mirrors vLLM's ``--generation-config``: ``"auto"`` reads
+ ``generation_config.json`` under ``model_path``, ``"vllm"`` ignores it and
+ uses the neutral defaults. The keyword overrides are the equivalent of
+ ``--override-generation-config`` and win over both.
+
+ ``model`` is the profile name the decode nodes serve (``--model``), used only
+ to decide whether a ``repetition_penalty`` in the config can be adopted.
+
+ Raises:
+ UnsupportedGenerationDefault: the config asks for a guarded field
+ (see :data:`GUARDED_FIELDS`) this deployment cannot apply on both
+ legs.
+ """
+ penalties_ok = penalties_supported_by(model)
+ config: dict = {}
+ if source == "auto" and model_path:
+ config = _read_config(model_path)
+ origin = os.path.join(model_path, "generation_config.json")
+ elif source == "auto":
+ logger.info("no --model-path given; using vLLM's neutral sampling " "defaults")
+ origin = "vllm neutral defaults"
+ else:
+ origin = "vllm neutral defaults"
+
+ if config:
+ _check_guarded(config, origin, model=model, penalties_ok=penalties_ok)
+
+ overrides = {
+ "temperature": temperature,
+ "top_p": top_p,
+ "top_k": top_k,
+ "repetition_penalty": repetition_penalty,
+ }
+ # The overrides win over the file, so they need the same guard -- otherwise
+ # --default-repetition-penalty on a family without the pre-pass starts the
+ # router with an adopted default the gate then refuses on EVERY request,
+ # which is a worse outcome than refusing to start.
+ _check_guarded(
+ {k: v for k, v in overrides.items() if v is not None},
+ "command-line overrides",
+ model=model,
+ penalties_ok=penalties_ok,
+ )
+ if not penalties_ok:
+ # Not executable here, so it is pinned to the no-op on BOTH legs rather
+ # than left for vLLM to resolve from the model config on one of them.
+ config = {k: v for k, v in config.items() if k != "repetition_penalty"}
+ resolved: dict[str, Any] = {}
+ for field, neutral in (
+ ("temperature", VLLM_NEUTRAL_TEMPERATURE),
+ ("top_p", VLLM_DEFAULT_TOP_P),
+ ("top_k", VLLM_NEUTRAL_TOP_K),
+ ("repetition_penalty", NEUTRAL_REPETITION_PENALTY),
+ ):
+ if overrides[field] is not None:
+ resolved[field] = overrides[field]
+ elif config.get(field) is not None:
+ resolved[field] = config[field]
+ else:
+ resolved[field] = neutral
+
+ if any(v is not None for v in overrides.values()):
+ origin = f"{origin} + command-line overrides"
+
+ defaults = GenerationDefaults(
+ temperature=float(resolved["temperature"]),
+ top_p=float(resolved["top_p"]),
+ top_k=int(resolved["top_k"]),
+ repetition_penalty=float(resolved["repetition_penalty"]),
+ source=origin,
+ )
+ # vLLM logs when the model's config displaces its neutral defaults, and an
+ # operator comparing the two stacks needs the same line from this side.
+ logger.info("sampling defaults for requests that omit a field: %s", defaults.describe())
+ return defaults
diff --git a/tilert/pd_vllm/grammar_backend.py b/tilert/pd_vllm/grammar_backend.py
new file mode 100644
index 0000000..81313b3
--- /dev/null
+++ b/tilert/pd_vllm/grammar_backend.py
@@ -0,0 +1,40 @@
+"""Where the engine keeps its xgrammar host wrapper.
+
+Engine builds have shipped the wrapper at two locations: ``tilert.grammar``
+(current) and, in older wheels, under a model package
+(``tilert.models.glm_5_2.grammar``). The model-package location decided which
+products could constrain anything -- a wheel that excluded that model package
+excluded the mask producer with it, so a ``response_format`` request failed
+naming a model the client was not serving.
+
+One serve version has to work against engine wheels from either side of that
+move, so the new path is tried first and the old one is the fallback.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def load_grammar_backend() -> tuple[Any, Any]:
+ """Return ``(GrammarEngine, GrammarSession)`` from the installed engine.
+
+ Raises ``ModuleNotFoundError`` when the engine ships neither path -- the
+ message names no model, because which model the caller is serving has
+ nothing to do with why the backend is missing.
+ """
+ try:
+ from tilert.grammar import GrammarEngine, GrammarSession
+
+ return GrammarEngine, GrammarSession
+ except ImportError:
+ pass
+ try:
+ # Engine wheels that predate ``tilert.grammar``.
+ from tilert.models.glm_5_2.grammar import GrammarEngine, GrammarSession
+
+ return GrammarEngine, GrammarSession
+ except ImportError as e:
+ raise ModuleNotFoundError(
+ "this engine build ships no xgrammar host backend (looked for tilert.grammar)"
+ ) from e
diff --git a/tilert/pd_vllm/grammar_spec.py b/tilert/pd_vllm/grammar_spec.py
new file mode 100644
index 0000000..5d91d90
--- /dev/null
+++ b/tilert/pd_vllm/grammar_spec.py
@@ -0,0 +1,191 @@
+"""Grammar/constrained-decoding serve-layer helpers for the pd_vllm path.
+
+Two concerns, both framework-agnostic (no vLLM / no tilert imports), so this
+module is unit-testable on a plain CPU box:
+
+1. ``extract_request_grammar_spec`` — translate an OpenAI-style request dict
+ into the engine's ``grammar_spec`` (``{"type", "value"}``). Selection
+ priority mirrors the reference sglang implementation's grammar_manager:
+ ``json_schema > regex > ebnf > structural_tag``.
+
+2. The serve-layer error hierarchy used for fail-closed classification:
+ - client sent a bad/unsupported spec -> ``InvalidGrammarError`` (HTTP 400)
+ - decode node lacks the xgrammar backend -> ``GrammarBackendUnavailable``
+ - an emitted token violated the grammar -> ``GrammarViolationError`` (400)
+
+ Each carries ``error_type`` + ``http_status`` and a ``to_payload()`` that
+ the decode server / router return verbatim, so the classification survives
+ the /pd/decode HTTP hop.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+# --------------------------------------------------------------------------- #
+# Error hierarchy (fail-closed classification)
+# --------------------------------------------------------------------------- #
+class GrammarError(Exception):
+ """Base for all grammar serve-layer errors. Never silently degrade."""
+
+ error_type = "grammar_error"
+ http_status = 500
+
+ def to_payload(self) -> dict[str, str]:
+ return {"error": str(self), "error_type": self.error_type}
+
+
+class InvalidGrammarError(GrammarError):
+ """Client sent a malformed/unsupported constraint -> HTTP 400.
+
+ Raised both by request parsing (:func:`extract_request_grammar_spec`) and
+ by grammar compilation (a schema/regex/EBNF xgrammar cannot compile).
+ """
+
+ error_type = "invalid_grammar"
+ http_status = 400
+
+
+class GrammarBackendUnavailable(GrammarError):
+ """xgrammar is not installed on the decode node -> HTTP 500.
+
+ A missing backend is a server-side deployment fault, NOT a client error: a
+ legitimate ``response_format`` request must never be told its grammar is
+ invalid just because the backend is absent.
+ """
+
+ error_type = "grammar_backend_unavailable"
+ http_status = 500
+
+
+class GrammarViolationError(GrammarError):
+ """An emitted token violated the grammar -> HTTP 400 / SSE error event.
+
+ Typically an unconstrained prefill first token that the matcher rejects.
+ We fail closed rather than serve the rest of the request unconstrained.
+ """
+
+ error_type = "grammar_violation"
+ http_status = 400
+
+
+class GrammarUnsupported(GrammarError):
+ """A constrained request hit a code path without grammar masking -> HTTP 500.
+
+ Fail loud instead of silently serving the request unconstrained (staged
+ rollout: e.g. the MTP path before it is wired).
+ """
+
+ error_type = "grammar_unsupported"
+ http_status = 500
+
+
+# --------------------------------------------------------------------------- #
+# Spec validation
+# --------------------------------------------------------------------------- #
+# Compilation runs inside the decode node's single-slot lock, so an unbounded
+# schema is downtime for everyone. Measured compile cost (150k-vocab tokenizer),
+# limits chosen for ~1s worst case:
+# depth 100 -> 0.2s 200 -> 1.7s 400 -> 13.3s
+# properties 1000 -> 1.2s 5000 -> 36.3s 10000 -> 170.9s
+# enum 10000 -> 0.1s 20000 -> 1.1s 50000 -> 6.2s
+# regex/EBNF need no cap: 200 000 chars of regex compiles in 0.7s.
+_MAX_SCHEMA_DEPTH = 64
+_MAX_SCHEMA_NODES = 1000
+_MAX_ENUM_VALUES = 10000
+
+
+def _walk_schema(node, depth, counts, path):
+ """Check one schema node, recursing into subschemas."""
+ where = path or ""
+ if depth > _MAX_SCHEMA_DEPTH:
+ raise InvalidGrammarError(f"json_schema nests deeper than {_MAX_SCHEMA_DEPTH} at {where}")
+
+ if isinstance(node, list):
+ for i, item in enumerate(node):
+ _walk_schema(item, depth, counts, f"{path}[{i}]")
+ return
+ if not isinstance(node, dict):
+ return
+
+ counts["nodes"] += 1
+ if counts["nodes"] > _MAX_SCHEMA_NODES:
+ raise InvalidGrammarError(f"json_schema has over {_MAX_SCHEMA_NODES} subschemas")
+
+ if isinstance(node.get("enum"), list):
+ counts["enums"] += len(node["enum"])
+ if counts["enums"] > _MAX_ENUM_VALUES:
+ raise InvalidGrammarError(f"json_schema has over {_MAX_ENUM_VALUES} enum values")
+
+ for key, child in node.items():
+ # Literal data, not subschemas: a 500-value enum is not 500 nodes.
+ if key in ("enum", "const", "examples", "default"):
+ continue
+ if isinstance(child, (dict, list)):
+ _walk_schema(child, depth + 1, counts, f"{path}.{key}" if path else key)
+
+
+def validate_grammar_spec(spec: dict[str, Any] | None) -> None:
+ """Reject a spec whose compilation would stall the decode node.
+
+ Walks whatever the spec carries: json_schema and structural_tag hold
+ subschemas, regex/ebnf hold a string and json_object nothing, both of which
+ the walk skips. Called from :func:`extract_request_grammar_spec`, i.e. in
+ the router before the prefill request runs, so an oversized schema never reaches
+ a decode node. A direct ``/pd/decode`` call is not covered.
+
+ Raises:
+ InvalidGrammarError: over a compile-cost limit (400).
+ """
+ if spec:
+ _walk_schema(spec.get("value"), 0, {"nodes": 0, "enums": 0}, "")
+
+
+# --------------------------------------------------------------------------- #
+# Request -> grammar_spec translation
+# --------------------------------------------------------------------------- #
+_SPEC_TYPES = ("json_schema", "json_object", "ebnf", "regex", "structural_tag")
+
+
+def extract_request_grammar_spec(
+ request: dict[str, Any],
+) -> dict[str, Any] | None:
+ """Return the engine grammar spec for a request, or None if unconstrained.
+
+ Priority: json_schema > regex > ebnf > structural_tag (``json_object``
+ returns early, before regex/ebnf, matching the reference).
+
+ The returned spec has already passed :func:`validate_grammar_spec`.
+
+ Raises:
+ InvalidGrammarError: malformed, or over a compile-cost limit (400).
+ """
+ response_format = request.get("response_format") or {}
+ if not isinstance(response_format, dict):
+ raise InvalidGrammarError("response_format must be an object")
+ rf_type = response_format.get("type")
+
+ json_schema = None
+ structural_tag = None
+ if rf_type == "json_schema":
+ json_schema = (response_format.get("json_schema") or {}).get("schema")
+ if json_schema is None:
+ raise InvalidGrammarError("response_format json_schema requires json_schema.schema")
+ elif rf_type == "json_object":
+ return {"type": "json_object", "value": None}
+ if rf_type == "structural_tag":
+ structural_tag = response_format
+
+ if json_schema is not None:
+ spec = {"type": "json_schema", "value": json_schema}
+ elif request.get("regex") is not None:
+ spec = {"type": "regex", "value": request["regex"]}
+ elif request.get("ebnf") is not None:
+ spec = {"type": "ebnf", "value": request["ebnf"]}
+ elif structural_tag is not None:
+ spec = {"type": "structural_tag", "value": structural_tag}
+ else:
+ return None
+ validate_grammar_spec(spec)
+ return spec
diff --git a/tilert/pd_vllm/logprobs.py b/tilert/pd_vllm/logprobs.py
new file mode 100644
index 0000000..93759ef
--- /dev/null
+++ b/tilert/pd_vllm/logprobs.py
@@ -0,0 +1,183 @@
+"""Chat-completions ``logprobs`` / ``top_logprobs``: request parsing and response assembly.
+
+Contract (OpenAI chat completions, narrowed to the stricter vendor reading):
+
+``logprobs``
+ boolean, default false. "Whether to return log probabilities of the output
+ tokens or not. If true, returns the log probabilities of each output token
+ returned in the ``content`` of ``message``."
+
+``top_logprobs``
+ integer, default 0, range ``[0, 5]``. The number of most likely tokens to
+ return at each position, each with a log probability. ``logprobs`` must be
+ true if it is used.
+
+Two consequences of that wording drive this module:
+
+* logprobs cover ``message.content`` only. A reasoning segment lives in a
+ different field (``reasoning_content``), so its tokens carry no logprobs.
+* ``top_logprobs`` is capped at 5 here, not OpenAI's current 20. The API
+ reference says 20 and the cookbook still says 5; 5 is the narrower of the
+ two, so a client written against either gets a consistent answer.
+
+Out of range is rejected rather than clamped, matching how vLLM handles a
+logprobs count above its own cap.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+__all__ = [
+ "GREEDY_TEMPERATURE",
+ "LOGPROB_UNAVAILABLE",
+ "MIN_LOGPROBS_TEMPERATURE",
+ "TOP_LOGPROBS_MAX",
+ "LogprobsRequest",
+ "LogprobsUnsupported",
+ "build_logprobs",
+ "resolve_logprobs_request",
+]
+
+TOP_LOGPROBS_MAX = 5
+
+# At or below this the request is greedy: the engine takes its top-1 kernel, and
+# temperature cannot change which token that picks.
+GREEDY_TEMPERATURE = 1e-5
+
+# Lowest temperature the decode sampler's log-probability export is valid at.
+# Set by the engine, not by this contract: see the note in
+# resolve_logprobs_request and RawPowSum in include/ops/deepseek_v3_2/top_p.cuh.
+MIN_LOGPROBS_TEMPERATURE = 0.2
+
+# OpenAI documents -9999.0 as the value standing in for "very unlikely", and
+# uses it where a real log probability is not available. JSON has no -inf, so a
+# missing or -inf entry must surface as this rather than null or Infinity.
+LOGPROB_UNAVAILABLE = -9999.0
+
+
+class LogprobsUnsupported(Exception):
+ """The client's logprobs request cannot be served -> HTTP 400.
+
+ Mirrors ``GrammarError``'s shape so the router can return it the same way.
+ """
+
+ error_type = "invalid_logprobs"
+ http_status = 400
+
+ def to_payload(self) -> dict[str, str]:
+ return {"error": str(self), "error_type": self.error_type}
+
+
+@dataclass(frozen=True)
+class LogprobsRequest:
+ """A validated request for logprobs. ``top_n`` is already in [0, 5]."""
+
+ top_n: int
+
+
+def resolve_logprobs_request(body: dict) -> LogprobsRequest | None:
+ """Parse and validate ``logprobs`` / ``top_logprobs``.
+
+ Returns ``None`` when the client did not ask for logprobs, so callers can
+ skip the whole path. Raises :class:`LogprobsUnsupported` (400) on a request
+ that is malformed or outside the supported range -- never degrades silently,
+ because a client that asked for logprobs and got none has no way to tell.
+ """
+ enabled = body.get("logprobs")
+ if enabled is not None and not isinstance(enabled, bool):
+ raise LogprobsUnsupported(f"logprobs must be a boolean, got {type(enabled).__name__}")
+
+ raw = body.get("top_logprobs")
+ top_n = 0
+ if raw is not None:
+ # bool is an int subclass in Python; `top_logprobs: true` is a type
+ # error, not a request for 1.
+ if isinstance(raw, bool) or not isinstance(raw, int):
+ raise LogprobsUnsupported(f"top_logprobs must be an integer, got {type(raw).__name__}")
+ if not 0 <= raw <= TOP_LOGPROBS_MAX:
+ raise LogprobsUnsupported(f"top_logprobs must be in [0, {TOP_LOGPROBS_MAX}], got {raw}")
+ if raw > 0 and not enabled:
+ raise LogprobsUnsupported("when using top_logprobs, logprobs must be set to true")
+ top_n = raw
+
+ if not enabled:
+ return None
+
+ # Greedy and the top-p sampler both export a chosen-token value and a
+ # candidate row, so both serve any top_logprobs up to TOP_LOGPROBS_MAX --
+ # which is the greedy kernel's row width, the narrower of the two.
+ #
+ # The band between them is refused rather than degraded, for the reason an
+ # out-of-range top_logprobs is refused: a wrong number is indistinguishable
+ # from a right one. There the temperature genuinely selects the
+ # distribution being sampled, so greedy's export cannot stand in for it,
+ # and the top-p path's raw denominator loses accuracy (order 0.01 nat at
+ # T = 0.1, 0.3 at T = 0.05). A non-numeric temperature is vLLM's to reject.
+ temp = body.get("temperature")
+ if isinstance(temp, (int, float)) and not isinstance(temp, bool):
+ t = float(temp)
+ if GREEDY_TEMPERATURE <= t < MIN_LOGPROBS_TEMPERATURE:
+ raise LogprobsUnsupported(
+ f"logprobs require temperature >= {MIN_LOGPROBS_TEMPERATURE} "
+ f"or greedy (temperature < {GREEDY_TEMPERATURE}), got {t}"
+ )
+ return LogprobsRequest(top_n=top_n)
+
+
+def _entry(token_id: int, logprob: float | None, decode_one: Callable[[int], str]) -> dict:
+ """One ``{token, logprob, bytes}`` object.
+
+ ``bytes`` carries the UTF-8 encoding of the token text, which is how a
+ caller reassembles text when a single token is not valid UTF-8 on its own
+ (byte-level BPE splits multi-byte characters across tokens).
+ """
+ text = decode_one(token_id)
+ return {
+ "token": text,
+ "logprob": LOGPROB_UNAVAILABLE if logprob is None else _finite(logprob),
+ "bytes": list(text.encode("utf-8")),
+ }
+
+
+def _finite(value: float) -> float:
+ """JSON has no infinities; -inf becomes the documented sentinel."""
+ return LOGPROB_UNAVAILABLE if value == float("-inf") else float(value)
+
+
+def build_logprobs(
+ token_ids: list[int],
+ token_logprobs: list[float | None],
+ top_logprobs: list[list[tuple[int, float]]] | None,
+ req: LogprobsRequest,
+ decode_one: Callable[[int], str],
+) -> dict:
+ """Assemble ``choices[].logprobs`` for the tokens of ``message.content``.
+
+ ``token_ids`` are the content tokens in order -- callers must already have
+ dropped any reasoning segment (see :func:`content_token_slice`).
+ ``token_logprobs[i]`` is the log probability of ``token_ids[i]``.
+ ``top_logprobs[i]`` is that position's candidate list, longest-first; it is
+ truncated to ``req.top_n`` here so a decode node may return more than asked.
+
+ ``refusal`` is always ``None``: this path produces no refusal channel, and
+ the field is documented as nullable.
+ """
+ if len(token_logprobs) != len(token_ids):
+ raise ValueError(
+ f"token_logprobs has {len(token_logprobs)} entries for " f"{len(token_ids)} tokens"
+ )
+ if top_logprobs is not None and len(top_logprobs) != len(token_ids):
+ raise ValueError(
+ f"top_logprobs has {len(top_logprobs)} entries for " f"{len(token_ids)} tokens"
+ )
+
+ content = []
+ for i, tid in enumerate(token_ids):
+ item = _entry(tid, token_logprobs[i], decode_one)
+ alts = [] if top_logprobs is None else top_logprobs[i][: req.top_n]
+ item["top_logprobs"] = [_entry(alt_id, alt_lp, decode_one) for alt_id, alt_lp in alts]
+ content.append(item)
+
+ return {"content": content, "refusal": None}
diff --git a/tilert/pd_vllm/oai_parser.py b/tilert/pd_vllm/oai_parser.py
index 931daa9..650e5ff 100644
--- a/tilert/pd_vllm/oai_parser.py
+++ b/tilert/pd_vllm/oai_parser.py
@@ -12,6 +12,11 @@
Runs in the ROUTER environment only — that env must have vllm installed
(CPU-only import is fine; verified with CUDA_VISIBLE_DEVICES=""). The decode
node never imports vllm.
+
+Verified against real AIME transcripts (prefilled- convention, incl. a
+135K-char truncated-thinking sample) and template-format tool calls with
+random-delta streaming fuzz. Key engine semantics (cost a bug to learn):
+``TOOL_NAME`` is an incremental chunk event — fragments must be concatenated.
"""
import logging
@@ -27,6 +32,11 @@ class ToolCall:
name: str
arguments: str # JSON string (OpenAI convention)
+ @property
+ def id(self) -> str: # noqa: A003
+ """Alias of ``call_id`` under the OpenAI field name."""
+ return self.call_id
+
def to_openai(self, index: int) -> dict:
return {
"index": index,
@@ -50,7 +60,7 @@ def _new_call_id() -> str:
# family -> (config-builder import path, arg-converter import path). The
# glm47_moe parser engine uses the vllm.parser API shape (a `*_config(thinking)`
# builder + a `_*_arg_converter(raw, partial)`); the adapter picks the engine
-# by family name.
+# by family name, so another family with the same shape is one table entry.
_FAMILIES = {
"glm47": ("vllm.parser.glm47_moe", "glm47_moe_config", "_glm47_arg_converter"),
}
@@ -203,25 +213,68 @@ class IncrementalDetok:
r"""Incremental token→text for byte-level BPE tokenizers.
Decodes a bounded trailing window; holds output while the window ends in
- a partial multi-byte sequence (\\ufffd). Window folding is safe for
+ a partial multi-byte sequence (\ufffd). Window folding is safe for
byte-level BPE: separate windows decode to concatenable byte streams.
- Specials are KEPT (skip_special_tokens=False) — the parser consumes
- etc.; the stop token never reaches the stream (engine adapter
- suppresses it).
+
+ ``finish`` releases what is held. A generation can end mid-character -- EOS
+ or ``max_tokens`` after the first byte of two -- and the ids are reported
+ either way, so without the release the reply would omit a character the
+ tokenizer makes of the ids it reports.
+
+ Specials are kept by default — a parser consumes and friends, and
+ the stop token never reaches the stream because the engine adapter
+ suppresses it. A caller with no parser passes True instead: nothing
+ downstream would consume a special, so one would surface as content. The
+ policy is a constructor argument rather than a per-call one so a request
+ cannot be matched against one spelling of its own output and shown another.
"""
_FOLD = 256
- def __init__(self, tokenizer):
+ def __init__(self, tokenizer, skip_special_tokens: bool = False):
self._tok = tokenizer
+ self._skip = skip_special_tokens
self._ids: list[int] = []
self._emitted = 0
+ self._holding = False
+
+ @property
+ def holding(self) -> bool:
+ """Whether the last push produced nothing because it was incomplete.
+
+ Distinguishes a byte fragment — whose text arrives with a later token —
+ from a token that genuinely decodes to nothing, such as a special the
+ caller strips. Both return an empty delta, and a caller pairing
+ per-token metadata against the text has to tell them apart.
+ """
+ return self._holding
+
+ def finish(self) -> str:
+ """Whatever ``push`` held back because the window ended mid-character.
+
+ Generation can stop between the byte-level tokens of one character -- EOS
+ or ``max_tokens`` arriving after its first byte -- and the tokenizer's own
+ decode of the complete id list then ends in the replacement character.
+ Dropping it loses a character the reply had, while ``token_ids``, usage
+ and the logprob entry all still count the token.
+ """
+ # Idempotent through `_emitted`, which the first call advances to the
+ # end; the `_holding` early-out is a shortcut, not the guarantee.
+ if not self._holding:
+ return ""
+ text = self._tok.decode(self._ids, skip_special_tokens=self._skip)
+ delta = text[self._emitted :]
+ self._emitted = len(text)
+ self._holding = False
+ return delta # noqa: R504 (self._emitted mutated after delta is computed)
def push(self, ids: list[int]) -> str:
self._ids.extend(ids)
- text = self._tok.decode(self._ids, skip_special_tokens=False)
- if text.endswith("�"):
+ text = self._tok.decode(self._ids, skip_special_tokens=self._skip)
+ if text.endswith("\ufffd"):
+ self._holding = True
return ""
+ self._holding = False
delta = text[self._emitted :]
self._emitted = len(text)
if len(self._ids) > self._FOLD:
diff --git a/tilert/pd_vllm/openai_params.py b/tilert/pd_vllm/openai_params.py
new file mode 100644
index 0000000..a92f149
--- /dev/null
+++ b/tilert/pd_vllm/openai_params.py
@@ -0,0 +1,86 @@
+"""Shared reading of the OpenAI request fields every backend accepts.
+
+All backends expose the same OpenAI-compatible ``/v1/chat/completions``
+surface (see ``docs/architecture.md``), so the rules for interpreting a client
+body belong in exactly one place rather than once per subpackage. Nothing here
+imports a backend; it is plain dict handling.
+"""
+
+from __future__ import annotations
+
+from pydantic import TypeAdapter, ValidationError
+
+__all__ = ["InvalidOutputLength", "resolve_max_tokens"]
+
+# Two orders, because vLLM applies two.
+#
+# Declaration order on ``ChatCompletionRequest``
+# (v0.24.0 protocol.py:202, 207).
+# Pydantic validates every field and reports in this order, so a body wrong in
+# both names is answered about ``max_tokens``.
+_DECLARED = ("max_tokens", "max_completion_tokens")
+# Precedence, applied after validation and on ``is not None`` -- NOT on
+# truthiness, so a preferred ``0`` does not defer to the other name
+# (v0.24.0 chat_completion/serving.py:302-306).
+_PRECEDENCE = ("max_completion_tokens", "max_tokens")
+
+_INT = TypeAdapter(int)
+
+
+class InvalidOutputLength(ValueError):
+ """The client's output length cannot be honoured; each edge maps it to 400.
+
+ A ``ValueError`` subclass so a caller that already guards this call keeps
+ working.
+ """
+
+
+def resolve_max_tokens(body: dict, default: int | None = None) -> int | None:
+ """The output length for this request, or ``default`` if none was asked for.
+
+ The single owner of this field. It used to own the precedence only, leaving
+ coercion to whichever ``int()`` call happened to run first and the refusal
+ to nobody -- so ``1.9`` was truncated to 1 on one path while ``"20.0"``
+ raised on another, after the prefill had already run.
+
+ Coercion is delegated to the request model's own validator rather than
+ restated: ``max_tokens`` is a plain ``int | None`` on vLLM's
+ ``ChatCompletionRequest``, so pydantic's lax rules ARE the specification.
+ ``True`` is 1, ``"20"`` / ``" 20 "`` / ``"20.0"`` coerce, ``"1e3"`` and
+ ``1.9`` and ``inf`` do not, and ``10**309`` is a valid integer that
+ ``float()`` cannot hold. That table has no guessable edges, which is why
+ every hand-written version of it has been wrong somewhere.
+
+ Both names are coerced but only the effective one is range-checked, because
+ vLLM's two layers have different reach: pydantic validates every declared
+ field before the precedence is applied, while ``_verify_args`` then sees
+ one resolved number. So ``{max_completion_tokens: 3, max_tokens: 1.9}`` is
+ refused and ``{max_completion_tokens: 16, max_tokens: 0}`` is served.
+
+ ``default`` is an argument because it is a property of the backend, not of
+ the request: the PD decode node uses 256, ``serve_native`` 4096. Omitting
+ it asks the question without answering it -- what the CLIENT requested,
+ which is what an edge validating before it has a backend in hand needs.
+
+ Pure and total, so calling it more than once per request is a repeated
+ computation and not a second opinion. That is the property the two call
+ sites lacked when each did its own coercion.
+
+ Raises:
+ InvalidOutputLength: unusable type or value; the edge answers 400.
+ """
+ seen: dict[str, int] = {}
+ for name in _DECLARED:
+ raw = body.get(name)
+ if raw is None:
+ continue
+ try:
+ seen[name] = _INT.validate_python(raw)
+ except ValidationError:
+ raise InvalidOutputLength(f"{name} must be an integer, got {raw!r}") from None
+ for name in _PRECEDENCE:
+ if name in seen:
+ if seen[name] < 1:
+ raise InvalidOutputLength(f"{name} must be at least 1, got {seen[name]}")
+ return seen[name]
+ return default
diff --git a/tilert/pd_vllm/pd_router.py b/tilert/pd_vllm/pd_router.py
index 87a61e5..2f98fa5 100644
--- a/tilert/pd_vllm/pd_router.py
+++ b/tilert/pd_vllm/pd_router.py
@@ -1,10 +1,11 @@
-"""PD router (W6): client-facing entry over vLLM prefill + TileRT decode.
+"""PD router: client-facing entry that orchestrates vLLM prefill and TileRT decode.
-Does OpenAI-semantics output parsing (reasoning + tool calls), streaming and
+OpenAI-semantics output parsing (reasoning + tool calls), streaming and
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
@@ -20,15 +21,15 @@
Run:
CUDA_VISIBLE_DEVICES= python -m tilert.pd_vllm.pd_router \
- --vllm-url http://prefill-node:8000 \
- --decode decode-node:5556:5557 --port 23333 \
- --model-path /path/to/GLM-5.1 --parser glm47
+ --vllm-url http://:8000 \
+ --decode :5556:5557 --port 23333 \
+ --model-path /path/to/model --model glm5 --parser glm47
"""
import argparse
-import json
+import contextlib
+import functools
import logging
-import threading
import time
import requests
@@ -36,39 +37,118 @@
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse
+from tilert.pd_vllm import generation_defaults
+from tilert.pd_vllm.capabilities import CapabilityError, validate_generation_request
+from tilert.pd_vllm.decode_pool import DecodeNode, Pool, acquire_lease
+from tilert.pd_vllm.decode_response import (
+ BUSY,
+ OK,
+ PROPAGATED_ERROR_TYPES,
+ REFUSED,
+ RETRY,
+ SERVER_ERROR,
+ TRUNCATED,
+ TYPED_ERROR,
+ UNTYPED_ERROR,
+ DecodeReader,
+ classify_decode_status,
+ decode_refusal,
+ terminal_verdict,
+)
+from tilert.pd_vllm.generation_defaults import (
+ GenerationDefaults,
+ UnsupportedGenerationDefault,
+)
+from tilert.pd_vllm.grammar_spec import GrammarError
+from tilert.pd_vllm.logprobs import LogprobsUnsupported
+from tilert.pd_vllm.openai_params import resolve_max_tokens
+from tilert.pd_vllm.presentation import (
+ SseWriter,
+ blocking_choice,
+ blocking_envelope,
+ collect,
+)
+from tilert.pd_vllm.presentation import finish_reason as reply_finish_reason
+from tilert.pd_vllm.presentation import (
+ sse_chunk,
+ textless_choice,
+ usage_chunk,
+)
+from tilert.pd_vllm.reply import ReplyStream
+from tilert.pd_vllm.request_gate import gate_request
from tilert.pd_vllm.wire import derive_rid
logger = logging.getLogger("pd_vllm.router")
+QUEUE_LOG_SECONDS = 0.1
-class DecodeNode:
- def __init__(self, host: str, ctrl_port: int, http_port: int):
- self.host = host
- self.ctrl_port = ctrl_port
- self.http_port = http_port
- self.busy = False
+# Output length when the client names none; carried over from the original
+# import rather than tuned. Decode only -- the prefill request is pinned to
+# max_tokens=1 with both client field names stripped (_PREFILL_DROP_FIELDS).
+_DECODE_MAX_TOKENS_DEFAULT = 256
- @property
- def http_base(self) -> str:
- return f"http://{self.host}:{self.http_port}"
+# error_types the decode node may return that carry client-meaningful HTTP
+# status (400/500/501) and must be propagated verbatim rather than masked as 502.
+#
+# The rule is "typed and classified by the decode node": those statuses are the
+# node's considered answer about THIS request, so flattening them into 502 tells
+# the operator to go restart a healthy component. An untyped failure stays 502,
+# which is what it is -- a component call that did not work.
+# A decode node answers 429 when its single slot is still held. The router's own
+# gating should make that rare, but not impossible: it frees its reservation as
+# soon as it stops reading, while the node's slot unwinds a little later (a
+# cancel handshake, an engine reset), so a request dispatched into that window
+# meets the node's admission. One short retry absorbs it; a node still busy after
+# that gets an honest 429, which is retryable, rather than a 502, which tells the
+# operator to go restart a healthy component.
+_DECODE_BUSY_ATTEMPTS = 2
+_DECODE_BUSY_RETRY_S = 0.5
+_RETRY_LOG = "decode node %s busy for %s, retrying in %.1fs"
-class Pool:
- def __init__(self, nodes: list[DecodeNode]):
- self.nodes = nodes
- self._lock = threading.Lock()
+def _log_refusal(verdict: str, node, rid: str, status: int, payload=None) -> None:
+ """One line per refused dispatch, at the level the verdict deserves."""
+ if verdict == BUSY:
+ logger.warning(
+ "decode node %s still busy for %s after %d attempts",
+ node.http_base,
+ rid,
+ _DECODE_BUSY_ATTEMPTS,
+ )
+ elif verdict == SERVER_ERROR:
+ logger.error(
+ "decode node %s returned %d for %s: %s", node.http_base, status, rid, str(payload)[:200]
+ )
- def acquire(self) -> DecodeNode | None:
- with self._lock:
- for n in self.nodes:
- if not n.busy:
- n.busy = True
- return n
- return None
- def release(self, node: DecodeNode) -> None:
- with self._lock:
- node.busy = False
+# How often the streaming preflight checks whether the client is still there
+# while the decode node holds its headers. Starlette offers no awaitable for a
+# disconnect before the response begins, so this is a poll.
+_DISCONNECT_POLL_S = 0.2
+
+
+# The status each typed error deserves, taken from the exception class that
+# raises it on the node so the two protocols cannot disagree.
+#
+# Needed because they carry the type differently. Over the BLOCKING protocol the
+# node answers an HTTP status and the router forwards it. Over the STREAMING one
+# the error arrives inside a 200 body -- the status is already spent -- so the
+# router has to reconstruct it, and mapping every propagated type to one status
+# is how `logprobs_unavailable` came back as 400 on one path and 501 on the
+# other for the same inability.
+class PrefillClientError(Exception):
+ """vLLM rejected the prefill request itself (4xx).
+
+ That is the client's fault, not a component failure, so it must reach the
+ client with vLLM's own status and body instead of being flattened into a 502.
+ Hit in practice by `response_format` specs vLLM validates before we ever get
+ to the decode node (e.g. an uncompilable json_schema).
+ """
+
+ def __init__(self, status: int, payload):
+ super().__init__(f"vLLM rejected the request with {status}")
+ self.status = status
+ self.payload = payload
def first_token_from_logprobs(resp: dict, is_chat: bool) -> int:
@@ -92,19 +172,202 @@ def first_token_from_logprobs(resp: dict, is_chat: bool) -> int:
)
-def _thinking_enabled(body: dict) -> bool:
- ctk = body.get("chat_template_kwargs") or {}
- return bool(ctk.get("enable_thinking", True))
+def _token_id_of(tok: object) -> int | None:
+ """The integer id behind a vLLM ``token_id:N`` string, or None."""
+ if isinstance(tok, str) and tok.startswith("token_id:"):
+ try:
+ return int(tok.split(":", 1)[1])
+ except ValueError:
+ return None
+ return None
+
+
+def first_token_logprob_from_prefill(resp: dict, top_n: int):
+ """Token 1's ``(logprob, candidates)`` out of the prefill response.
+
+ The decode node echoes ``first_token_id`` without sampling it, so the only
+ distribution for that position is the prefill instance's. The router already
+ reads this entry for the token id; this reads the numbers beside it.
+
+ ``(None, [])`` when there is no usable entry -- the caller surfaces the
+ documented sentinel for that one position rather than failing the request.
+
+ The value carries the prefill instance's ``logprobs_mode``, vLLM's default
+ being ``raw_logprobs``, so the decode side must be in its raw mode for the
+ array to sit on one scale (see the engine adapter's ``supports_logprobs``).
+ """
+ choice = (resp.get("choices") or [{}])[0]
+ content = (choice.get("logprobs") or {}).get("content") or []
+ if not content:
+ return None, []
+ entry = content[0]
+ lp = entry.get("logprob")
+ lp = float(lp) if isinstance(lp, (int, float)) else None
+ cands = []
+ for alt in (entry.get("top_logprobs") or [])[:top_n]:
+ alt_id = _token_id_of(alt.get("token"))
+ alt_lp = alt.get("logprob")
+ if alt_id is not None and isinstance(alt_lp, (int, float)):
+ cands.append((alt_id, float(alt_lp)))
+ return lp, cands
+
+
+# Client fields that must not survive into the prefill request. The body is
+# forwarded verbatim apart from the fields we set, so anything vLLM checks
+# against one of those overrides has to go first.
+#
+# stream_options -- we force stream=False, and vLLM rejects the pair
+# with 400 "Stream options can only be defined when
+# `stream=True`". Its validator is mode="before", so
+# the request dies before the model is looked at.
+# max_completion_tokens -- takes precedence over max_tokens in vLLM, so it
+# would override our max_tokens=1 and make the
+# prefill instance decode the client's whole output
+# length, defeating the split.
+# stop, include_stop_str_in_output
+# -- the router matches these itself, over the whole
+# reply. Left in, vLLM would match against the one
+# token it generates and could report
+# finish_reason="stop" for a prefill that succeeded.
+#
+# Any streaming client sends the first two, `vllm bench serve
+# --backend openai-chat` included -- which made the official benchmark fail
+# every request with 502.
+_PREFILL_DROP_FIELDS = (
+ "stream_options",
+ "max_completion_tokens",
+ "stop",
+ "include_stop_str_in_output",
+)
+
+
+def build_prefill_body(
+ path: str,
+ body: dict,
+ node: DecodeNode,
+ logprobs_req=None,
+ defaults: GenerationDefaults | None = None,
+) -> dict:
+ """The vLLM request that prefills only and hands the KV state to ``node``.
+
+ Lives outside ``build_app`` so the rewrite can be unit-tested without a
+ router process, a vLLM instance or a decode node.
+
+ ``logprobs`` is always requested, because the first token's id is recovered
+ from it. When the client asked for logprobs as well, ``top_logprobs`` is
+ raised to the count they asked for: token 1's candidate row can only come
+ from here, since the decode node never sampled that position.
+ """
+ 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
+ # 1 is the floor, not the default: the id extraction needs one entry.
+ prefill_body["top_logprobs"] = max(1, logprobs_req.top_n if logprobs_req is not None else 0)
+ else:
+ prefill_body["logprobs"] = 1
+ # Pin temperature/top_p/top_k explicitly on BOTH legs. Left absent, this leg
+ # would resolve them through vLLM's own chain (generation_config.json, then
+ # the neutral defaults) while the decode leg resolved its own -- the split
+ # that had token 1 sampled at temperature 0.6 / top_k 20 and tokens 2..N at
+ # 1.0 / uncapped. One resolution, written to both requests, so the
+ # prefill instance's own --generation-config can no longer move one leg
+ # without the other.
+ prefill_body.update((defaults or GenerationDefaults()).resolve(body))
+ prefill_body["kv_transfer_params"] = {
+ "tilert_host": node.host,
+ "tilert_ctrl_port": node.ctrl_port,
+ }
+ return prefill_body
+
+
+def build_usage(prompt_tokens, completion_tokens: int) -> dict:
+ """``usage``, shaped like vLLM's ``UsageInfo``.
+
+ All three fields are integers there (``prompt_tokens: int = 0``,
+ ``total_tokens: int = 0``, ``completion_tokens: int | None = 0``), and
+ ``total_tokens`` is always present. This endpoint used to omit it on the
+ non-streaming path while sending it on the streaming one, so a client reading
+ ``usage.total_tokens`` -- which every OpenAI SDK does -- got a KeyError from
+ one shape and a number from the other.
+
+ ``prompt_tokens`` is coerced because it is copied from the prefill response,
+ where it can be absent; ``None`` would violate the contract just as surely as
+ the missing key did.
+ """
+ prompt = int(prompt_tokens or 0)
+ completion = int(completion_tokens)
+ return {
+ "prompt_tokens": prompt,
+ "completion_tokens": completion,
+ "total_tokens": prompt + completion,
+ }
+
+
+def should_include_usage(body: dict, force: bool = False) -> bool:
+ """Whether this stream carries usage at all, per the OpenAI contract.
+
+ Absent ``include_usage`` means no usage anywhere in the stream -- not usage
+ relocated to another chunk. vLLM (``serve/utils/api_utils.py``) and SGLang
+ (``serving_chat.py``) both resolve it through a function of this name and
+ shape; this endpoint is judged against them.
+
+ Relocating it, which the router used to do, is worse than non-conformant:
+ a client written as ``if chunk.choices: ... elif chunk.usage:`` takes the
+ choices branch and never reads usage riding on the same chunk. That is how
+ the InferenceX bench reported `Total generated tokens: 0`.
+
+ ``force`` is the deployment escape hatch, mirroring vLLM's
+ ``enable_force_include_usage`` and SGLang's
+ ``stream_response_default_include_usage``. Off by default in all three.
+
+ Reads the client body only: vLLM never sees ``stream_options`` (stripped,
+ since we force stream=False and vLLM 400s the pair) and the decode node is
+ sent only sampling params.
+ """
+ if force:
+ return True
+ opts = body.get("stream_options")
+ return bool(isinstance(opts, dict) and opts.get("include_usage"))
+
+
+def _json_or_none(response) -> dict | None:
+ """The body as JSON, or None.
+
+ A node classifies its failures there; a proxy or a crash answers with something else.
+ """
+ try:
+ return response.json()
+ except ValueError:
+ return None
class RouterCtx:
"""Immutable per-process context (tokenizer, parser factory, config)."""
- def __init__(self, vllm_url: str, pool: Pool, tokenizer, parser_name: str):
+ def __init__(
+ self,
+ vllm_url: str,
+ pool: Pool,
+ tokenizer,
+ parser_name: str,
+ force_include_usage: bool = False,
+ gen_defaults: GenerationDefaults | None = None,
+ ):
self.vllm_url = vllm_url
self.pool = pool
self.tokenizer = tokenizer
self.parser_name = parser_name
+ # Resolved once at startup, exactly as vLLM resolves its own
+ # default_sampling_params. Both legs are handed THESE values, so they
+ # cannot disagree whatever they are; see generation_defaults.
+ self.gen_defaults = gen_defaults or GenerationDefaults()
+ # Defaults off, as it does in vLLM and SGLang: a deployment may opt its
+ # whole fleet into usage, but never by omission.
+ self.force_include_usage = force_include_usage
self._parsers = {}
if parser_name != "none":
if tokenizer is None:
@@ -123,6 +386,22 @@ def build_app(ctx: RouterCtx) -> FastAPI:
app = FastAPI()
pool = ctx.pool
+ def _parser_active(thinking: bool) -> bool:
+ """Whether an output parser would run, given the thinking flag.
+
+ A router-configuration answer the gate needs but cannot look up.
+ """
+ return ctx.parser(thinking) is not None
+
+ def _busy_response(waited: float) -> JSONResponse:
+ """429 body, telling a full pool apart from an exhausted 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)}
@@ -132,191 +411,458 @@ def pool_status():
return {"nodes": [{"host": n.host, "busy": n.busy} for n in pool.nodes]}
# ── 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,
- }
+ def _prefill(path, body, node, logprobs_req=None):
+ prefill_body = build_prefill_body(path, body, node, logprobs_req, ctx.gen_defaults)
r = requests.post(f"{ctx.vllm_url}{path}", json=prefill_body, timeout=600)
+ if 400 <= r.status_code < 500:
+ try:
+ payload = r.json()
+ except ValueError:
+ payload = {"error": r.text[:500]}
+ raise PrefillClientError(r.status_code, payload)
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}
+ # ignore_eos is the decode node's alone: the prefill request is pinned to
+ # max_tokens=1 so it never reaches a stop token; the decode loop owns the
+ # stop set and clears it for the flag (adapters declare this via
+ # /capabilities, and the router refuses the request if none can).
+ #
+ # Every OTHER field a client may send is either forwarded here or
+ # refused by validate_generation_request. A field that is neither would
+ # be applied by the vLLM prefill instance to token 1 and silently
+ # dropped for the rest of the reply -- so this whitelist and that gate
+ # must be kept in step (test_no_forwarded_field_is_left_ungated).
+ forwarded = (
+ "temperature",
+ "top_p",
+ "top_k",
+ "repetition_penalty",
+ "presence_penalty",
+ "ignore_eos",
+ )
+ sampling = {k: body[k] for k in forwarded if k in body}
+ # Resolved, never defaulted downstream: the same call on the same body
+ # produced the values already written into the prefill request.
+ sampling.update(ctx.gen_defaults.resolve(body))
+ return sampling
- def _max_tokens_of(body):
- return int(body.get("max_tokens") or body.get("max_completion_tokens") or 256)
+ def _decode_body(
+ rid, first_token_id, body, grammar_spec, *, stream=False, logprobs_req=None, thinking=True
+ ):
+ payload = {
+ "rid": rid,
+ "first_token_id": first_token_id,
+ # The whole output length is decoded here; the prefill
+ # request is pinned to max_tokens=1 and both client field names
+ # are stripped
+ # from it (see _PREFILL_DROP_FIELDS).
+ "max_tokens": resolve_max_tokens(body, _DECODE_MAX_TOKENS_DEFAULT),
+ "sampling": _sampling_of(body),
+ }
+ if stream:
+ payload["stream"] = True
+ if grammar_spec is not None:
+ # Only sent for constrained requests; plain requests stay byte-for-
+ # byte on the existing unconstrained path.
+ payload["grammar_spec"] = grammar_spec
+ payload["enable_thinking"] = thinking
+ if logprobs_req is not None:
+ # Omitted entirely when not requested, so a decode node that
+ # predates the field is unaffected.
+ payload["top_logprobs"] = logprobs_req.top_n
+ return payload
# ── non-streaming ────────────────────────────────────────────────────
def _handle(path: str, body: dict):
- is_chat = path.endswith("chat/completions")
- node = pool.acquire()
- if node is None:
- return JSONResponse({"error": "all decode nodes busy"}, status_code=429)
+ try:
+ # Request-only first, then the probe -- `request_gate` says why.
+ req = gate_request(path, body, tokenizer=ctx.tokenizer, parser_active=_parser_active)
+ validate_generation_request(body, pool.capabilities(), ctx.gen_defaults.resolve({}))
+ except (CapabilityError, GrammarError, LogprobsUnsupported) as e:
+ return JSONResponse(e.to_payload(), status_code=e.http_status)
+ is_chat, stop, include_stop = req.is_chat, req.stop, req.include_stop
+ logprobs_req, grammar_spec = req.logprobs_req, req.grammar_spec
+ lease, waited = acquire_lease(pool)
+ if lease is None:
+ return _busy_response(waited)
+ node = lease.node
t0 = time.time()
+ created = int(t0)
+ reader = None
try:
- prefill = _prefill(path, body, node)
+ prefill = _prefill(path, body, node, logprobs_req)
t_prefill = time.time()
- rid = derive_rid(prefill["id"])
+ rid = lease.rid = derive_rid(prefill["id"])
first_token_id = first_token_from_logprobs(prefill, is_chat)
+ first_lp = (
+ first_token_logprob_from_prefill(prefill, logprobs_req.top_n)
+ if logprobs_req is not None
+ else None
+ )
- dr = requests.post(
- f"{node.http_base}/pd/decode",
- json={
- "rid": rid,
- "first_token_id": first_token_id,
- "max_tokens": _max_tokens_of(body),
- "sampling": _sampling_of(body),
- },
- timeout=600,
+ parser = ctx.parser(req.thinking) if is_chat else None
+ # Same object, same arguments as the streaming path; this function
+ # only concatenates the emissions instead of framing them.
+ asm = (
+ ReplyStream(
+ ctx.tokenizer,
+ stop=stop,
+ include_stop_in_output=include_stop,
+ parser_session=parser.stream() if parser else None,
+ logprobs_req=logprobs_req,
+ first_token_logprob=first_lp,
+ )
+ if ctx.tokenizer is not None
+ else None
)
- dr.raise_for_status()
- decode = dr.json()
- token_ids = decode["token_ids"]
- timing = decode.get("timing_ms", {})
- finish = timing.get("finish_reason", "stop")
- if finish == "cancelled":
- finish = "stop"
-
- choice: dict = {"index": 0, "finish_reason": finish}
- parser = ctx.parser(_thinking_enabled(body)) if is_chat else None
- if parser is not None:
- text = ctx.tokenizer.decode(token_ids, skip_special_tokens=False)
- parsed = parser.parse_complete(text)
- msg = {"role": "assistant", "content": parsed.content or ""}
- if parsed.reasoning_content:
- msg["reasoning_content"] = parsed.reasoning_content
- if parsed.tool_calls:
- msg["tool_calls"] = [c.to_openai(i) for i, c in enumerate(parsed.tool_calls)]
- choice["finish_reason"] = "tool_calls"
- choice["message"] = msg
- else:
- text = (
- ctx.tokenizer.decode(token_ids, skip_special_tokens=True)
- if ctx.tokenizer
- else None
+
+ # Only the router can see a stop, and not before the tokens arrive.
+ # Asking for the whole sequence up front would make a stop merely
+ # trim the answer: the node would run to max_tokens and hold its slot
+ # for all of it. So a request with something to match reads the
+ # node's STREAMING protocol even though its own reply is not.
+ want_stream = bool(stop)
+ reader = DecodeReader(stream=asm, logprobs_req=logprobs_req, rid=rid)
+ decode_payload = _decode_body(
+ rid,
+ first_token_id,
+ body,
+ grammar_spec,
+ stream=want_stream,
+ logprobs_req=logprobs_req,
+ thinking=req.thinking,
+ )
+ # True from the moment a POST goes out. The node may have admitted
+ # the request before the call failed -- a timeout or a reset while
+ # waiting for headers -- and it then holds its slot until its own
+ # timeout. Cancelling an rid the node never saw is harmless; not
+ # cancelling one it did is a slot lost for `timeout_s`.
+ dr = None
+ for attempt in range(1, _DECODE_BUSY_ATTEMPTS + 1):
+ lease.dispatched = True
+ dr = requests.post(
+ f"{node.http_base}/pd/decode",
+ json=decode_payload,
+ timeout=600,
+ stream=want_stream,
+ )
+ if dr.status_code == 200:
+ break
+ verdict = classify_decode_status(
+ dr.status_code,
+ _json_or_none(dr),
+ attempts_left=attempt < _DECODE_BUSY_ATTEMPTS,
+ propagated_types=PROPAGATED_ERROR_TYPES,
)
- if is_chat:
- choice["message"] = {"role": "assistant", "content": text}
+ if want_stream:
+ # An unread streamed body holds the connection; a
+ # non-streamed one was already consumed by `post`.
+ dr.close()
+ if verdict == RETRY:
+ logger.info(_RETRY_LOG, node.http_base, rid, _DECODE_BUSY_RETRY_S)
+ time.sleep(_DECODE_BUSY_RETRY_S)
+ continue
+ _log_refusal(verdict, node, rid, dr.status_code)
+ body, status = decode_refusal(verdict, dr.status_code, _json_or_none(dr), rid)
+ return JSONResponse(body, status_code=status)
+ assert dr is not None # the loop always posts at least once
+
+ emissions = []
+ try:
+ if not want_stream:
+ emissions = reader.feed_blocking(dr.json())
else:
- choice["text"] = text
- choice["token_ids"] = token_ids
+ for line in dr.iter_lines(decode_unicode=True):
+ emissions += reader.feed(line)
+ if reader.finished:
+ break
+ finally:
+ if want_stream and dr is not None:
+ dr.close()
+
+ verdict, payload, status = terminal_verdict(reader)
+ if verdict == TRUNCATED:
+ logger.warning(
+ "decode stream for %s ended after %d tokens with " "no done/error message",
+ rid,
+ len(asm.token_ids) if asm else 0,
+ )
+ timing = reader.timing
+ finish = reader.finish_reason
+ token_ids = reader.token_ids
+ if verdict != OK:
+ # Every one of them is a status here; the streaming path has to
+ # say the same things inside a spent 200.
+ return JSONResponse(payload, status_code=status)
+ if asm is not None:
+ emissions += asm.finish()
+ # One source for the id list, the count and the entries, so
+ # they cannot disagree. The node's batching makes them disagree
+ # otherwise: a stop lands part-way into a batch, and the tokens
+ # behind it are not part of the reply.
+ token_ids = asm.token_ids
+ if asm is None:
+ choice, n_completion = textless_choice(
+ is_chat=is_chat, from_node=finish, token_ids=token_ids
+ )
+ else:
+ choice, n_completion = blocking_choice(
+ collect(emissions),
+ is_chat=is_chat,
+ stream=asm,
+ from_node=finish,
+ logprobs_asked=logprobs_req is not None,
+ token_ids=token_ids,
+ )
return JSONResponse(
- {
- "id": prefill["id"],
- "object": "chat.completion" if is_chat else "text_completion",
- "created": int(time.time()),
- "model": prefill.get("model"),
- "choices": [choice],
- "usage": {
- "prompt_tokens": (prefill.get("usage") or {}).get("prompt_tokens"),
- "completion_tokens": len(token_ids),
- },
- "pd_timing_ms": {
- "prefill": round(1000 * (t_prefill - t0), 1),
- **timing,
- },
- }
+ blocking_envelope(
+ choice,
+ is_chat=is_chat,
+ prefill=prefill,
+ created=created,
+ usage=build_usage(
+ (prefill.get("usage") or {}).get("prompt_tokens"), n_completion
+ ),
+ timing={"prefill": round(1000 * (t_prefill - t0), 1), **timing},
+ )
)
+ except PrefillClientError as e:
+ logger.info("vLLM rejected the request (%d): %s", e.status, str(e.payload)[:200])
+ return JSONResponse(e.payload, status_code=e.status)
except Exception as e:
logger.exception("pd request failed")
return JSONResponse({"error": str(e)}, status_code=502)
finally:
- pool.release(node)
+ # The lease owns both rules; the one fact it cannot know is whether
+ # the NODE said it was done.
+ lease.release(terminated=reader is not None and reader.node_terminated)
+
+ async def _send_watching_client(client, req, request):
+ """Open the decode stream, giving up if the client leaves first.
+
+ Returns the response, or None if the client disconnected while waiting.
+
+ Needed because the decode node holds its response headers through the
+ whole KV wire-wait (``/pd/decode`` streams only after the transfer
+ lands), so this await can sit for ``timeout_s`` -- 120 s by default. At
+ that point ``StreamingResponse`` does not exist yet, so neither the
+ generator's ``finally`` nor its ``is_disconnected`` poll is running: a
+ client that hangs up here would otherwise hold both the router's
+ reservation and the decode slot for the full wait, which is exactly the
+ failure this endpoint is meant to have stopped having.
+
+ The disconnect is polled rather than awaited because that is the only
+ signal Starlette offers for a request whose response has not begun.
+ """
+ import asyncio
+
+ send = asyncio.ensure_future(client.send(req, stream=True))
+
+ async def _watch():
+ while not await request.is_disconnected():
+ await asyncio.sleep(_DISCONNECT_POLL_S)
+
+ watch = asyncio.ensure_future(_watch())
+ try:
+ await asyncio.wait({send, watch}, return_when=asyncio.FIRST_COMPLETED)
+ if send.done():
+ return send.result()
+ # The watcher won: the client is gone. Abandon the send rather than
+ # waiting it out -- the finally below cancels the decode node.
+ send.cancel()
+ # CancelledError is a BaseException, so suppress(Exception) would
+ # let it escape and turn a clean 499 into a 500.
+ with contextlib.suppress(Exception, asyncio.CancelledError):
+ await send
+ return None
+ finally:
+ watch.cancel()
+ with contextlib.suppress(Exception, asyncio.CancelledError):
+ await watch
# ── 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)
-
try:
- prefill = await run_in_threadpool(_prefill, path, body, node)
- rid = derive_rid(prefill["id"])
+ req = gate_request(path, body, tokenizer=ctx.tokenizer, parser_active=_parser_active)
+ # capabilities() may probe over HTTP on a cache miss, so it goes
+ # off the event loop even though it is usually a dict lookup.
+ caps = await run_in_threadpool(pool.capabilities)
+ validate_generation_request(body, caps, ctx.gen_defaults.resolve({}))
+ except (CapabilityError, GrammarError, LogprobsUnsupported) as e:
+ return JSONResponse(e.to_payload(), status_code=e.http_status)
+ stop, include_stop = req.stop, req.include_stop
+ logprobs_req, grammar_spec = req.logprobs_req, req.grammar_spec
+ lease = None # taken inside the try: every handler below releases it
+ try:
+ # Shielded: a disconnect must not strand a reservation mid-acquire.
+ with anyio.CancelScope(shield=True):
+ lease, waited = await run_in_threadpool(acquire_lease, pool)
+ if lease is None:
+ return _busy_response(waited)
+ node = lease.node
+ prefill = await run_in_threadpool(_prefill, path, body, node, logprobs_req)
+ rid = lease.rid = derive_rid(prefill["id"])
first_token_id = first_token_from_logprobs(prefill, True)
+ first_lp = (
+ first_token_logprob_from_prefill(prefill, logprobs_req.top_n)
+ if logprobs_req is not None
+ else None
+ )
+ except PrefillClientError as e:
+ if lease is not None:
+ lease.release()
+ logger.info("vLLM rejected the stream request (%d): %s", e.status, str(e.payload)[:200])
+ return JSONResponse(e.payload, status_code=e.status)
except Exception as e:
- pool.release(node)
+ if lease is not None:
+ lease.release()
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 it stays busy.
+ if lease is not None:
+ lease.release()
+ raise
+
+ # Open the decode stream HERE, not inside the generator. Once the
+ # generator runs the response has begun and the status is spent, so a
+ # busy decode node could only be reported as an SSE error inside a 200 --
+ # or, as it was, a stream truncated with no terminator. Sending the
+ # request first keeps the status available for exactly the case that
+ # needs it.
+ import asyncio
+
+ import httpx
+
+ client = httpx.AsyncClient(timeout=httpx.Timeout(600, read=600))
+ decode_resp = None
+ try:
+ for attempt in range(1, _DECODE_BUSY_ATTEMPTS + 1):
+ lease.dispatched = True
+ decode_resp = await _send_watching_client(
+ client,
+ client.build_request(
+ "POST",
+ f"{node.http_base}/pd/decode",
+ json=_decode_body(
+ rid,
+ first_token_id,
+ body,
+ grammar_spec,
+ stream=True,
+ logprobs_req=logprobs_req,
+ thinking=req.thinking,
+ ),
+ ),
+ request,
+ )
+ if decode_resp is None: # client left; see the helper
+ logger.info(
+ "client disconnected while the decode node was "
+ "still holding headers for %s",
+ rid,
+ )
+ return JSONResponse(
+ {
+ "error": "client disconnected",
+ "error_type": "request_cancelled",
+ "rid": rid,
+ },
+ status_code=499,
+ )
+ if decode_resp.status_code != 429 or attempt == _DECODE_BUSY_ATTEMPTS:
+ break
+ await decode_resp.aclose()
+ logger.info(_RETRY_LOG, node.http_base, rid, _DECODE_BUSY_RETRY_S)
+ await asyncio.sleep(_DECODE_BUSY_RETRY_S)
+ assert decode_resp is not None # the loop always posts at least once
+
+ if decode_resp.status_code != 200:
+ payload = None
+ if decode_resp.status_code != 429:
+ try:
+ await decode_resp.aread()
+ payload = decode_resp.json()
+ except Exception:
+ payload = None
+ await decode_resp.aclose()
+ # Same classification and the same answers as the blocking
+ # path: only the transport above this line differs.
+ verdict = classify_decode_status(
+ decode_resp.status_code,
+ payload,
+ attempts_left=False,
+ propagated_types=PROPAGATED_ERROR_TYPES,
+ )
+ _log_refusal(verdict, node, rid, decode_resp.status_code, payload)
+ body, status = decode_refusal(verdict, decode_resp.status_code, payload, rid)
+ return JSONResponse(body, status_code=status)
+ except Exception as e:
+ logger.exception("opening the decode stream failed for %s", rid)
+ with contextlib.suppress(Exception):
+ if decode_resp is not None:
+ await decode_resp.aclose()
+ return JSONResponse({"error": str(e)}, status_code=502)
+ finally:
+ # Every early return above leaves the request unserved, so the node
+ # goes back to the pool, the decode node is told to stop, and the
+ # client is closed. The success path hands all of that to the
+ # generator instead.
+ if decode_resp is None or decode_resp.status_code != 200:
+ lease.release()
+ with contextlib.suppress(Exception):
+ await client.aclose()
chunk_id = prefill["id"]
model = prefill.get("model")
prompt_tokens = (prefill.get("usage") or {}).get("prompt_tokens")
- parser = ctx.parser(_thinking_enabled(body))
-
- def _chunk(delta: dict, finish=None, usage=None) -> str:
- payload = {
- "id": chunk_id,
- "object": "chat.completion.chunk",
- "created": int(time.time()),
- "model": model,
- "choices": [{"index": 0, "delta": delta, "finish_reason": finish}],
- }
- if usage is not None:
- payload["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"]}
- if ev["kind"] == "content":
- return {"content": ev["text"]}
- return {
- "tool_calls": [
- {
- "index": ev["index"],
- "id": ev["id"],
- "type": "function",
- "function": {"name": ev["name"], "arguments": ev["arguments"]},
- }
- ]
- }
-
- def _fire_cancel():
- try:
- requests.post(f"{node.http_base}/pd/cancel", json={"rid": rid}, timeout=5)
- except Exception:
- logger.warning("cancel POST failed for %s", rid)
+ parser = ctx.parser(req.thinking)
+ # Stamped ONCE for the whole response, as vLLM does (it threads a single
+ # `created_time` through every chunk it builds). Re-reading the clock per
+ # chunk gave one response several timestamps, which breaks a client that
+ # groups or de-duplicates by (id, created).
+ created = int(time.time())
+
+ # Every frame carries the same id / model / created, so they are bound
+ # once here and the generator passes only what varies.
+ _chunk = functools.partial(sse_chunk, chunk_id=chunk_id, model=model, created=created)
+ _usage_chunk = functools.partial(
+ usage_chunk, chunk_id=chunk_id, model=model, created=created
+ )
async def _gen():
import anyio
- import httpx
- from tilert.pd_vllm.oai_parser import IncrementalDetok
+ # Both paths drive this with the same arguments, which is why the
+ # two replies agree about text, channels, logprobs and count.
+ asm = ReplyStream(
+ ctx.tokenizer,
+ stop=stop,
+ include_stop_in_output=include_stop,
+ parser_session=parser.stream() if parser else None,
+ logprobs_req=logprobs_req,
+ first_token_logprob=first_lp,
+ )
- n_tokens = 0
- saw_tool = False
+ # Shared with the non-streaming path: this loop keeps only the
+ # transport (async) and the presentation (SSE).
+ reader = DecodeReader(stream=asm, logprobs_req=logprobs_req, rid=rid)
+ out = SseWriter(asm, _chunk)
finish_reason = "stop"
client_gone = False
- completed_ok = False
- detok = IncrementalDetok(ctx.tokenizer)
- sess = parser.stream() if parser else None
- client = httpx.AsyncClient(timeout=httpx.Timeout(600, read=600))
+ # Whether the decode node has told us it is done with this request.
+ # It owns its slot until then, so anything that leaves this loop
+ # early -- a client hanging up, a stop string, an exception -- has
+ # to cancel; a node that reported `done` must not be.
+ decode_done = False
try:
- yield _chunk({"role": "assistant"})
- async with client.stream(
- "POST",
- f"{node.http_base}/pd/decode",
- json={
- "rid": rid,
- "first_token_id": first_token_id,
- "max_tokens": _max_tokens_of(body),
- "sampling": _sampling_of(body),
- "stream": True,
- },
- ) as resp:
- resp.raise_for_status()
+ async with contextlib.aclosing(decode_resp) as resp:
async for line in resp.aiter_lines():
# Deterministic client-liveness check: writes to a
# dead socket do NOT raise (verified by drill), so
@@ -325,59 +871,92 @@ async def _gen():
client_gone = True
logger.info("client disconnected, cancelling %s", rid)
break
- if not line:
- continue
- msg = json.loads(line)
- if "t" in msg:
- n_tokens += len(msg["t"])
- text = detok.push(msg["t"])
- if not text:
- continue
- if sess is None:
- yield _chunk({"content": text})
- continue
- for ev in sess.feed(text):
- if ev["kind"] == "tool":
- saw_tool = True
- yield _chunk(_event_delta(ev))
- elif "done" in msg:
- finish_reason = msg.get("finish_reason", "stop")
- if finish_reason == "cancelled":
- finish_reason = "stop"
- elif "error" in msg:
- yield _chunk({"content": f"\n[decode error: {msg['error']}]"})
- finish_reason = "stop"
+ for frame in out.frames(reader.feed(line)):
+ yield frame
+ if reader.finished:
+ if reader.stop_hit:
+ # Complete. The node cannot see text and is
+ # still generating; the finally cancels it.
+ logger.info("stop string %r ended %s", asm.stop_reason, rid)
+ break
+ decode_done = reader.node_terminated
+ finish_reason = reader.finish_reason
+ # Same conditions and the same payloads as the blocking path,
+ # from the same function; only how they are SAID differs, and
+ # only because this 200 is already spent.
+ verdict, payload, _status = terminal_verdict(reader, client_gone=client_gone)
+ if verdict == REFUSED:
+ logger.warning(
+ "decode node %s sent an unusable logprobs " "line for %s",
+ node.http_base,
+ rid,
+ )
+ elif verdict == TRUNCATED:
+ logger.warning(
+ "decode stream for %s ended after %d tokens " "with no done/error message",
+ rid,
+ len(asm.token_ids) if asm else 0,
+ )
+ if verdict in (REFUSED, TRUNCATED, TYPED_ERROR):
+ # A typed error is the node's considered answer about this
+ # request: emitting it as a content marker and finishing
+ # normally would report a successful completion that broke
+ # the contract asked for -- unconstrained output for a
+ # grammar, or a reply without the logprobs requested.
+ for _c in out.fail_closed(payload):
+ yield _c
+ return
+ if verdict == UNTYPED_ERROR:
+ # No status left to carry it, so it goes in the text. Flush
+ # first, or the marker lands before text the reply earned:
+ # `prefix[decode error]suffix`.
+ for _c in out.flush_held():
+ yield _c
+ yield _chunk({"content": "\n[decode error: " f"{payload['error']}]"})
+ finish_reason = "stop"
if client_gone:
logger.info("client gone mid-stream for %s", rid)
return # finally fires the cancel
- if sess is not None:
- for ev in sess.finish():
- if ev["kind"] == "tool":
- saw_tool = True
- yield _chunk(_event_delta(ev))
- if saw_tool:
- finish_reason = "tool_calls"
+ for _c in out.flush_held():
+ yield _c
+ # The same decision the non-streaming path makes, from the
+ # same function.
yield _chunk(
{},
- finish=finish_reason,
- usage={
- "prompt_tokens": prompt_tokens,
- "completion_tokens": n_tokens,
- },
+ finish=reply_finish_reason(
+ saw_tool=out.saw_tool, from_node=finish_reason, stream=asm
+ ),
+ stop_reason=asm.stop_reason,
)
+ if should_include_usage(body, ctx.force_include_usage):
+ # From the reply stream, not a counter in this loop.
+ yield _usage_chunk(build_usage(prompt_tokens, asm.completion_tokens))
yield "data: [DONE]\n\n"
- completed_ok = True
- except Exception:
+ except Exception as exc:
+ # Malformed NDJSON, or a wrong-typed field that makes
+ # `reader.feed` raise. The 200 is spent, so exiting here would
+ # leave the client with a partial response and no terminator --
+ # the same failure the clean-EOF check above refuses, reached by
+ # a different route. Best effort: the client may already be gone,
+ # in which case yielding raises again and there is nothing left
+ # to say.
logger.exception("stream failed mid-flight for %s", rid)
+ with contextlib.suppress(Exception):
+ for _c in out.fail_closed(
+ {
+ "error": f"stream failed: {exc}",
+ "error_type": "decode_stream_failed",
+ "rid": rid,
+ }
+ ):
+ yield _c
finally:
# Runs under cancellation too (client disconnect cancels this
# task). Order matters: release first (sync, can't be
# cancelled), then best-effort cancel via a plain thread
# (an await here could be cancelled before firing), then a
# shielded aclose.
- pool.release(node)
- if not completed_ok:
- threading.Thread(target=_fire_cancel, daemon=True).start()
+ lease.release(terminated=decode_done)
with anyio.CancelScope(shield=True):
await client.aclose()
@@ -423,7 +1002,69 @@ def main() -> None:
"--parser",
choices=["glm47", "none"],
default="glm47",
- help="output parser (reasoning + tool calls)",
+ help="output parser (reasoning + tool calls); anything but 'none' loads "
+ "vLLM's parser engine and needs vllm importable",
+ )
+ ap.add_argument(
+ "--model",
+ default="",
+ help="model profile the decode nodes serve (glm5 / glm5_2 / "
+ "glm5_3 / dsv32); must match their --model. Used to decide "
+ "whether a repetition_penalty in the model's generation_config "
+ "can be adopted: only a profile whose decode runtime declares "
+ "penalties may adopt one (the GLM-5 / GLM-5.2 / DSV3.2 runtimes "
+ "do not). Omitted means unknown, and the conservative answer is "
+ "taken.",
+ )
+ ap.add_argument(
+ "--generation-config",
+ choices=["auto", "vllm"],
+ default="auto",
+ help="where sampling defaults come from, mirroring vLLM's "
+ "flag of the same name: 'auto' reads the model's "
+ "generation_config.json under --model-path, 'vllm' "
+ "ignores it and uses the neutral defaults. Whichever "
+ "is chosen, the resolved values are sent explicitly to "
+ "BOTH legs so they cannot disagree.",
+ )
+ ap.add_argument(
+ "--default-temperature",
+ type=float,
+ default=None,
+ help="override the resolved temperature default " "(vLLM: --override-generation-config)",
+ )
+ ap.add_argument(
+ "--default-top-p", type=float, default=None, help="override the resolved top_p default"
+ )
+ ap.add_argument(
+ "--default-top-k",
+ type=int,
+ default=None,
+ help="override the resolved top_k default; 0 disables the " "rank cut, as it does in vLLM",
+ )
+ ap.add_argument(
+ "--default-repetition-penalty",
+ type=float,
+ default=None,
+ help="override the resolved repetition_penalty default. "
+ "Only executable where --model names a family whose "
+ "decode runtime implements penalties; 1.0 is the "
+ "runtime no-op.",
+ )
+ 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)",
+ )
+ ap.add_argument(
+ "--force-include-usage",
+ action="store_true",
+ help="emit the trailing usage chunk on every stream, even "
+ "when the client omits stream_options.include_usage "
+ "(vLLM: enable_force_include_usage). Off by default; "
+ "note it makes every stream end with a choices:[] "
+ "chunk, which some clients cannot read.",
)
args = ap.parse_args()
@@ -440,7 +1081,29 @@ def main() -> None:
args.model_path, trust_remote_code=True
) # nosec B615
- ctx = RouterCtx(args.vllm_url, Pool(nodes), tokenizer, args.parser)
+ try:
+ gen_defaults = generation_defaults.load(
+ args.model_path,
+ args.generation_config,
+ model=args.model,
+ temperature=args.default_temperature,
+ top_p=args.default_top_p,
+ top_k=args.default_top_k,
+ repetition_penalty=args.default_repetition_penalty,
+ )
+ except UnsupportedGenerationDefault as e:
+ # Startup, not per request: a default the router cannot promise both legs
+ # will apply must be settled by the operator before traffic arrives.
+ raise SystemExit(f"cannot serve with these sampling defaults:\n{e}")
+
+ ctx = RouterCtx(
+ args.vllm_url,
+ Pool(nodes, args.queue_timeout),
+ tokenizer,
+ args.parser,
+ force_include_usage=args.force_include_usage,
+ gen_defaults=gen_defaults,
+ )
app = build_app(ctx)
logger.info(
"router on :%d -> vllm=%s, %d decode node(s), parser=%s",
diff --git a/tilert/pd_vllm/prefill_connector.py b/tilert/pd_vllm/prefill_connector.py
index abd4819..485c0ae 100644
--- a/tilert/pd_vllm/prefill_connector.py
+++ b/tilert/pd_vllm/prefill_connector.py
@@ -19,6 +19,11 @@
tracking, worker init, staging, background send, TCP handshake); all per-model
extraction / layout / RDMA planning is delegated to the selected model profile
(``tilert_model``, default ``glm5``).
+
+``tilert_sync_send`` (sending inside the forward window) is no longer
+supported: admission retries must be able to outlast the prefill response, so
+every send runs on the background sender thread. The key is ignored with a
+warning.
"""
import logging
@@ -38,6 +43,25 @@
logger = logging.getLogger("pd_vllm.connector")
+# _send outcomes.
+_SENT = "sent"
+_REJECTED_TRANSIENT = "rejected_transient"
+_REJECTED_PERMANENT = "rejected_permanent"
+
+# Reject reasons the receive slot leaves behind on its own: it is serving another
+# rid, or draining one it abandoned. Both end within the receiver's own socket
+# timeout, so coming back shortly is worth more than dropping the shard.
+# States the slot leaves on its own, so coming back is worth it. A decode node
+# also reports a tombstoned rid as `cancelling`: the same rid is reused when
+# vLLM reschedules a preempted request, and the node drops the tombstone as soon
+# as that retry's /pd/decode arrives. Adding a reason for that case instead
+# would be PERMANENT to every connector built before it, which is how a rolling
+# upgrade would start dropping shards.
+_TRANSIENT_REJECTS = frozenset({"busy", "cancelling"})
+
+_ADMISSION_ATTEMPTS = 5
+_ADMISSION_BACKOFF_S = 0.2
+
@dataclass
class _ReqMeta:
@@ -49,6 +73,9 @@ class _ReqMeta:
tilert_host: str
tilert_ctrl_port: int
sampling: dict | None = None
+ # Full prompt ids, so the decode node can seed its repetition-penalty prompt
+ # bitmap. Kept only when a penalty is actually requested (see _emit).
+ prompt_token_ids: list = field(default_factory=list)
@dataclass
@@ -75,9 +102,15 @@ def __init__(self, vllm_config, role, kv_cache_config=None):
extra = vllm_config.kv_transfer_config.kv_connector_extra_config or {}
self._default_host = extra.get("tilert_host")
self._default_port = int(extra.get("tilert_ctrl_port", 5556))
- self._sync_send = bool(extra.get("tilert_sync_send", False))
+ self._admission_attempts = int(extra.get("tilert_admission_attempts", _ADMISSION_ATTEMPTS))
self._max_seq = int(extra.get("tilert_max_seq_len", vllm_config.model_config.max_model_len))
self._profile = profiles.get_profile(extra.get("tilert_model", "glm5"))
+ if extra.get("tilert_sync_send"):
+ logger.warning(
+ "tilert_sync_send is no longer supported and is ignored: sends run on "
+ "the background sender thread so admission retries can outlast the "
+ "prefill response"
+ )
self._transport_name = extra.get("tilert_transport", "mooncake")
# scheduler-side
@@ -93,12 +126,11 @@ def __init__(self, vllm_config, role, kv_cache_config=None):
self._sender_thread: threading.Thread | None = None
logger.info(
- "TileRTConnector: role=%s profile=%s target=%s:%s sync=%s",
+ "TileRTConnector: role=%s profile=%s target=%s:%s",
role,
self._profile.name,
self._default_host,
self._default_port,
- self._sync_send,
)
# ══════════════════════ scheduler side ═════════════════════
@@ -114,7 +146,19 @@ def _params_of(self, new_req) -> dict | None:
sp = getattr(new_req, "sampling_params", None)
extra = getattr(sp, "extra_args", None) if sp is not None else None
if extra:
- return self._claim(extra.get("kv_transfer_params"))
+ claimed = self._claim(extra.get("kv_transfer_params"))
+ if claimed is not None:
+ # Stash whether this request needs its prompt ids shipped. It has
+ # to be read from vLLM's OWN SamplingParams: kv_transfer_params
+ # carries only {tilert_host, tilert_ctrl_port} (pd_router sets
+ # it), and the wire's `sampling` field is vestigial -- the real
+ # sampling params reach the decode node over the router's
+ # /pd/decode call, never through this connector.
+ claimed = dict(claimed)
+ claimed["_wants_prompt_ids"] = wire.wants_prompt_token_ids(
+ {"repetition_penalty": getattr(sp, "repetition_penalty", 1.0)}
+ )
+ return claimed
return None
def get_num_new_matched_tokens(self, request, num_computed_tokens):
@@ -178,6 +222,7 @@ def _emit(self, req_id, token_ids, groups, params) -> _ReqMeta:
rid=derive_rid(req_id),
num_tokens=len(token_ids),
last_prompt_token=int(token_ids[-1]),
+ prompt_token_ids=(list(token_ids) if params.get("_wants_prompt_ids") else []),
block_ids_per_group=groups,
tilert_host=host,
tilert_ctrl_port=int(params.get("tilert_ctrl_port", self._default_port)),
@@ -226,11 +271,10 @@ def _ensure_worker_ready(self) -> None:
self._transport.init(hostname)
self._transport.register(self._staging.data_ptr(), total, dev)
- if not self._sync_send:
- self._sender_thread = threading.Thread(
- target=self._sender_loop, name="tilert-pd-sender", daemon=True
- )
- self._sender_thread.start()
+ self._sender_thread = threading.Thread(
+ target=self._sender_loop, name="tilert-pd-sender", daemon=True
+ )
+ self._sender_thread.start()
logger.info(
"worker ready: rank=%d transport=%s staging=%.1f MB profile=%s",
self._tp_rank,
@@ -264,11 +308,13 @@ def wait_for_save(self):
except Exception:
logger.exception("extraction failed for %s", m.rid)
continue
- job = {"meta": m, "sections": sections, "seq": sections["seq"]}
- if self._sync_send:
- self._send(job)
- else:
- self._send_q.put(job)
+ # Always handed to the sender thread. Sending inside the forward
+ # window used to be selectable; it cannot work with admission,
+ # because the retries would then all run before the prefill
+ # response returns -- and the router cannot call /pd/decode, the
+ # only thing that clears a tombstone for a rescheduled rid, until
+ # it has.
+ self._send_q.put({"meta": m, "sections": sections, "seq": sections["seq"]})
def get_finished(self, finished_req_ids):
return None, None
@@ -279,11 +325,15 @@ def _sender_loop(self) -> None:
while True:
job = self._send_q.get()
try:
- self._send(job)
+ self._send_with_retry(job)
except Exception:
logger.exception("send failed for %s", job["meta"].rid)
- def _send(self, job: dict) -> None:
+ def _send(self, job: dict) -> str:
+ """One admission + RDMA attempt.
+
+ Returns ``_SENT``, ``_REJECTED_TRANSIENT`` or ``_REJECTED_PERMANENT``.
+ """
import socket as _socket
import time as _time
@@ -300,6 +350,17 @@ def _send(self, job: dict) -> None:
conn.connect((m.tilert_host, m.tilert_ctrl_port))
hello = wire.recv_msg(conn)
assert hello.get("magic") == wire.MAGIC, f"bad hello: {hello}"
+ # Control-plane version, checked separately from the buffer layout.
+ # A receiver that predates the admission step would never send an
+ # accept, so waiting for one would hang every request; a receiver
+ # that expects it must never be written to blind. Either mismatch is
+ # a deployment error, so it fails here rather than being guessed at.
+ remote_proto = hello.get("protocol_version", 1)
+ assert remote_proto == wire.PROTOCOL_VERSION, (
+ f"control-plane protocol mismatch: decode={remote_proto} "
+ f"vs prefill={wire.PROTOCOL_VERSION}; upgrade both ends "
+ f"together"
+ )
assert hello.get("layout_version") == self._profile.layout_version, (
f"layout version mismatch: {hello.get('layout_version')} "
f"vs {self._profile.layout_version}"
@@ -311,16 +372,49 @@ def _send(self, job: dict) -> None:
remote_max_seq = int(hello["max_seq_len"])
assert seq <= remote_max_seq, f"seq {seq} exceeds decode max_seq_len {remote_max_seq}"
- wire.send_msg(
- conn,
- {
- "rid": m.rid,
- "rank": self._tp_rank,
- "seq_len": seq,
- "last_prompt_token": m.last_prompt_token,
- "sampling": m.sampling,
- },
- )
+ msg = {
+ "rid": m.rid,
+ "rank": self._tp_rank,
+ "seq_len": seq,
+ "last_prompt_token": m.last_prompt_token,
+ "sampling": m.sampling,
+ "admission_window_s": self._admission_window(),
+ }
+ # Rank 0 only: every rank opens its own connection, and the decode
+ # side broadcasts the ids to all 8 devices itself, so sending them
+ # per rank would just multiply the payload by 8.
+ if self._tp_rank == 0 and m.prompt_token_ids:
+ msg["prompt_token_ids"] = m.prompt_token_ids
+ wire.send_msg(conn, msg)
+
+ # ADMISSION. Nothing may touch RDMA before this: the receive buffer
+ # holds one request at a time, and writing into it uninvited lands
+ # this request's KV inside whatever the decode node is currently
+ # serving. That corruption is undetectable downstream -- the victim
+ # decodes from a mix of two prompts and answers confidently -- so the
+ # reply is checked field by field rather than just for an `error`
+ # key. An unrecognised reply is a rejection.
+ ack = wire.recv_msg(conn)
+ if not ack.get("accepted"):
+ reason = ack.get("error")
+ logger.warning("decode node refused %s rank=%d: %s", m.rid, self._tp_rank, ack)
+ # busy / cancelling are states the slot leaves on its own, so
+ # the caller may come back. Anything else is about THIS request
+ # and will be refused again.
+ return _REJECTED_TRANSIENT if reason in _TRANSIENT_REJECTS else _REJECTED_PERMANENT
+ if (
+ ack.get("rid") != m.rid
+ or ack.get("rank") != self._tp_rank
+ or not isinstance(ack.get("generation"), int)
+ ):
+ logger.error(
+ "discarding %s rank=%d: admission does not match " "the request (%s)",
+ m.rid,
+ self._tp_rank,
+ ack,
+ )
+ return _REJECTED_PERMANENT
+ generation = ack["generation"]
base = self._staging.data_ptr()
srcs, dsts, lens = self._profile.rdma_plan(
@@ -328,14 +422,72 @@ def _send(self, job: dict) -> None:
)
self._transport.write(hello, srcs, dsts, lens)
- wire.send_msg(conn, {"done": True, "rid": m.rid, "rank": self._tp_rank})
+ wire.send_msg(conn, wire.done_msg(m.rid, self._tp_rank, generation))
logger.info(
- "sent %s: rank=%d seq=%d %.1f MB in %.1f ms",
+ "sent %s: rank=%d seq=%d gen=%d %.1f MB in %.1f ms",
m.rid,
self._tp_rank,
seq,
+ generation,
sum(lens) / 1e6,
1000 * (_time.time() - t0),
)
+ return _SENT
finally:
conn.close()
+
+ def _admission_window(self) -> float:
+ """Total backoff this sender will spend before giving up on a rid.
+
+ Sent with every request so the decode node can size a tombstone to
+ outlast it: each retry opens a NEW connection, so the socket timeout
+ bounds one attempt and says nothing about the sequence. Derived rather
+ than cached, so it cannot drift from the attempt count it describes.
+ """
+ return _ADMISSION_BACKOFF_S * (2 ** max(0, self._admission_attempts - 1) - 1)
+
+ def _send_with_retry(self, job: dict) -> None:
+ """Re-attempt admission while the slot is only transiently unavailable.
+
+ Without this a rank turned away -- a previous transfer still draining, a
+ router that lost its busy state over a restart -- drops its shard for
+ good, and nothing tells the router: the prefill response still succeeds
+ and `/pd/decode` then waits out its whole kv_transfer_timeout for shards
+ that will never arrive. Retrying absorbs the short races; a rank that is
+ still refused after the last attempt is logged with the consequence
+ named, because this connector has no path back to the router to fail the
+ request properly.
+
+ Runs on the sender thread, never inside a forward window: the sleeps
+ below must not stall a vLLM step, and admission has to be able to
+ outlast the prefill response, which is what clears a tombstone for a
+ rescheduled rid.
+ """
+ import time as _time
+
+ m = job["meta"]
+ delay = _ADMISSION_BACKOFF_S
+ for attempt in range(1, self._admission_attempts + 1):
+ outcome = self._send(job)
+ if outcome != _REJECTED_TRANSIENT:
+ return
+ if attempt == self._admission_attempts:
+ break
+ logger.info(
+ "retrying admission for %s rank=%d in %.1fs " "(attempt %d/%d)",
+ m.rid,
+ self._tp_rank,
+ delay,
+ attempt,
+ self._admission_attempts,
+ )
+ _time.sleep(delay)
+ delay *= 2
+ logger.error(
+ "gave up admitting %s rank=%d after %d attempts: its shard was "
+ "NOT transferred, so the decode node will wait out its "
+ "kv_transfer_timeout for this request",
+ m.rid,
+ self._tp_rank,
+ self._admission_attempts,
+ )
diff --git a/tilert/pd_vllm/presentation.py b/tilert/pd_vllm/presentation.py
new file mode 100644
index 0000000..5045a5a
--- /dev/null
+++ b/tilert/pd_vllm/presentation.py
@@ -0,0 +1,298 @@
+"""One reply, in two presentations.
+
+Both response paths consume the same :class:`reply.Emission` values, so they
+agree about text, channels, logprob entries and the token count by construction.
+What they did NOT share was the policy applied to those emissions -- the
+``tool_calls`` over ``stop`` precedence, ``stop_reason``, where ``usage`` comes
+from, and how a channel becomes a field. Each path spelled that out for itself
+and they agreed because a test drove one request through both and compared the
+whole reply.
+
+Here the policy is decided once (:func:`finish_reason`, :func:`collect`) and
+rendered twice: :func:`blocking_envelope` builds the single JSON body,
+:func:`sse_chunk` the frames. A difference between the two presentations is now
+a difference in this file.
+
+Free of HTTP and asyncio: dicts and strings out; the caller wraps them in
+whatever response class it uses.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any
+
+from tilert.pd_vllm.reply import (
+ CONTENT,
+ REASONING,
+ TOOL_CALL,
+ as_logprobs,
+)
+
+__all__ = [
+ "Collected",
+ "SseWriter",
+ "blocking_envelope",
+ "collect",
+ "finish_reason",
+ "sse_chunk",
+ "sse_delta",
+ "usage_chunk",
+]
+
+
+@dataclass
+class Collected:
+ """Emissions, gathered per channel, as the blocking body needs them."""
+
+ content: str = ""
+ reasoning: str = ""
+ entries: list = field(default_factory=list)
+ tool_calls: list = field(default_factory=list)
+
+
+def collect(emissions) -> Collected:
+ """Split emissions by channel.
+
+ Logprob entries follow content, which is what #22 means by ``logprobs`` covering
+ ``message.content``.
+ """
+ out = Collected()
+ for e in emissions:
+ if e.channel == CONTENT:
+ out.content += e.text
+ out.entries += e.logprobs
+ elif e.channel == REASONING:
+ out.reasoning += e.text
+ else:
+ out.tool_calls.append(e.tool_call)
+ return out
+
+
+def finish_reason(*, saw_tool: bool, from_node: str, stream) -> str:
+ """What ended the reply, in vLLM's precedence at this layer.
+
+ ``tool_calls`` outranks a matched stop, as it does in vLLM's own serving
+ path: it is set whenever the tool parser extracted calls, ahead of the
+ engine's reason. A client that needs to know the reply was cut short reads
+ ``stop_reason`` instead. A stop then outranks what the NODE said, since the
+ node cannot see text and cannot reach that conclusion.
+ """
+ if saw_tool:
+ return "tool_calls"
+ return stream.finish_reason(from_node)
+
+
+def _tool_calls_field(tool_calls: list) -> list:
+ return [
+ {
+ "index": c["index"],
+ "id": c["id"],
+ "type": "function",
+ "function": {"name": c["name"], "arguments": c["arguments"]},
+ }
+ for c in tool_calls
+ ]
+
+
+def blocking_choice(
+ got: Collected,
+ *,
+ is_chat: bool,
+ stream,
+ from_node: str,
+ logprobs_asked: bool,
+ token_ids: list[int],
+) -> tuple[dict, int]:
+ """The one ``choices[0]`` of a non-streamed reply, and its token count."""
+ choice: dict = {
+ "index": 0,
+ "finish_reason": finish_reason(
+ saw_tool=bool(got.tool_calls), from_node=from_node, stream=stream
+ ),
+ "logprobs": as_logprobs(got.entries) if logprobs_asked else None,
+ # Non-null when a stop string ended the reply. Names which one, which a
+ # client cannot recover from the text when the string was cut out of it.
+ "stop_reason": stream.stop_reason,
+ }
+ if is_chat:
+ msg: dict[str, Any] = {"role": "assistant", "content": got.content}
+ if got.reasoning:
+ msg["reasoning_content"] = got.reasoning
+ if got.tool_calls:
+ msg["tool_calls"] = _tool_calls_field(got.tool_calls)
+ choice["message"] = msg
+ else:
+ choice["text"] = got.content
+ # Every id that ran, including those whose text a stop removed:
+ # `token_ids` reports what ran and `text` what came back. vLLM's
+ # `CompletionOutput.token_ids` is untruncated too.
+ choice["token_ids"] = token_ids
+ return choice, stream.completion_tokens
+
+
+def textless_choice(*, is_chat: bool, from_node: str, token_ids: list[int]) -> tuple[dict, int]:
+ """``--parser none`` with no ``--model-path``: no tokenizer, so no text.
+
+ No stop matching, no channels and no logprobs are possible; token ids are
+ all the reply can carry, and only ``/v1/completions`` exposes them.
+ """
+ choice: dict = {"index": 0, "finish_reason": from_node, "logprobs": None, "stop_reason": None}
+ if is_chat:
+ choice["message"] = {"role": "assistant", "content": None}
+ else:
+ choice["text"] = None
+ choice["token_ids"] = token_ids
+ return choice, len(token_ids)
+
+
+def blocking_envelope(
+ choice: dict, *, is_chat: bool, prefill: dict, created: int, usage: dict, timing: dict
+) -> dict:
+ """The whole non-streamed body."""
+ return {
+ "id": prefill["id"],
+ "object": "chat.completion" if is_chat else "text_completion",
+ "created": created,
+ "model": prefill.get("model"),
+ "choices": [choice],
+ "usage": usage,
+ # Empty apart from `prefill` when a stop string ended the reply: the node
+ # reports its timings on the `done` line, which we stopped reading
+ # before.
+ "pd_timing_ms": timing,
+ }
+
+
+def sse_delta(e) -> dict:
+ """One emission as an OpenAI streaming delta."""
+ if e.channel == REASONING:
+ return {"reasoning_content": e.text}
+ if e.channel == CONTENT:
+ return {"content": e.text}
+ return {"tool_calls": _tool_calls_field([e.tool_call])}
+
+
+def sse_chunk(
+ delta: dict,
+ *,
+ chunk_id: str,
+ model: Any,
+ created: int,
+ finish: str | None = None,
+ usage: dict | None = None,
+ logprobs: dict | None = None,
+ stop_reason: str | None = None,
+) -> str:
+ """One ``chat.completion.chunk`` frame."""
+ choice: dict = {"index": 0, "delta": delta, "finish_reason": finish}
+ if logprobs is not None:
+ choice["logprobs"] = logprobs
+ if finish is not None:
+ # Only on the chunk that closes the choice, which is where a client
+ # reads it, and where vLLM puts it.
+ choice["stop_reason"] = stop_reason
+ payload = {
+ "id": chunk_id,
+ "object": "chat.completion.chunk",
+ "created": created,
+ "model": model,
+ "choices": [choice],
+ }
+ if usage is not None:
+ payload["usage"] = usage
+ return _frame(payload)
+
+
+def usage_chunk(usage: dict, *, chunk_id: str, model: Any, created: int) -> str:
+ """The trailing usage-only chunk, the shape vLLM and OpenAI emit.
+
+ ``choices: []`` on a chunk of its own, not on the one that closes the
+ choice: clients read usage off the final chunk and stop at the first one
+ bearing a finish_reason.
+ """
+ return _frame(
+ {
+ "id": chunk_id,
+ "object": "chat.completion.chunk",
+ "created": created,
+ "model": model,
+ "choices": [],
+ "usage": usage,
+ }
+ )
+
+
+def _frame(payload: dict) -> str:
+ return f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
+
+
+class SseWriter:
+ """The order an SSE reply has to come out in.
+
+ Every rule here is one a client noticed when it was missing:
+
+ * the opening ``delta.role`` chunk goes out ONCE, and only if something
+ follows it (#34: a request that emits nothing sends no role either);
+ * an emission carrying neither text nor entries produces no chunk -- the
+ parser consuming a tag, say -- because a chunk for it is noise;
+ * whatever the reply stream still holds is flushed before any marker, or the
+ marker lands before text the reply earned: ``prefix[decode error]suffix``;
+ * a stream that cannot be completed correctly still ends properly. The 200 is
+ spent, so the only honest ending is the reply's own text, a finish_reason, a
+ typed error event, ``[DONE]`` -- not a truncated stream, and not more
+ content.
+
+ ``saw_tool`` is a fact about the reply the caller needs afterwards: it
+ outranks a stop in the finish reason.
+
+ The methods are plain generators, not async ones: ``yield from`` is not
+ allowed inside an async generator, so the caller iterates them.
+ """
+
+ def __init__(self, stream, chunk):
+ self._stream = stream # the ReplyStream
+ self._chunk = chunk # a bound sse_chunk
+ self.role_sent = False
+ self.saw_tool = False
+
+ def role_once(self) -> str:
+ self.role_sent = True
+ return self._chunk({"role": "assistant"})
+
+ def emit(self, e) -> str | None:
+ """One emission as a frame, or None if it carries nothing."""
+ if e.channel == TOOL_CALL:
+ self.saw_tool = True
+ elif not e.text and not e.logprobs:
+ return None
+ return self._chunk(sse_delta(e), logprobs=(as_logprobs(e.logprobs) if e.logprobs else None))
+
+ def frames(self, emissions):
+ """Emissions as frames, with the role chunk first if it is still owed."""
+ for e in emissions:
+ frame = self.emit(e)
+ if frame is None:
+ continue
+ if not self.role_sent:
+ yield self.role_once()
+ yield frame
+
+ def flush_held(self):
+ """What the reply stream still holds, then the role if still owed.
+
+ Safe to call twice: ``finish()`` is idempotent and ``role_once`` sets its
+ own flag. With a stop the matcher may be holding up to ``len(stop) - 1``
+ characters the reply earned.
+ """
+ yield from self.frames(self._stream.finish())
+ if not self.role_sent:
+ yield self.role_once()
+
+ def fail_closed(self, payload: dict):
+ """End a stream that cannot be completed correctly."""
+ yield from self.flush_held()
+ yield self._chunk({}, finish="stop")
+ yield _frame({"error": payload})
+ yield "data: [DONE]\n\n"
diff --git a/tilert/pd_vllm/profiles/base.py b/tilert/pd_vllm/profiles/base.py
index f1231d7..068970e 100644
--- a/tilert/pd_vllm/profiles/base.py
+++ b/tilert/pd_vllm/profiles/base.py
@@ -4,8 +4,18 @@
decode server orchestration, router) is model-agnostic and calls into the
active profile for the parts that differ between models:
- GLM-5 : replicated MLA latent KV + NSA KI index + MTP draft
- DeepSeek-V3.2 : replicated MLA latent KV + NSA KI index + MTP draft
+ GLM-5 / GLM-5.2 / GLM-5.3 : replicated MLA latent KV + NSA KI index + MTP draft
+ DeepSeek-V3.2 : replicated MLA latent KV + NSA KI index + MTP draft
+
+A profile owns four concerns:
+ 1. receive-buffer layout (decode node) + hello fields advertising it
+ 2. prefill-side extraction (kv_caches -> staged bytes) + the RDMA plan
+ 3. decode-side convert (received bytes -> native tensors) + inject
+ 4. engine construction + decode loop (MTP style differs per model)
+
+Registration returned by ``classify_layers`` and ``sections`` returned by
+``extract`` are opaque profile-internal objects threaded back by the
+framework — they never cross the profile boundary interpreted.
"""
from __future__ import annotations
@@ -16,6 +26,9 @@
class ModelProfile(Protocol):
name: str
num_ranks: int
+ # TP ranks that actually RDMA-send (the framework counts `done` against
+ # this set and skips extraction on other ranks). MLA-family profiles: {0}
+ # (the MLA latent is replicated across TP).
sender_ranks: frozenset
@property
@@ -64,12 +77,33 @@ def build_engine(
) -> Any:
"""Construct the decode engine adapter (inject/decode/reset)."""
+ # ── optional hooks (checked with hasattr; not every profile has them) ─
+ # configure(kv_cache_dtype) : MLA family — cache dtype sizes the buffer
+ # configure_weights(weights_dir) : members whose depth comes from the
+ # checkpoint; called on the decode node
+ # BEFORE buffer_bytes/hello_layout.
+ # declares_penalties: bool : whether the decode runtime implements
+ # repetition/presence penalties. Read by
+ # the router at startup (no engine exists
+ # yet to ask); absent means "no".
+
_REGISTRY: dict[str, ModelProfile] = {}
_ALIASES = {
"glm5": "glm5",
"glm_5": "glm5",
"glm-5": "glm5",
+ "glm5_2": "glm5_2",
+ "glm_5_2": "glm5_2",
+ "glm-5.2": "glm5_2",
+ "glm5.2": "glm5_2",
+ "glm52": "glm5_2",
+ # GLM-5.3 shares GLM-5.2's base model, config and PD data plane.
+ "glm5_3": "glm5_2",
+ "glm_5_3": "glm5_2",
+ "glm-5.3": "glm5_2",
+ "glm5.3": "glm5_2",
+ "glm53": "glm5_2",
"dsv32": "dsv32",
"deepseek_v3_2": "dsv32",
"deepseek-v3.2": "dsv32",
@@ -88,6 +122,8 @@ def get_profile(name: str) -> ModelProfile:
# lazy import so a profile's heavy deps load only when selected
if canon == "glm5":
from tilert.pd_vllm.profiles import glm5 # noqa: F401
+ elif canon == "glm5_2":
+ from tilert.pd_vllm.profiles import glm5_2 # noqa: F401
elif canon == "dsv32":
from tilert.pd_vllm.profiles import dsv32 # noqa: F401
if canon not in _REGISTRY:
diff --git a/tilert/pd_vllm/profiles/glm5_2.py b/tilert/pd_vllm/profiles/glm5_2.py
new file mode 100644
index 0000000..68bdda8
--- /dev/null
+++ b/tilert/pd_vllm/profiles/glm5_2.py
@@ -0,0 +1,108 @@
+"""GLM-5.2 / GLM-5.3 profile — thin config over the shared MLA+NSA data plane.
+
+GLM-5.3 is a post-training update of GLM-5.2 (same base model, same config,
+same 79-layer plane); ``glm5_3`` and friends alias to this profile.
+
+GLM-5.2 has the same PD data plane as GLM-5 (MLA latent KV + NSA KI index +
+1 MTP draft layer) with ONE difference: the DSA indexer is sparsified — only
+the "full" layers carry an indexer/KI cache; the "shared" layers reuse the
+previous full layer's top-k at runtime and have NO KI cache.
+
+vLLM registers a KI (indexer) cache only on the full layers (HF
+``config.indexer_types``, mirrored by vLLM's ``_skip_topk`` in the DeepSeek-V3.2
+model code and by the engine's own full-layer rule):
+
+ full layer <=> max(L - 2, 0) % 4 == 0 -> {0,1,2,6,10,...,74} (21)
+ MTP tail (layer 78) is always full -> + {78} (= 22)
+
+The KV and PE planes still cover all 79 layers (every layer does MLA
+attention). The 22-vs-79 KI difference is absorbed entirely in
+``MlaNsaProfile.classify_layers``: the registered full-layer KI list is
+expanded to 79 entries (each shared layer references the previous full
+layer's KI cache tensor), so extract / rdma_plan / convert / inject /
+buffer_bytes stay layer-uniform and unchanged. ``ki_layer_ids`` below turns on
+a strict check that vLLM's registered KI set equals exactly this set (the MTP
+tail may be absent when prefill runs without --speculative-config).
+
+Engine selection: on a ROCm torch build (``torch.version.hip``), or when
+``TILERT_PD_ENGINE_BACKEND=rocm``, the decode engine is the ROCm adapter in
+``glm5_rocm_engine``; otherwise the CUDA engine ``tilert.models.glm_5_2`` is
+imported lazily and a clear ``ImportError`` names it when the installed tilert
+build does not ship it.
+"""
+
+from __future__ import annotations
+
+import os
+
+from tilert.pd_vllm.profiles import base
+from tilert.pd_vllm.profiles.glm5_rocm_engine import build_rocm_engine, is_rocm_torch
+from tilert.pd_vllm.profiles.mla_nsa import MlaNsaEngineAdapter, MlaNsaProfile
+
+NUM_LAYERS = 79 # 78 main + 1 MTP draft layer (same skeleton as GLM-5)
+LAYOUT_VERSION = 12 # glm5_2 wire family (distinct from glm5's 10, dsv32's 11)
+
+# Full (indexer-carrying) layers: dense (0,1,2) + every-4th MoE (6,10,...,74)
+# + MTP tail (78). Mirrors the engine's full-layer rule (max(L-2,0)%4==0).
+FULL_LAYERS = [L for L in range(78) if max(L - 2, 0) % 4 == 0] + [78]
+
+
+def _use_rocm_engine() -> bool:
+ """Select the ROCm adapter on a HIP torch build or when forced by env; CUDA otherwise."""
+ backend = os.environ.get("TILERT_PD_ENGINE_BACKEND", "").strip().lower()
+ if backend in ("rocm", "hip"):
+ return True
+ if backend == "cuda":
+ return False
+ if backend:
+ raise ValueError(f"TILERT_PD_ENGINE_BACKEND={backend!r}; want 'rocm' or 'cuda'")
+ return is_rocm_torch()
+
+
+def _build_cuda_engine(model_weights_dir, max_seq_len, with_mtp):
+ try:
+ import tilert
+
+ # multi-backend builds load the per-model .so on demand; single-backend
+ # builds auto-register on import and lack load_backend.
+ if hasattr(tilert, "load_backend"):
+ tilert.load_backend("glm5_2")
+ from tilert.models.glm_5_2.generator import GLM5_2Generator
+ from tilert.models.glm_5_2.model_args import ModelArgsGLM5_2
+ except (ImportError, ValueError) as e:
+ raise ImportError(
+ "the installed tilert build ships no CUDA GLM-5.2 engine "
+ "(tilert.models.glm_5_2); use a tilert build that includes it, or "
+ "a ROCm tilert build (selected automatically on a HIP torch, or with "
+ "TILERT_PD_ENGINE_BACKEND=rocm)"
+ ) from e
+
+ gen = GLM5_2Generator(
+ model_args=ModelArgsGLM5_2(),
+ max_new_tokens=max(max_seq_len - 256, 4096 - 256),
+ model_weights_dir=model_weights_dir,
+ with_mtp=with_mtp,
+ use_topp=True,
+ enable_thinking=False,
+ )
+ gen.from_pretrained()
+ return MlaNsaEngineAdapter(gen, with_mtp)
+
+
+def _build_engine(model_weights_dir, max_seq_len, with_mtp, ar_steps):
+ # The ROCm tilert build serves GLM-5.2/5.3 through a different engine API
+ # (no inject_cache); same PD data plane, so only the engine adapter differs.
+ if _use_rocm_engine():
+ return build_rocm_engine(model_weights_dir, max_seq_len, with_mtp, ar_steps)
+ return _build_cuda_engine(model_weights_dir, max_seq_len, with_mtp)
+
+
+base.register(
+ MlaNsaProfile(
+ name="glm5_2",
+ num_layers=NUM_LAYERS,
+ layout_version=LAYOUT_VERSION,
+ engine_factory=_build_engine,
+ ki_layer_ids=FULL_LAYERS,
+ )
+)
diff --git a/tilert/pd_vllm/profiles/glm5_rocm_engine.py b/tilert/pd_vllm/profiles/glm5_rocm_engine.py
new file mode 100644
index 0000000..f1b5e5e
--- /dev/null
+++ b/tilert/pd_vllm/profiles/glm5_rocm_engine.py
@@ -0,0 +1,307 @@
+"""ROCm (MI350X / MI355X) engine adapter for the GLM-5.2 / GLM-5.3 PD plane.
+
+The ROCm tilert build serves GLM-5.2 (and GLM-5.3, same base model and
+config) through ``tilert.models.glm_5``: ``Glm52Generator`` owns one
+``Glm52ShowHands`` -- a single-process TP8 e2e whose API is
+``prefill / step / decode_n / mtp_n / seed_draft / accepted_tokens /
+set_cur_pos / reset_sequence / update_sampling``. Unlike the CUDA engine it
+has no ``inject_cache``; its caches are plain per-rank tensors:
+
+ rank 0 : [kv, pe] x n_layers (pure-MLA TP8 mode only; the full-layer
+ pairs are never read) followed by one ki cache per FULL
+ indexer layer (+1 for the MTP block when num_mtp > 0)
+ ranks 1..7 : [kv, pe] x (n_layers + 1) -- the last pair is the MTP block's
+
+ kv : [B=1, max_seq_len, 512] bf16 pe : [B=1, max_seq_len, 64] bf16
+ ki : [B=1, max_seq_len, 128] bf16, or with TILERT_GLM5_FP8_KI=1 one flat u8
+ plane [B*L*128 fp8 e4m3 | B*L f32 per-token scales]
+
+``cur_pos`` is the NEXT write row: after a prompt of P tokens the engine sits
+at cur_pos == P, so after injecting P rows we ``set_cur_pos(P)`` and feed the
+prefill-sampled first token exactly where the engine's own generate() would
+be after its prompt prefill. The decode loop below is that generate() tail
+(MTP: seed_draft + chained mtp_n; plain: one forced step + decode_n), reading
+the emitted stream back from the AR flat buffer.
+
+GPU-VERIFY (not checkable on the build machine): the ki fp8 plane layout and
+e4m3 flavour, that reset_sequence leaves the caches intact, and that the MTP
+block tolerates a cold last_hidden on the first verify step (the CUDA adapter
+has the same warm-up situation and it is benign there).
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+
+import torch
+
+from tilert.pd_vllm.grammar_spec import GrammarBackendUnavailable
+from tilert.pd_vllm.sampling import resolve_top_p
+
+logger = logging.getLogger("pd_vllm.profile.glm5_rocm")
+
+_FP8_MAX = 448.0 # OCP e4m3fn range (CDNA4 fp8)
+_INDEX_HEAD_DIM = 128
+
+
+def is_rocm_torch() -> bool:
+ """True when the installed torch is a ROCm build (HIP behind torch.cuda)."""
+ return getattr(torch.version, "hip", None) is not None
+
+
+def build_rocm_engine(model_weights_dir, max_seq_len, with_mtp, ar_steps, num_mtp=3):
+ """Load the ROCm GLM-5.2/5.3 engine and wrap it as a PDEngine."""
+ # Shipped by the ROCm tilert build; absent from the CUDA package tree.
+ from tilert.models.glm_5.generator import Glm52Generator # type: ignore[attr-defined]
+ from tilert.models.glm_5.model_args import ModelArgsGlm52 # type: ignore[attr-defined]
+
+ if with_mtp and num_mtp not in (1, 3):
+ raise ValueError(f"the ROCm GLM-5.2 engine builds MTP at depth 1 or 3, " f"not {num_mtp}")
+ gen = Glm52Generator(
+ model_weights_dir=model_weights_dir,
+ model_args=ModelArgsGlm52(),
+ max_new_tokens=max(max_seq_len - 256, 4096 - 256),
+ use_topp=True,
+ num_mtp=num_mtp if with_mtp else 0,
+ max_seq_len=max_seq_len,
+ )
+ gen.from_pretrained()
+ return RocmGlm52EngineAdapter(gen, with_mtp, ar_steps=ar_steps)
+
+
+class RocmGlm52EngineAdapter:
+ """PDEngine over the ROCm ``Glm52Generator`` (see module docstring)."""
+
+ def __init__(
+ self,
+ generator,
+ with_mtp: bool,
+ ar_steps: int = 8,
+ *,
+ pure_tp8: bool | None = None,
+ fp8_ki: bool | None = None,
+ ):
+ self.gen = generator
+ self.dl = generator.decode_layer
+ self.with_mtp = bool(with_mtp) and self.dl.num_mtp > 0
+ if with_mtp and not self.with_mtp:
+ raise ValueError("--with-mtp requested but the engine was built " "with num_mtp=0")
+ self.max_seq_len = int(self.dl.args.max_seq_len)
+ self.n_layers = int(self.dl.n_layers)
+ self.npes = int(self.dl.npes)
+ self.mtp_seq_len = self.dl.num_mtp + 1
+ self.ar_steps = max(1, min(1024, int(os.environ.get("GLM5_AR_N", str(ar_steps)))))
+ self.stop_ids = {int(t) for t in generator.stop_token_ids}
+ self.last_stats: dict = {}
+ self._ignore_eos = False
+ self._seq_len = 0
+ # Cache geometry knobs. Read from the engine's own helpers so the
+ # adapter follows whatever the C++ side was configured with; the
+ # keyword overrides exist for the CPU tests.
+ from tilert.models.glm_5.model_args import ( # type: ignore[attr-defined]
+ full_layer_ordinals,
+ )
+ from tilert.models.glm_5.weight_converter import ( # type: ignore[import-not-found]
+ fp8_ki_enabled,
+ pure_tp8_enabled,
+ )
+
+ self._full_layers = list(full_layer_ordinals(self.n_layers))
+ self._pure_tp8 = pure_tp8_enabled() if pure_tp8 is None else pure_tp8
+ self._fp8_ki = fp8_ki_enabled() if fp8_ki is None else fp8_ki
+
+ # ── capabilities ────────────────────────────────────────────────────
+ def supports_logprobs(self) -> bool:
+ return False
+
+ def supports_penalties(self) -> bool:
+ return False
+
+ def supports_ignore_eos(self) -> bool:
+ return True
+
+ def prepare_grammar(self, grammar_spec, enable_thinking=True):
+ if grammar_spec is None:
+ return
+ raise GrammarBackendUnavailable(
+ "constrained decoding is not available on the ROCm GLM engine "
+ "(no grammar bitmask path in this build)"
+ )
+
+ # ── inject ──────────────────────────────────────────────────────────
+ def _write_ki(self, slot: torch.Tensor, ki: torch.Tensor, seq: int) -> None:
+ """Write ``ki`` [seq,128] bf16 into one rank-0 ki cache slot.
+
+ ``ki`` is already Hadamard-rotated by the profile's convert.
+ """
+ if not self._fp8_ki:
+ slot[0, :seq].copy_(ki.to(slot.device, non_blocking=True))
+ return
+ # Flat u8 plane: [L*128 fp8 bytes | L f32 scales] (B == 1).
+ L = self.max_seq_len
+ nbytes = L * _INDEX_HEAD_DIM
+ x = ki.to(slot.device).float()
+ amax = x.abs().amax(dim=-1).clamp_(min=1e-12)
+ scale = amax / _FP8_MAX # [seq] f32
+ q = (x / scale.unsqueeze(-1)).clamp_(-_FP8_MAX, _FP8_MAX)
+ q = q.to(torch.float8_e4m3fn)
+ slot[:nbytes].view(torch.float8_e4m3fn).view(L, _INDEX_HEAD_DIM)[:seq].copy_(q)
+ slot[nbytes : nbytes + L * 4].view(torch.float32)[:seq].copy_(scale)
+
+ def inject(self, req) -> None:
+ dl = self.dl
+ layers = req.layers
+ seq = int(req.seq_len)
+ n_extra = 1 if dl.num_mtp > 0 else 0
+ if len(layers) not in (self.n_layers, self.n_layers + 1):
+ raise RuntimeError(
+ f"glm5_rocm inject: got {len(layers)} layers, engine has "
+ f"{self.n_layers} (+1 MTP block)"
+ )
+ if seq <= 0 or seq > self.max_seq_len:
+ raise RuntimeError(
+ f"glm5_rocm inject: seq_len {seq} outside " f"(0, {self.max_seq_len}]"
+ )
+ if n_extra and len(layers) == self.n_layers:
+ raise RuntimeError(
+ "glm5_rocm inject: engine has an MTP block but "
+ "the prefill sent no MTP-layer KV (prefill must "
+ "run with --speculative-config mtp)"
+ )
+ # The profile always ships n_layers + 1 (the MTP tail); an engine
+ # built without MTP has no slot for it.
+ use = layers[: self.n_layers + n_extra]
+
+ dl.reset_sequence() # AR buffers + cur_pos; caches stay
+ for rank in range(self.npes):
+ if rank == 0 and not self._pure_tp8:
+ continue # TP7 arm: rank 0 holds no kv/pe
+ caches = dl._caches[rank]
+ n_pairs = self.n_layers + (n_extra if rank != 0 else 0)
+ for lid in range(min(len(use), n_pairs)):
+ _ki, kv, pe = use[lid]
+ caches[2 * lid][0, :seq].copy_(kv, non_blocking=True)
+ caches[2 * lid + 1][0, :seq].copy_(pe, non_blocking=True)
+ caches0 = dl._caches[0]
+ ki_base = 2 * self.n_layers if self._pure_tp8 else 0
+ ki_layers = self._full_layers + ([self.n_layers] if n_extra else [])
+ for ki_slot, lid in enumerate(ki_layers):
+ if lid >= len(use):
+ break
+ self._write_ki(caches0[ki_base + ki_slot], use[lid][0], seq)
+ torch.cuda.synchronize()
+ dl.set_cur_pos(seq)
+ self._seq_len = seq
+
+ # ── decode ──────────────────────────────────────────────────────────
+ def decode(
+ self,
+ first_token_id,
+ max_tokens,
+ sampling,
+ on_token=None,
+ cancel_event=None,
+ grammar_session=None,
+ top_logprobs=None,
+ ):
+ if grammar_session is not None:
+ raise GrammarBackendUnavailable(
+ "constrained decoding is not available on the ROCm GLM engine"
+ )
+ if top_logprobs:
+ raise NotImplementedError("logprobs are not available on the ROCm GLM engine")
+ sampling = sampling or {}
+ rep = float(sampling.get("repetition_penalty", 1.0) or 1.0)
+ presence = float(sampling.get("presence_penalty", 0.0) or 0.0)
+ if rep != 1.0 or presence != 0.0:
+ raise NotImplementedError(
+ "repetition/presence penalties are not supported by this " "model's decode runtime"
+ )
+ temp = float(sampling.get("temperature", 1.0))
+ if temp < 1e-5:
+ self.dl.update_sampling(False, 1.0, 1.0) # greedy arm
+ else:
+ self.dl.update_sampling(True, temp, resolve_top_p(sampling))
+ self._ignore_eos = bool(sampling.get("ignore_eos"))
+ first = int(first_token_id)
+ budget = min(int(max_tokens), self.max_seq_len - self._seq_len - 1)
+ if budget <= 0:
+ self.last_stats = {"finish_reason": "length"}
+ return [first]
+ stop_ids = set() if self._ignore_eos else self.stop_ids
+ if first in stop_ids:
+ self.last_stats = {"finish_reason": "stop"}
+ return []
+ tokens = [first]
+ if on_token:
+ on_token(first)
+ if self.with_mtp:
+ finish = self._decode_mtp(first, budget, tokens, stop_ids, on_token, cancel_event)
+ else:
+ finish = self._decode_plain(first, budget, tokens, stop_ids, on_token, cancel_event)
+ self.last_stats = {"finish_reason": finish}
+ return tokens
+
+ def _emit(self, new, tokens, budget, stop_ids, on_token):
+ """Append ``new`` to ``tokens`` honouring stop / budget.
+
+ Returns 'stop' | 'length' | None (keep going).
+ """
+ for tok in new:
+ tok = int(tok)
+ if tok in stop_ids:
+ return "stop"
+ if len(tokens) >= budget:
+ return "length"
+ tokens.append(tok)
+ if on_token:
+ on_token(tok)
+ return "length" if len(tokens) >= budget else None
+
+ def _decode_mtp(self, first, budget, tokens, stop_ids, on_token, cancel_event):
+ dl = self.dl
+ pos_limit = self.max_seq_len
+ mtp_seq = self.mtp_seq_len
+ chain_slack = max(0, dl.num_mtp - 1)
+ # First draft is unknown: seed it with the token itself; a wrong draft
+ # is simply rejected (one accepted token that step).
+ dl.seed_draft(first, first)
+ base = dl.accepted_count
+ produced = 0
+ while True:
+ if cancel_event is not None and cancel_event.is_set():
+ return "cancelled"
+ room = pos_limit - (self._seq_len + produced)
+ k = min(self.ar_steps, (room - chain_slack) // mtp_seq)
+ if k < 1:
+ return "length"
+ got = int(dl.mtp_n(k))
+ new = dl.accepted_tokens(base + produced)
+ produced += got
+ verdict = self._emit(new, tokens, budget, stop_ids, on_token)
+ if verdict:
+ return verdict
+
+ def _decode_plain(self, first, budget, tokens, stop_ids, on_token, cancel_event):
+ dl = self.dl
+ pos_limit = self.max_seq_len
+ base = dl.accepted_count
+ dl.step(first) # row seq_len <- first, samples t1
+ produced = 0
+ while True:
+ new = dl.accepted_tokens(base + produced)
+ produced += len(new)
+ verdict = self._emit(new, tokens, budget, stop_ids, on_token)
+ if verdict:
+ return verdict
+ if cancel_event is not None and cancel_event.is_set():
+ return "cancelled"
+ room = pos_limit - (self._seq_len + 1 + produced)
+ if room < 1:
+ return "length"
+ n = min(8, budget - len(tokens), room)
+ dl.decode_n(n)
+
+ def reset(self) -> None:
+ pass
diff --git a/tilert/pd_vllm/profiles/mla_nsa.py b/tilert/pd_vllm/profiles/mla_nsa.py
index a270f7c..c184f74 100644
--- a/tilert/pd_vllm/profiles/mla_nsa.py
+++ b/tilert/pd_vllm/profiles/mla_nsa.py
@@ -1,14 +1,59 @@
-"""Shared MLA + NSA-KI data plane for the DeepSeek-family models."""
+"""Shared MLA + NSA-KI data plane for the DeepSeek-family models.
+
+GLM-5 and DeepSeek-V3.2 have the same PD data plane — MLA latent KV
+(kv_lora_rank=512 + qk_rope_head_dim=64), an NSA KI index (index_head_dim=128,
+FP8+scale, 8448 B/page), and one MTP draft layer — differing only in layer
+count and the engine generator class. Both are expressed as thin configs of
+``MlaNsaProfile`` + ``MlaNsaEngineAdapter`` below.
+
+Data plane (rank-0 only — MLA latent is replicated across TP):
+ KV plane : num_layers x [max_seq][kv_bpt] u8 (fp8: 528, bf16: 1024)
+ PE plane : num_layers x [max_seq][128] u8 (64 bf16 k_pe, both dtypes)
+ KI plane : num_layers x [max_pages][8448] u8 (FP8 index + FP32 scale)
+
+The MLA cache dtype is a launch choice (vLLM ``--kv-cache-dtype``), independent
+of the (fp8) model weights — both paths are supported and selected at runtime:
+
+ fp8_ds_mla (recommended, SGLang-aligned): cache [nblk, page, 656] u8; per
+ token 512 fp8 kv_c + 16 B (4 fp32) scale + 128 B bf16 k_pe. Prefill splits
+ raw 656 -> 528-B kv_merged + 128-B pe; decode dequantizes kv_merged
+ fp8->bf16 (per-128 block scale).
+ bf16: cache [nblk, page, 576] bf16; per token 512 bf16 kv_c + 64 bf16 k_pe.
+ Prefill splits 1152 -> 1024-B kv + 128-B pe; decode copies bf16 as-is.
+
+In both, KI is FP8 and is dequantized fp8->bf16 + Hadamard-rotated on the
+decode side (vLLM's indexer omits the Hadamard TileRT expects); PE is bf16.
+The prefill side auto-detects the dtype from the cache stride; the decode side
+is told via ``--kv-cache-dtype`` (layout_version differs per dtype so a
+mismatched pairing is rejected at hello). Matches the validated
+serve/tilert_decode dequant + serve_vllm connector split.
+
+Engine inject uses ``inject_cache([(ki[seq,128], kv[seq,512], pe[seq,64])]
+x num_layers, start_pos=0)`` + ``set_cur_pos`` and the three-phase MTP loop
+(warm-up / override / normal) — valid for MLA (KV-cache replay is idempotent).
+
+GPU-validation TODO (per model, W2-style dump on vLLM 0.24): MLA/KI cache
+tensor shapes and per-token layout, exact num_layers incl. MTP exposure under
+--speculative-config, kv-cache group topology.
+"""
from __future__ import annotations
import logging
+import os
import re
from dataclasses import dataclass
import torch
from tilert.pd_vllm import wire
+from tilert.pd_vllm.grammar_backend import load_grammar_backend
+from tilert.pd_vllm.grammar_spec import (
+ GrammarBackendUnavailable,
+ GrammarViolationError,
+ InvalidGrammarError,
+)
+from tilert.pd_vllm.sampling import resolve_top_k, resolve_top_p
logger = logging.getLogger("pd_vllm.profile.mla_nsa")
@@ -81,21 +126,41 @@ class _Reg:
class MlaNsaProfile:
"""Config-driven MLA+NSA profile.
- ``engine_factory(weights, max_seq, with_mtp, ar_steps) -> adapter`` builds
- the model-specific engine.
+ ``engine_factory(weights, max_seq, with_mtp, ar_steps) -> adapter`` builds the
+ model-specific engine.
"""
num_ranks = wire.NUM_RANKS
sender_ranks = frozenset({0}) # MLA latent replicated across TP
+ # The MLA/NSA runtimes (GLM-5, GLM-5.2, DSV3.2) carry no penalty pre-pass,
+ # so MlaNsaEngineAdapter refuses a non-neutral penalty outright. Stated here
+ # too because the router needs the answer at startup, where no engine exists
+ # to ask. When the runtime gains the pre-pass, this and
+ # MlaNsaEngineAdapter.supports_penalties flip together.
+ declares_penalties = False
+
def __init__(
- self, name: str, num_layers: int, layout_version: int, engine_factory, mla_fp8: bool = True
+ self,
+ name: str,
+ num_layers: int,
+ layout_version: int,
+ engine_factory,
+ mla_fp8: bool = True,
+ ki_layer_ids: list[int] | None = None,
):
self.name = name
self.num_layers = num_layers
self._base_version = layout_version
self._engine_factory = engine_factory
self.mla_fp8 = mla_fp8 # fp8_ds_mla (True) vs bf16 (False) MLA cache
+ # Optional strict validation of the sparse-indexer (KI) layer set:
+ # GLM-5.2 registers a KI cache only on the "full" layers; if provided,
+ # classify asserts the KI layer ids vLLM registered equal this set
+ # (ignoring the MTP tail layer, which is present only under speculative).
+ # None = data-driven (accept whatever full-layer KI set vLLM registers
+ # and expand it). GLM-5.1/DSV3.2 leave this None (dense: every layer KI).
+ self.ki_layer_ids = ki_layer_ids
def configure(self, kv_cache_dtype: str) -> MlaNsaProfile:
"""Select the MLA cache dtype (decode side; prefill auto-detects)."""
@@ -185,7 +250,7 @@ def convert(self, buffer, base_ptr, max_seq_len, received, num_devices=1):
@staticmethod
def _dequant_kv(kv_raw: torch.Tensor, seq_len: int) -> torch.Tensor:
- """Dequantize kv_merged [seq,528] u8 (512 fp8 + 4 fp32 scale) -> bf16 [seq,512].
+ """kv_merged [seq,528] u8 (512 fp8 + 4 fp32 scale) -> bf16 [seq,512].
Per-128-block scale: kv[:, b*128:(b+1)*128] *= scale[:, b].
"""
@@ -249,12 +314,59 @@ def lid_of(name):
mla.append((lid_of(name), name, t, gi))
mla.sort(key=lambda x: x[0])
ki.sort(key=lambda x: x[0])
- if len(mla) != self.num_layers or len(ki) != self.num_layers:
+ if len(mla) != self.num_layers:
raise RuntimeError(
- f"{self.name} classify: {len(mla)} MLA + {len(ki)} KI layers "
- f"(expected {self.num_layers} each); check --speculative-config"
+ f"{self.name} classify: {len(mla)} MLA layers "
+ f"(expected {self.num_layers}); check --speculative-config"
f" and the vLLM layer naming"
)
+ # KI (sparse indexer) layer set. GLM-5.2 registers a KI cache only on the
+ # "full" layers ({0,1,2,6,10,...} + MTP tail); "shared" layers reuse the
+ # previous full layer's indexer at runtime and have NO KI cache in vLLM
+ # (verified: config.indexer_types == C++ moe_layer_is_full == vLLM
+ # deepseek_v2 _skip_topk == runtime classify dump). Expand the registered
+ # full-layer KI list to num_layers entries so extract/rdma/convert/inject
+ # stay layer-uniform: logical layer L points at the KI cache of the
+ # largest full-layer id <= L (its controlling full layer); shared layers
+ # thus replicate the previous full layer's KI. For a dense model
+ # (GLM-5.1/DSV3.2, every layer full) this is an identity no-op.
+ ki_ids = [x[0] for x in ki]
+ if not ki or ki[0][0] != 0:
+ raise RuntimeError(
+ f"{self.name} classify: KI layer 0 missing (ids={ki_ids}); "
+ f"cannot expand sparse indexer set"
+ )
+ if len(ki) > self.num_layers or ki_ids != sorted(set(ki_ids)):
+ raise RuntimeError(
+ f"{self.name} classify: bad KI layer set {ki_ids} "
+ f"(num_layers={self.num_layers})"
+ )
+ if self.ki_layer_ids is not None:
+ want = [layer for layer in self.ki_layer_ids if layer < self.num_layers]
+ want_no_mtp = [layer for layer in want if layer != self.num_layers - 1]
+ if ki_ids not in (want, want_no_mtp):
+ raise RuntimeError(
+ f"{self.name} classify: KI layer ids {ki_ids} != expected "
+ f"{want} (or {want_no_mtp} without the MTP tail)"
+ )
+ ki_expanded, cur, idx = [], None, 0
+ for L in range(self.num_layers):
+ while idx < len(ki) and ki[idx][0] <= L:
+ cur = ki[idx]
+ idx += 1
+ # logical layer L, KI tensor/name/group of its controlling full layer
+ assert cur is not None # ki[0][0] == 0 was checked above
+ ki_expanded.append((L, cur[1], cur[2], cur[3]))
+ if len(ki) < self.num_layers:
+ logger.info(
+ "%s: sparse KI %d full layers %s expanded to %d "
+ "(shared layers reuse previous full layer's indexer)",
+ self.name,
+ len(ki),
+ ki_ids,
+ self.num_layers,
+ )
+ ki = ki_expanded
# auto-detect MLA cache dtype from the actual cache stride (the prefill
# cache is ground truth; the decode side is told via --kv-cache-dtype)
t0 = mla[0][2]
@@ -364,6 +476,10 @@ 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)
+ # Per-request, set by decode() before it picks a decode path. The
+ # adapter serves one request at a time (the decode server holds the
+ # node for the whole generation), so a plain attribute is enough.
+ self._ignore_eos = False
@staticmethod
def _resolve_stop_ids(generator) -> set:
@@ -380,30 +496,107 @@ def inject(self, req) -> None:
self._last_prompt_token = req.last_prompt_token
self._seq_len = req.seq_len
- def decode(self, first_token_id, max_tokens, sampling, on_token=None, cancel_event=None):
+ def prepare_grammar(self, grammar_spec, enable_thinking=True):
+ """Compile a per-request GrammarSession (engine cached on first use).
+
+ Raises before any GPU/inject work: unsupported path -> 500, missing
+ xgrammar -> 500, malformed spec -> 400. Never returns a session that
+ would decode unconstrained.
+ """
+ if grammar_spec is None:
+ return None
+ gen = self.gen
+ try:
+ GrammarEngine, GrammarSession = load_grammar_backend()
+ except ImportError as e:
+ raise GrammarBackendUnavailable(f"xgrammar/grammar backend unavailable: {e}") from e
+ if getattr(gen, "_grammar_engine", None) is None:
+ try:
+ gen._grammar_engine = GrammarEngine(
+ gen.tokenizer,
+ padded_vocab_size=gen.config.vocab_size,
+ stop_token_ids=sorted(self.stop_ids),
+ )
+ except ImportError as e:
+ raise GrammarBackendUnavailable(f"xgrammar backend unavailable: {e}") from e
+ think_end_id = None
+ if enable_thinking:
+ tid = gen.tokenizer.convert_tokens_to_ids("")
+ think_end_id = tid if isinstance(tid, int) and tid >= 0 else None
+ # MTP verifies mtp_seq_len positions per step (one mask row each);
+ # non-MTP AR has a single verify position.
+ num_positions = self.mtp_seq_len if self.with_mtp else 1
+ try:
+ return GrammarSession(
+ gen._grammar_engine,
+ grammar_spec,
+ num_positions=num_positions,
+ think_end_id=think_end_id,
+ )
+ except (ImportError, OSError) as e:
+ raise GrammarBackendUnavailable(f"xgrammar backend unavailable: {e}") from e
+ except Exception as e:
+ # Bad schema / regex / EBNF the compiler rejects -> client 400.
+ raise InvalidGrammarError(f"failed to compile grammar spec: {e}") from e
+
+ def supports_penalties(self) -> bool:
+ """No: the MLA/NSA runtimes carry no penalty pre-pass.
+
+ Stated rather than left undefined, because the router reads this over
+ ``/capabilities`` to refuse a penalty request BEFORE the prefill runs.
+ When this runtime gains the pre-pass, this becomes a probed claim
+ rather than a constant.
+ """
+ return False
+
+ def supports_ignore_eos(self) -> bool:
+ return True
+
+ def decode(
+ self,
+ first_token_id,
+ max_tokens,
+ sampling,
+ on_token=None,
+ cancel_event=None,
+ grammar_session=None,
+ ):
sampling = sampling or {}
+ # Fail loud on a penalty this runtime cannot apply. The router normally
+ # refuses these upstream (capabilities), but /pd/decode is reachable
+ # directly and a request that reaches here must not be decoded
+ # UNPENALISED while the caller is told it succeeded.
+ rep = float(sampling.get("repetition_penalty", 1.0) or 1.0)
+ presence = float(sampling.get("presence_penalty", 0.0) or 0.0)
+ if rep != 1.0 or presence != 0.0:
+ raise NotImplementedError(
+ "repetition/presence penalties are not supported by this " "model's decode runtime"
+ )
temp = float(sampling.get("temperature", 1.0))
if temp < 1e-5:
self.gen.update_sampling_params(temperature=1.0, top_p=1.0, top_k=1, use_topp=False)
else:
self.gen.update_sampling_params(
temperature=temp,
- top_p=float(sampling.get("top_p", 0.95)),
- top_k=int(sampling.get("top_k", 256)),
+ top_p=resolve_top_p(sampling),
+ top_k=resolve_top_k(sampling),
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"}
return [int(first_token_id)]
if self.with_mtp:
- return self._decode_mtp(first_token_id, budget, on_token, cancel_event)
- return self._decode_standard(first_token_id, budget, on_token, cancel_event)
+ return self._decode_mtp(first_token_id, budget, on_token, cancel_event, grammar_session)
+ return self._decode_standard(
+ first_token_id, budget, on_token, cancel_event, grammar_session
+ )
- def _decode_mtp(self, first_token_id, budget, on_token, cancel_event):
+ def _decode_mtp(self, first_token_id, budget, on_token, cancel_event, grammar_session=None):
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:
@@ -411,37 +604,99 @@ def _decode_mtp(self, first_token_id, budget, on_token, cancel_event):
if int(first_token_id) in stop_ids:
self.last_stats = {"finish_reason": "stop"}
return []
+ # Feed the prefill-sampled first token to the matcher before any masked
+ # step (see _decode_standard); a violation here fails closed (400).
+ finished = False
+ if grammar_session is not None:
+ try:
+ if grammar_session.accept(int(first_token_id)) == "terminated":
+ finished = True
+ except RuntimeError as e:
+ raise GrammarViolationError(
+ f"prefill first token {first_token_id} violates the " f"grammar"
+ ) from e
dl.set_prefill_valid_tokens(0)
+ ar_steps = max(1, min(1024, int(os.environ.get("GLM5_AR_N", "8"))))
+ ar_ok = hasattr(dl, "ar_accepted_tokens") and hasattr(dl, "ar_num_accepted")
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:
- if cancel_event is not None and cancel_event.is_set():
- finish = "cancelled"
- break
- if fwd == 1:
- 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:
- 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"
+ accepted, finish, fwd = [], "length", 0
+ grammar_mask_written = False
+ try:
+ while not finished and len(tokens) < budget:
+ if cancel_event is not None and cancel_event.is_set():
+ finish = "cancelled"
break
- tokens.append(tok)
- if on_token:
- on_token(tok)
- dl.reset_sequence()
+ if fwd == 1:
+ 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)
+ # Publish this step's per-verify-position mask BEFORE forward
+ # (synchronous path; forward-型, so no show_hands overlap).
+ # fwd==0 is a discarded warmup -> never masked. The draft chain
+ # (row j+1 accepts draft[:j]) mirrors the verified tokens;
+ # at fwd==1 the all-first_token placeholder yields n_acc==1, so
+ # the deeper rows are don't-care. Dormant never follows an
+ # active step (activation is one-way), so no mid-stream reset.
+ if grammar_session is not None and fwd >= 1 and not grammar_session.terminated:
+ chain = draft[0, 1:].cpu().tolist()
+ masks = grammar_session.fill_step_masks(chain)
+ if masks is not None:
+ dl.update_grammar_bitmask(masks)
+ grammar_mask_written = True
+ if fwd == 0 or grammar_session is not None or not ar_ok:
+ steps = 1
+ else:
+ rem = budget - len(tokens)
+ steps = max(1, min(ar_steps, -(-rem // T)))
+ if ar_ok:
+ dl.show_hands(draft, steps)
+ acc = dl.ar_accepted_tokens(0).cpu()[0]
+ num = dl.ar_num_accepted(0).cpu()[0]
+ 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 = 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
+ fwd += 1
+ 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)
+ if (
+ grammar_session is not None
+ and grammar_session.accept(tok) == "terminated"
+ ):
+ finished = True
+ finish = "stop"
+ break
+ accepted.append(na)
+ if finished or len(tokens) >= budget:
+ break
+ finally:
+ if grammar_mask_written:
+ dl.reset_grammar_bitmask()
+ dl.reset_sequence()
self.last_stats = {
"finish_reason": finish,
"mtp_accept_mean": round(sum(accepted) / max(1, len(accepted)), 3),
@@ -449,11 +704,11 @@ 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
-
+ def _decode_standard(
+ self, first_token_id, budget, on_token, cancel_event, grammar_session=None
+ ):
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:
@@ -462,23 +717,91 @@ def _decode_standard(self, first_token_id, budget, on_token, cancel_event):
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:
- 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
- dl.reset_sequence()
+ finished = False
+ grammar_mask_written = False
+ ar_ok = hasattr(dl, "show_hands_no_mtp") and hasattr(dl, "ar_accepted_tokens_no_mtp")
+ ar_steps = max(1, min(1024, int(os.environ.get("GLM5_AR_N", "8"))))
+ cur = last_tok = None
+ prev = None
+ if ar_ok:
+ dl.set_prefill_valid_tokens(0, with_mtp=False)
+ last_tok = int(first_token_id)
+ prev = torch.tensor([last_tok], dtype=torch.int32, device="cuda:0")
+ else:
+ cur = torch.tensor(int(first_token_id), dtype=torch.long, device="cuda:0")
+ try:
+ # Feed the prefill-sampled first token to the matcher before the
+ # first masked step. Its origin is the UNCONSTRAINED prefill, so it
+ # may legitimately violate the grammar (thinking dormant never
+ # rejects) -> fail closed with a client 400.
+ if grammar_session is not None:
+ try:
+ if grammar_session.accept(int(first_token_id)) == "terminated":
+ finish, finished = "stop", True
+ except RuntimeError as e:
+ raise GrammarViolationError(
+ f"prefill first token {first_token_id} violates the " f"grammar"
+ ) from e
+ while not finished and len(tokens) < budget:
+ if cancel_event is not None and cancel_event.is_set():
+ finish = "cancelled"
+ break
+ # Publish the row-0 mask for the token this step will sample
+ # (matcher is at the post-committed state). Dormant/terminated
+ # steps return None -> device stays allow-all (no write).
+ if (
+ grammar_session is not None
+ and grammar_session.active
+ and not grammar_session.terminated
+ ):
+ masks = grammar_session.fill_step_masks([])
+ if masks is not None:
+ dl.update_grammar_bitmask(masks)
+ grammar_mask_written = True
+ if ar_ok:
+ steps = (
+ 1
+ if grammar_session is not None
+ else max(1, min(ar_steps, budget - len(tokens)))
+ )
+ dl.show_hands_no_mtp(prev, steps)
+ acc = dl.ar_accepted_tokens_no_mtp(0).cpu()[0]
+ n_tokens = int(acc[0].item())
+ emitted = acc[1 : 1 + n_tokens].tolist()
+ else:
+ from tilert.models.deepseek_v3_2.temp_var_indices import Idx
+
+ res = dl.forward(cur)
+ intermediates, *_ = res[0]
+ nxt = intermediates[Idx.TOKEN_OUT][0][0]
+ emitted = [int(nxt.item())]
+ cur = nxt
+ 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)
+ if grammar_session is not None and grammar_session.accept(tok) == "terminated":
+ finished = True
+ finish = "stop"
+ break
+ if ar_ok:
+ prev = torch.tensor([last_tok], dtype=torch.int32, device="cuda:0")
+ finally:
+ if grammar_mask_written:
+ # Restore the all-ones no-op mask so later unconstrained
+ # requests on this generator sample the full vocabulary.
+ dl.reset_grammar_bitmask()
+ # Always reset the decode-layer sequence, even on a fail-closed
+ # violation, so the next request starts from a clean state.
+ dl.reset_sequence()
self.last_stats = {"finish_reason": finish}
return tokens
diff --git a/tilert/pd_vllm/receive_server.py b/tilert/pd_vllm/receive_server.py
index 3ad06f2..1a9b22d 100644
--- a/tilert/pd_vllm/receive_server.py
+++ b/tilert/pd_vllm/receive_server.py
@@ -1,4 +1,27 @@
-"""Decode-side receive server (W4): Mooncake buffer + TCP control plane."""
+"""Decode-side receive server (W4): Mooncake buffer + TCP control plane.
+
+Owns one big cuda:0 receive buffer (region layout supplied by the model
+profile), registered with a local Mooncake TransferEngine. Runs a TCP
+server; each participating prefill rank connects per request, gets the hello
+(session id + buffer base addresses), sends its request metadata, **waits to be
+admitted**, then RDMA-writes its sections and reports ``done``. When all 8 ranks
+report, the assembled request is handed to the consumer via a queue.
+
+bs=1: one in-flight request, so the buffer has exactly one tenant at a time and
+admission is what enforces it. A rank asking for a different rid while the
+current tenant is still in use is rejected (``accepted: false``) and must not
+write; ``busy`` in the hello is advisory only, because the rid is not known yet.
+The router's gated dispatch should make a rejection rare, but it is the
+correctness fence, not an optimisation: a write that ignores it lands inside
+another request's KV.
+
+Tenancy is tracked explicitly (``FREE``/``RESERVED``/``TRANSFERRING``/
+``COMPLETE``/``CANCELLING``) with a monotonic ``generation`` and a count of
+ranks that may still be writing. The buffer is never handed to a different
+request while that count is non-zero -- the previous version inferred ownership
+from a wall-clock comparison and replaced timed-out requests whose ranks were
+still mid-RDMA.
+"""
import contextlib
import logging
@@ -15,6 +38,18 @@
logger = logging.getLogger("pd_vllm.receive")
+# Receive-buffer tenancy states. The buffer holds one request at a time, so
+# "who owns it right now" has to be explicit -- the previous code inferred it
+# from `t_complete == 0.0` plus a wall-clock comparison, and that inference is
+# what allowed a timed-out request to be replaced while its ranks were still
+# writing into the buffer.
+FREE = "free" # no tenant; admissible
+RESERVED = "reserved" # rid claimed, no rank writing yet
+TRANSFERRING = "transferring" # >=1 rank admitted and writing
+COMPLETE = "complete" # every sender rank reported done; queued
+CANCELLING = "cancelling" # abandoned; NOT reusable until writers stop
+
+
@dataclass
class ReceivedRequest:
rid: str
@@ -22,9 +57,30 @@ class ReceivedRequest:
last_prompt_token: int
first_token_id: int | None
sampling: dict | None
+ # Full prompt ids when the prefill connector sent them (rank 0, penalty requests
+ # only); empty otherwise -- the engine then leaves the prompt bitmap clear.
+ prompt_token_ids: list = field(default_factory=list)
done_ranks: set = field(default_factory=set)
t_first_conn: float = 0.0
t_complete: float = 0.0
+ # Monotonic tenancy id. Echoed on accept and checked on done, so a message
+ # from a previous tenant cannot be counted towards the current one.
+ generation: int = 0
+ state: str = RESERVED
+ # Ranks admitted but not yet known to have stopped writing. The buffer must
+ # not be handed to another request while this is non-zero.
+ active_writers: int = 0
+
+ @property
+ def has_live_writer(self) -> bool:
+ """Whether a rank may still be RDMA-writing into the buffer.
+
+ The hard half of the reuse rule: while this is true the buffer cannot be
+ handed to another request under any circumstance, timeout included. The
+ soft half (has this tenancy finished, or has it aged out?) is policy and
+ lives in :meth:`ReceiveServer._reusable`, which owns the timeout.
+ """
+ return self.active_writers > 0
class ReceiveServer:
@@ -71,6 +127,12 @@ def __init__(
self._lock = threading.Lock()
self._current: ReceivedRequest | None = None
+ # Tombstones: rid -> deadline. A rid whose consumer has let go must not
+ # be admitted again, and until now nothing recorded that. See release().
+ self._cancelled: dict[str, float] = {}
+ # Monotonic tenancy counter. Never reused, so a message from an earlier
+ # tenant is always distinguishable from the current one.
+ self._generation = 0
self.completed: queue.Queue[ReceivedRequest] = queue.Queue()
# dual-stack: accept IPv4 (v4-mapped) and IPv6, incl. link-local peers
@@ -90,11 +152,159 @@ def __init__(
# ── public ───────────────────────────────────────────────────────────
- def release(self) -> None:
- """Mark the single receive slot free (call after inject/decode)."""
+ def release(self, rid: str) -> None:
+ """Give up the receive slot held by ``rid``.
+
+ Scoped to a rid ON PURPOSE. A caller only ever knows about its own
+ request, and the slot it once held may since have been handed to another:
+ a transfer that arrives after its consumer gave up is enqueued with no
+ one waiting for it, and the NEXT request drains it as unmatched. An
+ unscoped release there would free whatever tenancy is current -- that
+ next request's own -- and it would then never complete, because its
+ ranks' `done` messages are dropped once the tenancy is cancelled.
+ Releasing a rid that no longer owns the slot is a no-op.
+
+ Only actually frees the buffer when nothing is writing into it. A caller
+ that gives up on a request whose ranks are still mid-RDMA (a rejected
+ request, a drained entry) moves it to ``CANCELLING`` instead: the slot
+ stays claimed until the writers stop, because freeing it there is
+ precisely how one request's KV lands inside another's buffer. Every
+ sender connection carries ``request_timeout`` as its socket timeout, so
+ the drain is bounded without needing to force it.
+
+ Always leaves a TOMBSTONE for ``rid``, whatever the slot turns out to
+ hold. The ``CANCELLING`` state above only refuses senders while a writer
+ is still live; it says nothing about a rank that has not connected YET.
+ A cancel that wins before any sender arrives leaves ``_current`` None,
+ so without a tombstone the late rank finds a free buffer, is admitted,
+ and holds it while the next request's ranks are turned away "busy" until
+ they exhaust their admission retries -- that request then waits out its
+ whole kv_transfer_timeout. The same hole is open after a normal
+ completion: a rank that never reported done is not in ``done_ranks``, so
+ it would open a fresh tenancy rather than be refused as a duplicate.
+
+ ``request_timeout`` is the right lifetime because it is the senders'
+ socket timeout: past it, no rank can still be trying to join this rid.
+ """
with self._lock:
+ self._tombstone(rid)
+ cur = self._current
+ if cur is None:
+ return
+ if cur.rid != rid:
+ logger.info("release(%s) ignored: the slot now holds %s", rid, cur.rid)
+ return
+ if cur.has_live_writer and cur.state != COMPLETE:
+ cur.state = CANCELLING
+ logger.warning(
+ "release(%s) with %d writer(s) still active: slot stays "
+ "claimed (cancelling) until they stop",
+ cur.rid,
+ cur.active_writers,
+ )
+ return
self._current = None
+ def _tombstone(self, rid: str) -> None:
+ """Record ``rid`` as done with, and drop the tombstones that expired.
+
+ Caller holds ``self._lock``. Pruning here rather than on a timer keeps
+ the map bounded without a second thread to reason about.
+ """
+ now = time.time()
+ self._cancelled = {r: t for r, t in self._cancelled.items() if t > now}
+ self._cancelled[rid] = now + self.request_timeout
+
+ def _extend_tombstone(self, rid: str, window_s) -> None:
+ """Push ``rid``'s tombstone out to cover ``window_s`` from now.
+
+ Caller holds ``self._lock``. Never shortens one: a second rank
+ declaring a smaller budget must not expose the request to the first.
+ """
+ try:
+ window = float(window_s)
+ except (TypeError, ValueError):
+ return # older sender: keep the default
+ if window <= 0:
+ return
+ deadline = time.time() + window
+ if deadline > self._cancelled.get(rid, 0.0):
+ self._cancelled[rid] = deadline
+
+ def expect(self, rid: str) -> None:
+ """Announce that a consumer is now waiting for ``rid``.
+
+ Clears any tombstone for it. Called from ``/pd/decode`` on admission,
+ which is precisely the event that distinguishes the two things a
+ tombstone cannot tell apart: a request everyone has given up on, and the
+ same request coming back because vLLM rescheduled it. The first is never
+ re-announced; the second always is, and its senders' admission retries
+ are what carry them across the gap.
+ """
+ with self._lock:
+ if self._cancelled.pop(rid, None) is not None:
+ logger.info("request %s re-announced; its tombstone is dropped", rid)
+
+ def _is_tombstoned(self, rid: str) -> bool:
+ """Caller holds ``self._lock``."""
+ deadline = self._cancelled.get(rid)
+ if deadline is None:
+ return False
+ if deadline <= time.time():
+ del self._cancelled[rid]
+ return False
+ return True
+
+ def _next_generation(self) -> int:
+ """Caller holds ``self._lock``."""
+ self._generation += 1
+ return self._generation
+
+ def _reusable(self, cur: ReceivedRequest) -> bool:
+ """Whether ``cur``'s buffer may be given to a different rid.
+
+ Three clauses, and the order matters:
+
+ 1. ``COMPLETE`` is reusable whatever ``active_writers`` says. A rank sends
+ ``done`` only AFTER its RDMA write has returned, so once every sender
+ rank is done nothing is writing; a non-zero count there is just
+ sockets that have not closed yet. ``release()`` makes the same
+ exception, and the two must agree -- if this clause moved below the
+ writer check, a request could be released by one rule and refused by
+ the other.
+ 2. **Otherwise, never under a live writer.** No timeout overrides this.
+ Freeing the buffer while a rank is mid-RDMA is what puts one request's
+ KV inside another's, and no elapsed time makes that safe.
+ 3. Reusable once abandoned (``CANCELLING``) or **aged out** -- a tenancy
+ still marked ``RESERVED``/``TRANSFERRING`` whose writers have all gone
+ without completing it. That happens when every sender dies
+ mid-transfer, and without the age-out the slot would be held for good:
+ nothing calls ``release()`` for a rid whose ``/pd/decode`` never
+ arrived (the router gave up on the prefill leg, say). The original code
+ aged out on wall clock ALONE, which is how a live writer got
+ overwritten; this keeps the bound and drops the hazard.
+ """
+ if cur.state == COMPLETE:
+ return True
+ if cur.has_live_writer:
+ return False
+ if cur.state in (FREE, CANCELLING):
+ return True
+ return (time.time() - cur.t_first_conn) >= self.request_timeout
+
+ def state_snapshot(self) -> dict:
+ """Current tenancy, for tests and the decode server's status endpoint."""
+ with self._lock:
+ cur = self._current
+ if cur is None:
+ return {"state": FREE, "rid": None, "generation": None, "active_writers": 0}
+ return {
+ "state": cur.state,
+ "rid": cur.rid,
+ "generation": cur.generation,
+ "active_writers": cur.active_writers,
+ }
+
def close(self) -> None:
self._stop.set()
with contextlib.suppress(OSError):
@@ -111,14 +321,131 @@ def _accept_loop(self) -> None:
t = threading.Thread(target=self._handle, args=(conn, addr), daemon=True)
t.start()
+ def _admit(self, req: dict, rid: str, rank: int) -> dict:
+ """Decide whether ``rank`` of ``rid`` may write, under the lock.
+
+ Returns the message to send back: an accept carrying the tenancy
+ generation, or a reject. Admission is granted only for a tenancy this
+ rank is actually joining, so the sender can key its RDMA on the reply.
+ """
+ with self._lock:
+ if self._is_tombstoned(rid):
+ # Checked BEFORE anything about the current tenancy: this rid is
+ # dead whatever the slot holds now, and admitting it would let a
+ # request nobody is waiting for occupy the buffer.
+ #
+ # RETRYABLE on purpose. A rid is not unique to a tenancy -- it
+ # is derived from the vLLM request id, so a preempted request
+ # rescheduled by vLLM sends again under the SAME rid. That retry
+ # must not lose its shard to a tombstone left by the attempt
+ # before it, and it announces itself by calling /pd/decode
+ # again, which is what clears the tombstone (see expect()).
+ # Refusing transiently lets the sender wait for that to happen;
+ # a rid nobody re-announces just exhausts its retries, which
+ # costs a dead request nothing.
+ #
+ # Reported as `cancelling`, a reason senders ALREADY retry,
+ # rather than a new one. A new reason is permanent to any
+ # connector built before it: during a rolling upgrade it would
+ # drop the shard, and the rescheduled request would then wait
+ # out its kv_transfer timeout. The distinction is only useful
+ # in a log, so it goes in `detail`.
+ #
+ # Extend to outlast THIS sender's remaining retries. The
+ # default lifetime is request_timeout, which is one
+ # connection's socket timeout -- but every retry opens a new
+ # connection, so a sender configured with enough attempts can
+ # still be trying after the tombstone has expired, and would
+ # then be admitted for a request nobody wants. Only the sender
+ # knows its budget, so it declares it.
+ self._extend_tombstone(rid, req.get("admission_window_s"))
+ logger.warning(
+ "refusing %s rank %d for now: no consumer is " "waiting for this request",
+ rid,
+ rank,
+ )
+ return wire.reject_msg("cancelling", rid=rid, detail="no_consumer")
+ cur = self._current
+ if cur is not None and cur.rid == rid:
+ if cur.state == CANCELLING:
+ return wire.reject_msg("cancelling", rid=rid)
+ if rank in cur.done_ranks:
+ # A duplicate for a rank that already finished would be
+ # counted twice and could complete the request early.
+ return wire.reject_msg("duplicate_rank", rid=rid, rank=rank)
+ elif cur is not None and not self._reusable(cur): # noqa: R505 (exclusive branches)
+ # Busy with a DIFFERENT rid whose buffer is still in use. This is
+ # the reply the sender used to ignore, writing anyway.
+ logger.warning(
+ "rejecting %s rank %d (busy with %s, state=%s, " "writers=%d)",
+ rid,
+ rank,
+ cur.rid,
+ cur.state,
+ cur.active_writers,
+ )
+ return wire.reject_msg("busy", busy_rid=cur.rid, busy_state=cur.state)
+ else:
+ # Free, or the previous tenant is finished/drained: new tenancy.
+ self._current = cur = ReceivedRequest(
+ rid=rid,
+ seq_len=int(req["seq_len"]),
+ last_prompt_token=int(req.get("last_prompt_token", 0)),
+ first_token_id=req.get("first_token_id"),
+ sampling=req.get("sampling"),
+ prompt_token_ids=list(req.get("prompt_token_ids") or []),
+ t_first_conn=time.time(),
+ generation=self._next_generation(),
+ state=RESERVED,
+ )
+ logger.info(
+ "request %s: seq_len=%d (generation %d)", rid, cur.seq_len, cur.generation
+ )
+
+ # Absorb the prompt ids from whichever rank carries them. Only
+ # rank 0 sends them (prefill_connector._send), but ranks connect
+ # in ARBITRARY order and only the first one to arrive builds the
+ # ReceivedRequest -- so keying this off the creation path would
+ # drop the ids ~7/8 of the time, non-deterministically.
+ if not cur.prompt_token_ids and req.get("prompt_token_ids"):
+ cur.prompt_token_ids = list(req["prompt_token_ids"])
+ logger.info(
+ "request %s: prompt bitmap seeded from rank %d " "(%d ids)",
+ rid,
+ rank,
+ len(cur.prompt_token_ids),
+ )
+
+ cur.state = TRANSFERRING
+ cur.active_writers += 1
+ return wire.accept_msg(rid, rank, cur.generation)
+
+ def _writer_left(self, rid: str, generation: int) -> None:
+ """One admitted rank has stopped writing (done, error, or disconnect).
+
+ Called from the connection's ``finally`` so a sender that dies mid-RDMA
+ still releases its claim. Bounded by the socket timeout, which is why the
+ cancelling drain needs no forced override.
+ """
+ with self._lock:
+ cur = self._current
+ if cur is None or cur.generation != generation:
+ return
+ cur.active_writers = max(0, cur.active_writers - 1)
+ if cur.state == CANCELLING and cur.active_writers == 0:
+ logger.info("request %s drained; receive slot free", cur.rid)
+ self._current = None
+
def _handle(self, conn: socket.socket, addr) -> None:
+ admitted: tuple[str, int] | None = None
try:
conn.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
conn.settimeout(self.request_timeout)
- # `busy` in hello is advisory (a same-rid rank must still proceed);
- # the authoritative accept/reject happens once the rid is known.
+ # `busy` in hello is advisory only (a same-rid rank must still
+ # proceed, and the rid is not known yet); the authoritative decision
+ # is the accept/reject below.
with self._lock:
- advisory_busy = self._current is not None and self._current.t_complete == 0.0
+ advisory_busy = self._current is not None and not self._reusable(self._current)
wire.send_msg(
conn,
wire.hello_msg(
@@ -134,30 +461,19 @@ def _handle(self, conn: socket.socket, addr) -> None:
req = wire.recv_msg(conn)
rid, rank = req["rid"], int(req["rank"])
if req.get("seq_len", 0) > self.max_seq_len:
- wire.send_msg(conn, {"error": "seq_len exceeds max_seq_len"})
+ wire.send_msg(
+ conn,
+ wire.reject_msg(
+ "seq_len exceeds max_seq_len", rid=rid, max_seq_len=self.max_seq_len
+ ),
+ )
return
- with self._lock:
- cur = self._current
- if cur is None or cur.rid != rid:
- if (
- cur is not None
- and cur.t_complete == 0.0
- and time.time() - cur.t_first_conn < self.request_timeout
- ):
- # busy with a different in-flight rid
- wire.send_msg(conn, {"error": "busy", "busy_rid": cur.rid})
- logger.warning("rejecting %s (busy with %s)", rid, cur.rid)
- return
- self._current = cur = ReceivedRequest(
- rid=rid,
- seq_len=int(req["seq_len"]),
- last_prompt_token=int(req.get("last_prompt_token", 0)),
- first_token_id=req.get("first_token_id"),
- sampling=req.get("sampling"),
- t_first_conn=time.time(),
- )
- logger.info("request %s: seq_len=%d", rid, cur.seq_len)
+ reply = self._admit(req, rid, rank)
+ wire.send_msg(conn, reply)
+ if not reply.get("accepted"):
+ return
+ admitted = (rid, reply["generation"])
# wait for this rank's done (RDMA happens meanwhile)
done = wire.recv_msg(conn)
@@ -166,7 +482,23 @@ def _handle(self, conn: socket.socket, addr) -> None:
return
with self._lock:
cur = self._current
- if cur is None or cur.rid != rid:
+ if cur is None or cur.rid != rid or cur.generation != admitted[1]:
+ # The tenancy this rank was admitted against is gone. Its
+ # bytes went into a buffer that has since been reassigned or
+ # abandoned, so counting the done would attribute them to
+ # whoever holds the slot now.
+ logger.warning(
+ "ignoring done from %s rank %d: it was admitted to "
+ "generation %d and the slot has moved on",
+ rid,
+ rank,
+ admitted[1],
+ )
+ return
+ if cur.state == CANCELLING:
+ logger.warning(
+ "ignoring done from %s rank %d: request was " "abandoned", rid, rank
+ )
return
cur.done_ranks.add(rank)
logger.info(
@@ -177,6 +509,7 @@ def _handle(self, conn: socket.socket, addr) -> None:
len(self.profile.sender_ranks),
)
if cur.done_ranks >= set(self.profile.sender_ranks):
+ cur.state = COMPLETE
cur.t_complete = time.time()
self.completed.put(cur)
logger.info(
@@ -187,4 +520,6 @@ def _handle(self, conn: socket.socket, addr) -> None:
except Exception:
logger.exception("connection from %s failed", addr)
finally:
+ if admitted is not None:
+ self._writer_left(*admitted)
conn.close()
diff --git a/tilert/pd_vllm/reply.py b/tilert/pd_vllm/reply.py
new file mode 100644
index 0000000..b81c87f
--- /dev/null
+++ b/tilert/pd_vllm/reply.py
@@ -0,0 +1,348 @@
+"""Token ids -> emissions, for both response channels.
+
+Three transformations stand between the node's ids and either reply: detokenise
+incrementally (``IncrementalDetok``), match stop strings and hold back what could
+still become one (``StopWindow``), route text to channels (the output parser).
+Both channels consume the same emissions, which is what makes them agree.
+
+Logprob entries ride along, and are the one place the transformations are not
+independent: an entry is keyed by TOKEN POSITION, the stop cut by CHARACTER
+OFFSET. What this module does about that is one-sided: it keeps an emission from
+ENDING inside a token whose entry is still due, so text and the entry describing
+it go out together. It does not attribute an entry to a CHANNEL --
+``refuse_unattributable_logprobs``
+refuses the only shape where the join has no answer, and the rest needs no offset
+arithmetic:
+
+ no parser multi-token chunk the reply IS content, so every
+ token is described
+ parser, no stop one token's text exact by construction
+ parser, no logprobs multi-token chunk nothing to attribute
+ parser + stop + lp -- REFUSED at the gate
+
+The refused row has no answer, not an expensive one: a parser reports no
+difference between buffering an ambiguous marker prefix and consuming a complete
+marker, so dropping the entry under-reports and keeping it corrupts
+``logprobs.content``.
+
+Free of HTTP, asyncio and vLLM: driven from a list of token ids.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterable
+from dataclasses import dataclass, field
+from typing import Any
+
+from tilert.pd_vllm.logprobs import LogprobsRequest, build_logprobs
+from tilert.pd_vllm.oai_parser import IncrementalDetok
+from tilert.pd_vllm.stop_strings import StopWindow
+
+__all__ = ["CONTENT", "REASONING", "TOOL_CALL", "Emission", "ReplyStream", "as_logprobs"]
+
+CONTENT = "content"
+REASONING = "reasoning"
+TOOL_CALL = "tool_call"
+
+
+def as_logprobs(entries: list[dict]) -> dict:
+ """``choices[].logprobs`` around a list of entries, as vLLM shapes it.
+
+ One envelope for both channels. ``refusal`` is always null: this path has no
+ refusal channel, and the field is nullable.
+ """
+ return {"content": entries, "refusal": None}
+
+
+@dataclass(frozen=True)
+class Emission:
+ """One piece of reply, on one channel, with the logprobs of its own tokens.
+
+ ``text`` for :data:`CONTENT` / :data:`REASONING`; ``tool_call`` instead for
+ :data:`TOOL_CALL`, already shaped for the OpenAI delta and always WHOLE --
+ the parser emits each index once, arguments complete.
+
+ ``logprobs`` is non-empty only on :data:`CONTENT`: the contract covers
+ ``message.content`` alone. Empty ``text`` with non-empty ``logprobs`` is
+ legitimate -- a token can run and produce no visible text, and vLLM reports
+ it.
+ """
+
+ channel: str
+ text: str = ""
+ logprobs: list[dict] = field(default_factory=list)
+ tool_call: dict | None = None
+
+
+@dataclass
+class _Entry:
+ """One token's logprob record, waiting for the emission carrying its text."""
+
+ token_id: int
+ logprob: float | None
+ candidates: list[tuple[int, float]]
+ # Character offset where this token's text ends, or -1 while unknown: a
+ # token can decode to nothing of its own -- the first half of a split
+ # multi-byte character -- and its text arrives with a later token.
+ ends_at: int
+
+
+class ReplyStream:
+ """Drive one request's tokens through the three transformations.
+
+ ``push`` may be called once with the whole sequence or per streamed message;
+ the emissions are the same either way, which is what the two channels'
+ agreement rests on. Then ``finish()``, then ``token_ids`` / ``stop_reason``.
+ """
+
+ def __init__(
+ self,
+ tokenizer,
+ *,
+ stop: Iterable[str] = (),
+ include_stop_in_output: bool = False,
+ parser_session: Any | None = None,
+ logprobs_req: LogprobsRequest | None = None,
+ first_token_logprob: tuple[float | None, list] | None = None,
+ ) -> None:
+ stop = list(stop)
+ if stop and parser_session is not None and logprobs_req is not None:
+ # Reaching here past the gate is a routing bug; 502 is the honest
+ # answer for one.
+ raise ValueError(
+ "stop with both an output parser and logprobs is not "
+ "attributable; the request gate must refuse it"
+ )
+
+ self._session = parser_session
+ self._logprobs_req = logprobs_req
+ self._first_lp = first_token_logprob
+
+ # Specials are kept only when a parser will consume them. One policy
+ # drives both the matcher and the output, or a stop spelled like a
+ # special ends one channel's reply and not the other's.
+ self._detok = IncrementalDetok(tokenizer, skip_special_tokens=parser_session is None)
+ self._decode_one: Callable[[int], str] = lambda t: tokenizer.decode(
+ [t], skip_special_tokens=False
+ )
+ self._window = StopWindow(stop, include_stop_in_output)
+
+ self._ids: list[int] = []
+ self._chars = 0 # characters the detokeniser produced
+ self._due: list[_Entry] = [] # entries not yet on an emission
+ self._held: list[dict] = [] # built entries with no channel yet
+ self._finished = False
+
+ # ── what the caller reports rather than computes ─────────────────────────
+
+ @property
+ def token_ids(self) -> list[int]:
+ """The tokens this reply is made of, in order.
+
+ The node delivers in batches, so a stop lands part-way into one; the
+ tokens behind it ran only because the node cannot see text. A caller
+ keeping its own list from the wire reported 30 ids against a
+ ``completion_tokens`` of 6 on a live pair.
+ """
+ return list(self._ids)
+
+ @property
+ def completion_tokens(self) -> int:
+ """Tokens generated for this reply, which is what the client is billed.
+
+ A stop truncates the text, not the count. Matches vLLM, whose
+ ``completion_tokens`` is its detokeniser's untruncated id list.
+ """
+ return len(self._ids)
+
+ @property
+ def stop_reason(self) -> str | None:
+ """The stop string that ended the sequence, or None."""
+ return self._window.stopped
+
+ def finish_reason(self, from_decode: str) -> str:
+ """What to report, given what the node said.
+
+ A stop overrides it: the node cannot see text, so it cannot reach that conclusion.
+ """
+ return "stop" if self._window.stopped is not None else from_decode
+
+ # ── the stream ──────────────────────────────────────────────────────────
+
+ def push(
+ self,
+ token_ids: Iterable[int],
+ logprobs: list[float | None] | None = None,
+ candidates: list[list[tuple[int, float]]] | None = None,
+ ) -> list[Emission]:
+ """Absorb tokens; return whatever became emittable.
+
+ One token at a time internally: batching the detokeniser is cheaper but
+ loses the boundaries the stop cut and the entries both need.
+ """
+ out: list[Emission] = []
+ for i, tid in enumerate(token_ids):
+ if self._window.stopped is not None:
+ break
+ self._ids.append(tid)
+ delta = self._detok.push([tid])
+ self._chars += len(delta)
+ self._queue(tid, logprobs, candidates, i, bool(delta))
+ self._window.push(delta)
+ out += self._route(self._window.take(limit=self._cap()))
+ return out
+
+ def finish(self) -> list[Emission]:
+ """Release the held tail, then flush the parser. Idempotent.
+
+ Everything owed after the last ``push`` comes from here, including
+ entries for tokens that produced no visible text.
+ """
+ if self._finished:
+ return []
+ self._finished = True
+
+ # The detokeniser may still hold a partial character. Its text belongs to
+ # the reply -- the tokens are already counted -- so it goes through the
+ # matcher like any other, before the window is drained.
+ tail = self._detok.finish()
+ if tail:
+ self._chars += len(tail)
+ for pending in self._due:
+ if pending.ends_at < 0:
+ pending.ends_at = self._chars
+ self._window.push(tail)
+ out = self._route(self._window.take(final=True))
+ if self._session is not None:
+ out += self._events(self._session.finish(), self._collect())
+ return out
+ leftover = self._collect()
+ if leftover:
+ # These tokens ran; the stop removed their text, or the caller
+ # strips them. vLLM reports them, so so do we.
+ out.append(Emission(CONTENT, "", leftover))
+ return out
+
+ # ── internals ───────────────────────────────────────────────────────────
+
+ def _queue(self, tid, logprobs, candidates, i, produced_text) -> None:
+ """Record this token's logprob, to be attached when its text goes out.
+
+ A token that produced nothing gets ``ends_at = -1`` and is backfilled by
+ whichever token completes its character. Marking it complete at the
+ current offset would make it due before the character it belongs to is
+ out, so its entry would ride an empty chunk while the chunk carrying the
+ character described one token too few.
+ """
+ if self._logprobs_req is None:
+ return
+ lp = logprobs[i] if logprobs and i < len(logprobs) else None
+ cands = list(candidates[i]) if candidates and i < len(candidates) else []
+ if len(self._ids) == 1 and lp is None and self._first_lp is not None:
+ # Token 1 was echoed, not sampled, by the node: its distribution
+ # exists only in the prefill reply.
+ lp, cands = self._first_lp[0], list(self._first_lp[1] or [])
+ if produced_text:
+ for pending in self._due:
+ if pending.ends_at < 0:
+ pending.ends_at = self._chars
+ self._due.append(_Entry(tid, lp, cands, ends_at=self._chars if produced_text else -1))
+
+ def _cap(self) -> int | None:
+ """Ceiling so an emission never ends inside a token whose entry is still due.
+
+ Text and the entry describing it then go out together.
+
+ The window holds text back by a CHARACTER count, which lands mid-token:
+ the visible part of a token would go out while its entry waited for the
+ offset to clear, and the entry then rode a later chunk. Only stop strings
+ create a holdback, so this binds for stop with logprobs and nothing else.
+ """
+ if self._logprobs_req is None:
+ return None
+ end = self._chars - self._window.hold
+ cap = self._window.visible
+ for pending in self._due:
+ if 0 <= pending.ends_at <= end:
+ cap = pending.ends_at
+ return cap
+
+ def _route(self, text: str) -> list[Emission]:
+ """Send text the client may see to its channel(s)."""
+ if self._session is None:
+ entries = self._ready()
+ if not text and not entries:
+ return []
+ return [Emission(CONTENT, text, entries)]
+ entries = self._held + self._ready()
+ if not text:
+ # No characters contributed, so no channel decided: the first half
+ # of a split multi-byte character waits for the second.
+ self._held = entries
+ return []
+ self._held = []
+ return self._events(self._session.feed(text), entries)
+
+ def _events(self, events: list[dict], entries: list[dict]) -> list[Emission]:
+ """Parser events as emissions, with the pending entries if content."""
+ out: list[Emission] = []
+ for ev in events:
+ kind = ev.get("kind")
+ if kind == "tool":
+ out.append(
+ Emission(
+ TOOL_CALL,
+ tool_call={
+ "index": ev["index"],
+ "id": ev["id"],
+ "name": ev["name"],
+ "arguments": ev["arguments"],
+ },
+ )
+ )
+ elif kind == "reasoning":
+ out.append(Emission(REASONING, ev.get("text", "")))
+ else:
+ out.append(Emission(CONTENT, ev.get("text", ""), entries))
+ entries = [] # the first content emission carries them
+ return out
+
+ def _ready(self) -> list[dict]:
+ """Entries whose token's text the client can now see.
+
+ Once the window has stopped no more text can arrive, so a token the cut
+ ran through is as visible as it will ever be: its surviving prefix is in
+ the emission being built, and its entry belongs there rather than on a
+ later empty chunk. `ends_at` still points past the untruncated token, so
+ the offset comparison alone would defer it.
+
+ A token that produced nothing yet (`ends_at < 0`) is never ready: its
+ text arrives with a later token, or at the detokeniser's final flush.
+ """
+ visible = self._window.visible
+ stopped = self._window.stopped is not None
+ n = 0
+ while (
+ n < len(self._due)
+ and self._due[n].ends_at >= 0
+ and (stopped or self._due[n].ends_at <= visible)
+ ):
+ n += 1
+ return self._build(n)
+
+ def _collect(self) -> list[dict]:
+ """Every remaining entry, at end of stream."""
+ return self._build(len(self._due))
+
+ def _build(self, n: int) -> list[dict]:
+ if not n:
+ return []
+ taken, self._due = self._due[:n], self._due[n:]
+ assert self._logprobs_req is not None # entries are queued only when asked for
+ return build_logprobs(
+ [e.token_id for e in taken],
+ [e.logprob for e in taken],
+ [[(int(c[0]), float(c[1])) for c in e.candidates] for e in taken],
+ self._logprobs_req,
+ self._decode_one,
+ )["content"]
diff --git a/tilert/pd_vllm/request_gate.py b/tilert/pd_vllm/request_gate.py
new file mode 100644
index 0000000..27e3823
--- /dev/null
+++ b/tilert/pd_vllm/request_gate.py
@@ -0,0 +1,222 @@
+"""What is decided about a request before anything observable happens.
+
+Called before the vLLM prefill request, before a decode node is acquired, before
+the connector claims anything, so a refused request costs nothing and holds
+nothing. Both response paths ran an identical copy of these six checks; the
+copies drifted twice under review.
+
+Deliberately NOT here: ``validate_generation_request``, which needs the pool's
+capabilities. That probe talks HTTP to every uncached node and blocks per wedged
+one, while nothing below depends on it -- so the caller runs it after this
+returns, on whichever thread discipline it has (a direct call, or a threadpool
+off the event loop). The order is the point: a 400 for a malformed request must
+not wait on an unreachable node.
+
+Free of HTTP and asyncio: a dict in, a :class:`GatedRequest` or an exception out.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from pydantic import TypeAdapter, ValidationError
+
+from tilert.pd_vllm.capabilities import (
+ CapabilityUnavailable,
+ InvalidParameter,
+)
+from tilert.pd_vllm.grammar_spec import extract_request_grammar_spec
+from tilert.pd_vllm.logprobs import (
+ LogprobsRequest,
+ LogprobsUnsupported,
+ resolve_logprobs_request,
+)
+from tilert.pd_vllm.openai_params import (
+ InvalidOutputLength,
+ resolve_max_tokens,
+)
+from tilert.pd_vllm.stop_strings import resolve_stop
+
+__all__ = [
+ "GatedRequest",
+ "gate_request",
+ "refuse_unattributable_logprobs",
+ "require_tokenizer_for_logprobs",
+ "resolve_stop_request",
+]
+
+
+@dataclass(frozen=True)
+class GatedRequest:
+ """An accepted request, as the handlers need it.
+
+ Every field is derived from the body, so a handler never re-reads it for the
+ same decision -- `enable_thinking` was read three times per request, and the
+ two paths disagreed about whether `/v1/completions` has a parser at all.
+ """
+
+ is_chat: bool
+ thinking: bool
+ stop: list[str] = field(default_factory=list)
+ include_stop: bool = False
+ logprobs_req: LogprobsRequest | None = None
+ grammar_spec: dict | None = None
+
+
+def gate_request(path: str, body: dict, *, tokenizer, parser_active) -> GatedRequest:
+ """Accept the request or raise, without touching a backend.
+
+ ``parser_active`` answers "would a parser run for this request", given the
+ thinking flag: it decides one refusal and is a router-configuration
+ question, not a request one.
+
+ Raises:
+ InvalidParameter: unusable type or value (400).
+ CapabilityUnavailable: valid request, nothing to execute it (501).
+ LogprobsUnsupported: logprobs asked for where they are not served (400).
+ GrammarError: unusable constrained-decoding spec.
+ """
+ is_chat = path.endswith("chat/completions")
+ grammar_spec = extract_request_grammar_spec(body)
+ # Resolved here only to fail early: the router pins the prefill leg to
+ # max_tokens=1 and strips both client names from it, so an unusable length
+ # would otherwise surface from `_decode_body` with the prefill already
+ # spent. The value itself is taken there, from the same function.
+ try:
+ resolve_max_tokens(body)
+ except InvalidOutputLength as e:
+ raise InvalidParameter(str(e)) from None
+ logprobs_req = _logprobs_of(path, body)
+ require_tokenizer_for_logprobs(logprobs_req, tokenizer)
+ stop, include_stop = resolve_stop_request(body, tokenizer)
+ # Chat only, and in this position, because both are vLLM's behaviour:
+ # `ChatCompletionRequest` types the field and answers 422 for a string,
+ # while `CompletionRequest` has no such field and `extra="allow"` accepts
+ # and ignores it (measured on 0.25.1). Validating it for /v1/completions
+ # would refuse a request vLLM serves; validating it before `stop` would
+ # change which of two bad fields is reported.
+ thinking = _thinking_enabled(body) if is_chat else True
+ refuse_unattributable_logprobs(logprobs_req, stop, is_chat and parser_active(thinking))
+ return GatedRequest(
+ is_chat=is_chat,
+ thinking=thinking,
+ stop=stop,
+ include_stop=include_stop,
+ logprobs_req=logprobs_req,
+ grammar_spec=grammar_spec,
+ )
+
+
+def _thinking_enabled(body: dict) -> bool:
+ """Whether the request asked for the model's thinking segment.
+
+ Raises:
+ InvalidParameter: ``chat_template_kwargs`` present and not an object
+ (400). vLLM declares it ``dict[str, Any] | None`` and answers 422 for
+ anything else; reading `.get` off a string here raised AttributeError
+ instead, which surfaced as a 500 for a client mistake.
+ """
+ ctk = body.get("chat_template_kwargs")
+ if ctk is None:
+ return True
+ if not isinstance(ctk, dict):
+ raise InvalidParameter(
+ f"chat_template_kwargs must be an object, got " f"{type(ctk).__name__}"
+ )
+ return bool(ctk.get("enable_thinking", True))
+
+
+def _logprobs_of(path: str, body: dict):
+ """Validated logprobs request for this path, or None.
+
+ Chat only. ``/v1/completions`` takes a differently-typed ``logprobs`` (a
+ count, not a flag) and is not served here, so asking for it there is
+ rejected with the same shape as the existing streaming rejection rather than
+ being ignored.
+ """
+ if path.endswith("chat/completions"):
+ return resolve_logprobs_request(body)
+ if body.get("logprobs") is not None:
+ raise LogprobsUnsupported("logprobs is supported on /v1/chat/completions")
+ return None
+
+
+def refuse_unattributable_logprobs(logprobs_req, stop, has_parser) -> None:
+ """Refuse the one request shape whose logprob attribution has no answer.
+
+ A ``stop`` makes the router hold text back, so what reaches the parser spans
+ token boundaries and some tokens' channel becomes undecidable. Every other
+ combination is served; ``reply``'s module docstring has the four cases and
+ why this one has no answer rather than an expensive one.
+
+ Raises:
+ CapabilityUnavailable: all three asked for at once (501).
+ """
+ if logprobs_req is not None and stop and has_parser:
+ raise CapabilityUnavailable(
+ "logprobs cannot be attributed to message.content when stop "
+ "strings and an output parser are both in play: the parser buffers "
+ "across the boundary the stop hold-back creates, so some tokens' "
+ "channel is undecidable. Drop one of stop / logprobs, or use a "
+ "router started with --parser none."
+ )
+
+
+def require_tokenizer_for_logprobs(logprobs_req, tokenizer) -> None:
+ """Refuse a logprobs request this router cannot name tokens for.
+
+ Every entry carries the token's text, and ``--parser none`` without
+ ``--model-path`` is a supported configuration with no tokenizer. 200 with
+ ``logprobs: null`` would report success for a field the client asked for and
+ did not get, after both backends computed the values.
+
+ Raises:
+ CapabilityUnavailable: 501.
+ """
+ if logprobs_req is not None and tokenizer is None:
+ raise CapabilityUnavailable(
+ "logprobs need a tokenizer to name each token, and this router has "
+ "none: start it with --model-path to serve them."
+ )
+
+
+def resolve_stop_request(body: dict, tokenizer) -> tuple[list[str], bool]:
+ """The request's stop strings and whether to keep them, or ``([], False)``.
+
+ Not in the capability gate, because executing them needs a tokenizer and
+ that is a router-side resource.
+
+ Raises:
+ InvalidParameter: unusable type or value (400).
+ CapabilityUnavailable: asked for with no tokenizer (501).
+ """
+ try:
+ stop = resolve_stop(body)
+ except ValueError as e:
+ raise InvalidParameter(str(e)) from e
+
+ include = False
+ if "include_stop_str_in_output" in body:
+ # pydantic is the validator vLLM's request model uses, so the accepted
+ # set is identical by construction. Measured on 2.13.4: `1`, `0` and
+ # `"true"` accepted (refusing them would turn a working request into a
+ # 400), explicit `null` refused (vLLM answers 422, and the router strips
+ # the field before prefill so vLLM never gets to).
+ try:
+ include = TypeAdapter(bool).validate_python(body["include_stop_str_in_output"])
+ except ValidationError:
+ raise InvalidParameter(
+ f"include_stop_str_in_output must be a boolean, got "
+ f"{body['include_stop_str_in_output']!r}"
+ ) from None
+ # No check for `include` without `stop`: with no stop strings the flag is a
+ # no-op, and vLLM serves such a request. Refusing it would turn a request
+ # that works against a native endpoint into a 400 here.
+
+ if stop and tokenizer is None:
+ raise CapabilityUnavailable(
+ "stop strings need a tokenizer to match against the reply text, "
+ "and this router has none: matching is text-level, so the decode "
+ "node's token-id stop set cannot serve it."
+ )
+ return stop, include
diff --git a/tilert/pd_vllm/sampling.py b/tilert/pd_vllm/sampling.py
new file mode 100644
index 0000000..53f2cf5
--- /dev/null
+++ b/tilert/pd_vllm/sampling.py
@@ -0,0 +1,160 @@
+"""Request sampling params -> engine sampling params, for the decode node.
+
+The engine adapters each take the router's forwarded `sampling` dict and call
+their generator's ``update_sampling_params``. The dict comes straight from the
+client body (``pd_router._sampling_of``), so the translation from "what an
+OpenAI/vLLM client may send" to "what the engine accepts" belongs here, once,
+rather than three times across the profiles.
+"""
+
+from __future__ import annotations
+
+__all__ = [
+ "GREEDY_LOGPROBS_TOP_P",
+ "TOP_K_DISABLED",
+ "VLLM_DEFAULT_TOP_P",
+ "resolve_top_k",
+ "resolve_top_p",
+]
+
+# Two unrelated top_p values live here; keep them apart.
+#
+# ``VLLM_DEFAULT_TOP_P`` / ``resolve_top_p`` are the CLIENT-FACING nucleus: what
+# the request asked for, or the deployment's default when it asked for nothing.
+# ``GREEDY_LOGPROBS_TOP_P`` is an INTERNAL mechanism that makes the top-p kernel
+# behave as argmax; it is not a nucleus a client can ask for and never passes
+# through ``resolve_top_p``.
+
+# top_p for a greedy request that also wants log probabilities.
+#
+# The engine has no greedy log-probability export: greedy takes a separate top-1
+# kernel whose logits buffer the MTP draft head overwrites later in the tape. So
+# a greedy logprobs request is decoded on the TOP-P path instead, where the
+# export lives inside the sampling op — and this cutoff is what makes that path
+# behave exactly like argmax.
+#
+# The kernel picks its nucleus as the first index where the cumulative
+# probability exceeds top_p (top_p.cuh, "Find cutoff point"). cum_probs[0] is the
+# largest full-vocab probability and so is at least 1/vocab (~4e-6 at 248320),
+# which always exceeds this value: the cutoff is 0, and the multinomial then
+# iterates a single candidate. Under MTP verify the accept probability becomes
+# topks[0]/cum_probs[0] == 1, so an argmax draft is always accepted and any other
+# draft is rejected back to the argmax -- greedy speculative decoding, at full
+# MTP speed.
+#
+# Two orders of margin below 1/vocab is deliberate; do not raise it toward
+# realistic top_p values, or the cutoff stops being 0 and greedy silently becomes
+# sampling. The paired temperature is 1.0, NOT 0 -- see _greedy_logprobs_params.
+GREEDY_LOGPROBS_TOP_P = 1e-9
+
+# vLLM's framework default, from ``_DEFAULT_SAMPLING_PARAMS`` in
+# ``vllm/entrypoints/openai/chat_completion/protocol.py``. The protocol field
+# itself defaults to None; ``to_sampling_params`` then resolves
+#
+# client explicit value > default_sampling_params (the deployed model's
+# generation_config.json) > this constant
+#
+# so this is the LAST link in vLLM's own chain, not the whole chain.
+VLLM_DEFAULT_TOP_P = 1.0
+
+
+def resolve_top_p(sampling: dict, default: float = VLLM_DEFAULT_TOP_P) -> float:
+ """The ``top_p`` for this request, resolved once for both PD legs.
+
+ Exists because the two legs used to disagree. The decode adapters each read
+ ``sampling.get("top_p", 0.95)`` while the vLLM prefill instance resolved its
+ own default through the chain above -- so a client that sent ``temperature``
+ but no ``top_p`` had token 1 sampled under one nucleus and tokens 2..N under
+ another, with nothing in the response to show it.
+
+ The fix is not a better constant: no constant can be right, because the
+ middle link is a property of the deployed checkpoint. It is resolving the
+ value in ONE place (the router) and sending it explicitly to BOTH legs, so
+ they agree by construction whatever the value is. ``default`` is the
+ deployment's stand-in for that middle link (``--default-top-p``); wiring
+ ``generation_config.json`` in behind it changes only what is passed here.
+
+ Adapters must not re-default: they receive an already-resolved number. A
+ greedy logprobs request does not come through here at all -- its top_p is
+ ``GREEDY_LOGPROBS_TOP_P``, chosen to defeat the nucleus rather than express
+ one.
+ """
+ raw = sampling.get("top_p")
+ if raw is None:
+ return float(default)
+ if isinstance(raw, bool):
+ raise ValueError("top_p must be a number, got bool")
+ return float(raw)
+
+
+# The sampler materialises kTopK=256 candidates per GPU (all-gathered to
+# kNumGpus*256), and its cutoff treats any top_k at or above that per-GPU bound
+# as "no rank cut":
+#
+# effective_top_k = (top_k > 0 && top_k < kTopK) ? top_k : kTotalTopKs
+# -- the engine's top_p sampler kernel
+#
+# So 256 is the value that disables the rank cut, and [1, 255] is the range the
+# kernel can actually apply. These two constants mirror that kernel's kTopK and
+# must be kept in step with it: if kTopK ever changes, TOP_K_DISABLED stops
+# meaning "disabled" and silently becomes a real rank cut.
+_KERNEL_TOP_K_POOL = 256
+TOP_K_DISABLED = _KERNEL_TOP_K_POOL
+_TOP_K_APPLIED_MAX = _KERNEL_TOP_K_POOL - 1
+
+
+def resolve_top_k(sampling: dict) -> int:
+ """The engine's top_k for this request, following vLLM's convention.
+
+ vLLM is the reference because the router sends ``top_k`` to the vLLM
+ prefill instance as well as here (``pd_router.build_prefill_body`` copies the
+ client body). If the two disagreed, one request would sample its first token
+ under vLLM's rules and the rest under ours.
+
+ vLLM (``SamplingParams``, ``gpu_input_batch.py``)::
+
+ top_k: int = 0 # "Set to 0 (or -1) to consider all tokens."
+ if 0 < top_k < vocab_size: applied
+ else: top_k = vocab_size # i.e. disabled
+
+ We follow that shape, with the kernel's candidate pool standing in for
+ vocab_size:
+
+ ================== ==========================================
+ request top_k result
+ ================== ==========================================
+ absent / null disabled
+ 0, -1 disabled -- vLLM's documented sentinels
+ [1, 255] applied
+ >= 256 disabled (pool bound; vLLM would apply it)
+ < -1 disabled here; the vLLM prefill instance
+ 400s it first, so the client sees an error
+ ================== ==========================================
+
+ The last two rows are the only divergences from vLLM and both are one-way
+ (we disable where vLLM would cut), so they can widen the sampled set but
+ never narrow it unexpectedly.
+
+ The 255 bound is the kernel's, not a policy choice: measured on 8xB200,
+ ``top_k`` of 300 and 2048 sample bit-identically to 256 over 200 seeds, so
+ the sampler already treats everything at or above 256 as no cut.
+
+ Known kernel behaviour an applied value inherits: the cutoff is inclusive,
+ so ``top_k=N`` keeps N+1 candidates (measured: 1 -> logit ranks {0,1},
+ 4 -> {0..4}). ``top_k=1`` therefore samples between the top two tokens
+ rather than being strict argmax. Greedy requests do not go through here --
+ ``temperature ~ 0`` selects a separate captured graph.
+
+ Note ``>= 256`` is *not* what every client expects: some vendor APIs
+ document "a value greater than 100 indicates that the top_k strategy is not
+ enabled". A generation_config ``top_k`` in the low tens is well inside the
+ applied range, so such a default is honoured either way; only explicit
+ values in [101, 255] differ, and there we follow vLLM.
+ """
+ raw = sampling.get("top_k")
+ if raw is None:
+ return TOP_K_DISABLED
+ k = int(raw)
+ if k < 1 or k > _TOP_K_APPLIED_MAX:
+ return TOP_K_DISABLED
+ return k
diff --git a/tilert/pd_vllm/stop_strings.py b/tilert/pd_vllm/stop_strings.py
new file mode 100644
index 0000000..a2a3fcb
--- /dev/null
+++ b/tilert/pd_vllm/stop_strings.py
@@ -0,0 +1,200 @@
+"""Text-level ``stop`` matching, on the router's side of the detokeniser.
+
+``stop`` is a property of the decoded TEXT, not of token ids: a stop string
+routinely spans two tokens ("Observation:" is often three), and byte-level
+BPE can split one character across tokens. The node emits ids; the router is
+where text exists.
+
+:func:`check_stop_strings` is a port of vLLM's ``v1/engine/detokenizer.py``, so
+the same request stops at the same character on both stacks. Two rules are easy
+to get wrong: the search starts at ``1 - new_char_count - len(stop_str)``, so a
+stop straddling the delta boundary is still found; and when several match in one
+step -- routine under MTP -- the one COMPLETING EARLIEST wins, so the result does
+not depend on the batch size. Ties go to stop-list order.
+
+:class:`StopWindow` splits accumulation from release as vLLM does: ``push``
+absorbs and matches, ``take`` is a read cursor. Neither then needs a queue of
+held pieces. One difference: vLLM keeps the whole reply because its non-streaming
+response needs it, so this class drops what has gone out -- keeping it cost
+0.75 s and 200 KB over a 200k-token generation.
+"""
+
+from __future__ import annotations
+
+import sys
+
+__all__ = ["StopWindow", "check_stop_strings", "resolve_stop"]
+
+
+def resolve_stop(body: dict) -> list[str]:
+ """The request's stop strings, normalised to a list.
+
+ ``str | list[str] | None``, as vLLM accepts. An empty string is REJECTED, not
+ dropped: it matches at position 0 of everything, so it is not a neutral value
+ the way ``[]`` is, and dropping it would serve an unrestricted completion to
+ a client who asked for a restricted one. vLLM raises the same from
+ ``SamplingParams._verify_args``, and the router strips ``stop`` from the
+ prefill request, so vLLM no longer gets the chance to.
+ """
+ raw = body.get("stop")
+ if raw is None:
+ return []
+ if isinstance(raw, str):
+ raw = [raw]
+ if not isinstance(raw, (list, tuple)):
+ raise ValueError(f"stop must be a string or list of strings, " f"got {type(raw).__name__}")
+ out = []
+ for s in raw:
+ if not isinstance(s, str):
+ raise ValueError(f"stop entries must be strings, " f"got {type(s).__name__}")
+ if not s:
+ raise ValueError("stop cannot contain an empty string")
+ out.append(s)
+ return out
+
+
+def check_stop_strings(
+ output_text: str,
+ new_char_count: int,
+ stop: list[str],
+ include_in_output: bool,
+) -> tuple[str, int] | None:
+ """``(stop_string, truncate_to)`` if one matched, else ``None``.
+
+ ``truncate_to`` is the length to cut ``output_text`` to, or ``-1`` for none.
+ vLLM port; see the module docstring for the two rules.
+ """
+ if not new_char_count or not stop:
+ return None
+
+ best_stop_str: str | None = None
+ best_stop_index = 0
+ best_end = sys.maxsize
+ for stop_str in stop:
+ stop_len = len(stop_str)
+ # Start before the new text so a stop spanning the boundary is found,
+ # without re-scanning text that was already checked.
+ stop_index = output_text.find(stop_str, 1 - new_char_count - stop_len)
+ if stop_index == -1:
+ continue
+ end = stop_index + stop_len
+ if end < best_end:
+ best_stop_str = stop_str
+ best_stop_index = stop_index
+ best_end = end
+
+ if best_stop_str is None:
+ return None
+ if include_in_output:
+ if best_end >= len(output_text):
+ return best_stop_str, -1
+ return best_stop_str, best_end
+ return best_stop_str, best_stop_index
+
+
+class StopWindow:
+ """Decoded text, matched against the stop strings, behind a read cursor.
+
+ ``take`` lags ``push`` by ``max(len(s) for s in stop) - 1`` characters: text
+ on the wire cannot be recalled, and the last few may begin a stop string.
+ Nothing lags without stop strings, or when the stop stays in the output --
+ nothing to remove then, the same condition vLLM uses for its
+ ``stop_buffer_length``.
+ """
+
+ def __init__(self, stop: list[str], include_in_output: bool = False):
+ self.stop = list(stop)
+ self.include_in_output = include_in_output
+ longest = max((len(s) for s in self.stop), default=1) - 1
+ # How far `take` stays behind the end of the text.
+ self._hold = 0 if include_in_output else longest
+ # Look-back kept behind the cursor: the matcher reaches at most
+ # len(stop)-1 characters back from the newest delta.
+ self._keep = longest
+ # Slack above what matching needs, so trimming is amortised. Zero
+ # without stop strings, which empties the window on every `take` --
+ # one code path, degenerating to a pass-through.
+ self._slack = 4096 if self.stop else 0
+
+ self._text = "" # a WINDOW of the stream, not all of it
+ self._base = 0 # absolute index of self._text[0]
+ self._taken = 0 # absolute count handed to `take`
+ self._stopped: str | None = None
+
+ # ── what the caller reports ─────────────────────────────────────────────
+
+ @property
+ def stopped(self) -> str | None:
+ """The stop string that ended the sequence, or None."""
+ return self._stopped
+
+ @property
+ def hold(self) -> int:
+ """How far `take` stays behind the newest text, in characters.
+
+ A caller pairing text with per-token metadata rounds this down to a
+ token boundary: a character count lands mid-token.
+ """
+ return self._hold
+
+ @property
+ def visible(self) -> int:
+ """How many characters `take` has handed out.
+
+ A caller holding per-token metadata pairs it against this: an entry is
+ due once the text it describes is past this point.
+ """
+ return self._taken
+
+ # ── the stream ──────────────────────────────────────────────────────────
+
+ def push(self, delta: str) -> None:
+ """Absorb newly decoded text and match. Releasing is ``take``'s job --
+
+ doing both here is what forced a released/unreleased split.
+ """
+ if self._stopped is not None or not delta:
+ return
+ self._text += delta
+ # Matcher offsets are window-relative; `_base` converts them.
+ hit = check_stop_strings(self._text, len(delta), self.stop, self.include_in_output)
+ if hit is None:
+ return
+ self._stopped, truncate_to = hit
+ if truncate_to != -1:
+ self._text = self._text[:truncate_to]
+
+ def take(self, *, final: bool = False, limit: int | None = None) -> str:
+ """The text the client may see now.
+
+ ``final`` releases the held tail: nothing more can arrive to turn it into
+ a stop. A matched stop is also final -- the text is already cut.
+
+ ``limit`` is an absolute ceiling a caller sets to keep the release on a
+ boundary of its own; it does not apply once the text is final or cut,
+ when everything must go out regardless.
+ """
+ end = self._base + len(self._text)
+ if not (final or self._stopped is not None):
+ end -= self._hold
+ if limit is not None and limit < end:
+ end = limit
+ if end <= self._taken:
+ return ""
+ out = self._text[self._taken - self._base : end - self._base]
+ self._taken = end
+ self._trim()
+ return out # noqa: R504 (_trim mutates the window after the slice)
+
+ def _trim(self) -> None:
+ """Drop text the matcher can no longer reach.
+
+ Keeps from ``_taken - _keep`` onward, and only runs once the window
+ exceeds that by ``_slack``, so the copy is amortised.
+ """
+ if len(self._text) <= self._keep + self._hold + self._slack:
+ return
+ cut = self._taken - self._keep - self._base
+ if cut > 0:
+ self._text = self._text[cut:]
+ self._base += cut
diff --git a/tilert/pd_vllm/wire.py b/tilert/pd_vllm/wire.py
index 284850c..0d03583 100644
--- a/tilert/pd_vllm/wire.py
+++ b/tilert/pd_vllm/wire.py
@@ -1,4 +1,31 @@
-"""Shared control-plane protocol for vLLM-prefill -> TileRT-decode PD."""
+"""Shared control-plane protocol for vLLM-prefill -> TileRT-decode PD.
+
+Model-agnostic. The per-model wire *layout* (which regions exist, their
+sizes and offsets) lives in the model profile (``profiles/``); this module
+owns only the framing everything shares:
+
+ - length-prefixed JSON messages (send_msg / recv_msg)
+ - the hello envelope (server -> client): common fields + a profile-supplied
+ ``layout`` dict of region base addresses
+ - request / done messages (client -> server)
+ - local_ip / derive_rid helpers
+
+Per-request control flow (one TCP connection per participating rank):
+ server -> client : hello {magic, protocol_version, layout_version,
+ session_id, max_seq_len, busy, **layout}
+ client -> server : request {rid, rank, seq_len, last_prompt_token, sampling?,
+ prompt_token_ids?} (rank 0 only, penalties only)
+ server -> client : accept {accepted: true, rid, rank, generation}
+ or reject {accepted: false, error, ...}
+ client -> server : done {done, rid, rank, generation} (after the RDMA write)
+
+The accept step is an ADMISSION step, not an acknowledgement: the receive buffer
+holds one request at a time, so a sender that writes before being admitted can
+land its KV inside a request the decode node is already serving. Nothing detects
+that afterwards -- the victim decodes from a mix of two prompts' state and
+returns a confident wrong answer. So the sender must not touch RDMA until it has
+an accept whose rid, rank and generation all match what it asked for.
+"""
import json
import socket
@@ -6,6 +33,14 @@
MAGIC = "tilert-pd"
+# Control-plane version. 1 = the original "write immediately after the request
+# message" flow; 2 adds the accept/reject admission step and carries a
+# generation on accept/done. Bumped together on both ends: a v1 sender paired
+# with a v2 receiver would write without being admitted, which is the exact
+# failure this version exists to remove, so the pairing is refused at hello
+# rather than tolerated.
+PROTOCOL_VERSION = 2
+
NUM_RANKS = 8
EXPECTED_RANKS = tuple(range(NUM_RANKS))
@@ -23,6 +58,28 @@ def local_ip(probe_addr: str | None = None) -> str:
s.close()
+def wants_prompt_token_ids(sampling: dict | None) -> bool:
+ """Should this request ship its full prompt id list to the decode node?
+
+ Only ``repetition_penalty`` is scoped over prompt UNION output, so only it
+ needs the prompt half of the decode-side bitmap. ``presence_penalty`` is
+ output-scoped and does not justify the payload on its own; a value of
+ exactly 1.0 is the kernel's no-op and needs nothing either. Mirrors vLLM's
+ ``needs_prompt_token_ids`` gate in
+ ``v1/worker/gpu_input_batch.py::make_sampling_metadata``, which likewise
+ skips the copy when no request in the batch has penalties.
+ """
+ if not sampling:
+ return False
+ rep = sampling.get("repetition_penalty")
+ if rep is None:
+ return False
+ try:
+ return float(rep) != 1.0
+ except (TypeError, ValueError):
+ return False
+
+
def derive_rid(request_id: str) -> str:
"""Map a vLLM request/response id to the client-visible rid.
@@ -74,15 +131,21 @@ def hello_msg(
layout: dict,
busy: bool,
) -> dict:
- """Build the common hello envelope.
+ """Common hello envelope.
+
+ ``transport`` names the RDMA backend and ``transport_meta`` carries its connection info
+ (mooncake: session_id; nixl: nixl_meta/nixl_dev). ``layout`` carries profile-specific region
+ base addresses (e.g. kv_base / pe_base / ki_base).
- ``transport`` names the RDMA backend and ``transport_meta`` carries its
- connection info (mooncake: session_id; nixl: nixl_meta/nixl_dev).
- ``layout`` carries profile-specific region base addresses (e.g. gdn_base /
- gqa_k_base / kv_base).
+ ``protocol_version`` is the CONTROL-plane version, independent of
+ ``layout_version`` (which versions the buffer geometry). It exists so a
+ sender that does not wait for :func:`accept_msg` cannot be paired with a
+ receiver that expects it: the mismatch fails at handshake instead of being
+ guessed at run time.
"""
return {
"magic": MAGIC,
+ "protocol_version": PROTOCOL_VERSION,
"layout_version": layout_version,
"transport": transport,
"max_seq_len": max_seq_len,
@@ -90,3 +153,28 @@ def hello_msg(
**transport_meta,
**layout,
}
+
+
+def accept_msg(rid: str, rank: int, generation: int) -> dict:
+ """Receiver -> sender: this rank may now RDMA-write for ``rid``.
+
+ ``generation`` identifies the receive-buffer tenancy the write is authorised
+ against. It is echoed in :func:`done_msg` so a ``done`` arriving after the
+ buffer has been handed to a later request is recognisable as stale rather
+ than counted towards the current one.
+ """
+ return {"accepted": True, "rid": rid, "rank": rank, "generation": generation}
+
+
+def reject_msg(reason: str, **extra) -> dict:
+ """Receiver -> sender: do NOT write. ``reason`` is for the sender's log.
+
+ The sender must treat anything that is not a matching accept as a rejection,
+ so an unrecognised reason is still safe.
+ """
+ return {"accepted": False, "error": reason, **extra}
+
+
+def done_msg(rid: str, rank: int, generation: int) -> dict:
+ """Sender -> receiver: this rank's RDMA write has completed."""
+ return {"done": True, "rid": rid, "rank": rank, "generation": generation}