From 054ff0199c66e919f9eae1a4ab533b704293a9ca Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Mon, 3 Aug 2026 18:05:38 -0700 Subject: [PATCH 01/16] test(e2e): cover the ZMQ direct backend for vLLM and TokenSpeed Run the existing local (non-PD/EPD) e2e cases over the ZMQ direct backend, in addition to gRPC/HTTP, for both vLLM and TokenSpeed. - Add ConnectionMode.ZMQ and an E2E_CONNECTION_MODE per-lane override (mirrors how E2E_RUNTIME picks the engine): a grpc/http case runs over ZMQ without a separate parametrize value. - Worker: derive the ipc:// base_url from serve._zmq_ipc_url, skip the worker health wait (readiness is gateway-gated), and build the headless vLLM/TokenSpeed commands by delegating to the serve launchers so engine flags and the FNV-1a handshake port stay in lockstep with production. - Gateway: pass --router-backend explicitly (the router cannot probe the wire from an ipc:// URL). - Worker pool: do not cache ZMQ workers (an engine dials one gateway's handshake sockets and cannot be reused); stop them with the gateway. - Exclude the gRPC-only families from ZMQ lanes (PD, EPD, multi-worker) and collapse grpc/http twins onto a single ZMQ run, via pytest_collection_modifyitems using only public API. - CI: add an e2e-1gpu-chat-zmq lane (vllm, tokenspeed) and thread the connection_mode input through the gpu job. TokenSpeed keeps its pinned install commit. - Add no-GPU unit tests for the ZMQ command builders and base_url. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- .github/workflows/e2e-gpu-job.yml | 6 ++ .github/workflows/pr-test-rust.yml | 39 ++++++++++++- e2e_test/fixtures/hooks.py | 78 ++++++++++++++++++++++++- e2e_test/fixtures/setup_backend.py | 39 +++++++++++-- e2e_test/infra/__init__.py | 4 ++ e2e_test/infra/constants.py | 18 +++++- e2e_test/infra/gateway.py | 8 ++- e2e_test/infra/test_zmq_cmd_builders.py | 67 +++++++++++++++++++++ e2e_test/infra/worker.py | 62 ++++++++++++++++++-- e2e_test/infra/worker_pool.py | 6 +- 10 files changed, 311 insertions(+), 16 deletions(-) create mode 100644 e2e_test/infra/test_zmq_cmd_builders.py diff --git a/.github/workflows/e2e-gpu-job.yml b/.github/workflows/e2e-gpu-job.yml index b5abed138..8ed684791 100644 --- a/.github/workflows/e2e-gpu-job.yml +++ b/.github/workflows/e2e-gpu-job.yml @@ -54,6 +54,11 @@ on: type: string default: "nixl" description: "KV transfer backend for vLLM PD workers: nixl or mooncake" + connection_mode: + required: false + type: string + default: "" + description: "Wire override for local backends: zmq runs the local cases over ZMQ" jobs: run: @@ -66,6 +71,7 @@ jobs: E2E_RUNTIME: ${{ inputs.engine }} E2E_GPU_TIER: ${{ inputs.gpu_tier }} E2E_VLLM_KV_BACKEND: ${{ inputs.vllm_kv_backend }} + E2E_CONNECTION_MODE: ${{ inputs.connection_mode }} ROUTER_LOCAL_MODEL_PATH: /models steps: - name: Checkout code diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml index 9230953fc..8aae5d13f 100644 --- a/.github/workflows/pr-test-rust.yml +++ b/.github/workflows/pr-test-rust.yml @@ -543,6 +543,42 @@ jobs: test_dirs: ${{ matrix.test_dirs || 'e2e_test/chat_completions' }} secrets: inherit + e2e-1gpu-chat-zmq: + name: e2e-1gpu-chat-zmq (${{ matrix.engine }}) + needs: [build-wheel, detect-changes] + if: >- + always() + && !cancelled() + && needs.build-wheel.result == 'success' + && (github.event_name != 'pull_request' + || (needs.detect-changes.result == 'success' + && (needs.detect-changes.outputs.common == 'true' + || needs.detect-changes.outputs.chat-completions == 'true'))) + # Same single-worker chat suite as e2e-1gpu-chat, driven over the ZMQ + # direct-backend wire. The collection hook deselects the gRPC-only + # families (PD, EPD, multi-worker) for a ZMQ lane, so this covers the + # local cases only. + strategy: + fail-fast: false + matrix: + include: + - engine: vllm + timeout: 24 + test_timeout: 18 + - engine: tokenspeed + timeout: 50 + test_timeout: 18 + uses: ./.github/workflows/e2e-gpu-job.yml + with: + engine: ${{ matrix.engine }} + gpu_tier: "1" + runner: 1-gpu-h100 + timeout: ${{ matrix.timeout }} + test_timeout: ${{ matrix.test_timeout }} + test_dirs: e2e_test/chat_completions + connection_mode: zmq + secrets: inherit + e2e-1gpu-completions: name: e2e-1gpu-completions (${{ matrix.engine }}) needs: [build-wheel, detect-changes] @@ -1049,7 +1085,7 @@ jobs: path: benchmark_go_bindings/ finish: - needs: [pre-commit, python-lint, grpc-proto-build-check, build-wheel, python-unit-tests, unit-tests, benchmarks, e2e-1gpu-chat, e2e-1gpu-completions, e2e-1gpu-embeddings, e2e-1gpu-gateway, e2e-1gpu-responses, e2e-2gpu-pd, e2e-4gpu-chat, e2e-4gpu-gateway, e2e-4gpu-epd, e2e-vendor, go-unit-tests, go-bindings-e2e] + needs: [pre-commit, python-lint, grpc-proto-build-check, build-wheel, python-unit-tests, unit-tests, benchmarks, e2e-1gpu-chat, e2e-1gpu-chat-zmq, e2e-1gpu-completions, e2e-1gpu-embeddings, e2e-1gpu-gateway, e2e-1gpu-responses, e2e-2gpu-pd, e2e-4gpu-chat, e2e-4gpu-gateway, e2e-4gpu-epd, e2e-vendor, go-unit-tests, go-bindings-e2e] if: always() runs-on: k8s-runner-cpu permissions: {} @@ -1064,6 +1100,7 @@ jobs: "${{ needs.unit-tests.result }}" == "failure" || \ "${{ needs.benchmarks.result }}" == "failure" || \ "${{ needs.e2e-1gpu-chat.result }}" == "failure" || \ + "${{ needs.e2e-1gpu-chat-zmq.result }}" == "failure" || \ "${{ needs.e2e-1gpu-completions.result }}" == "failure" || \ "${{ needs.e2e-1gpu-embeddings.result }}" == "failure" || \ "${{ needs.e2e-1gpu-gateway.result }}" == "failure" || \ diff --git a/e2e_test/fixtures/hooks.py b/e2e_test/fixtures/hooks.py index 3eb440864..8986252e1 100644 --- a/e2e_test/fixtures/hooks.py +++ b/e2e_test/fixtures/hooks.py @@ -14,10 +14,16 @@ import os import pytest -from infra import cleanup_pool, get_runtime +from infra import ConnectionMode, cleanup_pool, get_connection_mode_override, get_runtime from .markers import resolve_class_marker +# Local wires a plain (non-PD/EPD) test case can be authored with; these are the +# only backends a ZMQ lane reuses (mapped onto ZMQ by the ``setup_backend`` +# fixture). Anything else — ``pd_*``/``epd_*``, EPD topology tuples, cloud +# vendor names — is out of scope for ZMQ. +_ZMQ_LOCAL_WIRES = frozenset({"grpc", "http"}) + # --------------------------------------------------------------------------- # Marker registration # --------------------------------------------------------------------------- @@ -117,6 +123,66 @@ def _get_marker(item: pytest.Item, name: str): return resolve_class_marker(item, name) +def _setup_backend_param(item: pytest.Item): + """Return the ``setup_backend`` parametrize value for an item, or None.""" + callspec = getattr(item, "callspec", None) + if callspec is None: + return None + return getattr(callspec, "params", {}).get("setup_backend") + + +def _is_multi_worker(item: pytest.Item) -> bool: + """True when the item's ``workers`` marker asks for a PD/multi-worker topology.""" + marker = resolve_class_marker(item, "workers") + if marker is None: + return False + kwargs = marker.kwargs + if (kwargs.get("count") or 1) > 1: + return True + return bool(kwargs.get("prefill") or kwargs.get("decode")) + + +def _zmq_dedup_key(item: pytest.Item) -> tuple: + """Group key ignoring the ``setup_backend`` value. + + Lets us collapse a case parametrized on both ``grpc`` and ``http`` into a + single ZMQ run (they map onto the same wire) while keeping distinct + ``api_client`` (or other) parametrizations apart. + """ + callspec = getattr(item, "callspec", None) + params = getattr(callspec, "params", {}) or {} + others = tuple(sorted((k, repr(v)) for k, v in params.items() if k != "setup_backend")) + return (item.nodeid.split("[", 1)[0], others) + + +def _filter_zmq_items(items: list[pytest.Item]) -> tuple[list[pytest.Item], list[pytest.Item]]: + """Split items into (kept, deselected) for a ZMQ lane. + + Keeps single-worker local cases (``grpc``/``http`` map onto ZMQ) and drops + the gRPC-only families: PD (``pd_*``), EPD (``epd_*`` and topology tuples), + and multiple-worker topologies. A case authored for both ``grpc`` and + ``http`` is collapsed to one ZMQ run. Items without a ``setup_backend`` + parametrization are left untouched. + """ + kept: list[pytest.Item] = [] + deselected: list[pytest.Item] = [] + groups_with_grpc = {_zmq_dedup_key(it) for it in items if _setup_backend_param(it) == "grpc"} + for item in items: + param = _setup_backend_param(item) + if param is None: + kept.append(item) + continue + if param not in _ZMQ_LOCAL_WIRES or _is_multi_worker(item): + deselected.append(item) + continue + # Collapse the http twin when a grpc one covers the same ZMQ run. + if param == "http" and _zmq_dedup_key(item) in groups_with_grpc: + deselected.append(item) + continue + kept.append(item) + return kept, deselected + + def pytest_collection_modifyitems( config: pytest.Config, items: list[pytest.Item], @@ -124,7 +190,9 @@ def pytest_collection_modifyitems( """Filter + order collected tests. Filtering: env vars ``E2E_ENGINE``, ``E2E_VENDOR``, ``E2E_GPU_TIER`` - select the matching slice when set. + select the matching slice when set. When ``E2E_CONNECTION_MODE=zmq`` the + lane additionally drops the gRPC-only families (PD, EPD, multi-worker) and + collapses ``grpc``/``http`` twins onto a single ZMQ run. Ordering: items are sorted by ``(backend, model)`` so consecutive classes that share a backend cluster together. This is what lets @@ -157,6 +225,12 @@ def pytest_collection_modifyitems( selected.append(item) items[:] = selected + if get_connection_mode_override() == ConnectionMode.ZMQ: + kept, deselected = _filter_zmq_items(items) + if deselected: + config.hook.pytest_deselected(items=deselected) + items[:] = kept + items.sort(key=_pool_sort_key) diff --git a/e2e_test/fixtures/setup_backend.py b/e2e_test/fixtures/setup_backend.py index aedce60fa..654f1aa14 100644 --- a/e2e_test/fixtures/setup_backend.py +++ b/e2e_test/fixtures/setup_backend.py @@ -30,6 +30,7 @@ ConnectionMode, Gateway, WorkerType, + get_connection_mode_override, get_runtime, launch_cloud_gateway, ) @@ -145,6 +146,11 @@ def setup_backend(request: pytest.FixtureRequest): is_pd = backend_name.startswith("pd_") protocol = backend_name.replace("epd_", "").replace("pd_", "") connection_mode = ConnectionMode(protocol) + # A lane can override the local wire (e.g. run grpc/http cases over ZMQ); + # PD/EPD keep their own wire since they are excluded from those lanes. + mode_override = get_connection_mode_override() + if mode_override is not None and not is_pd and not is_epd: + connection_mode = mode_override engine = get_runtime() model_path = get_model_spec(model_id)["model"] workers_config = get_marker_kwargs(request, "workers", defaults=_WORKER_DEFAULTS) @@ -232,18 +238,30 @@ def _setup_local( gpus=workers_config.get("gpus"), extra_engine_args=workers_config.get("extra_engine_args"), ) + # ZMQ engines dial this gateway's handshake sockets, so they cannot be + # reused by a later class's gateway — the pool starts them fresh and the + # caller owns their teardown (like the PD path). gRPC/HTTP workers stay + # in the pool and outlive the gateway. + is_zmq = connection_mode == ConnectionMode.ZMQ try: _start_gateway( gateway, gateway_config, worker_urls=[w.base_url for w in workers], model_path=model_path, + backend=engine if is_zmq else None, ) logger.info("%s backend ready at %s", backend_name, gateway.base_url) yield backend_name, model_path, _make_openai_client(gateway), gateway finally: - logger.info("Tearing down %s backend (workers stay in pool)", backend_name) + logger.info( + "Tearing down %s backend (%s)", + backend_name, + "stopping ZMQ workers" if is_zmq else "workers stay in pool", + ) gateway.shutdown() + if is_zmq: + stop_workers(workers) # --------------------------------------------------------------------------- @@ -463,20 +481,31 @@ def test_router_state(backend_router): backend_name = request.param model_id = os.environ.get(ENV_MODEL, DEFAULT_MODEL) connection_mode = ConnectionMode(backend_name) + mode_override = get_connection_mode_override() + if mode_override is not None: + connection_mode = mode_override + engine = get_runtime() model_path = get_model_spec(model_id)["model"] + is_zmq = connection_mode == ConnectionMode.ZMQ # Route through the pool so we evict any cached class-scope worker - # holding the GPUs we need. The pool retains ownership; we don't stop - # the workers ourselves. + # holding the GPUs we need. The pool retains ownership of gRPC/HTTP + # workers; ZMQ engines are bound to this gateway, so we stop them here. workers = get_pool().acquire( model_id=model_id, - engine=get_runtime(), + engine=engine, mode=connection_mode, count=1, ) gateway = Gateway() try: - gateway.start(worker_urls=[w.base_url for w in workers], model_path=model_path) + gateway.start( + worker_urls=[w.base_url for w in workers], + model_path=model_path, + backend=engine if is_zmq else None, + ) yield gateway finally: gateway.shutdown() + if is_zmq: + stop_workers(workers) diff --git a/e2e_test/infra/__init__.py b/e2e_test/infra/__init__.py index 89044d746..bde475869 100644 --- a/e2e_test/infra/__init__.py +++ b/e2e_test/infra/__init__.py @@ -11,6 +11,7 @@ DEFAULT_RUNTIME, DEFAULT_STARTUP_TIMEOUT, ENV_BACKENDS, + ENV_CONNECTION_MODE, ENV_MODEL, ENV_MODELS, ENV_RUNTIME, @@ -32,6 +33,7 @@ ConnectionMode, Runtime, WorkerType, + get_connection_mode_override, get_runtime, is_mlx, is_sglang, @@ -102,6 +104,7 @@ "ENV_BACKENDS", "ENV_MODEL", "ENV_RUNTIME", + "ENV_CONNECTION_MODE", "ENV_STARTUP_TIMEOUT", "ENV_SKIP_MODEL_POOL", "ENV_SKIP_BACKEND_SETUP", @@ -109,6 +112,7 @@ "ENV_SHOW_WORKER_LOGS", # Runtime helpers "get_runtime", + "get_connection_mode_override", "is_vllm", "is_sglang", "is_trtllm", diff --git a/e2e_test/infra/constants.py b/e2e_test/infra/constants.py index ce4b523c6..ea0db97d7 100644 --- a/e2e_test/infra/constants.py +++ b/e2e_test/infra/constants.py @@ -9,6 +9,7 @@ class ConnectionMode(StrEnum): HTTP = "http" GRPC = "grpc" + ZMQ = "zmq" class WorkerType(StrEnum): @@ -35,7 +36,7 @@ class Runtime(StrEnum): # Convenience sets -LOCAL_MODES = frozenset({ConnectionMode.HTTP, ConnectionMode.GRPC}) +LOCAL_MODES = frozenset({ConnectionMode.HTTP, ConnectionMode.GRPC, ConnectionMode.ZMQ}) LOCAL_RUNTIMES = frozenset( {Runtime.SGLANG, Runtime.VLLM, Runtime.TRTLLM, Runtime.MLX, Runtime.TOKENSPEED} ) @@ -59,6 +60,9 @@ class Runtime(StrEnum): ENV_RUNTIME = ( "E2E_RUNTIME" # Runtime for gRPC tests — one of Runtime.{SGLANG,VLLM,TRTLLM,TOKENSPEED} ) +ENV_CONNECTION_MODE = ( + "E2E_CONNECTION_MODE" # Per-lane wire override — see get_connection_mode_override +) ENV_STARTUP_TIMEOUT = "E2E_STARTUP_TIMEOUT" ENV_SKIP_MODEL_POOL = "SKIP_MODEL_POOL" ENV_SKIP_BACKEND_SETUP = "SKIP_BACKEND_SETUP" @@ -125,6 +129,18 @@ def is_tokenspeed() -> bool: return get_runtime() == "tokenspeed" +def get_connection_mode_override() -> "ConnectionMode | None": + """Per-lane wire-protocol override for local backends. + + Set ``E2E_CONNECTION_MODE`` to run the existing local test cases over a + different wire (like ``E2E_RUNTIME`` picks the engine): a ``grpc``/``http`` + case then runs over that mode without a separate parametrize value. PD/EPD + backends keep their own wire. Returns ``None`` when unset. + """ + value = os.environ.get(ENV_CONNECTION_MODE) + return ConnectionMode(value.lower()) if value else None + + ENV_VLLM_KV_BACKEND = "E2E_VLLM_KV_BACKEND" diff --git a/e2e_test/infra/gateway.py b/e2e_test/infra/gateway.py index 145e9d272..66488ad99 100644 --- a/e2e_test/infra/gateway.py +++ b/e2e_test/infra/gateway.py @@ -122,6 +122,7 @@ def start( igw_mode: bool = False, cloud_backend: str | None = None, history_backend: str = "memory", + backend: str | None = None, policy: str = "round_robin", timeout: float = DEFAULT_ROUTER_TIMEOUT, show_output: bool | None = None, @@ -229,8 +230,13 @@ def start( self.model_path = model_path self.pd_mode = False self.igw_mode = False + mode_args = ["--model-path", model_path, "--worker-urls", *worker_urls] + # ZMQ workers share one wire across engine runtimes, so the router + # cannot probe the backend from the ipc:// URL — pin it explicitly. + if backend is not None: + mode_args += ["--router-backend", backend] self._launch( - mode_args=["--model-path", model_path, "--worker-urls", *worker_urls], + mode_args=mode_args, timeout=timeout, show_output=show_output, extra_args=extra_args, diff --git a/e2e_test/infra/test_zmq_cmd_builders.py b/e2e_test/infra/test_zmq_cmd_builders.py new file mode 100644 index 000000000..820593364 --- /dev/null +++ b/e2e_test/infra/test_zmq_cmd_builders.py @@ -0,0 +1,67 @@ +"""Unit tests for the ZMQ direct-backend command builders (no GPU).""" + +from __future__ import annotations + +import pytest + +# The builders delegate to smg.serve, so the wheel must be importable. +serve = pytest.importorskip("smg.serve") + +from infra.constants import ConnectionMode # noqa: E402 +from infra.worker import Worker # noqa: E402 + +_VLLM_MODEL = "meta-llama/Llama-3.2-1B-Instruct" +_TS_MODEL = "Qwen/Qwen3.5-9B" + + +def _worker(engine, port=50111): + return Worker( + model_id=_TS_MODEL if engine == "tokenspeed" else _VLLM_MODEL, + engine=engine, + port=port, + gpu_ids=[0], + mode=ConnectionMode.ZMQ, + ) + + +def test_zmq_base_url_matches_serve_helper(): + w = _worker("vllm", port=50123) + assert w.base_url == serve._zmq_ipc_url(50123) + assert w.base_url.startswith("ipc://") + + +def test_vllm_zmq_cmd_is_headless_with_derived_handshake_port(): + w = _worker("vllm") + cmd = w._build_vllm_zmq_cmd("/models/llama", 1, {"vllm_args": ["--max-model-len", "2048"]}) + assert "serve" in cmd + assert "--headless" in cmd + assert "/models/llama" in cmd + # The engine dials the same tcp port SMG derives from the ipc path. + expected_port = serve._zmq_handshake_port(serve._zmq_ipc_url(w.port)) + assert cmd[cmd.index("--data-parallel-rpc-port") + 1] == str(expected_port) + # Model-spec engine args ride through. + assert cmd[cmd.index("--max-model-len") + 1] == "2048" + + +def test_tokenspeed_zmq_cmd_is_headless_with_derived_handshake_port(): + w = _worker("tokenspeed") + cmd = w._build_tokenspeed_zmq_cmd( + "/models/qwen", 1, {"tokenspeed_args": ["--attention-backend", "fa3"]} + ) + assert "serve" in cmd + assert "--headless" in cmd + assert "/models/qwen" in cmd + expected_port = serve._zmq_handshake_port(serve._zmq_ipc_url(w.port)) + assert cmd[cmd.index("--data-parallel-rpc-port") + 1] == str(expected_port) + assert cmd[cmd.index("--attention-backend") + 1] == "fa3" + + +def test_grpc_worker_still_uses_grpc_url(): + w = Worker( + model_id=_VLLM_MODEL, + engine="vllm", + port=50111, + gpu_ids=[0], + mode=ConnectionMode.GRPC, + ) + assert w.base_url == "grpc://127.0.0.1:50111" diff --git a/e2e_test/infra/worker.py b/e2e_test/infra/worker.py index e1bdb5c7e..f7b2346aa 100644 --- a/e2e_test/infra/worker.py +++ b/e2e_test/infra/worker.py @@ -2,6 +2,7 @@ from __future__ import annotations +import argparse import json import logging import os @@ -54,6 +55,13 @@ class Worker: @property def base_url(self) -> str: """Base URL for this worker.""" + if self.mode == ConnectionMode.ZMQ: + # ipc:// worker URL the router binds; the engine dials the tcp + # handshake port SMG derives from it. Reuse serve's helper so the + # path format stays in lockstep with the launcher and the router. + from smg.serve import _zmq_ipc_url + + return _zmq_ipc_url(self.port) if self.mode == ConnectionMode.GRPC: return f"grpc://{DEFAULT_HOST}:{self.port}" return f"http://{DEFAULT_HOST}:{self.port}" @@ -105,6 +113,17 @@ def start( return # Wait for health check + if self.mode == ConnectionMode.ZMQ: + # SMG (the router) binds the ZMQ sockets and this engine dials in; + # there is no worker port to probe. The gateway's readiness gate + # (wait_for_workers_ready) covers the engine, so just proceed. + logger.info( + "Worker %s spawned at %s (PID %d) — ZMQ readiness gated by the gateway", + self.model_id, + self.base_url, + self.process.pid, + ) + return if self.mode == ConnectionMode.GRPC: self._wait_grpc_healthy(timeout) else: @@ -184,7 +203,9 @@ def _build_cmd(self) -> list[str]: if self.engine == "sglang": cmd = self._build_sglang_cmd(model_path, tp_size, features, spec) elif self.engine == "vllm": - if self.mode == ConnectionMode.GRPC: + if self.mode == ConnectionMode.ZMQ: + cmd = self._build_vllm_zmq_cmd(model_path, tp_size, spec) + elif self.mode == ConnectionMode.GRPC: cmd = self._build_vllm_grpc_cmd(model_path, tp_size, spec) else: cmd = self._build_vllm_http_cmd(model_path, tp_size, spec) @@ -193,12 +214,15 @@ def _build_cmd(self) -> list[str]: elif self.engine == "mlx": cmd = self._build_mlx_cmd(model_path, spec) elif self.engine == "tokenspeed": - if self.mode != ConnectionMode.GRPC: + if self.mode == ConnectionMode.ZMQ: + cmd = self._build_tokenspeed_zmq_cmd(model_path, tp_size, spec) + elif self.mode == ConnectionMode.GRPC: + cmd = self._build_tokenspeed_grpc_cmd(model_path, tp_size, spec) + else: raise ValueError( - "TokenSpeed e2e workers only support gRPC mode; " + "TokenSpeed e2e workers only support gRPC or ZMQ mode; " "HTTP mode would go through the existing OpenAI frontend." ) - cmd = self._build_tokenspeed_grpc_cmd(model_path, tp_size, spec) else: raise ValueError(f"Unsupported engine: {self.engine}") @@ -252,6 +276,21 @@ def _build_sglang_cmd( return cmd + def _build_vllm_zmq_cmd(self, model_path: str, tp_size: int, spec: dict) -> list[str]: + """Build the headless vLLM EngineCore command for the ZMQ direct backend. + + Delegates to the ``smg serve`` launcher so the engine flags and the + FNV-1a handshake port stay identical to the production launch path. + """ + from smg.serve import VllmWorkerLauncher + + args = argparse.Namespace( + connection_mode="zmq", model=model_path, tensor_parallel_size=tp_size + ) + return VllmWorkerLauncher().build_command( + args, list(spec.get("vllm_args", [])), DEFAULT_HOST, self.port + ) + def _build_vllm_grpc_cmd(self, model_path: str, tp_size: int, spec: dict) -> list[str]: """Build vLLM gRPC server command.""" return self._build_vllm_base_cmd("vllm.entrypoints.grpc_server", model_path, tp_size, spec) @@ -316,6 +355,21 @@ def _build_mlx_cmd(self, model_path: str, spec: dict) -> list[str]: cmd.extend(extra) return cmd + def _build_tokenspeed_zmq_cmd(self, model_path: str, tp_size: int, spec: dict) -> list[str]: + """Build the headless TokenSpeed command for the ZMQ direct backend. + + Delegates to the ``smg serve`` launcher so the engine flags and the + FNV-1a handshake port stay identical to the production launch path. + """ + from smg.serve import TokenspeedWorkerLauncher + + args = argparse.Namespace( + connection_mode="zmq", model=model_path, tensor_parallel_size=tp_size + ) + return TokenspeedWorkerLauncher().build_command( + args, list(spec.get("tokenspeed_args", [])), DEFAULT_HOST, self.port + ) + def _build_tokenspeed_grpc_cmd(self, model_path: str, tp_size: int, spec: dict) -> list[str]: """Build TokenSpeed gRPC server command. diff --git a/e2e_test/infra/worker_pool.py b/e2e_test/infra/worker_pool.py index 75f58c4a3..7070cd943 100644 --- a/e2e_test/infra/worker_pool.py +++ b/e2e_test/infra/worker_pool.py @@ -90,8 +90,10 @@ def acquire( # Non-REGULAR workers (PD prefill/decode) aren't cached, but we # still have to release any cached regular worker first — it - # holds the GPUs the caller is about to claim. - if worker_type != WorkerType.REGULAR: + # holds the GPUs the caller is about to claim. ZMQ workers are + # likewise uncached: the engine dials one gateway's handshake + # sockets and cannot be reused by the next class's gateway. + if worker_type != WorkerType.REGULAR or mode == ConnectionMode.ZMQ: if self._key is not None: logger.info( "WorkerPool: evicting %s to free GPUs for non-REGULAR %s/%s", From 4e49efa490aba8248cf066ce46e904de9f545dc0 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Mon, 3 Aug 2026 19:32:54 -0700 Subject: [PATCH 02/16] fix(e2e): correct ZMQ router flag and address review feedback The ZMQ chat lane failed CI because gateway.py passed --router-backend, which launch_router.py does not accept (the flag is --backend); the router crashed on every launch and the lane timed out. Switch to --backend. Also address review feedback: - disambiguate worker-log artifact names by connection mode so the gRPC and ZMQ vllm chat legs no longer collide - reject ConnectionMode.ZMQ for engines that can't speak it via a shared validation helper, failing fast instead of timing out - distinguish an unset E2E_CONNECTION_MODE from a set-but-empty/invalid value, which now raises - scope the smg.serve importorskip to the ZMQ command-builder tests so the gRPC test runs without the wheel - add unit tests for connection-mode parsing, engine validation, and the ZMQ collection filter Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- .github/workflows/e2e-gpu-job.yml | 6 + e2e_test/fixtures/setup_backend.py | 21 +++ .../test_connection_mode_validation.py | 34 +++++ e2e_test/fixtures/test_hooks_zmq_filter.py | 127 ++++++++++++++++++ e2e_test/infra/constants.py | 16 ++- e2e_test/infra/gateway.py | 2 +- e2e_test/infra/test_connection_mode.py | 43 ++++++ e2e_test/infra/test_zmq_cmd_builders.py | 22 +-- 8 files changed, 259 insertions(+), 12 deletions(-) create mode 100644 e2e_test/fixtures/test_connection_mode_validation.py create mode 100644 e2e_test/fixtures/test_hooks_zmq_filter.py create mode 100644 e2e_test/infra/test_connection_mode.py diff --git a/.github/workflows/e2e-gpu-job.yml b/.github/workflows/e2e-gpu-job.yml index 8ed684791..b97a2995e 100644 --- a/.github/workflows/e2e-gpu-job.yml +++ b/.github/workflows/e2e-gpu-job.yml @@ -172,6 +172,12 @@ jobs: if [ "${{ inputs.engine }}" = "vllm" ]; then SUFFIX="-${{ inputs.vllm_kv_backend }}" fi + # gRPC and ZMQ chat legs share engine/GPU/test-dir and differ only by + # wire mode; fold it in so their log artifacts don't collide. + MODE="${{ inputs.connection_mode }}" + if [ -n "$MODE" ]; then + SUFFIX="${SUFFIX}-${MODE}" + fi ARTIFACT="e2e-worker-logs-${{ inputs.engine }}-gpu${{ inputs.gpu_tier }}-${LABEL}${SUFFIX}" bash scripts/ci_dump_worker_logs.sh e2e-logs "$ARTIFACT" echo "artifact=${ARTIFACT}" >> "$GITHUB_OUTPUT" diff --git a/e2e_test/fixtures/setup_backend.py b/e2e_test/fixtures/setup_backend.py index 654f1aa14..44bf855a6 100644 --- a/e2e_test/fixtures/setup_backend.py +++ b/e2e_test/fixtures/setup_backend.py @@ -29,6 +29,7 @@ THIRD_PARTY_MODELS, ConnectionMode, Gateway, + Runtime, WorkerType, get_connection_mode_override, get_runtime, @@ -62,6 +63,24 @@ _worker_start_failures: dict[str, int] = {} # engine -> count _MAX_WORKER_START_FAILURES = 3 # fail fast after this many failures (matches --reruns 2) +# Engines that speak the direct-ZMQ backend wire in e2e. Pairing ZMQ with any +# other engine can't work, so we reject it up front instead of timing out on a +# worker that never becomes ready. +ZMQ_CAPABLE_ENGINES = frozenset({Runtime.VLLM.value, Runtime.TOKENSPEED.value}) + + +def _validate_connection_mode(connection_mode: ConnectionMode, engine: str) -> None: + """Reject connection-mode/engine pairings that cannot start. + + Raises ``ValueError`` when a lane selects ZMQ for an engine that does not + support the direct-ZMQ backend. + """ + if connection_mode == ConnectionMode.ZMQ and engine not in ZMQ_CAPABLE_ENGINES: + raise ValueError( + f"ConnectionMode.ZMQ is only supported for engines " + f"{sorted(ZMQ_CAPABLE_ENGINES)}, not {engine!r}" + ) + def _start_workers_tracked(**kwargs) -> list: """Start workers via the session pool and track failures for fail-fast. @@ -152,6 +171,7 @@ def setup_backend(request: pytest.FixtureRequest): if mode_override is not None and not is_pd and not is_epd: connection_mode = mode_override engine = get_runtime() + _validate_connection_mode(connection_mode, engine) model_path = get_model_spec(model_id)["model"] workers_config = get_marker_kwargs(request, "workers", defaults=_WORKER_DEFAULTS) log_dir = os.environ.get("E2E_LOG_DIR") or gateway_config.get("log_dir") @@ -485,6 +505,7 @@ def test_router_state(backend_router): if mode_override is not None: connection_mode = mode_override engine = get_runtime() + _validate_connection_mode(connection_mode, engine) model_path = get_model_spec(model_id)["model"] is_zmq = connection_mode == ConnectionMode.ZMQ diff --git a/e2e_test/fixtures/test_connection_mode_validation.py b/e2e_test/fixtures/test_connection_mode_validation.py new file mode 100644 index 000000000..54ab0d782 --- /dev/null +++ b/e2e_test/fixtures/test_connection_mode_validation.py @@ -0,0 +1,34 @@ +"""Unit tests for connection-mode/engine validation (no GPU).""" + +from __future__ import annotations + +import pytest +from infra.constants import ConnectionMode, Runtime + +# setup_backend pulls in the cloud SDKs; skip if the env lacks them. +setup_backend = pytest.importorskip("fixtures.setup_backend") + + +@pytest.mark.parametrize("engine", [Runtime.VLLM.value, Runtime.TOKENSPEED.value]) +def test_zmq_allowed_for_capable_engines(engine): + # Does not raise. + setup_backend._validate_connection_mode(ConnectionMode.ZMQ, engine) + + +@pytest.mark.parametrize( + "engine", + [Runtime.SGLANG.value, Runtime.TRTLLM.value, Runtime.MLX.value], +) +def test_zmq_rejected_for_incapable_engines(engine): + with pytest.raises(ValueError, match="ConnectionMode.ZMQ is only supported"): + setup_backend._validate_connection_mode(ConnectionMode.ZMQ, engine) + + +@pytest.mark.parametrize("mode", [ConnectionMode.GRPC, ConnectionMode.HTTP]) +@pytest.mark.parametrize( + "engine", + [Runtime.SGLANG.value, Runtime.VLLM.value, Runtime.TOKENSPEED.value], +) +def test_non_zmq_modes_accept_any_engine(mode, engine): + # Non-ZMQ wires impose no engine restriction. + setup_backend._validate_connection_mode(mode, engine) diff --git a/e2e_test/fixtures/test_hooks_zmq_filter.py b/e2e_test/fixtures/test_hooks_zmq_filter.py new file mode 100644 index 000000000..b3237e67e --- /dev/null +++ b/e2e_test/fixtures/test_hooks_zmq_filter.py @@ -0,0 +1,127 @@ +"""Unit tests for the ZMQ-lane collection filter (no GPU).""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from fixtures import hooks + + +class _FakeItem: + """Minimal stand-in for a pytest ``Item`` for the collection helpers. + + Exposes only what ``_filter_zmq_items`` / ``_zmq_dedup_key`` touch: + ``nodeid``, ``callspec.params``, ``cls`` (always None so marker resolution + falls back to ``get_closest_marker``), and a ``workers`` marker. + """ + + def __init__(self, nodeid, params=None, workers=None): + self.nodeid = nodeid + self.cls = None + self.callspec = SimpleNamespace(params=params) if params is not None else None + self._workers = workers + + def get_closest_marker(self, name): + if name == "workers": + return self._workers + return None + + +def _item(nodeid, setup_backend=None, extra_params=None, workers=None): + params = None + if setup_backend is not None or extra_params is not None: + params = {} + if setup_backend is not None: + params["setup_backend"] = setup_backend + if extra_params: + params.update(extra_params) + return _FakeItem(nodeid, params=params, workers=workers) + + +def _workers_marker(**kwargs): + return pytest.mark.workers(**kwargs).mark + + +# --------------------------------------------------------------------------- +# _zmq_dedup_key +# --------------------------------------------------------------------------- + + +def test_dedup_key_ignores_setup_backend_value(): + grpc = _item("t.py::test_x[grpc]", setup_backend="grpc") + http = _item("t.py::test_x[http]", setup_backend="http") + assert hooks._zmq_dedup_key(grpc) == hooks._zmq_dedup_key(http) + + +def test_dedup_key_keeps_other_params_apart(): + a = _item("t.py::test_x[grpc-a]", setup_backend="grpc", extra_params={"api_client": "a"}) + b = _item("t.py::test_x[grpc-b]", setup_backend="grpc", extra_params={"api_client": "b"}) + assert hooks._zmq_dedup_key(a) != hooks._zmq_dedup_key(b) + + +def test_dedup_key_splits_nodeid_at_bracket(): + key, _others = hooks._zmq_dedup_key(_item("t.py::test_x[grpc]", setup_backend="grpc")) + assert key == "t.py::test_x" + + +# --------------------------------------------------------------------------- +# _filter_zmq_items +# --------------------------------------------------------------------------- + + +def test_grpc_http_twins_collapse_to_grpc(): + grpc = _item("t.py::test_x[grpc]", setup_backend="grpc") + http = _item("t.py::test_x[http]", setup_backend="http") + kept, deselected = hooks._filter_zmq_items([grpc, http]) + assert kept == [grpc] + assert deselected == [http] + + +def test_http_only_case_is_retained(): + http = _item("t.py::test_x[http]", setup_backend="http") + kept, deselected = hooks._filter_zmq_items([http]) + assert kept == [http] + assert deselected == [] + + +def test_distinct_other_params_keep_both_wires(): + # A grpc/http pair that differs on another param is NOT a twin. + grpc = _item("t.py::test_x[grpc-a]", setup_backend="grpc", extra_params={"api_client": "a"}) + http = _item("t.py::test_x[http-b]", setup_backend="http", extra_params={"api_client": "b"}) + kept, deselected = hooks._filter_zmq_items([grpc, http]) + assert kept == [grpc, http] + assert deselected == [] + + +@pytest.mark.parametrize("param", ["pd_grpc", "epd_grpc", ("epd_grpc", (1, 1, 1)), ()]) +def test_non_local_wire_families_are_deselected(param): + item = _item("t.py::test_x[p]", setup_backend=param) + kept, deselected = hooks._filter_zmq_items([item]) + assert kept == [] + assert deselected == [item] + + +def test_multi_worker_case_is_deselected(): + item = _item("t.py::test_x[grpc]", setup_backend="grpc", workers=_workers_marker(count=2)) + kept, deselected = hooks._filter_zmq_items([item]) + assert kept == [] + assert deselected == [item] + + +def test_pd_worker_topology_is_deselected(): + item = _item( + "t.py::test_x[grpc]", + setup_backend="grpc", + workers=_workers_marker(prefill=1, decode=1), + ) + kept, deselected = hooks._filter_zmq_items([item]) + assert kept == [] + assert deselected == [item] + + +def test_items_without_setup_backend_are_untouched(): + item = _item("t.py::test_plain") # no callspec / params + kept, deselected = hooks._filter_zmq_items([item]) + assert kept == [item] + assert deselected == [] diff --git a/e2e_test/infra/constants.py b/e2e_test/infra/constants.py index ea0db97d7..7b9df2335 100644 --- a/e2e_test/infra/constants.py +++ b/e2e_test/infra/constants.py @@ -135,10 +135,22 @@ def get_connection_mode_override() -> "ConnectionMode | None": Set ``E2E_CONNECTION_MODE`` to run the existing local test cases over a different wire (like ``E2E_RUNTIME`` picks the engine): a ``grpc``/``http`` case then runs over that mode without a separate parametrize value. PD/EPD - backends keep their own wire. Returns ``None`` when unset. + backends keep their own wire. Returns ``None`` when the var is unset; a + set-but-empty or unrecognized value is a misconfiguration and raises. """ value = os.environ.get(ENV_CONNECTION_MODE) - return ConnectionMode(value.lower()) if value else None + if value is None: + return None + value = value.strip() + valid = [mode.value for mode in ConnectionMode] + if not value: + raise ValueError(f"{ENV_CONNECTION_MODE} is set but empty; unset it or use one of {valid}") + try: + return ConnectionMode(value.lower()) + except ValueError: + raise ValueError( + f"{ENV_CONNECTION_MODE}={value!r} is not a valid connection mode; use one of {valid}" + ) from None ENV_VLLM_KV_BACKEND = "E2E_VLLM_KV_BACKEND" diff --git a/e2e_test/infra/gateway.py b/e2e_test/infra/gateway.py index 66488ad99..68a123772 100644 --- a/e2e_test/infra/gateway.py +++ b/e2e_test/infra/gateway.py @@ -234,7 +234,7 @@ def start( # ZMQ workers share one wire across engine runtimes, so the router # cannot probe the backend from the ipc:// URL — pin it explicitly. if backend is not None: - mode_args += ["--router-backend", backend] + mode_args += ["--backend", backend] self._launch( mode_args=mode_args, timeout=timeout, diff --git a/e2e_test/infra/test_connection_mode.py b/e2e_test/infra/test_connection_mode.py new file mode 100644 index 000000000..9d25f4468 --- /dev/null +++ b/e2e_test/infra/test_connection_mode.py @@ -0,0 +1,43 @@ +"""Unit tests for E2E connection-mode env parsing (no GPU).""" + +from __future__ import annotations + +import pytest +from infra.constants import ( + ENV_CONNECTION_MODE, + ConnectionMode, + get_connection_mode_override, +) + + +def test_unset_returns_none(monkeypatch): + monkeypatch.delenv(ENV_CONNECTION_MODE, raising=False) + assert get_connection_mode_override() is None + + +@pytest.mark.parametrize( + "value,expected", + [ + ("zmq", ConnectionMode.ZMQ), + ("ZMQ", ConnectionMode.ZMQ), + ("Grpc", ConnectionMode.GRPC), + (" http ", ConnectionMode.HTTP), + ], +) +def test_valid_values_are_case_insensitive(monkeypatch, value, expected): + monkeypatch.setenv(ENV_CONNECTION_MODE, value) + assert get_connection_mode_override() == expected + + +@pytest.mark.parametrize("value", ["", " "]) +def test_set_but_empty_raises(monkeypatch, value): + # A set-but-blank var is a misconfiguration, not "unset". + monkeypatch.setenv(ENV_CONNECTION_MODE, value) + with pytest.raises(ValueError, match="set but empty"): + get_connection_mode_override() + + +def test_invalid_value_raises(monkeypatch): + monkeypatch.setenv(ENV_CONNECTION_MODE, "bogus") + with pytest.raises(ValueError, match="not a valid connection mode"): + get_connection_mode_override() diff --git a/e2e_test/infra/test_zmq_cmd_builders.py b/e2e_test/infra/test_zmq_cmd_builders.py index 820593364..e12229e6a 100644 --- a/e2e_test/infra/test_zmq_cmd_builders.py +++ b/e2e_test/infra/test_zmq_cmd_builders.py @@ -3,17 +3,21 @@ from __future__ import annotations import pytest - -# The builders delegate to smg.serve, so the wheel must be importable. -serve = pytest.importorskip("smg.serve") - -from infra.constants import ConnectionMode # noqa: E402 -from infra.worker import Worker # noqa: E402 +from infra.constants import ConnectionMode +from infra.worker import Worker _VLLM_MODEL = "meta-llama/Llama-3.2-1B-Instruct" _TS_MODEL = "Qwen/Qwen3.5-9B" +@pytest.fixture +def serve(): + # The ZMQ builders delegate to smg.serve, so the wheel must be importable. + # Scoped to a fixture (not module import) so the gRPC test below still runs + # when the wheel is absent. + return pytest.importorskip("smg.serve") + + def _worker(engine, port=50111): return Worker( model_id=_TS_MODEL if engine == "tokenspeed" else _VLLM_MODEL, @@ -24,13 +28,13 @@ def _worker(engine, port=50111): ) -def test_zmq_base_url_matches_serve_helper(): +def test_zmq_base_url_matches_serve_helper(serve): w = _worker("vllm", port=50123) assert w.base_url == serve._zmq_ipc_url(50123) assert w.base_url.startswith("ipc://") -def test_vllm_zmq_cmd_is_headless_with_derived_handshake_port(): +def test_vllm_zmq_cmd_is_headless_with_derived_handshake_port(serve): w = _worker("vllm") cmd = w._build_vllm_zmq_cmd("/models/llama", 1, {"vllm_args": ["--max-model-len", "2048"]}) assert "serve" in cmd @@ -43,7 +47,7 @@ def test_vllm_zmq_cmd_is_headless_with_derived_handshake_port(): assert cmd[cmd.index("--max-model-len") + 1] == "2048" -def test_tokenspeed_zmq_cmd_is_headless_with_derived_handshake_port(): +def test_tokenspeed_zmq_cmd_is_headless_with_derived_handshake_port(serve): w = _worker("tokenspeed") cmd = w._build_tokenspeed_zmq_cmd( "/models/qwen", 1, {"tokenspeed_args": ["--attention-backend", "fa3"]} From 72a287b72a64d79770aea634a8b38b3f0c4bf809 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Mon, 3 Aug 2026 19:57:07 -0700 Subject: [PATCH 03/16] fix(e2e): treat blank E2E_CONNECTION_MODE as no override The workflow always exports E2E_CONNECTION_MODE and leaves it empty for the non-override lanes, so raising on a set-but-blank value broke every regular e2e lane at collection time. Map blank to None (the default) while still rejecting genuinely-unrecognized values. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- e2e_test/infra/constants.py | 9 +++++---- e2e_test/infra/test_connection_mode.py | 8 ++++---- 2 files changed, 9 insertions(+), 8 deletions(-) diff --git a/e2e_test/infra/constants.py b/e2e_test/infra/constants.py index 7b9df2335..caac79b26 100644 --- a/e2e_test/infra/constants.py +++ b/e2e_test/infra/constants.py @@ -135,16 +135,17 @@ def get_connection_mode_override() -> "ConnectionMode | None": Set ``E2E_CONNECTION_MODE`` to run the existing local test cases over a different wire (like ``E2E_RUNTIME`` picks the engine): a ``grpc``/``http`` case then runs over that mode without a separate parametrize value. PD/EPD - backends keep their own wire. Returns ``None`` when the var is unset; a - set-but-empty or unrecognized value is a misconfiguration and raises. + backends keep their own wire. Returns ``None`` when the var is unset or + blank (the workflow always exports it and leaves it empty for non-override + lanes); a set-but-unrecognized value is a misconfiguration and raises. """ value = os.environ.get(ENV_CONNECTION_MODE) if value is None: return None value = value.strip() - valid = [mode.value for mode in ConnectionMode] if not value: - raise ValueError(f"{ENV_CONNECTION_MODE} is set but empty; unset it or use one of {valid}") + return None + valid = [mode.value for mode in ConnectionMode] try: return ConnectionMode(value.lower()) except ValueError: diff --git a/e2e_test/infra/test_connection_mode.py b/e2e_test/infra/test_connection_mode.py index 9d25f4468..44262b6c3 100644 --- a/e2e_test/infra/test_connection_mode.py +++ b/e2e_test/infra/test_connection_mode.py @@ -30,11 +30,11 @@ def test_valid_values_are_case_insensitive(monkeypatch, value, expected): @pytest.mark.parametrize("value", ["", " "]) -def test_set_but_empty_raises(monkeypatch, value): - # A set-but-blank var is a misconfiguration, not "unset". +def test_set_but_blank_returns_none(monkeypatch, value): + # The workflow always exports E2E_CONNECTION_MODE and leaves it empty for + # non-override lanes, so a blank value must mean "no override", not an error. monkeypatch.setenv(ENV_CONNECTION_MODE, value) - with pytest.raises(ValueError, match="set but empty"): - get_connection_mode_override() + assert get_connection_mode_override() is None def test_invalid_value_raises(monkeypatch): From f62c3e18965a816cfa188838665b1a0caeff6592 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Mon, 3 Aug 2026 23:42:21 -0700 Subject: [PATCH 04/16] feat(router): translate structured-output constraints over the ZMQ backend Port the engine-core structured-outputs params (guided decoding) into the engine ZMQ client and wire the proto constraint oneof onto the EngineCore wire, so ZMQ-backed vLLM workers honor json/regex/grammar/choice/ json_object/structural_tag constraints instead of returning unconstrained text. The client exposes an enum-backed StructuredOutputsParams and serializes through a private sparse wire struct (omit-defaults, skip-if-false) that matches Python engine-core; the backend is always lowered to guidance. translate_sampling maps json_object=false to no constraint. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- crates/engine_zmq_client/src/error.rs | 2 + .../src/protocol/vllm/mod.rs | 9 +- .../src/protocol/vllm/sampling.rs | 11 +- .../src/protocol/vllm/structured_outputs.rs | 315 ++++++++++++++++++ model_gateway/src/routers/grpc/zmq_client.rs | 98 +++++- 5 files changed, 407 insertions(+), 28 deletions(-) create mode 100644 crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs diff --git a/crates/engine_zmq_client/src/error.rs b/crates/engine_zmq_client/src/error.rs index ce8691983..77f3a6153 100644 --- a/crates/engine_zmq_client/src/error.rs +++ b/crates/engine_zmq_client/src/error.rs @@ -53,6 +53,8 @@ pub enum Error { context: &'static str, field: &'static str, }, + #[error("invalid structured outputs params: {message}")] + InvalidStructuredOutputsParams { message: String }, #[error("request `{request_id}` is already in flight")] DuplicateRequestId { request_id: String }, #[error("data parallel rank {rank} is out of range for {num_engines} engine(s)")] diff --git a/crates/engine_zmq_client/src/protocol/vllm/mod.rs b/crates/engine_zmq_client/src/protocol/vllm/mod.rs index 44d96abdd..0d8cf742f 100644 --- a/crates/engine_zmq_client/src/protocol/vllm/mod.rs +++ b/crates/engine_zmq_client/src/protocol/vllm/mod.rs @@ -7,10 +7,10 @@ //! shapes, field order, and `array_like` positional-tuple encoding are the wire //! contract with Python `EngineCoreProc` — do not reorder fields. //! -//! Text generation is typed fully. Multimodal features, structured outputs -//! (guided decoding), and pooling params are carried as [`crate::codec::OpaqueValue`] -//! for now — they serialize as `nil` on the text path and get strongly typed in -//! the multimodal phase. +//! Text generation and structured outputs (guided decoding) are typed fully. +//! Remaining multimodal features and pooling params are carried as +//! [`crate::codec::OpaqueValue`] for now — they serialize as `nil` on the text +//! path and get strongly typed in the multimodal phase. // The startup handshake is engine-neutral (TokenSpeed speaks the same // protocol); re-exported here so existing `vllm::handshake` paths keep working. @@ -21,6 +21,7 @@ pub mod output; pub mod request; pub mod sampling; pub mod stats; +pub mod structured_outputs; use bytes::Bytes; diff --git a/crates/engine_zmq_client/src/protocol/vllm/sampling.rs b/crates/engine_zmq_client/src/protocol/vllm/sampling.rs index 13ed51a9a..ac89ce7cd 100644 --- a/crates/engine_zmq_client/src/protocol/vllm/sampling.rs +++ b/crates/engine_zmq_client/src/protocol/vllm/sampling.rs @@ -2,16 +2,13 @@ // // Ported from the Apache-2.0 reference `vllm-engine-core-client` // (vllm-project/vllm): protocol/sampling.rs. -// -// `structured_outputs` (guided decoding) is carried as OpaqueValue for now; it -// serializes as `nil` when absent and gets strongly typed in a later phase. use std::collections::{BTreeSet, HashMap}; use serde::{Deserialize, Serialize}; use serde_default::DefaultFromSerde; -use crate::codec::OpaqueValue; +use super::structured_outputs::StructuredOutputsParams; fn default_top_p() -> f32 { 1.0 @@ -115,9 +112,9 @@ pub struct EngineCoreSamplingParams { /// Tokenized bad words to avoid during generation. #[serde(rename = "_bad_words_token_ids")] pub bad_words_token_ids: Option>>, - /// Structured outputs (guided decoding). Carried untyped for now; `nil` on - /// the text path. - pub structured_outputs: Option, + /// Structured outputs (guided decoding). `None` (serialized `nil`) on the + /// unconstrained text path. + pub structured_outputs: Option, /// Specific token IDs for which log probabilities should be returned at each /// position, in addition to the sampled/scored token. pub logprob_token_ids: Option>, diff --git a/crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs b/crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs new file mode 100644 index 000000000..3d9ad264b --- /dev/null +++ b/crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs @@ -0,0 +1,315 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Ported from the Apache-2.0 reference `vllm-engine-core-client` +// (vllm-project/vllm): protocol/structured_outputs.rs. + +use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use serde_json::Value; + +use crate::error::{Error, Result}; + +/// Structured-output backend selected for EngineCore grammar compilation. +/// +/// Python stores this in `StructuredOutputsParams._backend` after request +/// validation. This client currently always lowers structured-output requests +/// to guidance, ignoring any peer-supplied `_backend` value. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "kebab-case")] +pub enum StructuredOutputBackend { + Xgrammar, + #[default] + Guidance, + Outlines, + LmFormatEnforcer, +} + +/// The single structured-output constraint selected for a request. +#[derive(Debug, Clone, PartialEq)] +pub enum StructuredOutputConstraint { + /// JSON schema (as a dict/object or JSON string) constraining the output. + Json(Value), + /// Regular expression the output must match. + Regex(String), + /// List of allowed output strings (the model must produce one of these). + Choice(Vec), + /// Context-free grammar (in EBNF-like notation) the output must conform to. + Grammar(String), + /// Output must be valid JSON (free-form, no schema). + JsonObject, + /// Structural tag configuration (JSON-encoded string). + StructuralTag(String), +} + +/// Additional structured-output options that do not select the constraint mode. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct StructuredOutputOptions { + /// Disable any additional whitespace in guided JSON output. + pub disable_any_whitespace: bool, + /// Disable `additionalProperties` in JSON schema output. + pub disable_additional_properties: bool, + /// Custom whitespace pattern for guided JSON output. + pub whitespace_pattern: Option, +} + +/// Parameters for configuring structured outputs (guided decoding). +/// +/// This is the semantic Rust representation: exactly one constraint mode is +/// always selected. The Python/msgpack product-shaped representation is kept in +/// the private wire type below and used only at serde boundaries. +#[derive(Debug, Clone, PartialEq)] +pub struct StructuredOutputsParams { + pub constraint: StructuredOutputConstraint, + pub options: StructuredOutputOptions, + /// Structured-output backend, mirroring Python's internal `_backend`. + /// + /// Peer-supplied values are ignored during deserialization. This matches the + /// Python request boundary, where `_backend` is set by validation rather + /// than accepted as a request-level backend selector. + pub backend: StructuredOutputBackend, +} + +impl StructuredOutputsParams { + pub fn json(json: Value) -> Self { + Self::from_constraint(StructuredOutputConstraint::Json(json)) + } + + pub fn regex(regex: impl Into) -> Self { + Self::from_constraint(StructuredOutputConstraint::Regex(regex.into())) + } + + pub fn choice(choice: Vec) -> Self { + Self::from_constraint(StructuredOutputConstraint::Choice(choice)) + } + + pub fn grammar(grammar: impl Into) -> Self { + Self::from_constraint(StructuredOutputConstraint::Grammar(grammar.into())) + } + + pub fn json_object() -> Self { + Self::from_constraint(StructuredOutputConstraint::JsonObject) + } + + pub fn structural_tag(structural_tag: impl Into) -> Self { + Self::from_constraint(StructuredOutputConstraint::StructuralTag( + structural_tag.into(), + )) + } + + fn from_constraint(constraint: StructuredOutputConstraint) -> Self { + Self { + constraint, + options: StructuredOutputOptions::default(), + backend: StructuredOutputBackend::default(), + } + } +} + +/// `true` when a boolean is `false`; used to drop default-`false` flags from the +/// serialized map to match the sparse `omit_defaults` wire shape. Takes `&bool` +/// because serde's `skip_serializing_if` requires a by-reference predicate. +#[allow(clippy::trivially_copy_pass_by_ref)] +fn is_false(v: &bool) -> bool { + !*v +} + +/// Wire-compatible structured-output payload used by Python engine-core. +/// +/// Python models `StructuredOutputsParams` as a product-shaped dataclass with +/// several optional constraint fields, then validates that exactly one of those +/// fields is present. This client exposes [`StructuredOutputsParams`] as an +/// enum-backed domain type instead, while using this private wire type for +/// ser/de. +#[serde_with::skip_serializing_none] +#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)] +#[serde(default)] +struct WireStructuredOutputsParams { + json: Option, + regex: Option, + choice: Option>, + grammar: Option, + json_object: Option, + #[serde(skip_serializing_if = "is_false")] + disable_any_whitespace: bool, + #[serde(skip_serializing_if = "is_false")] + disable_additional_properties: bool, + whitespace_pattern: Option, + structural_tag: Option, + #[serde( + default, + rename = "_backend", + deserialize_with = "serde_with::rust::deserialize_ignore_any" + )] + backend: StructuredOutputBackend, +} + +impl TryFrom for StructuredOutputsParams { + type Error = Error; + + fn try_from(raw: WireStructuredOutputsParams) -> Result { + use StructuredOutputConstraint::*; + + let mut constraint = None; + + macro_rules! insert_constraint { + ($name:literal, $value:expr) => { + if let Some(value) = $value { + if let Some((existing, _)) = constraint { + return Err(Error::InvalidStructuredOutputsParams { + message: format!( + "multiple structured output constraints specified: {existing}, {}", + $name + ), + }); + } + constraint = Some(($name, value)); + } + }; + } + + insert_constraint!("json", raw.json.map(Json)); + insert_constraint!("regex", raw.regex.map(Regex)); + insert_constraint!("choice", raw.choice.map(Choice)); + insert_constraint!("grammar", raw.grammar.map(Grammar)); + match raw.json_object { + Some(true) => { + insert_constraint!("json_object", Some(JsonObject)); + } + Some(false) => { + return Err(Error::InvalidStructuredOutputsParams { + message: "structured_outputs.json_object must be true if set; omit structured_outputs to disable structured outputs".to_string(), + }); + } + None => {} + } + insert_constraint!("structural_tag", raw.structural_tag.map(StructuralTag)); + + Ok(Self { + constraint: constraint.map(|(_, c)| c).ok_or_else(|| { + Error::InvalidStructuredOutputsParams { + message: "missing structured output constraint".to_string(), + } + })?, + options: StructuredOutputOptions { + disable_any_whitespace: raw.disable_any_whitespace, + disable_additional_properties: raw.disable_additional_properties, + whitespace_pattern: raw.whitespace_pattern, + }, + backend: raw.backend, + }) + } +} + +impl From for WireStructuredOutputsParams { + fn from(params: StructuredOutputsParams) -> Self { + let mut raw = Self { + disable_any_whitespace: params.options.disable_any_whitespace, + disable_additional_properties: params.options.disable_additional_properties, + whitespace_pattern: params.options.whitespace_pattern, + backend: params.backend, + ..Self::default() + }; + + match params.constraint { + StructuredOutputConstraint::Json(json) => raw.json = Some(json), + StructuredOutputConstraint::Regex(regex) => raw.regex = Some(regex), + StructuredOutputConstraint::Choice(choice) => raw.choice = Some(choice), + StructuredOutputConstraint::Grammar(grammar) => raw.grammar = Some(grammar), + StructuredOutputConstraint::JsonObject => raw.json_object = Some(true), + StructuredOutputConstraint::StructuralTag(structural_tag) => { + raw.structural_tag = Some(structural_tag); + } + } + + raw + } +} + +impl Serialize for StructuredOutputsParams { + fn serialize(&self, serializer: S) -> std::result::Result + where + S: Serializer, + { + WireStructuredOutputsParams::from(self.clone()).serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for StructuredOutputsParams { + fn deserialize(deserializer: D) -> std::result::Result + where + D: Deserializer<'de>, + { + WireStructuredOutputsParams::deserialize(deserializer)? + .try_into() + .map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn structured_outputs_backend_ignores_deserialized_value() { + let params: StructuredOutputsParams = serde_json::from_value(serde_json::json!({ + "json_object": true, + "_backend": "xgrammar", + })) + .unwrap(); + + assert_eq!(params.backend, StructuredOutputBackend::Guidance); + assert_eq!(params.constraint, StructuredOutputConstraint::JsonObject); + + let value = serde_json::to_value(params).unwrap(); + assert_eq!(value["_backend"], "guidance"); + } + + #[test] + fn structured_outputs_rejects_missing_constraint() { + let error = + serde_json::from_value::(serde_json::json!({})).unwrap_err(); + + assert!(error.to_string().contains("missing structured output constraint")); + } + + #[test] + fn structured_outputs_rejects_multiple_constraints() { + let error = serde_json::from_value::(serde_json::json!({ + "json": {"type": "object"}, + "regex": ".*", + })) + .unwrap_err(); + + assert!( + error + .to_string() + .contains("multiple structured output constraints specified: json, regex") + ); + } + + #[test] + fn structured_outputs_rejects_json_object_false() { + let error = serde_json::from_value::(serde_json::json!({ + "json_object": false, + })) + .unwrap_err(); + + assert!(error.to_string().contains("json_object must be true")); + } + + #[test] + fn structured_outputs_serializes_through_raw_shape() { + let params = StructuredOutputsParams { + constraint: StructuredOutputConstraint::StructuralTag( + r#"{"type":"structural_tag"}"#.to_string(), + ), + options: StructuredOutputOptions::default(), + backend: StructuredOutputBackend::Xgrammar, + }; + + let value = serde_json::to_value(params).unwrap(); + + assert_eq!(value["structural_tag"], r#"{"type":"structural_tag"}"#); + assert_eq!(value["_backend"], "xgrammar"); + assert!(value.get("json").is_none()); + } +} diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index e3fc0b692..f4cabc3c2 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -30,6 +30,7 @@ use engine_zmq_client::{ output::{EngineCoreFinishReason, EngineCoreOutput, StopReason}, request::EngineCoreRequest, sampling::EngineCoreSamplingParams, + structured_outputs::StructuredOutputsParams, }, EngineLoad, }, @@ -786,15 +787,6 @@ fn translate_request(req: vllm::GenerateRequest) -> Result1 is fanned out into // single-sample sub-requests by `generate` before translation. // The ZMQ renderer path has no prompt-logprob merge, so the engine's @@ -848,10 +840,31 @@ fn translate_sampling(sp: vllm::SamplingParams) -> EngineCoreSamplingParams { // prompt_logprobs is rejected in `translate_request` (no renderer // support on the ZMQ path), so it is never forwarded. logit_bias, + structured_outputs: sp.constraint.and_then(translate_constraint), ..EngineCoreSamplingParams::default() } } +/// Map the proto `constraint` oneof onto typed structured-output params. The +/// backend defaults to guidance engine-side; `json_object=false` selects no +/// constraint (the caller opted out), so it maps to `None`. +fn translate_constraint(constraint: vllm::sampling_params::Constraint) -> Option { + use vllm::sampling_params::Constraint; + match constraint { + Constraint::JsonSchema(schema) => Some(StructuredOutputsParams::json( + // The engine accepts a JSON schema object or a schema string; parse + // to preserve object shape, falling back to the raw string. + serde_json::from_str(&schema).unwrap_or(serde_json::Value::String(schema)), + )), + Constraint::Regex(regex) => Some(StructuredOutputsParams::regex(regex)), + Constraint::Grammar(grammar) => Some(StructuredOutputsParams::grammar(grammar)), + Constraint::StructuralTag(tag) => Some(StructuredOutputsParams::structural_tag(tag)), + Constraint::JsonObject(true) => Some(StructuredOutputsParams::json_object()), + Constraint::JsonObject(false) => None, + Constraint::Choice(choice) => Some(StructuredOutputsParams::choice(choice.choices)), + } +} + fn map_matched_stop(reason: StopReason) -> vllm::generate_complete::MatchedStop { match reason { StopReason::TokenId(id) => vllm::generate_complete::MatchedStop::MatchedTokenId(id), @@ -1372,14 +1385,6 @@ mod tests { #[test] fn vllm_rejects_unsupported_sampling_features() { - // Structured-output constraints are not translated onto the wire. - let err = translate_request(tokenized_req(vllm::SamplingParams { - constraint: Some(vllm::sampling_params::Constraint::JsonObject(true)), - ..Default::default() - })) - .expect_err("constraint rejected"); - assert!(err.contains("structured output"), "{err}"); - // Prompt logprobs have no renderer merge on the ZMQ path. let err = translate_request(tokenized_req(vllm::SamplingParams { prompt_logprobs: Some(1), @@ -1389,6 +1394,65 @@ mod tests { assert!(err.contains("prompt logprobs"), "{err}"); } + #[test] + fn vllm_translates_structured_output_constraints() { + use engine_zmq_client::protocol::vllm::structured_outputs::{ + StructuredOutputBackend, StructuredOutputConstraint, + }; + + let translate = |constraint| { + translate_request(tokenized_req(vllm::SamplingParams { + constraint: Some(constraint), + ..Default::default() + })) + .expect("constraint translated") + .sampling_params + .expect("sampling params present") + .structured_outputs + }; + + // Each constraint mode maps onto its typed counterpart, always lowering + // to the guidance backend engine-side. + let json_object = translate(vllm::sampling_params::Constraint::JsonObject(true)) + .expect("json_object translated"); + assert_eq!( + json_object.constraint, + StructuredOutputConstraint::JsonObject + ); + assert_eq!(json_object.backend, StructuredOutputBackend::Guidance); + + let regex = translate(vllm::sampling_params::Constraint::Regex("a.*".to_string())) + .expect("regex translated"); + assert_eq!( + regex.constraint, + StructuredOutputConstraint::Regex("a.*".to_string()) + ); + + let choice = translate(vllm::sampling_params::Constraint::Choice( + vllm::ChoiceConstraint { + choices: vec!["yes".to_string(), "no".to_string()], + }, + )) + .expect("choice translated"); + assert_eq!( + choice.constraint, + StructuredOutputConstraint::Choice(vec!["yes".to_string(), "no".to_string()]) + ); + + // A JSON schema string is parsed to preserve object shape. + let json = translate(vllm::sampling_params::Constraint::JsonSchema( + r#"{"type":"object"}"#.to_string(), + )) + .expect("json schema translated"); + assert_eq!( + json.constraint, + StructuredOutputConstraint::Json(serde_json::json!({"type": "object"})) + ); + + // json_object=false means the caller opted out: no constraint. + assert!(translate(vllm::sampling_params::Constraint::JsonObject(false)).is_none()); + } + /// n=3 fans out into 3 single-sample wire requests with unique sub-rids. /// An explicit seed derives per-sub seeds (`seed + i`) so the samples /// differ deterministically; no seed stays `None` per sub (the engine From baf9332b81c4f5dfabeb86a1e9b2b327a40ada87 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Mon, 3 Aug 2026 23:42:21 -0700 Subject: [PATCH 05/16] fix(e2e): cover model load in the ZMQ gateway readiness gate ZMQ workers spawn and return immediately (readiness is gated by the gateway, not the worker), so model load happens inside the gateway's readiness wait. That wait used DEFAULT_ROUTER_TIMEOUT (60s), too short for large models. Use the model's startup_timeout for ZMQ gateways (never shrinking an explicitly larger gateway timeout); gRPC/HTTP unchanged. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- e2e_test/fixtures/setup_backend.py | 32 ++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/e2e_test/fixtures/setup_backend.py b/e2e_test/fixtures/setup_backend.py index 44bf855a6..3c38a76cf 100644 --- a/e2e_test/fixtures/setup_backend.py +++ b/e2e_test/fixtures/setup_backend.py @@ -23,6 +23,7 @@ from infra import ( DEFAULT_MODEL, DEFAULT_ROUTER_TIMEOUT, + DEFAULT_STARTUP_TIMEOUT, ENV_MODEL, ENV_SKIP_BACKEND_SETUP, RUNTIME_LABELS, @@ -111,6 +112,25 @@ def _start_gateway(gateway: Gateway, gateway_config: dict, **mode_kwargs) -> Non ) +def _gateway_readiness_timeout( + connection_mode: ConnectionMode, model_id: str, base_timeout: float +) -> float: + """Effective gateway readiness timeout for the given connection mode. + + gRPC/HTTP workers are health-checked (model fully loaded) by the pool + before the gateway starts, so the gateway only has to connect — the short + router timeout suffices. ZMQ engines instead spawn and return immediately; + their model load happens *inside* the gateway's readiness gate, so that gate + must cover model load too. Use the model's ``startup_timeout`` (what the + worker gate would have applied), never shrinking an explicitly larger + gateway timeout. + """ + if connection_mode != ConnectionMode.ZMQ: + return base_timeout + startup_timeout = get_model_spec(model_id).get("startup_timeout", DEFAULT_STARTUP_TIMEOUT) + return max(base_timeout, startup_timeout) + + def _make_openai_client(gateway: Gateway) -> openai.OpenAI: return openai.OpenAI(base_url=f"{gateway.base_url}/v1", api_key="not-used") @@ -263,6 +283,14 @@ def _setup_local( # caller owns their teardown (like the PD path). gRPC/HTTP workers stay # in the pool and outlive the gateway. is_zmq = connection_mode == ConnectionMode.ZMQ + # ZMQ engines load the model inside the gateway's readiness gate (the worker + # spawn returned immediately), so that gate must cover model load. + gateway_config = { + **gateway_config, + "timeout": _gateway_readiness_timeout( + connection_mode, model_id, gateway_config["timeout"] + ), + } try: _start_gateway( gateway, @@ -524,6 +552,10 @@ def test_router_state(backend_router): worker_urls=[w.base_url for w in workers], model_path=model_path, backend=engine if is_zmq else None, + # ZMQ loads the model inside the gateway's readiness gate; cover it. + timeout=_gateway_readiness_timeout( + connection_mode, model_id, DEFAULT_ROUTER_TIMEOUT + ), ) yield gateway finally: From 17b7bdb054dfe1786862274a3e3ab38efc55ebf7 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Mon, 3 Aug 2026 23:57:25 -0700 Subject: [PATCH 06/16] fix(zmq): satisfy fmt/clippy for structured-outputs port Use #[expect] instead of #[allow] for the trivially_copy_pass_by_ref predicate (clippy::allow_attributes is deny-by-default), apply nightly rustfmt to the ported module and the constraint translator, and let ruff-format normalize the readiness-timeout helper calls. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- .../src/protocol/vllm/structured_outputs.rs | 14 +++++++------- e2e_test/fixtures/setup_backend.py | 8 ++------ model_gateway/src/routers/grpc/zmq_client.rs | 4 +++- 3 files changed, 12 insertions(+), 14 deletions(-) diff --git a/crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs b/crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs index 3d9ad264b..6e93deb0d 100644 --- a/crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs +++ b/crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs @@ -107,7 +107,7 @@ impl StructuredOutputsParams { /// `true` when a boolean is `false`; used to drop default-`false` flags from the /// serialized map to match the sparse `omit_defaults` wire shape. Takes `&bool` /// because serde's `skip_serializing_if` requires a by-reference predicate. -#[allow(clippy::trivially_copy_pass_by_ref)] +#[expect(clippy::trivially_copy_pass_by_ref)] fn is_false(v: &bool) -> bool { !*v } @@ -268,7 +268,9 @@ mod tests { let error = serde_json::from_value::(serde_json::json!({})).unwrap_err(); - assert!(error.to_string().contains("missing structured output constraint")); + assert!(error + .to_string() + .contains("missing structured output constraint")); } #[test] @@ -279,11 +281,9 @@ mod tests { })) .unwrap_err(); - assert!( - error - .to_string() - .contains("multiple structured output constraints specified: json, regex") - ); + assert!(error + .to_string() + .contains("multiple structured output constraints specified: json, regex")); } #[test] diff --git a/e2e_test/fixtures/setup_backend.py b/e2e_test/fixtures/setup_backend.py index 3c38a76cf..8f63e3814 100644 --- a/e2e_test/fixtures/setup_backend.py +++ b/e2e_test/fixtures/setup_backend.py @@ -287,9 +287,7 @@ def _setup_local( # spawn returned immediately), so that gate must cover model load. gateway_config = { **gateway_config, - "timeout": _gateway_readiness_timeout( - connection_mode, model_id, gateway_config["timeout"] - ), + "timeout": _gateway_readiness_timeout(connection_mode, model_id, gateway_config["timeout"]), } try: _start_gateway( @@ -553,9 +551,7 @@ def test_router_state(backend_router): model_path=model_path, backend=engine if is_zmq else None, # ZMQ loads the model inside the gateway's readiness gate; cover it. - timeout=_gateway_readiness_timeout( - connection_mode, model_id, DEFAULT_ROUTER_TIMEOUT - ), + timeout=_gateway_readiness_timeout(connection_mode, model_id, DEFAULT_ROUTER_TIMEOUT), ) yield gateway finally: diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index f4cabc3c2..50766be78 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -848,7 +848,9 @@ fn translate_sampling(sp: vllm::SamplingParams) -> EngineCoreSamplingParams { /// Map the proto `constraint` oneof onto typed structured-output params. The /// backend defaults to guidance engine-side; `json_object=false` selects no /// constraint (the caller opted out), so it maps to `None`. -fn translate_constraint(constraint: vllm::sampling_params::Constraint) -> Option { +fn translate_constraint( + constraint: vllm::sampling_params::Constraint, +) -> Option { use vllm::sampling_params::Constraint; match constraint { Constraint::JsonSchema(schema) => Some(StructuredOutputsParams::json( From 0a5de615e24d57390da31b5b52cd98c5d37a8c89 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Tue, 4 Aug 2026 12:41:20 -0700 Subject: [PATCH 07/16] fix(zmq): default unset max_tokens to remaining context EngineCore requires a concrete max_tokens (the wire field is non-optional and serde-defaults to 16). The ZMQ path bypasses vLLM's OpenAI frontend, which would otherwise default an unset max_tokens to max_model_len - prompt_len. Without that, requests fell back to 16 tokens and truncated generation -- e.g. streaming reasoning never reached content past the transition. Take max_model_len from the connected engine's ready handshake and apply the same default; a connected engine is now required to translate a request. Validated on Qwen3-30B-A3B over ZMQ: reasoning and content both stream with finish delivered. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- model_gateway/src/routers/grpc/zmq_client.rs | 82 ++++++++++++++++---- 1 file changed, 68 insertions(+), 14 deletions(-) diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index 50766be78..3179294c5 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -166,10 +166,19 @@ impl ZmqEngineClient { // error, which auto-aborts their engine-side requests. match &self.backend { ZmqBackend::Vllm(client) => { + // EngineCore needs a concrete `max_tokens`; vLLM's OpenAI frontend + // (which the ZMQ path bypasses) defaults an unset value to + // `max_model_len - prompt_len`. The context length comes from the + // engine's ready handshake, so a connected engine is required. + let max_model_len = client + .engines() + .first() + .map(|e| e.ready_response.max_model_len) + .ok_or_else(|| tonic::Status::unavailable("no connected ZMQ engine"))?; let mut streams = SelectAll::new(); for (index, sub) in subs.into_iter().enumerate() { - let request = - translate_request(sub).map_err(tonic::Status::invalid_argument)?; + let request = translate_request(sub, max_model_len) + .map_err(tonic::Status::invalid_argument)?; let stream = client.submit(request).await.map_err(zmq_status)?; streams.push(VllmGenerateStream::new(stream, index as u32)); } @@ -772,7 +781,10 @@ fn translate_sampling_tokenspeed(sp: vllm::SamplingParams) -> TokenSpeedSampling /// Translate a vLLM-proto generate request into an `EngineCoreRequest`. ZMQ mode /// requires pre-tokenized input (SMG tokenizes upstream). -fn translate_request(req: vllm::GenerateRequest) -> Result { +fn translate_request( + req: vllm::GenerateRequest, + max_model_len: u64, +) -> Result { let prompt_token_ids = match req.input { Some(vllm::generate_request::Input::Tokenized(tokenized)) => Some(tokenized.input_ids), Some(vllm::generate_request::Input::Text(_)) => { @@ -795,17 +807,27 @@ fn translate_request(req: vllm::GenerateRequest) -> Result EngineCoreSamplingParams { +fn translate_sampling( + sp: vllm::SamplingParams, + default_max_tokens: u32, +) -> EngineCoreSamplingParams { let logit_bias = if sp.logit_bias.is_empty() { None } else { @@ -832,7 +854,7 @@ fn translate_sampling(sp: vllm::SamplingParams) -> EngineCoreSamplingParams { frequency_penalty: sp.frequency_penalty, presence_penalty: sp.presence_penalty, repetition_penalty: sp.repetition_penalty, - max_tokens: sp.max_tokens.unwrap_or(16), + max_tokens: sp.max_tokens.unwrap_or(default_max_tokens), min_tokens: sp.min_tokens, stop_token_ids: sp.stop_token_ids, seed: sp.seed.map(i64::from), @@ -1388,14 +1410,43 @@ mod tests { #[test] fn vllm_rejects_unsupported_sampling_features() { // Prompt logprobs have no renderer merge on the ZMQ path. - let err = translate_request(tokenized_req(vllm::SamplingParams { - prompt_logprobs: Some(1), - ..Default::default() - })) + let err = translate_request( + tokenized_req(vllm::SamplingParams { + prompt_logprobs: Some(1), + ..Default::default() + }), + 4096, + ) .expect_err("prompt logprobs rejected"); assert!(err.contains("prompt logprobs"), "{err}"); } + #[test] + fn vllm_defaults_unset_max_tokens_to_remaining_context() { + let max_tokens = |sampling, max_model_len| { + translate_request(tokenized_req(sampling), max_model_len) + .expect("request translated") + .sampling_params + .expect("sampling params present") + .max_tokens + }; + + // Unset max_tokens defaults to `max_model_len - prompt_len` (prompt is + // 3 tokens), mirroring vLLM's bypassed OpenAI frontend. + assert_eq!(max_tokens(vllm::SamplingParams::default(), 100), 97); + // An explicit value is always honored. + assert_eq!( + max_tokens( + vllm::SamplingParams { + max_tokens: Some(8), + ..Default::default() + }, + 100, + ), + 8, + ); + } + #[test] fn vllm_translates_structured_output_constraints() { use engine_zmq_client::protocol::vllm::structured_outputs::{ @@ -1403,10 +1454,13 @@ mod tests { }; let translate = |constraint| { - translate_request(tokenized_req(vllm::SamplingParams { - constraint: Some(constraint), - ..Default::default() - })) + translate_request( + tokenized_req(vllm::SamplingParams { + constraint: Some(constraint), + ..Default::default() + }), + 4096, + ) .expect("constraint translated") .sampling_params .expect("sampling params present") From 0a4f74813679aebe5552c345747253372f5d9dce Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Tue, 4 Aug 2026 13:20:58 -0700 Subject: [PATCH 08/16] fix(zmq): shape top_logprobs to the requested count over the ZMQ path The ZMQ direct path only forwarded per-token logprobs, dropping the top_logprobs alternatives that clients request via `logprobs: N`. The engine returns PositionLogprobs.entries as [sampled, top1..topk] (k+1 columns); take the first N entries per position to match what the gRPC servicer produces, and accumulate them for the completion aggregate. Adds a unit test asserting the requested-count shaping. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- model_gateway/src/routers/grpc/zmq_client.rs | 186 ++++++++++++++++++- 1 file changed, 176 insertions(+), 10 deletions(-) diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index 3179294c5..ae7bb09b9 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -177,10 +177,19 @@ impl ZmqEngineClient { .ok_or_else(|| tonic::Status::unavailable("no connected ZMQ engine"))?; let mut streams = SelectAll::new(); for (index, sub) in subs.into_iter().enumerate() { + // The engine returns the sampled logprob plus up to + // `logprobs` ranked candidates per position; carry the + // requested count so the stream can shape `top_logprobs`. + let top_logprobs = sub + .sampling_params + .as_ref() + .and_then(|sp| sp.logprobs) + .filter(|&n| n > 0) + .map_or(0, |n| n as usize); let request = translate_request(sub, max_model_len) .map_err(tonic::Status::invalid_argument)?; let stream = client.submit(request).await.map_err(zmq_status)?; - streams.push(VllmGenerateStream::new(stream, index as u32)); + streams.push(VllmGenerateStream::new(stream, index as u32, top_logprobs)); } Ok(ZmqGenerateStream::Vllm(streams)) } @@ -337,6 +346,9 @@ struct StreamState { /// `Complete`, so accumulate here and drain into `Complete`. output_logprobs_val: Vec, output_logprobs_idx: Vec, + /// Cumulative per-position ranked candidates (`top_logprobs`), accumulated + /// alongside the sampled logprobs and drained into the terminal `Complete`. + output_top_logprobs: Vec, } impl StreamState { @@ -346,7 +358,7 @@ impl StreamState { (!self.output_logprobs_val.is_empty()).then(|| vllm::OutputLogProbs { token_logprobs: std::mem::take(&mut self.output_logprobs_val), token_ids: std::mem::take(&mut self.output_logprobs_idx), - ..Default::default() + top_logprobs: std::mem::take(&mut self.output_top_logprobs), }) } } @@ -360,6 +372,10 @@ pub struct VllmGenerateStream { /// Choice index stamped on every chunk/complete (0 for n=1; the fan-out /// position for n>1) — the proto field the pipeline demuxes choices by. index: u32, + /// Number of ranked candidates the client requested per position; `0` when + /// only the sampled logprob (or nothing) was asked for, in which case no + /// `top_logprobs` are emitted. + top_logprobs: usize, /// Terminal `Complete` held back when the finish tick also carried new /// tokens: streaming frontends decode text/logprobs from chunks only, so /// the tick's delta goes out as a `Chunk` first. @@ -367,11 +383,12 @@ pub struct VllmGenerateStream { } impl VllmGenerateStream { - fn new(inner: EngineCoreStream, index: u32) -> Self { + fn new(inner: EngineCoreStream, index: u32, top_logprobs: usize) -> Self { Self { inner, state: StreamState::default(), index, + top_logprobs, pending: None, } } @@ -380,6 +397,7 @@ impl VllmGenerateStream { &mut self, output: EngineCoreOutput, ) -> Result { + let top_k = self.top_logprobs; let state = &mut self.state; if let Some(stats) = &output.prefill_stats { state.prompt_tokens = stats.num_prompt_tokens; @@ -389,11 +407,13 @@ impl VllmGenerateStream { state.completion_tokens += token_ids.len() as u32; state.output_ids.extend(token_ids.iter().copied()); - // Sampled-token logprobs (entry 0 per position), if requested. Chunks - // carry this tick's increment; the terminal `Complete` carries the - // cumulative set, so accumulate into `state` and drain it on finish. + // Sampled-token logprobs (entry 0 per position) plus the requested + // ranked candidates (`top_logprobs`). Chunks carry this tick's + // increment; the terminal `Complete` carries the cumulative set, so + // accumulate into `state` and drain it on finish. let mut tick_logprobs_val = Vec::new(); let mut tick_logprobs_idx = Vec::new(); + let mut tick_top_logprobs = Vec::new(); if let Some(logprobs) = &output.new_logprobs { let decoded = logprobs.as_direct().ok_or_else(|| { // The protocol layer resolves wire logprobs during decode, so @@ -401,19 +421,31 @@ impl VllmGenerateStream { tonic::Status::internal("unresolved wire logprobs in engine output") })?; for position in &decoded.positions { - if let Some(sampled) = position.entries.first() { - tick_logprobs_val.push(sampled.logprob); - tick_logprobs_idx.push(sampled.token_id); + let Some(sampled) = position.entries.first() else { + continue; + }; + tick_logprobs_val.push(sampled.logprob); + tick_logprobs_idx.push(sampled.token_id); + // The entries arrive sampled-first then rank-ordered; take the + // requested count so one ranked list lands per sampled token. + if top_k > 0 { + let mut top = vllm::TopLogProbs::default(); + for entry in position.entries.iter().take(top_k) { + top.values.push(entry.logprob); + top.token_ids.push(entry.token_id); + } + tick_top_logprobs.push(top); } } } let chunk_logprobs = (!tick_logprobs_val.is_empty()).then(|| vllm::OutputLogProbs { token_logprobs: tick_logprobs_val.clone(), token_ids: tick_logprobs_idx.clone(), - ..Default::default() + top_logprobs: tick_top_logprobs.clone(), }); state.output_logprobs_val.extend(tick_logprobs_val); state.output_logprobs_idx.extend(tick_logprobs_idx); + state.output_top_logprobs.extend(tick_top_logprobs); let response = match output.finish_reason { Some(reason) => { @@ -1102,6 +1134,140 @@ mod tests { engine_task.await.unwrap(); } + /// With `logprobs=k`, each position's ranked candidates are shaped into + /// `top_logprobs`, taking the sampled entry plus the leading candidates up + /// to the requested count (matching the gRPC servicer's `islice` behaviour). + #[tokio::test] + async fn generate_shapes_top_logprobs_to_requested_count() { + let dir = tempfile::tempdir().unwrap(); + let ep = |name: &str| format!("ipc://{}", dir.path().join(name).display()); + let (handshake, input, output) = (ep("hs.sock"), ep("in.sock"), ep("out.sock")); + + let (client, engine) = tokio::join!( + ZmqEngineClient::connect( + &handshake, + &input, + &output, + 1, + "m".to_string(), + RuntimeType::Vllm, + Duration::from_secs(10) + ), + connect_to_frontend( + &handshake, + EngineId::from_engine_index(0), + default_ready_response() + ), + ); + let client = client.expect("adapter connect"); + let engine = engine.expect("mock engine"); + + // One position with the sampled token (actual vocab rank) first, then + // the engine's ranked candidates. The wire carries `k + 1` entries. + let position = PositionLogprobs { + entries: vec![ + TokenLogprob { + token_id: 10, + logprob: -0.5, + rank: 5, + }, + TokenLogprob { + token_id: 20, + logprob: -0.1, + rank: 1, + }, + TokenLogprob { + token_id: 30, + logprob: -0.3, + rank: 2, + }, + ], + }; + let outputs = EngineCoreOutputs::RequestBatch(RequestBatchOutputs { + engine_index: 0, + outputs: vec![EngineCoreOutput { + request_id: "r1".to_string(), + new_token_ids: vec![10], + new_logprobs: Some(MaybeWireLogprobs::Direct(Logprobs { + positions: vec![position], + })), + finish_reason: Some(EngineCoreFinishReason::Length), + ..Default::default() + }], + finished_requests: Some(std::collections::BTreeSet::from(["r1".to_string()])), + ..Default::default() + }); + + #[expect( + clippy::disallowed_methods, + reason = "engine task ends after responding" + )] + let engine_task = tokio::spawn(async move { + let (mut input, mut output) = engine.split(); + let inbound = input.recv().await.unwrap(); + let request = match inbound { + EngineInbound::Add(request) => request, + other => panic!("expected Add, got {other:?}"), + }; + assert_eq!(request.sampling_params.as_ref().unwrap().logprobs, Some(2)); + output.send_outputs(&outputs).await.unwrap(); + }); + + let req = vllm::GenerateRequest { + request_id: "r1".to_string(), + input: Some(vllm::generate_request::Input::Tokenized( + vllm::TokenizedInput { + original_text: String::new(), + input_ids: vec![1, 2, 3], + }, + )), + sampling_params: Some(vllm::SamplingParams { + max_tokens: Some(1), + logprobs: Some(2), + ..Default::default() + }), + stream: true, + ..Default::default() + }; + let mut stream = client.generate(req).await.expect("generate"); + + // The requested count is 2, so `top_logprobs` keeps the sampled entry + // plus the leading candidate (the third entry is dropped). + let expected_top = vec![vllm::TopLogProbs { + values: vec![-0.5, -0.1], + token_ids: vec![10, 20], + }]; + + // The finish tick carried a token, so the delta streams as a chunk. + let chunk = stream.next().await.expect("chunk item").expect("chunk ok"); + match chunk.response { + Some(vllm::generate_response::Response::Chunk(chunk)) => { + let logprobs = chunk.output_logprobs.expect("chunk logprobs"); + assert_eq!(logprobs.token_logprobs, vec![-0.5]); + assert_eq!(logprobs.token_ids, vec![10]); + assert_eq!(logprobs.top_logprobs, expected_top); + } + other => panic!("expected chunk, got {other:?}"), + } + let complete = stream + .next() + .await + .expect("complete item") + .expect("complete ok"); + match complete.response { + Some(vllm::generate_response::Response::Complete(complete)) => { + let logprobs = complete.output_logprobs.expect("complete logprobs"); + assert_eq!(logprobs.token_logprobs, vec![-0.5]); + assert_eq!(logprobs.token_ids, vec![10]); + assert_eq!(logprobs.top_logprobs, expected_top); + } + other => panic!("expected complete, got {other:?}"), + } + assert!(stream.next().await.is_none()); + + engine_task.await.unwrap(); + } + /// End-to-end over ipc:// for a TokenSpeed backend: the adapter frames a /// tagged `TokenizedGenerateReqInput`, and maps `BatchTokenIDOutSlim` /// batches back to vLLM-proto responses. The mock engine speaks the shared From ead557e1459496d7b165c715752f99fc6a197350 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Tue, 4 Aug 2026 13:21:07 -0700 Subject: [PATCH 09/16] fix(zmq): report finish_reason=stop for stop-string matches Stop-string detection needs detokenized text, which vLLM does in its frontend OutputProcessor. Over gRPC the Python servicer forwards the stop strings to the engine, so vLLM trims the output and returns finish_reason=stop itself. The ZMQ direct path talks to EngineCore, which has no stop-string detection, so it ran to max_tokens and reported "length" even though the gateway's own StopSequenceDecoder had already trimmed the text correctly. Synthesize finish_reason=stop from the local decoder signal across the OpenAI chat (stream + non-stream), completions (non-stream), and Anthropic Messages (stream + non-stream) paths, mirroring the existing completions-streaming precedent. This is a no-op for gRPC: vLLM trims the stop string before returning, so the local decoder never fires, and matched_stop falls back to the engine value via or_else so token-stop and length results are preserved. Adds StopSequenceDecoder::matched_stop() so the matched string flows into the OpenAI matched_stop field and the Anthropic stop_reason= stop_sequence / stop_sequence field, matching gRPC parity. Unit tests cover the string-match and token-stop cases. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- crates/tokenizer/src/stop.rs | 46 ++++++++++++ .../src/routers/grpc/regular/processor.rs | 71 +++++++++++++++---- .../src/routers/grpc/regular/streaming.rs | 67 +++++++++++++++-- 3 files changed, 166 insertions(+), 18 deletions(-) diff --git a/crates/tokenizer/src/stop.rs b/crates/tokenizer/src/stop.rs index 820620cec..3b0ed3e6f 100644 --- a/crates/tokenizer/src/stop.rs +++ b/crates/tokenizer/src/stop.rs @@ -77,6 +77,9 @@ pub struct StopSequenceDecoder { jail_max_bytes: usize, /// Whether we've stopped stopped: bool, + /// The string stop sequence that triggered the stop, if any. Set only for + /// string-sequence matches; token-level stops leave this `None`. + matched_stop: Option, /// True when there are no string stop sequences (only token-level stops). /// In this mode the jail buffer is bypassed entirely for lower overhead. token_only: bool, @@ -139,6 +142,7 @@ impl StopSequenceDecoder { jail_buffer: String::new(), jail_max_bytes, stopped: false, + matched_stop: None, token_only, } } @@ -208,6 +212,7 @@ impl StopSequenceDecoder { let input = Input::new(&self.jail_buffer).span(search_start..self.jail_buffer.len()); if let Some(mat) = ac.find(input) { self.stopped = true; + self.matched_stop = Some(self.jail_buffer[mat.start()..mat.end()].to_string()); let is_visible = mat.pattern().as_usize() >= self.visible_boundary_idx; if is_visible { @@ -287,11 +292,18 @@ impl StopSequenceDecoder { self.stopped } + /// The string stop sequence that triggered the stop, if a string sequence + /// matched. `None` for token-level stops or when no stop has fired. + pub fn matched_stop(&self) -> Option<&str> { + self.matched_stop.as_deref() + } + /// Reset the decoder state pub fn reset(&mut self) { self.jail_buffer.clear(); self.sequence.clear(); self.stopped = false; + self.matched_stop = None; } } @@ -455,6 +467,40 @@ mod tests { )); } + #[test] + fn test_matched_stop_reports_matched_string() { + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default().with_stop_sequence("test"); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + // No match yet. + assert_eq!(decoder.matched_stop(), None); + + decoder.process_token(1).unwrap(); // "Hello" + decoder.process_token(2).unwrap(); // "world" + assert_eq!(decoder.matched_stop(), None); + + // "test" triggers the string stop; the matched string is captured. + decoder.process_token(3).unwrap(); // "test" + assert_eq!(decoder.matched_stop(), Some("test")); + + // Reset clears it. + decoder.reset(); + assert_eq!(decoder.matched_stop(), None); + } + + #[test] + fn test_matched_stop_none_for_token_stop() { + // A token-id stop is not a string match, so matched_stop stays None. + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default().with_stop_token(999); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + let result = decoder.process_token(999).unwrap(); + assert_eq!(result, SequenceDecoderOutput::Stopped); + assert_eq!(decoder.matched_stop(), None); + } + #[test] fn test_flush_after_partial() { let tokenizer = Arc::new(MockTokenizer::new()); diff --git a/model_gateway/src/routers/grpc/regular/processor.rs b/model_gateway/src/routers/grpc/regular/processor.rs index 6d1d8d9e0..2275f3eb9 100644 --- a/model_gateway/src/routers/grpc/regular/processor.rs +++ b/model_gateway/src/routers/grpc/regular/processor.rs @@ -78,14 +78,19 @@ impl ResponseProcessor { // Accumulate text with early breaks let mut final_text = String::new(); + let mut stopped = false; for output in outputs { match output { SequenceDecoderOutput::Text(t) => final_text.push_str(&t), SequenceDecoderOutput::StoppedWithText(t) => { final_text.push_str(&t); + stopped = true; + break; + } + SequenceDecoderOutput::Stopped => { + stopped = true; break; } - SequenceDecoderOutput::Stopped => break, SequenceDecoderOutput::Held => {} } } @@ -177,8 +182,14 @@ impl ResponseProcessor { } } - // Step 3: Use finish reason directly from proto (already OpenAI-compatible string) - let finish_reason_str = complete.finish_reason(); + // Step 3: Determine finish reason. A local stop-decoder match takes + // precedence over the engine's reason (which is "length" when stop + // strings are enforced gateway-side rather than by the backend). + let finish_reason_str = if stopped { + "stop" + } else { + complete.finish_reason() + }; // Override finish reason if we have tool calls let final_finish_reason_str = if tool_calls.is_some() { @@ -187,7 +198,12 @@ impl ResponseProcessor { finish_reason_str }; - let matched_stop = complete.matched_stop_json(); + // When the local decoder matched a stop string, surface it (the engine + // reports no stop_reason over the ZMQ path); otherwise use the engine's. + let matched_stop = stop_decoder + .matched_stop() + .map(|s| serde_json::Value::String(s.to_string())) + .or_else(|| complete.matched_stop_json()); // Step 4: Convert output logprobs if present let logprobs = complete.output_logprobs().map(|ref proto_logprobs| { @@ -579,14 +595,19 @@ impl ResponseProcessor { })?; let mut final_text = String::new(); + let mut stopped = false; for output in outputs { match output { SequenceDecoderOutput::Text(t) => final_text.push_str(&t), SequenceDecoderOutput::StoppedWithText(t) => { final_text.push_str(&t); + stopped = true; + break; + } + SequenceDecoderOutput::Stopped => { + stopped = true; break; } - SequenceDecoderOutput::Stopped => break, SequenceDecoderOutput::Held => {} } } @@ -726,10 +747,20 @@ impl ResponseProcessor { } } - // Step 4: Determine stop_reason and stop_sequence (derived from same conditions) - let finish_reason_str = complete.finish_reason(); - let matched_stop = complete.matched_stop_json(); - let stop_sequence = matched_stop.and_then(|v| v.as_str().map(String::from)); + // Step 4: Determine stop_reason and stop_sequence (derived from same conditions). + // A local stop-decoder match takes precedence over the engine's reason + // (the backend has no stop-string detection over ZMQ), surfacing the + // matched sequence for a StopSequence result. + let finish_reason_str = if stopped { + "stop" + } else { + complete.finish_reason() + }; + let stop_sequence = stop_decoder.matched_stop().map(String::from).or_else(|| { + complete + .matched_stop_json() + .and_then(|v| v.as_str().map(String::from)) + }); let stop_reason = if tool_calls.is_some() || finish_reason_str == "tool_calls" { Some(messages::StopReason::ToolUse) @@ -832,14 +863,19 @@ impl ResponseProcessor { }; let mut decoded_text = String::new(); + let mut stopped = false; for output in outputs { match output { SequenceDecoderOutput::Text(t) => decoded_text.push_str(&t), SequenceDecoderOutput::StoppedWithText(t) => { decoded_text.push_str(&t); + stopped = true; + break; + } + SequenceDecoderOutput::Stopped => { + stopped = true; break; } - SequenceDecoderOutput::Stopped => break, SequenceDecoderOutput::Held => {} } } @@ -851,7 +887,12 @@ impl ResponseProcessor { prompt_tokens = prompt_tokens.max(complete.prompt_tokens()); total_completion += complete.completion_tokens(); - let finish_reason = { + // A local stop-decoder match takes precedence over the engine's + // reason (which is "length" when stop strings are enforced + // gateway-side rather than by the backend). + let finish_reason = if stopped { + Some("stop".to_string()) + } else { let reason = complete.finish_reason(); if reason.is_empty() { None @@ -868,7 +909,13 @@ impl ResponseProcessor { } }; - let matched_stop = complete.matched_stop_json(); + // When the local decoder matched a stop string, surface it (the + // engine reports no stop_reason over the ZMQ path); otherwise use + // the engine's. + let matched_stop = stop_decoder + .matched_stop() + .map(|s| serde_json::Value::String(s.to_string())) + .or_else(|| complete.matched_stop_json()); let suffix_len = completion_req.suffix.as_ref().map_or(0, |s| s.len()); let echo_len = if completion_req.echo { diff --git a/model_gateway/src/routers/grpc/regular/streaming.rs b/model_gateway/src/routers/grpc/regular/streaming.rs index 505a8dc3b..b9d888f0c 100644 --- a/model_gateway/src/routers/grpc/regular/streaming.rs +++ b/model_gateway/src/routers/grpc/regular/streaming.rs @@ -271,6 +271,9 @@ impl StreamingProcessor { let mut stream_buffers: HashMap = HashMap::new(); let mut finish_reasons: HashMap = HashMap::new(); let mut matched_stops: HashMap> = HashMap::new(); + // Indices whose local stop decoder fired: their finish reason is pinned + // to "stop" and later engine output for the index is ignored. + let mut stopped_indices: HashSet = HashSet::new(); let mut prompt_tokens: HashMap = HashMap::new(); let mut completion_tokens = CompletionTokenTracker::new(); let mut cached_tokens: HashMap = HashMap::new(); @@ -384,6 +387,12 @@ impl StreamingProcessor { let index = chunk.index(); + // Once the local stop decoder has fired for an index, ignore + // any further engine output the backend emits for it. + if stopped_indices.contains(&index) { + continue; + } + completion_tokens.record_chunk(&chunk); // Get or create stop decoder for this index @@ -406,9 +415,26 @@ impl StreamingProcessor { }); // Process tokens through stop decoder - let (chunk_text, _should_stop) = + let (chunk_text, should_stop) = Self::process_chunk_tokens(stop_decoder, chunk.token_ids()); + if should_stop { + // Stop-decoder match takes precedence: pin "stop" even if + // the backend's eventual Complete carries "length" (the + // local stop sequence fired first). Any pre-stop text in + // `chunk_text` is still emitted below before the finish + // reason is flushed in Phase 4. + finish_reasons + .entry(index) + .or_insert_with(|| "stop".to_string()); + matched_stops.entry(index).or_insert_with(|| { + stop_decoder + .matched_stop() + .map(|s| Value::String(s.to_string())) + }); + stopped_indices.insert(index); + } + if chunk_text.is_empty() { continue; } @@ -568,9 +594,13 @@ impl StreamingProcessor { cached_tokens.insert(index, complete.cached_tokens()); reasoning_tokens.insert(index, complete.reasoning_tokens()); - finish_reasons.insert(index, complete.finish_reason().to_string()); - matched_stops.insert(index, complete.matched_stop_json()); + // A local stop-decoder match already pinned "stop" for this + // index; don't let the engine's finish reason overwrite it. + if !stopped_indices.contains(&index) { + finish_reasons.insert(index, complete.finish_reason().to_string()); + matched_stops.insert(index, complete.matched_stop_json()); + } // Don't break - continue reading all Complete messages for n>1 } @@ -1740,6 +1770,9 @@ impl StreamingProcessor { let mut prompt_tokens: u32 = 0; let mut finish_reason_str = String::new(); let mut matched_stop: Option = None; + // Set once the local stop decoder fires: pins "stop" and ignores later + // engine output (the backend has no stop-string detection over ZMQ). + let mut stopped = false; // Check parser availability once upfront. Run parser when the user explicitly // enabled thinking, or when the selected parser needs structural special tokens. @@ -1867,11 +1900,29 @@ impl StreamingProcessor { first_token_time = Some(Instant::now()); } + // Once the local stop decoder has fired, ignore further + // engine output for this (single-choice) request. + if stopped { + continue; + } + completion_tokens.record_chunk(&chunk); - let (chunk_text, _should_stop) = + let (chunk_text, should_stop) = Self::process_chunk_tokens(&mut stop_decoder, chunk.token_ids()); + if should_stop { + // Stop-decoder match takes precedence over the engine's + // eventual finish reason (the local stop sequence fired + // first). Pre-stop text in `chunk_text` is still emitted + // below; Phase 4 derives StopSequence from `matched_stop`. + stopped = true; + finish_reason_str = "stop".to_string(); + matched_stop = stop_decoder + .matched_stop() + .map(|s| Value::String(s.to_string())); + } + if chunk_text.is_empty() { continue; } @@ -2150,8 +2201,12 @@ impl StreamingProcessor { prompt_tokens = complete.prompt_tokens(); completion_tokens.record_complete(&complete); - finish_reason_str = complete.finish_reason().to_string(); - matched_stop = complete.matched_stop_json(); + // A local stop-decoder match already pinned "stop"; don't let + // the engine's finish reason overwrite it. + if !stopped { + finish_reason_str = complete.finish_reason().to_string(); + matched_stop = complete.matched_stop_json(); + } } ProtoResponseVariant::None => continue, } From c3f0209e9cc51036677527d8f047f6b78f215631 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Tue, 4 Aug 2026 14:05:13 -0700 Subject: [PATCH 10/16] feat(grpc): unify router-authoritative string-stop resolution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Engines that can't match string `stop` sequences server-side previously handled them in three divergent ways: SGLang gRPC (skip_tokenizer_init) had a dedicated resolver that dropped the strings and converted single-token stops to stop_token_ids; the ZMQ vLLM path silently dropped strings with no early-stop conversion; and the ZMQ TokenSpeed path rejected any request carrying `stop` with a 400. Consolidate all of them onto one helper, resolve_string_stops, applied in the chat/completion/generate/messages request-building stages. For backends whose engine sees token ids only — SGLang on any transport, and every direct-ZMQ backend — it drops the string `stop` list (the router-side StopSequenceDecoder trims the text) and forwards single-token stops as stop_token_ids for early stopping; multi-token stops fall to the decoder. gRPC vLLM (servicer detokenizes), TRT-LLM (server-side), and MLX keep their strings untouched, gated by the backend's is_zmq() flag. This fixes the TokenSpeed 400 (stop strings now work over ZMQ) and gives the ZMQ vLLM path the single-token early-stop optimization it lacked. The finish_reason=stop synthesis added earlier is already backend-neutral, so the two halves together give consistent stop-string behavior across gRPC-SGLang, ZMQ-vLLM, and ZMQ-TokenSpeed. Supersedes #1877 (SGLang-only input-side fix). Adds unit tests for the shared resolver and for TokenSpeed accepting residual stop strings. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- .../src/routers/grpc/backend_client.rs | 6 + .../src/routers/grpc/common/stages/helpers.rs | 229 ++++++++++++++++++ .../regular/stages/chat/request_building.rs | 7 + .../stages/completion/request_building.rs | 23 +- .../stages/generate/request_building.rs | 7 + .../stages/messages/request_building.rs | 7 + model_gateway/src/routers/grpc/zmq_client.rs | 38 +-- 7 files changed, 291 insertions(+), 26 deletions(-) diff --git a/model_gateway/src/routers/grpc/backend_client.rs b/model_gateway/src/routers/grpc/backend_client.rs index 908e0e989..4fd2916b5 100644 --- a/model_gateway/src/routers/grpc/backend_client.rs +++ b/model_gateway/src/routers/grpc/backend_client.rs @@ -50,6 +50,12 @@ impl BackendClient { } } + /// True if this is a direct-ZMQ backend (the engine receives token ids only + /// and cannot match string stops itself). + pub fn is_zmq(&self) -> bool { + matches!(self, Self::Zmq(_)) + } + /// Local liveness. gRPC has no cheap local flag (it uses a health RPC), so /// this reports `true` for gRPC; ZMQ reflects its connection liveness. pub fn is_alive(&self) -> bool { diff --git a/model_gateway/src/routers/grpc/common/stages/helpers.rs b/model_gateway/src/routers/grpc/common/stages/helpers.rs index 3cecfe4a6..701505a86 100644 --- a/model_gateway/src/routers/grpc/common/stages/helpers.rs +++ b/model_gateway/src/routers/grpc/common/stages/helpers.rs @@ -2,6 +2,7 @@ use std::sync::Arc; +use llm_tokenizer::traits::Tokenizer; use rand::RngExt; use smg_grpc_client::{ mlx_proto, @@ -291,6 +292,102 @@ fn apply_tokenspeed_sampling_defaults( apply_opt!(repetition_penalty); } +/// Convert single-token stop strings into `stop_token_ids` entries so the engine +/// can halt generation early for the common case (e.g. `["."]`, `["\n"]`). +/// +/// The proto `stop_token_ids` field is a flat list of single token ids, so a +/// multi-token stop string cannot be represented there — pushing its sub-tokens +/// would stop far too eagerly (on any one of them). Multi-token, empty, and +/// unknown stops are therefore left to the router-side `StopSequenceDecoder`, +/// which detokenizes worker output and trims the stop text. Existing +/// `stop_token_ids` are preserved and deduped. +fn encode_single_token_stops( + stops: Vec, + stop_token_ids: &mut Vec, + tokenizer: Option<&Arc>, +) { + // Without a tokenizer we cannot encode (not expected on paths that resolve + // one to tokenize the prompt). Safe: the strings are already dropped by the + // caller, so the router-side decoder remains the source of truth. + let Some(tokenizer) = tokenizer else { + if !stops.is_empty() { + warn!( + "No tokenizer available to encode string stop sequences; \ + relying on router-side stop decoder only" + ); + } + return; + }; + + for stop in stops { + if stop.is_empty() { + continue; + } + // add_special_tokens=false: we want the literal token(s) for the stop + // string, not a BOS/EOS-wrapped encoding. + match tokenizer.encode(&stop, false) { + Ok(encoding) => match encoding.token_ids() { + [id] => { + if !stop_token_ids.contains(id) { + stop_token_ids.push(*id); + } + } + ids => debug!( + stop = %stop, + token_count = ids.len(), + "string stop is not single-token; handled by router-side stop decoder" + ), + }, + Err(e) => warn!( + stop = %stop, + error = %e, + "Failed to encode string stop sequence; relying on router-side stop decoder" + ), + } + } +} + +/// Router-authoritative string-`stop` resolution for backends whose engine +/// cannot match string stops itself. +/// +/// vLLM over gRPC detokenizes server-side (`detokenize=bool(stop)`), TRT-LLM +/// tokenizes stop words server-side, and MLX has no string-`stop` field — those +/// keep their strings untouched. Two paths cannot: +/// - SGLang gRPC workers run with `skip_tokenizer_init=True` and reject string +/// stops outright (a 400 for any request carrying `stop`); and +/// - every direct-ZMQ backend (vLLM EngineCore, TokenSpeed) receives token ids +/// only, so the engine never sees — and cannot match — a stop string. +/// +/// For both, the router owns the tokenizer and already matches string stops via +/// `StopSequenceDecoder` (it detokenizes worker output and trims), so the worker +/// never needs the raw strings. This drops the string `stop` list and forwards +/// any single-token stop as a `stop_token_ids` entry for early stopping; the +/// router-side decoder handles the rest. This is the single resolution point +/// shared by SGLang gRPC and every ZMQ backend. +pub(crate) fn resolve_string_stops( + request: &mut ProtoGenerateRequest, + tokenizer: Option<&Arc>, + is_zmq: bool, +) { + // SGLang always needs it; other protos only when talking to a ZMQ backend + // (which, on this path, always carries the vLLM proto). + match request { + ProtoGenerateRequest::Sglang(req) => { + if let Some(params) = req.sampling_params.as_mut() { + let stops = std::mem::take(&mut params.stop); + encode_single_token_stops(stops, &mut params.stop_token_ids, tokenizer); + } + } + ProtoGenerateRequest::Vllm(req) if is_zmq => { + if let Some(params) = req.sampling_params.as_mut() { + let stops = std::mem::take(&mut params.stop); + encode_single_token_stops(stops, &mut params.stop_token_ids, tokenizer); + } + } + _ => {} + } +} + /// Inject PD bootstrap metadata for SGLang if needed. /// /// SGLang uses DisaggregatedParams with bootstrap host/port/room. @@ -434,3 +531,135 @@ mod request_id_tests { assert!(id.starts_with("chatcmpl-")); } } + +#[cfg(test)] +mod stop_resolution_tests { + use std::sync::Arc; + + use llm_tokenizer::{mock::MockTokenizer, traits::Tokenizer}; + use smg_grpc_client::{sglang_proto, vllm_proto}; + + use super::{resolve_string_stops, ProtoGenerateRequest}; + + fn mock_tokenizer() -> Arc { + // MockTokenizer vocab: "." => 6, "Hello" => 1, "world" => 2. `encode` + // splits on whitespace, so "." => [6] (single) and "Hello world" => + // [1, 2] (multi); unknown words encode to []. + Arc::new(MockTokenizer::new()) + } + + fn sglang_request(stop: Vec<&str>, stop_token_ids: Vec) -> ProtoGenerateRequest { + ProtoGenerateRequest::Sglang(Box::new(sglang_proto::GenerateRequest { + sampling_params: Some(sglang_proto::SamplingParams { + stop: stop.into_iter().map(str::to_string).collect(), + stop_token_ids, + ..Default::default() + }), + ..Default::default() + })) + } + + fn vllm_request(stop: Vec<&str>, stop_token_ids: Vec) -> ProtoGenerateRequest { + ProtoGenerateRequest::Vllm(Box::new(vllm_proto::GenerateRequest { + sampling_params: Some(vllm_proto::SamplingParams { + stop: stop.into_iter().map(str::to_string).collect(), + stop_token_ids, + ..Default::default() + }), + ..Default::default() + })) + } + + fn sglang_params(req: &ProtoGenerateRequest) -> &sglang_proto::SamplingParams { + match req { + ProtoGenerateRequest::Sglang(r) => r.sampling_params.as_ref().unwrap(), + _ => panic!("expected SGLang request"), + } + } + + fn vllm_params(req: &ProtoGenerateRequest) -> &vllm_proto::SamplingParams { + match req { + ProtoGenerateRequest::Vllm(r) => r.sampling_params.as_ref().unwrap(), + _ => panic!("expected vLLM request"), + } + } + + #[test] + fn sglang_single_token_becomes_stop_token_id() { + let mut req = sglang_request(vec!["."], vec![]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), false); + + let params = sglang_params(&req); + assert!(params.stop.is_empty(), "string stop should be cleared"); + assert_eq!(params.stop_token_ids, vec![6]); + } + + #[test] + fn sglang_multi_token_relies_on_router_decoder() { + // "Hello world" => [1, 2]: can't be a flat stop_token_id, so it must not + // be forwarded (would over-eagerly stop on any subtoken). + let mut req = sglang_request(vec!["Hello world"], vec![]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), false); + + let params = sglang_params(&req); + assert!(params.stop.is_empty()); + assert!(params.stop_token_ids.is_empty()); + } + + #[test] + fn sglang_mixed_only_single_token_forwarded_and_dedups() { + let mut req = sglang_request(vec![".", "Hello world"], vec![6, 42]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), false); + + let params = sglang_params(&req); + assert!(params.stop.is_empty()); + assert_eq!( + params.stop_token_ids, + vec![6, 42], + "existing ids kept, no dup" + ); + } + + #[test] + fn sglang_without_tokenizer_still_clears_strings() { + let mut req = sglang_request(vec!["."], vec![]); + resolve_string_stops(&mut req, None, false); + + let params = sglang_params(&req); + assert!( + params.stop.is_empty(), + "strings dropped so worker won't 400" + ); + assert!(params.stop_token_ids.is_empty()); + } + + #[test] + fn vllm_resolved_only_over_zmq() { + // gRPC vLLM keeps its strings (the servicer detokenizes engine-side). + let mut grpc = vllm_request(vec!["."], vec![]); + resolve_string_stops(&mut grpc, Some(&mock_tokenizer()), false); + let params = vllm_params(&grpc); + assert_eq!( + params.stop, + vec![".".to_string()], + "gRPC vLLM stop preserved" + ); + assert!(params.stop_token_ids.is_empty()); + + // ZMQ vLLM (EngineCore sees token ids only) resolves like SGLang. + let mut zmq = vllm_request(vec!["."], vec![]); + resolve_string_stops(&mut zmq, Some(&mock_tokenizer()), true); + let params = vllm_params(&zmq); + assert!(params.stop.is_empty(), "ZMQ vLLM stop cleared"); + assert_eq!(params.stop_token_ids, vec![6]); + } + + #[test] + fn noop_when_no_string_stops() { + let mut req = sglang_request(vec![], vec![7]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), false); + + let params = sglang_params(&req); + assert_eq!(params.stop_token_ids, vec![7], "unrelated ids untouched"); + } +} diff --git a/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs index 4967044e9..3dd71e948 100644 --- a/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs @@ -147,6 +147,13 @@ impl PipelineStage for ChatRequestBuildingStage { ctx.state.workers.as_ref(), ); + // Resolve string `stop` sequences for engines that can't match them + // server-side (SGLang skip_tokenizer_init, and every direct-ZMQ + // backend): drop the strings, convert single-token stops to + // stop_token_ids; the router-side StopSequenceDecoder trims the text. + let is_zmq = builder_client.is_zmq(); + helpers::resolve_string_stops(&mut proto_request, ctx.tokenizer_arc().as_ref(), is_zmq); + if self.inject_pd_metadata { if let Some(workers) = ctx.state.workers.as_ref() { helpers::maybe_inject_pd_metadata(&mut proto_request, workers); diff --git a/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs index 6477ccecc..8bd9b63da 100644 --- a/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs @@ -141,6 +141,12 @@ impl PipelineStage for CompletionRequestBuildingStage { let disaggregated = matches!(clients, ClientSelection::Disaggregated { .. }); let request_type = &ctx.input.request_type; let workers = ctx.state.workers.as_ref(); + // Resolve string `stop` sequences for engines that can't match them + // server-side (SGLang skip_tokenizer_init, and every direct-ZMQ + // backend): drop the strings, convert single-token stops to + // stop_token_ids; the router-side StopSequenceDecoder trims the text. + let is_zmq = builder_client.is_zmq(); + let tokenizer = ctx.tokenizer_arc(); let plan = match items.as_slice() { [] => { @@ -149,9 +155,8 @@ impl PipelineStage for CompletionRequestBuildingStage { "No prompts prepared", )) } - [item] => ExecutionPlan::generate( - self.plan_kind, - self.build_proto_request( + [item] => { + let mut proto_request = self.build_proto_request( builder_client, helpers::resolve_request_id( request_type, @@ -163,8 +168,10 @@ impl PipelineStage for CompletionRequestBuildingStage { &completion_request, request_type, workers, - )?, - ), + )?; + helpers::resolve_string_stops(&mut proto_request, tokenizer.as_ref(), is_zmq); + ExecutionPlan::generate(self.plan_kind, proto_request) + } batch_items => { // The shared id (client rid or middleware request id) stays // clean for the response; per-sub engine ids get a uniqueness @@ -186,14 +193,16 @@ impl PipelineStage for CompletionRequestBuildingStage { } else { format!("{shared_request_id}-p{i}") }; - requests.push(self.build_proto_request( + let mut proto_request = self.build_proto_request( builder_client, sub_request_id, item, &completion_request, request_type, workers, - )?); + )?; + helpers::resolve_string_stops(&mut proto_request, tokenizer.as_ref(), is_zmq); + requests.push(proto_request); } ExecutionPlan::Batch { kind: self.plan_kind, diff --git a/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs index 3be932e43..2f7634632 100644 --- a/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs @@ -86,6 +86,13 @@ impl PipelineStage for GenerateRequestBuildingStage { ctx.state.workers.as_ref(), ); + // Resolve string `stop` sequences for engines that can't match them + // server-side (SGLang skip_tokenizer_init, and every direct-ZMQ + // backend): drop the strings, convert single-token stops to + // stop_token_ids; the router-side StopSequenceDecoder trims the text. + let is_zmq = builder_client.is_zmq(); + helpers::resolve_string_stops(&mut proto_request, ctx.tokenizer_arc().as_ref(), is_zmq); + if self.inject_pd_metadata { if let Some(workers) = ctx.state.workers.as_ref() { helpers::maybe_inject_pd_metadata(&mut proto_request, workers); diff --git a/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs index 23d91b89a..aa73ff82a 100644 --- a/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs @@ -148,6 +148,13 @@ impl PipelineStage for MessageRequestBuildingStage { ctx.state.workers.as_ref(), ); + // Resolve string `stop` sequences for engines that can't match them + // server-side (SGLang skip_tokenizer_init, and every direct-ZMQ + // backend): drop the strings, convert single-token stops to + // stop_token_ids; the router-side StopSequenceDecoder trims the text. + let is_zmq = builder_client.is_zmq(); + helpers::resolve_string_stops(&mut proto_request, ctx.tokenizer_arc().as_ref(), is_zmq); + if self.inject_pd_metadata { if let Some(workers) = ctx.state.workers.as_ref() { helpers::maybe_inject_pd_metadata(&mut proto_request, workers); diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index ae7bb09b9..02f2437ce 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -736,17 +736,11 @@ fn translate_request_tokenspeed( } // The TokenSpeed wire has no per-sample demux; n>1 is fanned out into // single-sample sub-requests by `generate` before translation. - // Stop strings are not forwarded: the direct ZMQ path sends token ids - // only (the engine-side transport normalizes without a tokenizer) and - // the gateway's stop decoder does not enforce them, so they would be - // ignored and the request would run to max_tokens. - if !sp.stop.is_empty() { - return Err( - "stop strings are not supported over the TokenSpeed ZMQ backend yet; \ - use stop_token_ids" - .to_string(), - ); - } + // String `stop` sequences are resolved upstream in request building + // (`resolve_string_stops`): the strings are dropped and single-token + // stops become `stop_token_ids`, while the router-side stop decoder + // trims the text. Any residual strings here are harmless — the decoder + // still enforces them — so this path forwards token ids only. // logit_bias is not translated onto the TokenSpeed wire either. if !sp.logit_bias.is_empty() { return Err("logit_bias is not supported over the TokenSpeed ZMQ backend".to_string()); @@ -1544,14 +1538,6 @@ mod tests { .expect_err("constraint rejected"); assert!(err.contains("structured output"), "{err}"); - // Stop strings have no wire slot and are not enforced by the gateway. - let err = translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { - stop: vec!["".to_string()], - ..Default::default() - })) - .expect_err("stop strings rejected"); - assert!(err.contains("stop_token_ids"), "{err}"); - // logit_bias has no wire slot. let err = translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { logit_bias: HashMap::from([(7, 1.0)]), @@ -1561,6 +1547,20 @@ mod tests { assert!(err.contains("logit_bias"), "{err}"); } + #[test] + fn tokenspeed_ignores_residual_stop_strings() { + // String stops are resolved upstream (`resolve_string_stops`); any that + // reach here are harmless — the translator forwards token ids only and + // the router-side decoder trims the text. + let req = translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { + stop: vec!["".to_string()], + stop_token_ids: vec![13], + ..Default::default() + })) + .expect("residual stop strings must not be rejected"); + assert_eq!(req.sampling_params.stop_token_ids, Some(vec![13])); + } + #[test] fn tokenspeed_rejects_nonzero_dp_rank() { // Single-engine backend: only rank 0 (or none) is valid. From 7d19ab7ed675e1123d7055b4a61ef1962109b160 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Tue, 4 Aug 2026 16:29:42 -0700 Subject: [PATCH 11/16] fix(zmq): carry the model's EOS ids on EngineCore requests The engine has no tokenizer or model config, so stopping at EOS is the frontend's job: without ids on the request, generation only ends at max_tokens. This was masked while the wire default truncated everything at 16 tokens; the remaining-context default exposed it as never-ending generations (the e2e chat-zmq hang). Resolve the EOS set from the worker's local model directory at connect (config.json + generation_config.json) and attach it in translation, mirroring the reference frontend: the primary id rides _eos_token_id and extra ids merge into stop_token_ids, both gated on ignore_eos, with the union in _all_stop_token_ids for engine-side min_tokens masking. Live-verified on GB300 + Qwen3-0.6B: the previously-hanging no-max_tokens thinking request now finishes in 2.6s with finish_reason=stop; streaming terminates; ignore_eos still runs to the cap. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- model_gateway/src/routers/grpc/zmq_client.rs | 166 +++++++++++++++++-- model_gateway/src/worker/worker.rs | 17 +- 2 files changed, 172 insertions(+), 11 deletions(-) diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index 02f2437ce..66cbf6b6f 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -13,7 +13,8 @@ // the request-execution stage is reused unchanged. use std::{ - collections::HashMap, + collections::{BTreeSet, HashMap}, + path::Path, sync::Arc, time::{Duration, SystemTime, UNIX_EPOCH}, }; @@ -57,6 +58,58 @@ enum ZmqBackend { TokenSpeed(Arc), } +/// The model's EOS stop set, resolved from its local directory. EngineCore +/// has no tokenizer or model config — stopping at EOS is the frontend's job +/// (the ids ride each request), and without them generation only ends at +/// `max_tokens`. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct EosTokenIds { + /// Primary EOS id, carried as the request's `_eos_token_id`. + primary: Option, + /// Extra EOS ids (multi-EOS models), merged into `stop_token_ids`. + extra: Vec, +} + +impl EosTokenIds { + pub fn new(primary: Option, extra: Vec) -> Self { + Self { primary, extra } + } + + /// Resolve from `config.json` + `generation_config.json` in a local model + /// directory: primary = the model config's first id, extras = every other + /// listed id. Missing files or fields degrade to fewer ids. + pub fn from_model_dir(dir: &Path) -> Self { + let model_ids = eos_ids_from_file(&dir.join("config.json")); + let gen_ids = eos_ids_from_file(&dir.join("generation_config.json")); + let primary = model_ids.first().copied().or(gen_ids.first().copied()); + let mut extra = Vec::new(); + for id in model_ids.into_iter().chain(gen_ids) { + if Some(id) != primary && !extra.contains(&id) { + extra.push(id); + } + } + Self { primary, extra } + } +} + +/// Read a config file's `eos_token_id`, which is a single id or a list. +fn eos_ids_from_file(path: &Path) -> Vec { + std::fs::read_to_string(path) + .ok() + .and_then(|text| serde_json::from_str::(&text).ok()) + .map(|config| eos_ids_from_value(config.get("eos_token_id"))) + .unwrap_or_default() +} + +fn eos_ids_from_value(value: Option<&serde_json::Value>) -> Vec { + let as_id = |v: &serde_json::Value| v.as_u64().and_then(|id| u32::try_from(id).ok()); + match value { + Some(serde_json::Value::Array(ids)) => ids.iter().filter_map(as_id).collect(), + Some(id) => as_id(id).into_iter().collect(), + None => Vec::new(), + } +} + /// Direct ZMQ connection to a same-host engine (vLLM EngineCore or TokenSpeed), /// presented behind the vLLM gRPC client surface. #[derive(Clone)] @@ -65,6 +118,9 @@ pub struct ZmqEngineClient { /// Model id advertised for metadata (the engine does not report it on the /// wire; it is configured at worker registration). model_id: String, + /// EOS ids attached to every vLLM request (the engine can't stop at EOS + /// without them). + eos: EosTokenIds, } impl ZmqEngineClient { @@ -81,6 +137,7 @@ impl ZmqEngineClient { output_address: &str, engine_count: usize, model_id: String, + eos: EosTokenIds, runtime: RuntimeType, timeout: Duration, ) -> Result> { @@ -127,7 +184,11 @@ impl ZmqEngineClient { // runtimes were rejected before the handshake. _ => ZmqBackend::Vllm(Arc::new(EngineCoreClient::new(transport))), }; - Ok(Self { backend, model_id }) + Ok(Self { + backend, + model_id, + eos, + }) } /// The engine runtime behind this connection (the wire protocol chosen at @@ -186,7 +247,7 @@ impl ZmqEngineClient { .and_then(|sp| sp.logprobs) .filter(|&n| n > 0) .map_or(0, |n| n as usize); - let request = translate_request(sub, max_model_len) + let request = translate_request(sub, max_model_len, &self.eos) .map_err(tonic::Status::invalid_argument)?; let stream = client.submit(request).await.map_err(zmq_status)?; streams.push(VllmGenerateStream::new(stream, index as u32, top_logprobs)); @@ -810,6 +871,7 @@ fn translate_sampling_tokenspeed(sp: vllm::SamplingParams) -> TokenSpeedSampling fn translate_request( req: vllm::GenerateRequest, max_model_len: u64, + eos: &EosTokenIds, ) -> Result { let prompt_token_ids = match req.input { Some(vllm::generate_request::Input::Tokenized(tokenized)) => Some(tokenized.input_ids), @@ -843,7 +905,7 @@ fn translate_request( prompt_token_ids, sampling_params: req .sampling_params - .map(|sp| translate_sampling(sp, default_max_tokens)), + .map(|sp| translate_sampling(sp, default_max_tokens, eos)), arrival_time: now_secs(), data_parallel_rank, ..EngineCoreRequest::default() @@ -853,7 +915,23 @@ fn translate_request( fn translate_sampling( sp: vllm::SamplingParams, default_max_tokens: u32, + eos: &EosTokenIds, ) -> EngineCoreSamplingParams { + // Stopping at EOS is the frontend's duty here: the primary id rides + // `_eos_token_id`, extra ids merge into `stop_token_ids`, and the union + // feeds `_all_stop_token_ids` (engine-side `min_tokens` masking, built + // regardless of `ignore_eos`). + let mut stop_token_ids = sp.stop_token_ids; + if !sp.ignore_eos { + for id in &eos.extra { + if !stop_token_ids.contains(id) { + stop_token_ids.push(*id); + } + } + } + let mut all_stop_token_ids: BTreeSet = stop_token_ids.iter().copied().collect(); + all_stop_token_ids.extend(eos.primary); + all_stop_token_ids.extend(eos.extra.iter().copied()); let logit_bias = if sp.logit_bias.is_empty() { None } else { @@ -882,7 +960,9 @@ fn translate_sampling( repetition_penalty: sp.repetition_penalty, max_tokens: sp.max_tokens.unwrap_or(default_max_tokens), min_tokens: sp.min_tokens, - stop_token_ids: sp.stop_token_ids, + stop_token_ids, + eos_token_id: (!sp.ignore_eos).then_some(eos.primary).flatten(), + all_stop_token_ids, seed: sp.seed.map(i64::from), logprobs: sp.logprobs, // prompt_logprobs is rejected in `translate_request` (no renderer @@ -1026,6 +1106,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::Vllm, Duration::from_secs(10) ), @@ -1144,6 +1225,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::Vllm, Duration::from_secs(10) ), @@ -1288,6 +1370,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::TokenSpeed, Duration::from_secs(10) ), @@ -1582,6 +1665,7 @@ mod tests { ..Default::default() }), 4096, + &EosTokenIds::default(), ) .expect_err("prompt logprobs rejected"); assert!(err.contains("prompt logprobs"), "{err}"); @@ -1590,11 +1674,15 @@ mod tests { #[test] fn vllm_defaults_unset_max_tokens_to_remaining_context() { let max_tokens = |sampling, max_model_len| { - translate_request(tokenized_req(sampling), max_model_len) - .expect("request translated") - .sampling_params - .expect("sampling params present") - .max_tokens + translate_request( + tokenized_req(sampling), + max_model_len, + &EosTokenIds::default(), + ) + .expect("request translated") + .sampling_params + .expect("sampling params present") + .max_tokens }; // Unset max_tokens defaults to `max_model_len - prompt_len` (prompt is @@ -1613,6 +1701,60 @@ mod tests { ); } + #[test] + fn vllm_attaches_eos_stop_ids() { + let eos = EosTokenIds::new(Some(5), vec![7]); + let sampling = |sp| { + translate_request(tokenized_req(sp), 4096, &eos) + .expect("request translated") + .sampling_params + .expect("sampling params present") + }; + + // Primary rides `_eos_token_id`, extras merge into `stop_token_ids` + // without duplicating, and the union lands in `_all_stop_token_ids`. + let sp = sampling(vllm::SamplingParams { + stop_token_ids: vec![7, 9], + ..Default::default() + }); + assert_eq!(sp.eos_token_id, Some(5)); + assert_eq!(sp.stop_token_ids, vec![7, 9]); + assert_eq!(sp.all_stop_token_ids, BTreeSet::from([5, 7, 9])); + + // ignore_eos drops the EOS stops from the wire but keeps the + // bookkeeping set (mirrors the reference frontend). + let sp = sampling(vllm::SamplingParams { + stop_token_ids: vec![9], + ignore_eos: true, + ..Default::default() + }); + assert_eq!(sp.eos_token_id, None); + assert_eq!(sp.stop_token_ids, vec![9]); + assert_eq!(sp.all_stop_token_ids, BTreeSet::from([5, 7, 9])); + } + + #[test] + fn eos_token_ids_resolve_from_model_dir() { + let dir = tempfile::tempdir().expect("tempdir"); + std::fs::write(dir.path().join("config.json"), r#"{"eos_token_id": 5}"#).unwrap(); + std::fs::write( + dir.path().join("generation_config.json"), + r#"{"eos_token_id": [5, 7, 9]}"#, + ) + .unwrap(); + assert_eq!( + EosTokenIds::from_model_dir(dir.path()), + EosTokenIds::new(Some(5), vec![7, 9]), + ); + + // Missing files degrade to no ids, not an error. + let empty = tempfile::tempdir().expect("tempdir"); + assert_eq!( + EosTokenIds::from_model_dir(empty.path()), + EosTokenIds::default(), + ); + } + #[test] fn vllm_translates_structured_output_constraints() { use engine_zmq_client::protocol::vllm::structured_outputs::{ @@ -1626,6 +1768,7 @@ mod tests { ..Default::default() }), 4096, + &EosTokenIds::default(), ) .expect("constraint translated") .sampling_params @@ -1739,6 +1882,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::Vllm, Duration::from_secs(10) ), @@ -1836,6 +1980,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::TokenSpeed, Duration::from_secs(10) ), @@ -1934,6 +2079,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::Vllm, Duration::from_secs(10) ), diff --git a/model_gateway/src/worker/worker.rs b/model_gateway/src/worker/worker.rs index a117fc545..d34b9a4e3 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -35,7 +35,7 @@ use crate::{ grpc::{ backend_client::BackendClient, client::GrpcClient, - zmq_client::{ZmqEngineClient, ZMQ_LOOPBACK_HOST}, + zmq_client::{EosTokenIds, ZmqEngineClient, ZMQ_LOOPBACK_HOST}, }, }, }; @@ -200,6 +200,19 @@ async fn connect_zmq_backend( let (handshake, input, output) = zmq_socket_addresses(&base_url, handshake_override.as_deref())?; ensure_ipc_socket_dir(&base_url).await?; + // The engine can't stop at EOS on its own (it has no tokenizer or model + // config); resolve the EOS ids from the local model dir so every request + // carries them. + let model_dir = std::path::Path::new(&model_id); + let eos = if model_dir.is_dir() { + EosTokenIds::from_model_dir(model_dir) + } else { + tracing::warn!( + "ZMQ worker model id '{model_id}' is not a local model directory; EOS ids \ + unavailable — generation stops only at max_tokens or explicit stops" + ); + EosTokenIds::default() + }; tracing::info!("Binding ZMQ client for worker {base_url} (handshake={handshake})"); match ZmqEngineClient::connect( &handshake, @@ -207,6 +220,7 @@ async fn connect_zmq_backend( &output, 1, model_id, + eos, runtime, ZMQ_CONNECT_TIMEOUT, ) @@ -2826,6 +2840,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::Vllm, Duration::from_secs(10) ), From b1d4ce73e8a305481a2f230cdfdb550acf0f19da Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Tue, 4 Aug 2026 16:30:07 -0700 Subject: [PATCH 12/16] ci(e2e): bump the TokenSpeed pin to the msgpack ZMQ wire merge The chat-zmq (tokenspeed) lane needs the headless engine mode that merged upstream; the previous pin predates it by 67 commits, so the installed CLI fell back to the orchestrator path and rejected the launcher flags (unrecognized arguments, rc=2). Pin to the merge commit itself rather than current main: the KV-pool cache-recipe refactor that landed after it breaks the gRPC servicer's launch path (MHATokenToKVPool publishes no PagedCacheRuntimeContract), confirmed on a CI run against latest main. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- scripts/ci_install_tokenspeed.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scripts/ci_install_tokenspeed.sh b/scripts/ci_install_tokenspeed.sh index 1b1d4f5aa..352f94d5e 100755 --- a/scripts/ci_install_tokenspeed.sh +++ b/scripts/ci_install_tokenspeed.sh @@ -22,7 +22,7 @@ fi # engine-watch workflow files an issue when this drifts) rather than # floating against ``main`` — upstream has renamed APIs before and the # gRPC servicer broke until we caught up. -TOKENSPEED_REF="${TOKENSPEED_REF:-0f68676069141857a605b8805c13131d9f53e901}" +TOKENSPEED_REF="${TOKENSPEED_REF:-57b49061a1aa258a7c27d8e25f32beb91a4ed896}" TOKENSPEED_REPO="${TOKENSPEED_REPO:-https://github.com/lightseekorg/tokenspeed.git}" TOKENSPEED_DIR="${TOKENSPEED_DIR:-/tmp/tokenspeed-src}" From 703865b2c1e566796f4d2dd9a3474d6df931062e Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Tue, 4 Aug 2026 18:58:42 -0700 Subject: [PATCH 13/16] feat(zmq): carry multimodal inputs to the engine over the direct backend MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The direct-ZMQ path bypasses the Python servicer that converted batched proto tensors into the engine's per-item multimodal structures, so image and video requests were rejected outright. Type the mm_features slot in the wire protocol (feature specs, kwargs items, batched/flat/shared field factories, placeholder ranges) matching the engine's msgspec encoding, and perform the same per-item split in the gateway: batched keys index row i, flat keys slice by the cumulative sizes tensor, everything else is shared. Floating tensors are cast to the model dtype from the ready handshake — the engine applies no cast on this path — and placeholder embed masks are derived from the image token id. Tensors ride the request frame as inline ext-3 raw views (the engine's aux-frame split is an encoder-side size optimization; inline decodes at any size). Assembly reuses the existing vLLM arm with shm/RDMA payloads forced off, and the TokenSpeed translate rejects multimodal loudly until its wire grows a slot. Live-verified on GB300 + Qwen3-VL-8B over ipc://: single image, two images with correct per-position attribution, and streaming. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- crates/engine_zmq_client/src/codec/tensor.rs | 36 ++ .../src/protocol/vllm/mod.rs | 8 +- .../src/protocol/vllm/multimodal.rs | 324 +++++++++++ .../src/protocol/vllm/request.rs | 16 +- .../src/routers/grpc/backend_client.rs | 23 +- model_gateway/src/routers/grpc/mod.rs | 1 + .../src/routers/grpc/multimodal/assemble.rs | 36 +- model_gateway/src/routers/grpc/zmq_client.rs | 36 +- .../src/routers/grpc/zmq_multimodal.rs | 518 ++++++++++++++++++ 9 files changed, 958 insertions(+), 40 deletions(-) create mode 100644 crates/engine_zmq_client/src/protocol/vllm/multimodal.rs create mode 100644 model_gateway/src/routers/grpc/zmq_multimodal.rs diff --git a/crates/engine_zmq_client/src/codec/tensor.rs b/crates/engine_zmq_client/src/codec/tensor.rs index 4f3e959ea..dadba06f0 100644 --- a/crates/engine_zmq_client/src/codec/tensor.rs +++ b/crates/engine_zmq_client/src/codec/tensor.rs @@ -136,6 +136,42 @@ impl WireNdArray { Self::from_raw_bytes(dtype, shape, Bytes::from(data)) } + /// Build from little-endian `float32` bytes, casting each element to + /// `dtype`. Mirrors the model-dtype cast the engine's own frontend applies + /// to floating multimodal tensors before they reach the model. + pub fn from_f32_bytes_cast( + dtype: super::dtype::ModelDtype, + shape: Vec, + data: &[u8], + ) -> std::result::Result { + use super::dtype::ModelDtype; + if !data.len().is_multiple_of(4) { + return Err(format!( + "float32 buffer length {} is not a multiple of 4", + data.len() + )); + } + validate_element_count(&shape, data.len() / 4)?; + let floats = data + .as_chunks::<4>() + .0 + .iter() + .map(|c| f32::from_le_bytes(*c)); + Ok(match dtype { + ModelDtype::Float32 => Self::from_raw("float32", shape, data.to_vec()), + ModelDtype::Float16 => Self::from_raw_bytes( + "float16", + shape, + bytes_from_pod_vec(floats.map(f16::from_f32).collect::>()), + ), + ModelDtype::BFloat16 => Self::from_raw_bytes( + "bfloat16", + shape, + bytes_from_pod_vec(floats.map(bf16::from_f32).collect::>()), + ), + }) + } + /// Build from an owned immutable raw-view buffer. pub fn from_raw_bytes(dtype: impl Into, shape: Vec, data: Bytes) -> Self { Self { diff --git a/crates/engine_zmq_client/src/protocol/vllm/mod.rs b/crates/engine_zmq_client/src/protocol/vllm/mod.rs index 0d8cf742f..cec63b4db 100644 --- a/crates/engine_zmq_client/src/protocol/vllm/mod.rs +++ b/crates/engine_zmq_client/src/protocol/vllm/mod.rs @@ -7,16 +7,16 @@ //! shapes, field order, and `array_like` positional-tuple encoding are the wire //! contract with Python `EngineCoreProc` — do not reorder fields. //! -//! Text generation and structured outputs (guided decoding) are typed fully. -//! Remaining multimodal features and pooling params are carried as -//! [`crate::codec::OpaqueValue`] for now — they serialize as `nil` on the text -//! path and get strongly typed in the multimodal phase. +//! Text generation, structured outputs (guided decoding), and multimodal +//! features are typed fully. Pooling params and prompt embeds are carried as +//! [`crate::codec::OpaqueValue`] — they serialize as `nil` on supported paths. // The startup handshake is engine-neutral (TokenSpeed speaks the same // protocol); re-exported here so existing `vllm::handshake` paths keep working. pub use crate::protocol::handshake; pub mod logprobs; pub mod lora; +pub mod multimodal; pub mod output; pub mod request; pub mod sampling; diff --git a/crates/engine_zmq_client/src/protocol/vllm/multimodal.rs b/crates/engine_zmq_client/src/protocol/vllm/multimodal.rs new file mode 100644 index 000000000..50c65a757 --- /dev/null +++ b/crates/engine_zmq_client/src/protocol/vllm/multimodal.rs @@ -0,0 +1,324 @@ +// SPDX-License-Identifier: Apache-2.0 +// +// Ported from the Apache-2.0 reference `vllm-engine-core-client` +// (vllm-project/vllm): protocol/multimodal.rs. + +use std::collections::BTreeMap; + +use serde::{Deserialize, Serialize}; +use serde_tuple::{Deserialize_tuple, Serialize_tuple}; + +use crate::codec::tensor::WireTensor; + +/// Multimodal feature payload carried at `EngineCoreRequest.mm_features`. +/// +/// Python: `list[MultiModalFeatureSpec] | None` (`vllm/v1/engine/__init__.py`). +pub type MmFeatures = Vec; + +/// A single multimodal input with its processed data and metadata. A request +/// containing multiple multimodal items carries one `MmFeatureSpec` per item. +/// +/// Python: `MultiModalFeatureSpec` (`vllm/multimodal/inputs.py`), a dataclass — +/// encodes as a string-keyed msgpack map. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct MmFeatureSpec { + /// Processed multimodal data for this item. `None` only when the engine's + /// receiver cache already holds `identifier` — an external frontend that + /// does not mirror that cache protocol must always send the full data. + pub data: Option, + + /// The input modality, e.g. `"image"`, `"audio"`, `"video"`. + pub modality: String, + + /// The hash for caching encoder outputs (with LoRA prefix if applicable). + pub identifier: String, + + /// The location of the `modality` tokens corresponding to this item in + /// the prompt. + pub mm_position: PlaceholderRange, + + /// The hash for caching processor outputs (without LoRA prefix). + #[serde(default)] + pub mm_hash: Option, +} + +/// Placeholder location information for one multimodal item. +/// +/// Python: `PlaceholderRange` (`vllm/multimodal/inputs.py`), a dataclass — +/// encodes as a string-keyed msgpack map. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct PlaceholderRange { + /// The start index of the placeholder in the prompt. + pub offset: usize, + + /// The length of the placeholder. + pub length: usize, + + /// A boolean mask of shape `(length,)` indicating which positions between + /// `offset` and `offset + length` receive embeddings. `None` means all. + #[serde(default)] + pub is_embed: Option, +} + +/// Processed keyword arguments for a single multimodal item, keyed by model +/// kwarg name (e.g. `pixel_values`). +/// +/// Python: `MultiModalKwargsItem` (`vllm/multimodal/inputs.py`) — encoded by +/// the serializer hooks as a string-keyed map. +pub type MmKwargsItem = BTreeMap; + +/// One processed keyword argument of a `MmKwargsItem`. +/// +/// Python: `MultiModalFieldElem` (`vllm/multimodal/inputs.py`). +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct MmFieldElem { + /// The keyword argument value passed to the model. `None` only when the + /// item is cached engine-side (see [`MmFeatureSpec::data`]). + pub data: Option, + + /// How this field's values combine with other items' for batching. + pub field: MmField, +} + +/// Processed multimodal keyword argument value (Python `NestedTensors`). +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum MmKwargValue { + Tensor(WireTensor), + Int(i64), + Float(f64), + List(Vec), +} + +/// How to interpret tensor data belonging to a keyword argument. +/// +/// Wire form is a 2-tuple `(factory_name, kwargs_map)` with factory names +/// `"batched"`, `"flat"`, `"shared"` — the serializer's +/// `MMF_CLASS_TO_FACTORY` encoding. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(try_from = "MmFieldWire", into = "MmFieldWire")] +pub enum MmField { + Batched(MmBatchedField), + Flat(MmFlatField), + Shared(MmSharedField), +} + +/// Python `MultiModalFieldConfig.batched`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MmBatchedField { + /// If `true`, this field is excluded from being moved to the accelerator + /// when multimodal items are grouped and batched. + pub keep_on_cpu: bool, +} + +/// Python `MultiModalFieldConfig.flat` / `flat_from_sizes`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MmFlatField { + /// For each multimodal item, a slice (`dim=0`) or a tuple of slices + /// (`dim>0`) that extracts the data corresponding to it. + pub slices: Vec, + + /// The dimension to extract data from, default 0. + pub dim: i32, + + /// If `true`, this field is excluded from being moved to the accelerator + /// when multimodal items are grouped and batched. + pub keep_on_cpu: bool, +} + +/// Python `MultiModalFieldConfig.shared`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct MmSharedField { + pub batch_size: usize, + + /// If `true`, this field is excluded from being moved to the accelerator + /// when multimodal items are grouped and batched. + pub keep_on_cpu: bool, +} + +/// Python slice encoded as `(start, stop, step)`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize_tuple, Deserialize_tuple)] +pub struct SliceSpec { + pub start: Option, + pub stop: Option, + pub step: Option, +} + +/// A single slice or a tuple of slices used by [`MmFlatField`]. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum MmSlice { + Slice(SliceSpec), + Slices(Vec), +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize_tuple, Deserialize_tuple)] +struct MmFieldWire { + name: String, + inner: MmFieldWireInner, +} + +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(untagged)] +enum MmFieldWireInner { + Batched(MmBatchedField), + Flat(MmFlatField), + Shared(MmSharedField), +} + +impl TryFrom for MmField { + type Error = String; + + fn try_from(value: MmFieldWire) -> Result { + match (value.name.as_str(), value.inner) { + ("batched", MmFieldWireInner::Batched(kwargs)) => Ok(Self::Batched(kwargs)), + ("flat", MmFieldWireInner::Flat(kwargs)) => Ok(Self::Flat(kwargs)), + ("shared", MmFieldWireInner::Shared(kwargs)) => Ok(Self::Shared(kwargs)), + (name, _) => Err(format!( + "mismatched or unknown multimodal field factory {name:?}" + )), + } + } +} + +impl From for MmFieldWire { + fn from(value: MmField) -> Self { + match value { + MmField::Batched(kwargs) => Self { + name: "batched".to_string(), + inner: MmFieldWireInner::Batched(kwargs), + }, + MmField::Flat(kwargs) => Self { + name: "flat".to_string(), + inner: MmFieldWireInner::Flat(kwargs), + }, + MmField::Shared(kwargs) => Self { + name: "shared".to_string(), + inner: MmFieldWireInner::Shared(kwargs), + }, + } + } +} + +#[cfg(test)] +mod tests { + use std::io::Cursor; + + use rmpv::Value; + + use super::*; + use crate::codec::encode_msgpack; + + fn encode_value(value: &T) -> Value { + let bytes = encode_msgpack(value).expect("encode value"); + rmpv::decode::read_value(&mut Cursor::new(bytes)).expect("decode value") + } + + #[test] + fn field_serializes_to_python_factory_tuple() { + let field = MmField::Flat(MmFlatField { + slices: vec![MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some(1200), + step: None, + })], + dim: 0, + keep_on_cpu: false, + }); + + let value = encode_value(&field); + let Value::Array(items) = value else { + panic!("field should encode as a 2-tuple array"); + }; + assert_eq!(items.len(), 2); + assert_eq!(items[0].as_str(), Some("flat")); + let Value::Map(kwargs) = &items[1] else { + panic!("field kwargs should encode as a map"); + }; + for key in ["slices", "dim", "keep_on_cpu"] { + assert!( + kwargs.iter().any(|(k, _)| k.as_str() == Some(key)), + "missing kwarg {key}" + ); + } + } + + #[test] + fn field_round_trips_python_factory_tuple() { + for field in [ + MmField::Batched(MmBatchedField { keep_on_cpu: true }), + MmField::Shared(MmSharedField { + batch_size: 4, + keep_on_cpu: false, + }), + ] { + let encoded = encode_msgpack(&field).expect("encode field"); + let decoded: MmField = rmp_serde::from_slice(&encoded).expect("decode field"); + assert_eq!(decoded, field); + } + } + + #[test] + fn feature_spec_serializes_as_named_map_with_tensor_ext() { + let mut item = MmKwargsItem::new(); + item.insert( + "pixel_values".to_string(), + MmFieldElem { + data: Some(MmKwargValue::Tensor( + WireTensor::from_f32(vec![2, 3], vec![0.0, 1.0, 2.0, 3.0, 4.0, 5.0]) + .expect("tensor built"), + )), + field: MmField::Batched(MmBatchedField { keep_on_cpu: false }), + }, + ); + let spec = MmFeatureSpec { + data: Some(item), + modality: "image".to_string(), + identifier: "abc123".to_string(), + mm_position: PlaceholderRange { + offset: 5, + length: 6, + is_embed: None, + }, + mm_hash: Some("abc123".to_string()), + }; + + let value = encode_value(&spec); + let Value::Map(entries) = value else { + panic!("feature spec should encode as a map"); + }; + for key in ["data", "modality", "identifier", "mm_position", "mm_hash"] { + assert!( + entries.iter().any(|(k, _)| k.as_str() == Some(key)), + "missing key {key}" + ); + } + + // The tensor payload must reach the wire as the 3-tuple + // (dtype, shape, ext-3 raw view). + let data = entries + .iter() + .find(|(k, _)| k.as_str() == Some("data")) + .map(|(_, v)| v) + .expect("data present"); + let Value::Map(kwargs) = data else { + panic!("kwargs item should encode as a map"); + }; + let Value::Map(elem) = &kwargs[0].1 else { + panic!("field elem should encode as a map"); + }; + let tensor = elem + .iter() + .find(|(k, _)| k.as_str() == Some("data")) + .map(|(_, v)| v) + .expect("elem data present"); + let Value::Array(tuple) = tensor else { + panic!("tensor should encode as (dtype, shape, data)"); + }; + assert_eq!(tuple[0].as_str(), Some("float32")); + assert!(matches!(&tuple[2], Value::Ext(3, _))); + } +} diff --git a/crates/engine_zmq_client/src/protocol/vllm/request.rs b/crates/engine_zmq_client/src/protocol/vllm/request.rs index de7546f5e..8d3b0b0d1 100644 --- a/crates/engine_zmq_client/src/protocol/vllm/request.rs +++ b/crates/engine_zmq_client/src/protocol/vllm/request.rs @@ -2,9 +2,6 @@ // // Ported from the Apache-2.0 reference `vllm-engine-core-client` // (vllm-project/vllm): protocol/request.rs. -// -// `mm_features` is carried as OpaqueValue for now (typed in the multimodal -// phase); on the text path it is `nil` and carries no aux-frame tensors. use std::collections::{BTreeMap, HashMap}; @@ -15,7 +12,7 @@ use serde_tuple::{Deserialize_tuple, Serialize_tuple}; use crate::{ codec::OpaqueValue, - protocol::vllm::{lora, sampling::EngineCoreSamplingParams}, + protocol::vllm::{lora, multimodal::MmFeatures, sampling::EngineCoreSamplingParams}, Error, Result, }; @@ -72,8 +69,8 @@ pub struct ReasoningParserKwargs { pub struct EngineCoreRequest { pub request_id: String, pub prompt_token_ids: Option>, - /// Multimodal features (untyped for now; `nil` on the text path). - pub mm_features: Option, + /// Multimodal features, one per input item, sorted by placeholder offset. + pub mm_features: Option, pub sampling_params: Option, /// Pooling parameters, preserved in the schema but not yet strongly typed. pub pooling_params: Option, @@ -131,9 +128,10 @@ impl EngineCoreRequest { Ok(()) } - // NOTE: send-side aux-frame extraction (walking `mm_features` for large - // tensors) is added with the typed multimodal module. Text requests carry - // no tensors, so the transport send path appends no aux frames for now. + // NOTE: multimodal tensors are sent as inline ext-3 raw views in the + // request frame — valid at any size (the engine's aux-frame split is an + // encoder-side optimization only). Send-side aux extraction is a perf + // follow-up. } #[cfg(test)] diff --git a/model_gateway/src/routers/grpc/backend_client.rs b/model_gateway/src/routers/grpc/backend_client.rs index 4fd2916b5..0492d762c 100644 --- a/model_gateway/src/routers/grpc/backend_client.rs +++ b/model_gateway/src/routers/grpc/backend_client.rs @@ -23,6 +23,7 @@ use crate::routers::grpc::{ ProtoStream, }, zmq_client::ZmqEngineClient, + MultimodalData, }; /// A backend connection: gRPC (any engine) or direct ZMQ (vLLM EngineCore or @@ -210,8 +211,11 @@ impl BackendClient { client.build_chat_request(request_id, body, processed_text, token_ids, options) } Self::Zmq(_) => { - reject_zmq_multimodal(&options)?; - finish_vllm_request(None, |mm| { + let vllm_mm = options.multimodal_inputs.map(|mm| match mm { + MultimodalData::Vllm(data) => data.into_proto(), + _ => unreachable!("caller guarantees matching variant"), + }); + finish_vllm_request(vllm_mm, |mm| { VllmEngineClient::build_generate_request_from_chat( request_id, body, @@ -238,8 +242,11 @@ impl BackendClient { client.build_messages_request(request_id, body, processed_text, token_ids, options) } Self::Zmq(_) => { - reject_zmq_multimodal(&options)?; - finish_vllm_request(None, |mm| { + let vllm_mm = options.multimodal_inputs.map(|mm| match mm { + MultimodalData::Vllm(data) => data.into_proto(), + _ => unreachable!("caller guarantees matching variant"), + }); + finish_vllm_request(vllm_mm, |mm| { VllmEngineClient::build_generate_request_from_messages( request_id, body, @@ -299,11 +306,3 @@ impl BackendClient { } } } - -/// ZMQ text path does not carry multimodal inputs yet. -fn reject_zmq_multimodal(options: &GenerateRequestBuildOptions) -> Result<(), String> { - if options.multimodal_inputs.is_some() { - return Err("ZMQ backend does not support multimodal inputs yet".to_string()); - } - Ok(()) -} diff --git a/model_gateway/src/routers/grpc/mod.rs b/model_gateway/src/routers/grpc/mod.rs index 6014c7b2f..6f374ba86 100644 --- a/model_gateway/src/routers/grpc/mod.rs +++ b/model_gateway/src/routers/grpc/mod.rs @@ -18,6 +18,7 @@ pub(crate) mod regular; pub(crate) mod router; // Used by routers/factory pub mod utils; // Used by routers/http and bindings/golang pub mod zmq_client; // ZMQ backend adapter behind the vLLM client surface +pub(crate) mod zmq_multimodal; // Proto mm inputs → EngineCore mm_features // Re-export for convenience pub use proto_wrapper::{MultimodalData, TensorBytes}; diff --git a/model_gateway/src/routers/grpc/multimodal/assemble.rs b/model_gateway/src/routers/grpc/multimodal/assemble.rs index ff839eec7..d08abcbb8 100644 --- a/model_gateway/src/routers/grpc/multimodal/assemble.rs +++ b/model_gateway/src/routers/grpc/multimodal/assemble.rs @@ -28,16 +28,19 @@ use super::{ transport::{mm_encoder_input_dtype, resolve_mm_shm_enabled, resolve_mm_shm_min_bytes}, MediaBatch, MultimodalIntermediate, PrecomputedMultimodalIntermediate, PromptBinding, }; -use crate::routers::grpc::{ - backend_client::BackendClient, - client::GrpcClient, - context::WorkerSelection, - proto_wrapper::{ - cleanup_tokenspeed_items_encoder_shm, SglangMultimodalData, TensorBytes, - TokenSpeedModality, TokenSpeedMultimodalData, TokenSpeedMultimodalItem, TokenSpeedTensor, - TrtllmMultimodalData, VllmMultimodalData, +use crate::{ + routers::grpc::{ + backend_client::BackendClient, + client::GrpcClient, + context::WorkerSelection, + proto_wrapper::{ + cleanup_tokenspeed_items_encoder_shm, SglangMultimodalData, TensorBytes, + TokenSpeedModality, TokenSpeedMultimodalData, TokenSpeedMultimodalItem, + TokenSpeedTensor, TrtllmMultimodalData, VllmMultimodalData, + }, + MultimodalData, }, - MultimodalData, + worker::RuntimeType, }; /// Assemble backend-specific multimodal data from the intermediate. @@ -108,7 +111,20 @@ async fn assemble_multimodal_data_impl( BackendClient::Grpc(GrpcClient::Mlx(_)) => { anyhow::bail!("MLX does not support multimodal inputs") } - BackendClient::Zmq(_) => anyhow::bail!("ZMQ backend does not support multimodal inputs"), + BackendClient::Zmq(client) => match client.runtime() { + RuntimeType::Vllm | RuntimeType::Unspecified => { + let batch = into_single_batch(intermediate, "vLLM")?; + let mut data = assemble_vllm(batch, workers)?; + // The ZMQ translate reads tensor bytes inline; this wire has no + // /dev/shm or RDMA pull on the engine side. + data.shm_enabled = false; + data.rdma_enabled = false; + Ok(MultimodalData::Vllm(data)) + } + runtime => anyhow::bail!( + "multimodal inputs are not supported over the {runtime} ZMQ backend yet" + ), + }, } } diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index 66cbf6b6f..9a2c4b433 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -20,6 +20,7 @@ use std::{ }; use engine_zmq_client::{ + codec::dtype::ModelDtype, connect_handshake, connector::{EngineCoreClient, EngineCoreStream, TokenSpeedClient, TokenSpeedStream}, protocol::{ @@ -41,7 +42,7 @@ use futures::{stream::SelectAll, Stream, StreamExt}; use openai_protocol::worker::{SchedulerLoadSnapshot, WorkerLoadResponse}; use smg_grpc_client::vllm_proto as vllm; -use crate::worker::RuntimeType; +use crate::{routers::grpc::zmq_multimodal, worker::RuntimeType}; /// Loopback host for the same-host ZMQ transport (TCP handshake and local /// binds). Shared with the worker-side socket derivation. @@ -231,10 +232,10 @@ impl ZmqEngineClient { // (which the ZMQ path bypasses) defaults an unset value to // `max_model_len - prompt_len`. The context length comes from the // engine's ready handshake, so a connected engine is required. - let max_model_len = client + let (max_model_len, model_dtype) = client .engines() .first() - .map(|e| e.ready_response.max_model_len) + .map(|e| (e.ready_response.max_model_len, e.ready_response.dtype)) .ok_or_else(|| tonic::Status::unavailable("no connected ZMQ engine"))?; let mut streams = SelectAll::new(); for (index, sub) in subs.into_iter().enumerate() { @@ -247,7 +248,7 @@ impl ZmqEngineClient { .and_then(|sp| sp.logprobs) .filter(|&n| n > 0) .map_or(0, |n| n as usize); - let request = translate_request(sub, max_model_len, &self.eos) + let request = translate_request(sub, max_model_len, model_dtype, &self.eos) .map_err(tonic::Status::invalid_argument)?; let stream = client.submit(request).await.map_err(zmq_status)?; streams.push(VllmGenerateStream::new(stream, index as u32, top_logprobs)); @@ -748,6 +749,13 @@ fn fan_out_requests(req: vllm::GenerateRequest) -> Vec { fn translate_request_tokenspeed( req: vllm::GenerateRequest, ) -> Result { + // The TokenSpeed wire has no multimodal slot yet; reject loudly rather + // than silently dropping pixels (assembly also refuses upstream). + if req.mm_inputs.is_some() { + return Err( + "multimodal inputs are not supported over the TokenSpeed ZMQ backend".to_string(), + ); + } let input_ids = match req.input { Some(vllm::generate_request::Input::Tokenized(tokenized)) => tokenized.input_ids, Some(vllm::generate_request::Input::Text(_)) => { @@ -871,6 +879,7 @@ fn translate_sampling_tokenspeed(sp: vllm::SamplingParams) -> TokenSpeedSampling fn translate_request( req: vllm::GenerateRequest, max_model_len: u64, + model_dtype: ModelDtype, eos: &EosTokenIds, ) -> Result { let prompt_token_ids = match req.input { @@ -882,6 +891,19 @@ fn translate_request( return Err("ZMQ mode requires pre-tokenized input; no input provided".to_string()); } }; + // Per-item mm features: the split the Python servicer performs before the + // engine happens here instead (the ZMQ path bypasses it). + let mm_features = req + .mm_inputs + .map(|mm| { + zmq_multimodal::build_mm_features( + mm, + prompt_token_ids.as_deref().unwrap_or(&[]), + model_dtype, + ) + }) + .transpose()? + .filter(|features| !features.is_empty()); let data_parallel_rank = req .data_parallel_rank .map(|rank| u32::try_from(rank).map_err(|_| format!("invalid data_parallel_rank: {rank}"))) @@ -903,6 +925,7 @@ fn translate_request( Ok(EngineCoreRequest { request_id: req.request_id, prompt_token_ids, + mm_features, sampling_params: req .sampling_params .map(|sp| translate_sampling(sp, default_max_tokens, eos)), @@ -1665,6 +1688,7 @@ mod tests { ..Default::default() }), 4096, + ModelDtype::BFloat16, &EosTokenIds::default(), ) .expect_err("prompt logprobs rejected"); @@ -1677,6 +1701,7 @@ mod tests { translate_request( tokenized_req(sampling), max_model_len, + ModelDtype::BFloat16, &EosTokenIds::default(), ) .expect("request translated") @@ -1705,7 +1730,7 @@ mod tests { fn vllm_attaches_eos_stop_ids() { let eos = EosTokenIds::new(Some(5), vec![7]); let sampling = |sp| { - translate_request(tokenized_req(sp), 4096, &eos) + translate_request(tokenized_req(sp), 4096, ModelDtype::BFloat16, &eos) .expect("request translated") .sampling_params .expect("sampling params present") @@ -1768,6 +1793,7 @@ mod tests { ..Default::default() }), 4096, + ModelDtype::BFloat16, &EosTokenIds::default(), ) .expect("constraint translated") diff --git a/model_gateway/src/routers/grpc/zmq_multimodal.rs b/model_gateway/src/routers/grpc/zmq_multimodal.rs new file mode 100644 index 000000000..f20bd3c08 --- /dev/null +++ b/model_gateway/src/routers/grpc/zmq_multimodal.rs @@ -0,0 +1,518 @@ +//! Proto multimodal inputs → EngineCore `mm_features` for the direct-ZMQ path. +//! +//! The gRPC servicer converts the batched proto tensors into per-item engine +//! structures Python-side (`_build_preprocessed_mm_inputs` + the engine's +//! `from_hf_inputs` split). The ZMQ path bypasses that process, so the same +//! split happens here: batched keys index row `i`, flat keys slice by the +//! cumulative sizes tensor, everything else is shared (replicated per item). +//! Floating tensors are cast to the model dtype — the engine applies no cast +//! on this path. + +use std::collections::{BTreeMap, HashMap, HashSet}; + +use bytes::Bytes; +use engine_zmq_client::{ + codec::{ + dtype::ModelDtype, + tensor::{WireArrayData, WireTensor}, + }, + protocol::vllm::multimodal::{ + MmBatchedField, MmFeatureSpec, MmFeatures, MmField, MmFieldElem, MmFlatField, MmKwargValue, + MmKwargsItem, MmSharedField, MmSlice, PlaceholderRange, SliceSpec, + }, +}; +use smg_grpc_client::{common_proto as common, vllm_proto as vllm}; + +/// A decoded (and dtype-cast) proto tensor ready for per-item slicing. +struct Decoded { + dtype: String, + shape: Vec, + bytes: Bytes, +} + +impl Decoded { + fn elem_size(&self) -> Result { + match self.dtype.as_str() { + "bool" => Ok(1), + "float16" | "bfloat16" => Ok(2), + "float32" | "uint32" | "int32" => Ok(4), + "int64" | "float64" => Ok(8), + other => Err(format!("unsupported multimodal tensor dtype {other:?}")), + } + } + + /// Bytes per index step along dim 0. + fn row_nbytes(&self) -> Result { + let inner: usize = self.shape.iter().skip(1).product(); + Ok(inner * self.elem_size()?) + } + + /// Zero-copy view of rows `[start, stop)` along dim 0. + fn slice_rows(&self, start: usize, stop: usize) -> Result { + let row = self.row_nbytes()?; + let (lo, hi) = (start * row, stop * row); + if hi > self.bytes.len() || start > stop { + return Err(format!( + "row slice {start}..{stop} out of bounds for tensor of {} bytes", + self.bytes.len() + )); + } + let mut shape = self.shape.clone(); + shape[0] = stop - start; + Ok(WireTensor::from_raw_bytes( + self.dtype.clone(), + shape, + self.bytes.slice(lo..hi), + )) + } + + fn whole(&self) -> WireTensor { + WireTensor::from_raw_bytes(self.dtype.clone(), self.shape.clone(), self.bytes.clone()) + } + + /// Flattened values as widened i64 (sizes tensors are int64 or uint32). + fn flat_i64(&self) -> Result, String> { + match self.dtype.as_str() { + "int64" => Ok((self.bytes.as_chunks::<8>().0.iter()) + .map(|c| i64::from_le_bytes(*c)) + .collect()), + "uint32" => Ok((self.bytes.as_chunks::<4>().0.iter()) + .map(|c| i64::from(u32::from_le_bytes(*c))) + .collect()), + other => Err(format!("flat sizes tensor has unsupported dtype {other:?}")), + } + } +} + +fn decode_tensor( + name: &str, + tensor: vllm::TensorData, + model_dtype: ModelDtype, +) -> Result { + let shape: Vec = tensor.shape.iter().map(|&d| d as usize).collect(); + let data = match tensor.payload { + Some(vllm::tensor_data::Payload::Inline(data)) => data, + Some(_) => { + return Err(format!( + "multimodal tensor {name:?} uses a non-inline payload; the ZMQ wire carries \ + tensors inline" + )); + } + None => return Err(format!("multimodal tensor {name:?} has no payload")), + }; + // Floating tensors arrive as float32 and are cast to the model dtype, + // mirroring the cast the engine's own frontend applies. + if tensor.dtype == "float32" { + let cast = WireTensor::from_f32_bytes_cast(model_dtype, shape.clone(), &data)?; + let WireArrayData::RawView(bytes) = cast.data else { + return Err(format!("cast tensor {name:?} lost its raw view")); + }; + return Ok(Decoded { + dtype: cast.dtype, + shape, + bytes, + }); + } + Ok(Decoded { + dtype: tensor.dtype, + shape, + bytes: Bytes::from(data), + }) +} + +/// Rename generic keys for video inputs, mirroring the servicer's `mm_key`. +fn mm_key(key: &str, is_video: bool) -> String { + if is_video && key == "pixel_values" { + "pixel_values_videos".to_string() + } else { + key.to_string() + } +} + +/// Build per-item `mm_features` from batched proto multimodal inputs. +pub(crate) fn build_mm_features( + mm: vllm::MultimodalInputs, + prompt_token_ids: &[u32], + model_dtype: ModelDtype, +) -> Result { + let num_items = mm.mm_placeholders.len(); + if num_items == 0 { + return Ok(Vec::new()); + } + if mm.mm_hashes.len() != num_items { + return Err(format!( + "multimodal hash count {} does not match placeholder count {num_items}", + mm.mm_hashes.len() + )); + } + let is_video = mm.modality == common::Modality::Video as i32; + let modality = if is_video { "video" } else { "image" }; + + // Decode every tensor once, applying the video key rename. + let mut tensors: BTreeMap = BTreeMap::new(); + if let Some(pixel_values) = mm.pixel_values { + tensors.insert( + mm_key("pixel_values", is_video), + decode_tensor("pixel_values", pixel_values, model_dtype)?, + ); + } + for (key, tensor) in mm.model_specific_tensors { + let decoded = decode_tensor(&key, tensor, model_dtype)?; + tensors.insert(mm_key(&key, is_video), decoded); + } + + let batched: HashSet = mm + .batched_keys + .iter() + .map(|k| mm_key(k, is_video)) + .collect(); + let flat: HashMap = mm + .flat_keys + .iter() + .map(|(k, v)| (mm_key(k, is_video), mm_key(v, is_video))) + .collect(); + let keep_on_cpu: HashSet = mm + .keep_on_cpu_keys + .iter() + .map(|k| mm_key(k, is_video)) + .collect(); + + // Split every kwarg into per-item elems. + let mut items: Vec = vec![MmKwargsItem::new(); num_items]; + for (key, decoded) in &tensors { + let on_cpu = keep_on_cpu.contains(key); + if batched.contains(key) { + if decoded.shape.first() != Some(&num_items) { + return Err(format!( + "batched tensor {key:?} has leading dim {:?}, expected {num_items} items", + decoded.shape.first() + )); + } + for (i, item) in items.iter_mut().enumerate() { + let mut tensor = decoded.slice_rows(i, i + 1)?; + tensor.shape.remove(0); + item.insert( + key.clone(), + MmFieldElem { + data: Some(MmKwargValue::Tensor(tensor)), + field: MmField::Batched(MmBatchedField { + keep_on_cpu: on_cpu, + }), + }, + ); + } + } else if let Some(sizes_key) = flat.get(key) { + let sizes = tensors + .get(sizes_key) + .ok_or_else(|| format!("flat sizes tensor {sizes_key:?} missing for {key:?}"))? + .flat_i64()?; + if sizes.len() != num_items { + return Err(format!( + "flat sizes tensor {sizes_key:?} has {} entries, expected {num_items}", + sizes.len() + )); + } + // Cumulative row offsets, and the full per-item slice list every + // elem carries (the engine's flat field serializes all slices). + let mut bounds = Vec::with_capacity(num_items + 1); + let mut total = 0usize; + bounds.push(total); + for size in &sizes { + let size = usize::try_from(*size) + .map_err(|_| format!("negative size in flat sizes tensor {sizes_key:?}"))?; + total += size; + bounds.push(total); + } + if decoded.shape.first() != Some(&total) { + return Err(format!( + "flat tensor {key:?} has leading dim {:?}, expected {total} total rows", + decoded.shape.first(), + )); + } + let slices: Vec = bounds + .windows(2) + .map(|w| { + MmSlice::Slice(SliceSpec { + start: Some(w[0] as isize), + stop: Some(w[1] as isize), + step: None, + }) + }) + .collect(); + for (i, item) in items.iter_mut().enumerate() { + item.insert( + key.clone(), + MmFieldElem { + data: Some(MmKwargValue::Tensor( + decoded.slice_rows(bounds[i], bounds[i + 1])?, + )), + field: MmField::Flat(MmFlatField { + slices: slices.clone(), + dim: 0, + keep_on_cpu: on_cpu, + }), + }, + ); + } + } else { + // Shared: the full tensor replicated per item (the servicer's + // fallback for keys in neither batched nor flat sets). + for item in &mut items { + item.insert( + key.clone(), + MmFieldElem { + data: Some(MmKwargValue::Tensor(decoded.whole())), + field: MmField::Shared(MmSharedField { + batch_size: num_items, + keep_on_cpu: false, + }), + }, + ); + } + } + } + + // One feature per placeholder, in prompt-offset order. + let mut features: MmFeatures = Vec::with_capacity(num_items); + for ((placeholder, item), hash) in mm + .mm_placeholders + .iter() + .zip(items) + .zip(mm.mm_hashes.iter()) + { + let offset = placeholder.offset as usize; + let length = placeholder.length as usize; + features.push(MmFeatureSpec { + data: Some(item), + modality: modality.to_string(), + identifier: hash.clone(), + mm_position: PlaceholderRange { + offset, + length, + is_embed: is_embed_mask(prompt_token_ids, offset, length, mm.im_token_id)?, + }, + mm_hash: Some(hash.clone()), + }); + } + features.sort_by_key(|f| f.mm_position.offset); + Ok(features) +} + +/// Boolean embed mask over a placeholder range: `true` where the prompt token +/// is the image token, excluding structural tokens (vision start/end markers) +/// from the embedding scatter. `None` when every position is an embed slot. +fn is_embed_mask( + prompt_token_ids: &[u32], + offset: usize, + length: usize, + im_token_id: Option, +) -> Result, String> { + let Some(im_token_id) = im_token_id else { + return Ok(None); + }; + let end = offset + .checked_add(length) + .filter(|&end| end <= prompt_token_ids.len()) + .ok_or_else(|| { + format!( + "placeholder range {offset}+{length} exceeds prompt of {} tokens", + prompt_token_ids.len() + ) + })?; + let mask: Vec = prompt_token_ids[offset..end] + .iter() + .map(|&id| id == im_token_id) + .collect(); + if mask.iter().all(|&m| m) { + return Ok(None); + } + Ok(Some(WireTensor::from_bool(vec![length], mask)?)) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn inline_tensor(shape: Vec, dtype: &str, data: Vec) -> vllm::TensorData { + vllm::TensorData { + shape, + dtype: dtype.to_string(), + payload: Some(vllm::tensor_data::Payload::Inline(data)), + } + } + + fn f32_bytes(values: &[f32]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() + } + + fn i64_bytes(values: &[i64]) -> Vec { + values.iter().flat_map(|v| v.to_le_bytes()).collect() + } + + fn placeholders(ranges: &[(u32, u32)]) -> Vec { + ranges + .iter() + .map(|&(offset, length)| vllm::PlaceholderRange { offset, length }) + .collect() + } + + fn base_inputs() -> vllm::MultimodalInputs { + vllm::MultimodalInputs { + pixel_values: Some(inline_tensor( + vec![2, 4], + "float32", + f32_bytes(&[0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0]), + )), + model_specific_tensors: Default::default(), + im_token_id: None, + mm_placeholders: placeholders(&[(1, 3), (6, 3)]), + mm_hashes: vec!["h0".to_string(), "h1".to_string()], + batched_keys: vec!["pixel_values".to_string()], + flat_keys: Default::default(), + keep_on_cpu_keys: vec![], + modality: common::Modality::Image as i32, + } + } + + fn tensor_of(elem: &MmFieldElem) -> &WireTensor { + match elem.data.as_ref().expect("data present") { + MmKwargValue::Tensor(tensor) => tensor, + other => panic!("expected tensor, got {other:?}"), + } + } + + #[test] + fn batched_keys_split_per_row_and_cast_to_model_dtype() { + let features = build_mm_features(base_inputs(), &[], ModelDtype::BFloat16).expect("built"); + assert_eq!(features.len(), 2); + + for (i, feature) in features.iter().enumerate() { + assert_eq!(feature.modality, "image"); + assert_eq!(feature.identifier, format!("h{i}")); + assert_eq!(feature.mm_hash.as_deref(), Some(format!("h{i}").as_str())); + let item = feature.data.as_ref().expect("item present"); + let tensor = tensor_of(&item["pixel_values"]); + // Row i of the [2, 4] float32 batch, cast to bfloat16. + assert_eq!(tensor.dtype, "bfloat16"); + assert_eq!(tensor.shape, vec![4]); + assert!(matches!( + item["pixel_values"].field, + MmField::Batched(MmBatchedField { keep_on_cpu: false }) + )); + } + assert_eq!(features[0].mm_position.offset, 1); + assert_eq!(features[1].mm_position.offset, 6); + } + + #[test] + fn flat_keys_slice_by_cumulative_sizes() { + let mut mm = base_inputs(); + mm.pixel_values = Some(inline_tensor(vec![5, 2], "float32", f32_bytes(&[0.0; 10]))); + mm.batched_keys = vec!["patches_per_image".to_string()]; + mm.flat_keys = [("pixel_values".to_string(), "patches_per_image".to_string())].into(); + mm.model_specific_tensors = [( + "patches_per_image".to_string(), + inline_tensor(vec![2], "int64", i64_bytes(&[2, 3])), + )] + .into(); + + let features = build_mm_features(mm, &[], ModelDtype::Float32).expect("built"); + let item0 = features[0].data.as_ref().expect("item 0"); + let item1 = features[1].data.as_ref().expect("item 1"); + assert_eq!(tensor_of(&item0["pixel_values"]).shape, vec![2, 2]); + assert_eq!(tensor_of(&item1["pixel_values"]).shape, vec![3, 2]); + + // Every elem carries the full per-item slice list. + let expected_slices = vec![ + MmSlice::Slice(SliceSpec { + start: Some(0), + stop: Some(2), + step: None, + }), + MmSlice::Slice(SliceSpec { + start: Some(2), + stop: Some(5), + step: None, + }), + ]; + for item in [item0, item1] { + let MmField::Flat(flat) = &item["pixel_values"].field else { + panic!("expected flat field"); + }; + assert_eq!(flat.slices, expected_slices); + assert_eq!(flat.dim, 0); + } + } + + #[test] + fn unlisted_keys_are_shared_and_replicated() { + let mut mm = base_inputs(); + mm.model_specific_tensors = [( + "video_second_per_grid".to_string(), + inline_tensor(vec![2], "int64", i64_bytes(&[1, 1])), + )] + .into(); + + let features = build_mm_features(mm, &[], ModelDtype::BFloat16).expect("built"); + for feature in &features { + let item = feature.data.as_ref().expect("item present"); + let elem = &item["video_second_per_grid"]; + assert_eq!(tensor_of(elem).shape, vec![2]); + assert!(matches!( + elem.field, + MmField::Shared(MmSharedField { + batch_size: 2, + keep_on_cpu: false, + }) + )); + } + } + + #[test] + fn is_embed_masks_structural_tokens() { + let mut mm = base_inputs(); + mm.im_token_id = Some(7); + // Placeholder 0 covers tokens [7, 7, 5] (mixed); placeholder 1 covers + // [7, 7, 7] (all image tokens). + let prompt = [9, 7, 7, 5, 9, 9, 7, 7, 7]; + + let features = build_mm_features(mm, &prompt, ModelDtype::BFloat16).expect("built"); + let mask = features[0] + .mm_position + .is_embed + .as_ref() + .expect("mixed range keeps a mask"); + assert_eq!(mask.dtype, "bool"); + assert_eq!(mask.shape, vec![3]); + assert!(features[1].mm_position.is_embed.is_none()); + } + + #[test] + fn video_renames_pixel_values() { + let mut mm = base_inputs(); + mm.modality = common::Modality::Video as i32; + + let features = build_mm_features(mm, &[], ModelDtype::BFloat16).expect("built"); + let item = features[0].data.as_ref().expect("item present"); + assert!(item.contains_key("pixel_values_videos")); + assert!(!item.contains_key("pixel_values")); + assert_eq!(features[0].modality, "video"); + } + + #[test] + fn rejects_hash_mismatch_and_non_inline_payloads() { + let mut mm = base_inputs(); + mm.mm_hashes.pop(); + let err = build_mm_features(mm, &[], ModelDtype::BFloat16).expect_err("hash mismatch"); + assert!(err.contains("hash count"), "{err}"); + + let mut mm = base_inputs(); + mm.pixel_values = Some(vllm::TensorData { + shape: vec![2, 4], + dtype: "float32".to_string(), + payload: Some(vllm::tensor_data::Payload::Shm(Default::default())), + }); + let err = build_mm_features(mm, &[], ModelDtype::BFloat16).expect_err("shm rejected"); + assert!(err.contains("inline"), "{err}"); + } +} From a597911f8f9bb53c49004131f7f1bc6320770326 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Tue, 4 Aug 2026 21:13:55 -0700 Subject: [PATCH 14/16] ci(e2e): give the chat-zmq lanes a realistic test budget MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The lane restarts the engine per model group (66-90s of load and warmup each) instead of amortizing bring-up across a pool like the gRPC lane, and the latest run hit the 18-minute pytest cap with 53 passed and zero failures — the suite simply no longer fits now that the multimodal and stop-string tests do real GPU work. Raise the pytest step to 40 minutes with matching job timeouts. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- .github/workflows/pr-test-rust.yml | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/.github/workflows/pr-test-rust.yml b/.github/workflows/pr-test-rust.yml index 8aae5d13f..fa8260a39 100644 --- a/.github/workflows/pr-test-rust.yml +++ b/.github/workflows/pr-test-rust.yml @@ -562,12 +562,15 @@ jobs: fail-fast: false matrix: include: + # The ZMQ lane restarts the engine per model group (66-90s of + # load+warmup each); the suite needs more headroom than the gRPC + # lane, whose pool amortizes worker bring-up. - engine: vllm - timeout: 24 - test_timeout: 18 + timeout: 46 + test_timeout: 40 - engine: tokenspeed timeout: 50 - test_timeout: 18 + test_timeout: 40 uses: ./.github/workflows/e2e-gpu-job.yml with: engine: ${{ matrix.engine }} From ac01a77aea5fda5897c6a3f71b376651839fb8b3 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Tue, 4 Aug 2026 22:25:43 -0700 Subject: [PATCH 15/16] fix(zmq): source EOS stops from the tokenizer, not just the model dir MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The connect-time EOS resolution reads config.json from the worker's model id, which only works when that id is a local directory. In CI the ZMQ workers register under HF repo ids, so every model except the locally-mounted one ran with no EOS ids — an uncapped request generated to the full context window, wedging the chat-zmq lane at the first unbounded test (53 passed, zero failed, then 30 minutes of silence). The gateway's tokenizer already carries the merged EOS set from config.json + generation_config.json regardless of how the model was resolved, so fold those ids into the request's stop tokens during request building (ZMQ vLLM only, skipped when ignore_eos is set). The model-dir path stays as the higher-fidelity source when it resolves. Live-verified on GB300: uncapped chat request finishes in 2.8s with finish_reason=stop on the rebuilt gateway. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- crates/tokenizer/src/mock.rs | 5 ++ .../src/routers/grpc/common/stages/helpers.rs | 47 ++++++++++++++++++- 2 files changed, 50 insertions(+), 2 deletions(-) diff --git a/crates/tokenizer/src/mock.rs b/crates/tokenizer/src/mock.rs index f057aabe2..245e1968b 100644 --- a/crates/tokenizer/src/mock.rs +++ b/crates/tokenizer/src/mock.rs @@ -122,6 +122,11 @@ impl TokenizerTrait for MockTokenizer { self.reverse_vocab.get(&id).cloned() } + fn eos_token_ids(&self) -> &[u32] { + // `` in the mock vocab. + &[999] + } + fn as_any(&self) -> &dyn std::any::Any { self } diff --git a/model_gateway/src/routers/grpc/common/stages/helpers.rs b/model_gateway/src/routers/grpc/common/stages/helpers.rs index 701505a86..d8f5331fd 100644 --- a/model_gateway/src/routers/grpc/common/stages/helpers.rs +++ b/model_gateway/src/routers/grpc/common/stages/helpers.rs @@ -382,6 +382,22 @@ pub(crate) fn resolve_string_stops( if let Some(params) = req.sampling_params.as_mut() { let stops = std::mem::take(&mut params.stop); encode_single_token_stops(stops, &mut params.stop_token_ids, tokenizer); + // The engine only stops at EOS when the frontend supplies the + // ids, and the connect-time model-dir resolution has nothing + // to read when the worker's model id is a repo id rather than + // a local path. The tokenizer carries the merged EOS set, so + // fold it into the stop tokens as the always-available + // backstop — without it an uncapped request generates to the + // full context window. + if !params.ignore_eos { + if let Some(tokenizer) = tokenizer { + for &id in tokenizer.eos_token_ids() { + if !params.stop_token_ids.contains(&id) { + params.stop_token_ids.push(id); + } + } + } + } } } _ => {} @@ -646,12 +662,39 @@ mod stop_resolution_tests { ); assert!(params.stop_token_ids.is_empty()); - // ZMQ vLLM (EngineCore sees token ids only) resolves like SGLang. + // ZMQ vLLM (EngineCore sees token ids only) resolves like SGLang, and + // gains the tokenizer's EOS ids so generation always terminates. let mut zmq = vllm_request(vec!["."], vec![]); resolve_string_stops(&mut zmq, Some(&mock_tokenizer()), true); let params = vllm_params(&zmq); assert!(params.stop.is_empty(), "ZMQ vLLM stop cleared"); - assert_eq!(params.stop_token_ids, vec![6]); + assert_eq!(params.stop_token_ids, vec![6, 999]); + } + + #[test] + fn vllm_zmq_appends_tokenizer_eos_ids() { + let mut req = vllm_request(vec![], vec![7]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), true); + assert_eq!(vllm_params(&req).stop_token_ids, vec![7, 999]); + + // Already-present EOS ids are not duplicated. + let mut req = vllm_request(vec![], vec![999]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), true); + assert_eq!(vllm_params(&req).stop_token_ids, vec![999]); + } + + #[test] + fn vllm_zmq_ignore_eos_skips_injection() { + let mut req = ProtoGenerateRequest::Vllm(Box::new(vllm_proto::GenerateRequest { + sampling_params: Some(vllm_proto::SamplingParams { + stop_token_ids: vec![7], + ignore_eos: true, + ..Default::default() + }), + ..Default::default() + })); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), true); + assert_eq!(vllm_params(&req).stop_token_ids, vec![7]); } #[test] From a10ba4dd94b5f1df006255485ded2d7a64137b59 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Wed, 5 Aug 2026 08:17:54 -0700 Subject: [PATCH 16/16] feat(zmq): native per-engine dispatch + TokenSpeed structured outputs Make the direct-ZMQ backend a true sibling of gRPC instead of forcing every request through a vLLM-shaped path. BackendClient::Zmq now builds the native proto for its runtime (vLLM EngineCore or TokenSpeed) and ZmqEngineClient::generate dispatches on the ProtoGenerateRequest variant, translating native proto -> wire directly. This mirrors the per-engine dispatch already used on the gRPC path; the shared execution pipeline (ProtoStream / ProtoGenerateRequest) is untouched. To enable the symmetry, the TokenSpeed request builders are made static (as the vLLM ones already were) so both the gRPC-TS and ZMQ-TS paths call the same builders. Close the TokenSpeed structured-output gap: the proto sampling constraint oneof (regex / json_schema / ebnf / structural_tag) now maps 1:1 onto the wire SamplingParams fields, wired end-to-end. Previously these were rejected as unsupported over the ZMQ backend. Also lands harmony-over-ZMQ: stop synthesis and constrained-content handling for the token-only backend, with a text-based stop scanner. Migrate the affected tests to the native proto types, add coverage for the constraint mapping and multimodal rejection, and satisfy clippy under -D warnings. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- .../src/protocol/tokenspeed/sampling.rs | 9 +- .../src/protocol/vllm/structured_outputs.rs | 27 +- .../grpc_client/src/tokenspeed_scheduler.rs | 25 - .../src/routers/grpc/backend_client.rs | 216 ++++--- model_gateway/src/routers/grpc/client.rs | 16 +- model_gateway/src/routers/grpc/harmony/mod.rs | 1 + .../src/routers/grpc/harmony/processor.rs | 47 +- .../grpc/harmony/stages/request_building.rs | 86 ++- .../harmony/stages/response_processing.rs | 37 +- .../src/routers/grpc/harmony/stop.rs | 163 ++++++ .../src/routers/grpc/harmony/streaming.rs | 139 ++++- model_gateway/src/routers/grpc/zmq_client.rs | 529 ++++++++++-------- 12 files changed, 919 insertions(+), 376 deletions(-) create mode 100644 model_gateway/src/routers/grpc/harmony/stop.rs diff --git a/crates/engine_zmq_client/src/protocol/tokenspeed/sampling.rs b/crates/engine_zmq_client/src/protocol/tokenspeed/sampling.rs index 81ac5e8ea..7fe448588 100644 --- a/crates/engine_zmq_client/src/protocol/tokenspeed/sampling.rs +++ b/crates/engine_zmq_client/src/protocol/tokenspeed/sampling.rs @@ -54,13 +54,14 @@ pub struct SamplingParams { pub repetition_penalty: f64, /// Minimum number of tokens to generate before EOS / stop handling. pub min_new_tokens: u32, - /// Structured-output JSON schema. SMG rejects constraints upstream. + /// Structured-output JSON schema. Set from the request's `constraint`; + /// at most one of these four structured-output fields is populated. pub json_schema: Option, - /// Structured-output regex. SMG rejects constraints upstream. + /// Structured-output regex. Set from the request's `constraint`. pub regex: Option, - /// Structured-output EBNF grammar. SMG rejects constraints upstream. + /// Structured-output EBNF grammar. Set from the request's `constraint`. pub ebnf: Option, - /// Structured-output structural tag. SMG rejects constraints upstream. + /// Structured-output structural tag. Set from the request's `constraint`. pub structural_tag: Option, /// Ignore the EOS token and keep generating until another stop condition. pub ignore_eos: bool, diff --git a/crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs b/crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs index 6e93deb0d..ffef05a19 100644 --- a/crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs +++ b/crates/engine_zmq_client/src/protocol/vllm/structured_outputs.rs @@ -11,8 +11,10 @@ use crate::error::{Error, Result}; /// Structured-output backend selected for EngineCore grammar compilation. /// /// Python stores this in `StructuredOutputsParams._backend` after request -/// validation. This client currently always lowers structured-output requests -/// to guidance, ignoring any peer-supplied `_backend` value. +/// validation. This client selects the backend per constraint: structural +/// tags require xgrammar (the triggered-tags format is not understood by +/// guidance's legacy structures/triggers parser); everything else lowers to +/// guidance. Peer-supplied `_backend` values are ignored. #[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, Deserialize)] #[serde(rename_all = "kebab-case")] pub enum StructuredOutputBackend { @@ -96,10 +98,17 @@ impl StructuredOutputsParams { } fn from_constraint(constraint: StructuredOutputConstraint) -> Self { + // Structural tags use the triggered-tags format that only xgrammar + // compiles; guidance's parser expects the legacy structures/triggers + // shape and fails the request at grammar build. + let backend = match &constraint { + StructuredOutputConstraint::StructuralTag(_) => StructuredOutputBackend::Xgrammar, + _ => StructuredOutputBackend::default(), + }; Self { constraint, options: StructuredOutputOptions::default(), - backend: StructuredOutputBackend::default(), + backend, } } } @@ -263,6 +272,18 @@ mod tests { assert_eq!(value["_backend"], "guidance"); } + #[test] + fn structural_tag_selects_xgrammar_backend() { + // The triggered-tags format only compiles under xgrammar; guidance's + // legacy parser fails the request at grammar build. + let params = StructuredOutputsParams::structural_tag(r#"{"format":{}}"#); + assert_eq!(params.backend, StructuredOutputBackend::Xgrammar); + + let value = serde_json::to_value(params).unwrap(); + assert_eq!(value["_backend"], "xgrammar"); + assert_eq!(value["structural_tag"], r#"{"format":{}}"#); + } + #[test] fn structured_outputs_rejects_missing_constraint() { let error = diff --git a/crates/grpc_client/src/tokenspeed_scheduler.rs b/crates/grpc_client/src/tokenspeed_scheduler.rs index bd428dd6b..9a543f947 100644 --- a/crates/grpc_client/src/tokenspeed_scheduler.rs +++ b/crates/grpc_client/src/tokenspeed_scheduler.rs @@ -178,12 +178,7 @@ impl TokenSpeedSchedulerClient { // ── Request builders ────────────────────────────────────────────── - #[expect( - clippy::unused_self, - reason = "receiver kept for API parity with the other engine clients" - )] pub fn build_generate_request_from_chat( - &self, request_id: String, body: &ChatCompletionRequest, processed_text: String, @@ -208,12 +203,7 @@ impl TokenSpeedSchedulerClient { }) } - #[expect( - clippy::unused_self, - reason = "receiver kept for API parity with the other engine clients" - )] pub fn build_plain_generate_request( - &self, request_id: String, body: &GenerateRequest, original_text: Option, @@ -242,12 +232,7 @@ impl TokenSpeedSchedulerClient { }) } - #[expect( - clippy::unused_self, - reason = "receiver kept for API parity with the other engine clients" - )] pub fn build_generate_request_from_responses( - &self, request_id: String, body: &ResponsesRequest, processed_text: String, @@ -267,12 +252,7 @@ impl TokenSpeedSchedulerClient { }) } - #[expect( - clippy::unused_self, - reason = "receiver kept for API parity with the other engine clients" - )] pub fn build_generate_request_from_messages( - &self, request_id: String, body: &CreateMessageRequest, processed_text: String, @@ -295,12 +275,7 @@ impl TokenSpeedSchedulerClient { }) } - #[expect( - clippy::unused_self, - reason = "receiver kept for API parity with the other engine clients" - )] pub fn build_generate_request_from_completion( - &self, request_id: String, body: &CompletionRequest, original_text: String, diff --git a/model_gateway/src/routers/grpc/backend_client.rs b/model_gateway/src/routers/grpc/backend_client.rs index 0492d762c..226556ae9 100644 --- a/model_gateway/src/routers/grpc/backend_client.rs +++ b/model_gateway/src/routers/grpc/backend_client.rs @@ -13,17 +13,23 @@ use openai_protocol::{ messages::CreateMessageRequest, worker::WorkerLoadResponse, }; use smg_grpc_client::{ - common_proto, tokenizer_bundle::StreamBundle, SglangSchedulerClient, VllmEngineClient, + common_proto, tokenizer_bundle::StreamBundle, SglangSchedulerClient, TokenSpeedSchedulerClient, + VllmEngineClient, }; -use crate::routers::grpc::{ - client::{GenerateRequestBuildOptions, GrpcClient, HealthCheckResponse, ModelInfo, ServerInfo}, - proto_wrapper::{ - finish_vllm_request, ProtoEmbedComplete, ProtoEmbedRequest, ProtoGenerateRequest, - ProtoStream, +use crate::{ + routers::grpc::{ + client::{ + GenerateRequestBuildOptions, GrpcClient, HealthCheckResponse, ModelInfo, ServerInfo, + }, + proto_wrapper::{ + finish_tokenspeed_request, finish_vllm_request, ProtoEmbedComplete, ProtoEmbedRequest, + ProtoGenerateRequest, ProtoStream, + }, + zmq_client::ZmqEngineClient, + MultimodalData, }, - zmq_client::ZmqEngineClient, - MultimodalData, + worker::RuntimeType, }; /// A backend connection: gRPC (any engine) or direct ZMQ (vLLM EngineCore or @@ -36,21 +42,13 @@ pub enum BackendClient { impl BackendClient { /// Runtime type backing this client. - pub fn runtime_type(&self) -> crate::worker::RuntimeType { + pub fn runtime_type(&self) -> RuntimeType { match self { Self::Grpc(client) => client.runtime_type(), Self::Zmq(client) => client.runtime(), } } - /// True if this backend speaks the vLLM protocol (gRPC-vLLM or ZMQ). - pub fn is_vllm(&self) -> bool { - match self { - Self::Grpc(client) => client.is_vllm(), - Self::Zmq(_) => true, - } - } - /// True if this is a direct-ZMQ backend (the engine receives token ids only /// and cannot match string stops itself). pub fn is_zmq(&self) -> bool { @@ -95,14 +93,14 @@ impl BackendClient { pub async fn get_model_info(&self) -> Result { match self { Self::Grpc(client) => client.get_model_info().await, - Self::Zmq(client) => Ok(ModelInfo::Vllm(client.get_model_info())), + Self::Zmq(client) => Ok(client.get_model_info()), } } pub async fn get_server_info(&self) -> Result { match self { Self::Grpc(client) => client.get_server_info().await, - Self::Zmq(client) => Ok(ServerInfo::Vllm(client.get_server_info())), + Self::Zmq(client) => Ok(client.get_server_info()), } } @@ -175,14 +173,7 @@ impl BackendClient { ) -> Result { match self { Self::Grpc(client) => client.generate(req).await, - Self::Zmq(client) => match req { - ProtoGenerateRequest::Vllm(boxed_req) => { - Ok(ProtoStream::Zmq(client.generate(*boxed_req).await?)) - } - _ => Err(tonic::Status::internal( - "ZMQ backend expects a vLLM generate request", - )), - }, + Self::Zmq(client) => Ok(ProtoStream::Zmq(client.generate(req).await?)), } } @@ -198,6 +189,10 @@ impl BackendClient { } } + #[expect( + clippy::unreachable, + reason = "assembly stage guarantees matching MultimodalData variant for each backend" + )] pub fn build_chat_request( &self, request_id: String, @@ -210,25 +205,50 @@ impl BackendClient { Self::Grpc(client) => { client.build_chat_request(request_id, body, processed_text, token_ids, options) } - Self::Zmq(_) => { - let vllm_mm = options.multimodal_inputs.map(|mm| match mm { - MultimodalData::Vllm(data) => data.into_proto(), - _ => unreachable!("caller guarantees matching variant"), - }); - finish_vllm_request(vllm_mm, |mm| { - VllmEngineClient::build_generate_request_from_chat( - request_id, - body, - processed_text, - token_ids, - mm, - options.tool_constraints, - ) - }) - } + // A ZMQ backend speaks vLLM EngineCore or TokenSpeed directly; build + // the native request for its runtime, mirroring the gRPC per-engine + // dispatch in `GrpcClient::build_chat_request`. + Self::Zmq(client) => match client.runtime() { + RuntimeType::TokenSpeed => { + let tokenspeed_mm = options.multimodal_inputs.map(|mm| match mm { + MultimodalData::TokenSpeed(data) => data.into_proto(true), + _ => unreachable!("caller guarantees matching variant"), + }); + finish_tokenspeed_request(tokenspeed_mm, |mm| { + TokenSpeedSchedulerClient::build_generate_request_from_chat( + request_id, + body, + processed_text, + token_ids, + mm, + options.tool_constraints, + ) + }) + } + _ => { + let vllm_mm = options.multimodal_inputs.map(|mm| match mm { + MultimodalData::Vllm(data) => data.into_proto(), + _ => unreachable!("caller guarantees matching variant"), + }); + finish_vllm_request(vllm_mm, |mm| { + VllmEngineClient::build_generate_request_from_chat( + request_id, + body, + processed_text, + token_ids, + mm, + options.tool_constraints, + ) + }) + } + }, } } + #[expect( + clippy::unreachable, + reason = "assembly stage guarantees matching MultimodalData variant for each backend" + )] pub fn build_messages_request( &self, request_id: String, @@ -241,22 +261,42 @@ impl BackendClient { Self::Grpc(client) => { client.build_messages_request(request_id, body, processed_text, token_ids, options) } - Self::Zmq(_) => { - let vllm_mm = options.multimodal_inputs.map(|mm| match mm { - MultimodalData::Vllm(data) => data.into_proto(), - _ => unreachable!("caller guarantees matching variant"), - }); - finish_vllm_request(vllm_mm, |mm| { - VllmEngineClient::build_generate_request_from_messages( - request_id, - body, - processed_text, - token_ids, - mm, - options.tool_constraints, - ) - }) - } + // Mirrors the gRPC per-engine dispatch: build the request natively for + // the ZMQ backend's runtime (vLLM EngineCore or TokenSpeed). + Self::Zmq(client) => match client.runtime() { + RuntimeType::TokenSpeed => { + let tokenspeed_mm = options.multimodal_inputs.map(|mm| match mm { + MultimodalData::TokenSpeed(data) => data.into_proto(true), + _ => unreachable!("caller guarantees matching variant"), + }); + finish_tokenspeed_request(tokenspeed_mm, |mm| { + TokenSpeedSchedulerClient::build_generate_request_from_messages( + request_id, + body, + processed_text, + token_ids, + mm, + options.tool_constraints, + ) + }) + } + _ => { + let vllm_mm = options.multimodal_inputs.map(|mm| match mm { + MultimodalData::Vllm(data) => data.into_proto(), + _ => unreachable!("caller guarantees matching variant"), + }); + finish_vllm_request(vllm_mm, |mm| { + VllmEngineClient::build_generate_request_from_messages( + request_id, + body, + processed_text, + token_ids, + mm, + options.tool_constraints, + ) + }) + } + }, } } @@ -271,15 +311,26 @@ impl BackendClient { Self::Grpc(client) => { client.build_completion_request(request_id, body, original_text, token_ids) } - Self::Zmq(_) => { - let req = VllmEngineClient::build_generate_request_from_completion( - request_id, - body, - original_text, - token_ids, - )?; - Ok(ProtoGenerateRequest::Vllm(Box::new(req))) - } + Self::Zmq(client) => match client.runtime() { + RuntimeType::TokenSpeed => { + let req = TokenSpeedSchedulerClient::build_generate_request_from_completion( + request_id, + body, + original_text, + token_ids, + )?; + Ok(ProtoGenerateRequest::TokenSpeed(Box::new(req))) + } + _ => { + let req = VllmEngineClient::build_generate_request_from_completion( + request_id, + body, + original_text, + token_ids, + )?; + Ok(ProtoGenerateRequest::Vllm(Box::new(req))) + } + }, } } @@ -294,15 +345,26 @@ impl BackendClient { Self::Grpc(client) => { client.build_generate_request(request_id, body, original_text, token_ids) } - Self::Zmq(_) => { - let req = VllmEngineClient::build_plain_generate_request( - request_id, - body, - original_text, - token_ids, - )?; - Ok(ProtoGenerateRequest::Vllm(Box::new(req))) - } + Self::Zmq(client) => match client.runtime() { + RuntimeType::TokenSpeed => { + let req = TokenSpeedSchedulerClient::build_plain_generate_request( + request_id, + body, + original_text, + token_ids, + )?; + Ok(ProtoGenerateRequest::TokenSpeed(Box::new(req))) + } + _ => { + let req = VllmEngineClient::build_plain_generate_request( + request_id, + body, + original_text, + token_ids, + )?; + Ok(ProtoGenerateRequest::Vllm(Box::new(req))) + } + }, } } } diff --git a/model_gateway/src/routers/grpc/client.rs b/model_gateway/src/routers/grpc/client.rs index b420173ab..2bd1221b5 100644 --- a/model_gateway/src/routers/grpc/client.rs +++ b/model_gateway/src/routers/grpc/client.rs @@ -541,13 +541,13 @@ impl GrpcClient { )?; Ok(ProtoGenerateRequest::Mlx(Box::new(req))) } - Self::TokenSpeed(client) => { + Self::TokenSpeed(_) => { let tokenspeed_mm = options.multimodal_inputs.map(|mm| match mm { MultimodalData::TokenSpeed(data) => data.into_proto(true), _ => unreachable!("caller guarantees matching variant"), }); finish_tokenspeed_request(tokenspeed_mm, |mm| { - client.build_generate_request_from_chat( + TokenSpeedSchedulerClient::build_generate_request_from_chat( request_id, body, processed_text, @@ -633,13 +633,13 @@ impl GrpcClient { )?; Ok(ProtoGenerateRequest::Mlx(Box::new(req))) } - Self::TokenSpeed(client) => { + Self::TokenSpeed(_) => { let tokenspeed_mm = options.multimodal_inputs.map(|mm| match mm { MultimodalData::TokenSpeed(data) => data.into_proto(true), _ => unreachable!("caller guarantees matching variant"), }); finish_tokenspeed_request(tokenspeed_mm, |mm| { - client.build_generate_request_from_messages( + TokenSpeedSchedulerClient::build_generate_request_from_messages( request_id, body, processed_text, @@ -696,8 +696,8 @@ impl GrpcClient { )?; Ok(ProtoGenerateRequest::Mlx(Box::new(req))) } - Self::TokenSpeed(client) => { - let req = client.build_generate_request_from_completion( + Self::TokenSpeed(_) => { + let req = TokenSpeedSchedulerClient::build_generate_request_from_completion( request_id, body, original_text, @@ -752,8 +752,8 @@ impl GrpcClient { )?; Ok(ProtoGenerateRequest::Mlx(Box::new(req))) } - Self::TokenSpeed(client) => { - let req = client.build_plain_generate_request( + Self::TokenSpeed(_) => { + let req = TokenSpeedSchedulerClient::build_plain_generate_request( request_id, body, original_text, diff --git a/model_gateway/src/routers/grpc/harmony/mod.rs b/model_gateway/src/routers/grpc/harmony/mod.rs index 59896abfe..dd0c3641c 100644 --- a/model_gateway/src/routers/grpc/harmony/mod.rs +++ b/model_gateway/src/routers/grpc/harmony/mod.rs @@ -35,6 +35,7 @@ pub(crate) mod parser; pub(crate) mod processor; pub(crate) mod responses; pub(crate) mod stages; +pub(crate) mod stop; pub(crate) mod streaming; pub(crate) mod types; diff --git a/model_gateway/src/routers/grpc/harmony/processor.rs b/model_gateway/src/routers/grpc/harmony/processor.rs index 6a74a1bec..df975093a 100644 --- a/model_gateway/src/routers/grpc/harmony/processor.rs +++ b/model_gateway/src/routers/grpc/harmony/processor.rs @@ -16,6 +16,7 @@ use tracing::error; use super::{ builder::{convert_harmony_logprobs, try_harmony_encoding}, + stop::TextStopScanner, HarmonyParserAdapter, }; use crate::routers::{ @@ -39,11 +40,16 @@ impl HarmonyResponseProcessor { } /// Process a non-streaming Harmony chat response + /// + /// `router_stop_strings` is non-empty only when the router must enforce + /// string `stop` sequences itself (direct-ZMQ backends: the engine sees + /// token ids only). pub async fn process_non_streaming_chat_response( &self, execution_result: ExecutionResult, chat_request: Arc, dispatch: DispatchMetadata, + router_stop_strings: &[String], ) -> Result { let request_logprobs = chat_request.logprobs; @@ -103,16 +109,47 @@ impl HarmonyResponseProcessor { None }; + let mut analysis = parsed.analysis; + let mut final_text = parsed.final_text; + let mut tool_calls = parsed.commentary; + let mut finish_reason = parsed.finish_reason; + let mut matched_stop = matched_stop; + + // Router-enforced string stops: scan channel text in emission + // order (analysis before final), truncate after the first match + // (retaining the stop as suffix, like the token-forwarding gRPC + // path), and drop everything generated past it. + if !router_stop_strings.is_empty() { + let mut scanner = TextStopScanner::new(router_stop_strings.to_vec()); + if let Some(text) = analysis.take() { + let (out, stopped) = scanner.scan_complete(&text); + analysis = (!out.is_empty()).then_some(out); + if stopped { + final_text.clear(); + tool_calls = None; + } + } + if scanner.matched().is_none() && !final_text.is_empty() { + let (out, stopped) = scanner.scan_complete(&final_text); + final_text = out; + if stopped { + tool_calls = None; + } + } + if let Some(stop) = scanner.matched() { + finish_reason = "stop".to_string(); + matched_stop = Some(serde_json::Value::String(stop.to_string())); + } + } + // Build response message (assistant) let message = ChatCompletionMessage { role: "assistant".to_string(), - content: (!parsed.final_text.is_empty()).then_some(parsed.final_text), - tool_calls: parsed.commentary, - reasoning_content: parsed.analysis, + content: (!final_text.is_empty()).then_some(final_text), + tool_calls, + reasoning_content: analysis, }; - let finish_reason = parsed.finish_reason; - // Accumulate reasoning tokens across all responses total_reasoning_tokens += parsed.reasoning_token_count; diff --git a/model_gateway/src/routers/grpc/harmony/stages/request_building.rs b/model_gateway/src/routers/grpc/harmony/stages/request_building.rs index dda6e4426..568afb603 100644 --- a/model_gateway/src/routers/grpc/harmony/stages/request_building.rs +++ b/model_gateway/src/routers/grpc/harmony/stages/request_building.rs @@ -2,21 +2,24 @@ use async_trait::async_trait; use axum::response::Response; -use smg_grpc_client::{SglangGenerateRequestOptions, VllmEngineClient}; +use smg_grpc_client::{SglangGenerateRequestOptions, TokenSpeedSchedulerClient, VllmEngineClient}; use tracing::{debug, error}; -use crate::routers::{ - error, - grpc::{ - backend_client::BackendClient, - client::GrpcClient, - common::stages::{helpers, PipelineStage}, - context::{ - ClientSelection, ExecutionPlan, ExecutionPlanKind, PreparationOutput, RequestContext, - RequestType, +use crate::{ + routers::{ + error, + grpc::{ + backend_client::BackendClient, + client::GrpcClient, + common::stages::{helpers, PipelineStage}, + context::{ + ClientSelection, ExecutionPlan, ExecutionPlanKind, PreparationOutput, + RequestContext, RequestType, + }, + proto_wrapper::ProtoGenerateRequest, }, - proto_wrapper::ProtoGenerateRequest, }, + worker::RuntimeType, }; /// Harmony Request Building stage: Convert Harmony tokens to gRPC request @@ -299,12 +302,11 @@ impl PipelineStage for HarmonyRequestBuildingStage { }; ProtoGenerateRequest::Mlx(Box::new(req)) } - BackendClient::Grpc(GrpcClient::TokenSpeed(tokenspeed_client)) => { + BackendClient::Grpc(GrpcClient::TokenSpeed(_)) => { let req = match &ctx.input.request_type { RequestType::Chat(request) => { let body = modified_request.as_deref().unwrap_or_else(|| request.as_ref()); - tokenspeed_client - .build_generate_request_from_chat( + TokenSpeedSchedulerClient::build_generate_request_from_chat( request_id, body, placeholder_processed_text, @@ -317,8 +319,7 @@ impl PipelineStage for HarmonyRequestBuildingStage { error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) })? } - RequestType::Responses(request) => tokenspeed_client - .build_generate_request_from_responses( + RequestType::Responses(request) => TokenSpeedSchedulerClient::build_generate_request_from_responses( request_id, request.as_ref(), placeholder_processed_text, @@ -344,9 +345,58 @@ impl PipelineStage for HarmonyRequestBuildingStage { }; ProtoGenerateRequest::TokenSpeed(Box::new(req)) } + // A ZMQ worker speaks vLLM EngineCore or TokenSpeed directly; build the + // request natively for its runtime, mirroring the gRPC per-engine + // dispatch above. Both support the Harmony request types (Chat + + // Responses). + BackendClient::Zmq(zmq_client) if zmq_client.runtime() == RuntimeType::TokenSpeed => { + let req = match &ctx.input.request_type { + RequestType::Chat(request) => { + let body = modified_request + .as_deref() + .unwrap_or_else(|| request.as_ref()); + TokenSpeedSchedulerClient::build_generate_request_from_chat( + request_id, + body, + placeholder_processed_text, + token_ids, + None, // Harmony path: multimodal not yet wired + tool_constraints, + ) + .map_err(|e| { + error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build TokenSpeed ZMQ generate request"); + error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) + })? + } + RequestType::Responses(request) => { + TokenSpeedSchedulerClient::build_generate_request_from_responses( + request_id, + request.as_ref(), + placeholder_processed_text, + token_ids, + tool_constraints, + ) + .map_err(|e| { + error!(function = "HarmonyRequestBuildingStage::execute", error = %e, "Failed to build TokenSpeed ZMQ generate request from responses"); + error::bad_request("invalid_request_parameters", format!("Invalid request parameters: {e}")) + })? + } + RequestType::Embedding(_) => { + return Err(error::bad_request( + "harmony_embedding_not_supported", + "Embedding requests are not supported with Harmony models".to_string(), + )); + } + _ => { + return Err(error::bad_request( + "unsupported_request_type", + "Unsupported request type for Harmony models".to_string(), + )); + } + }; + ProtoGenerateRequest::TokenSpeed(Box::new(req)) + } BackendClient::Zmq(_) => { - // A ZMQ worker is a vLLM engine, so it uses the same vLLM request - // builders and supports the same request types (Chat + Responses). let req = match &ctx.input.request_type { RequestType::Chat(request) => { let body = modified_request diff --git a/model_gateway/src/routers/grpc/harmony/stages/response_processing.rs b/model_gateway/src/routers/grpc/harmony/stages/response_processing.rs index f745ca589..488c16b4d 100644 --- a/model_gateway/src/routers/grpc/harmony/stages/response_processing.rs +++ b/model_gateway/src/routers/grpc/harmony/stages/response_processing.rs @@ -12,12 +12,38 @@ use crate::{ error, grpc::{ common::stages::PipelineStage, - context::{FinalResponse, RequestContext, RequestType}, + context::{ClientSelection, FinalResponse, RequestContext, RequestType}, }, }, worker::AttachedBody, }; +/// String `stop` sequences the ROUTER must enforce: only for direct-ZMQ +/// backends, where the engine receives token ids and never sees the strings. +/// Empty for gRPC backends (the engine matches stops itself). +fn router_stop_strings(ctx: &RequestContext) -> Vec { + let is_zmq = ctx + .state + .clients + .as_ref() + .is_some_and(|clients| match clients { + ClientSelection::Single { client } => client.is_zmq(), + ClientSelection::Disaggregated { decode, .. } => decode.is_zmq(), + }); + if !is_zmq { + return Vec::new(); + } + match &ctx.input.request_type { + RequestType::Chat(_) => ctx + .chat_request_arc() + .stop + .as_ref() + .map(|stop| stop.to_vec()) + .unwrap_or_default(), + _ => Vec::new(), + } +} + /// Harmony Response Processing stage: Parse and format Harmony responses /// /// Takes output tokens from execution and parses them using HarmonyParserAdapter @@ -80,6 +106,7 @@ impl PipelineStage for HarmonyResponseProcessingStage { execution_result, ctx.chat_request_arc(), dispatch, + router_stop_strings(ctx), ); // Attach load guards to response body for proper RAII lifecycle @@ -93,9 +120,15 @@ impl PipelineStage for HarmonyResponseProcessingStage { // For non-streaming, delegate to Harmony response processor to build ChatCompletionResponse let chat_request = ctx.chat_request_arc(); + let stops = router_stop_strings(ctx); let response = self .processor - .process_non_streaming_chat_response(execution_result, chat_request, dispatch) + .process_non_streaming_chat_response( + execution_result, + chat_request, + dispatch, + &stops, + ) .await?; ctx.state.response.final_response = Some(FinalResponse::Chat(response)); diff --git a/model_gateway/src/routers/grpc/harmony/stop.rs b/model_gateway/src/routers/grpc/harmony/stop.rs new file mode 100644 index 000000000..f6f49a72d --- /dev/null +++ b/model_gateway/src/routers/grpc/harmony/stop.rs @@ -0,0 +1,163 @@ +//! Incremental stop-string scanning for harmony channel text. +//! +//! The direct-ZMQ engine receives token ids only and cannot match string +//! `stop` sequences; the regular pipeline covers this with the router-side +//! `StopSequenceDecoder`, but harmony emits channel-parsed text rather than +//! raw tokens, so stops are matched here on the decoded text instead. The +//! matched stop stays in the output as its suffix, mirroring what the +//! token-forwarding gRPC path yields for harmony models. + +/// Result of feeding text through the scanner. +pub(crate) struct StopScan { + /// Text safe to emit: everything up to and including a match, or the + /// prefix that cannot be part of a future match. + pub emit: String, + /// A stop sequence completed inside this push. + pub stopped: bool, +} + +pub(crate) struct TextStopScanner { + stops: Vec, + /// Tail held back because it could still grow into a match. + pending: String, + matched: Option, +} + +impl TextStopScanner { + pub fn new(stops: Vec) -> Self { + Self { + stops, + pending: String::new(), + matched: None, + } + } + + pub fn matched(&self) -> Option<&str> { + self.matched.as_deref() + } + + /// Feed a text delta. Once a stop has matched, further pushes emit nothing. + pub fn push(&mut self, text: &str) -> StopScan { + if self.matched.is_some() { + return StopScan { + emit: String::new(), + stopped: true, + }; + } + self.pending.push_str(text); + + // Earliest match across all stops wins; ties prefer the longer stop. + let mut best: Option<(usize, &str)> = None; + for stop in &self.stops { + if let Some(at) = self.pending.find(stop.as_str()) { + let better = match best { + None => true, + Some((best_at, best_stop)) => { + at < best_at || (at == best_at && stop.len() > best_stop.len()) + } + }; + if better { + best = Some((at, stop)); + } + } + } + if let Some((at, stop)) = best { + let emit = self.pending[..at + stop.len()].to_string(); + self.matched = Some(stop.to_string()); + self.pending.clear(); + return StopScan { + emit, + stopped: true, + }; + } + + // Hold back the longest suffix that is a prefix of some stop; emit the + // rest (it can never become part of a match). + let max_hold = (self.stops.iter().map(String::len).max()) + .unwrap_or(1) + .saturating_sub(1); + let mut hold = 0; + for len in (1..=max_hold.min(self.pending.len())).rev() { + let Some(start) = self.pending.len().checked_sub(len) else { + continue; + }; + if !self.pending.is_char_boundary(start) { + continue; + } + let tail = &self.pending[start..]; + if self.stops.iter().any(|s| s.starts_with(tail)) { + hold = len; + break; + } + } + let emit = self.pending[..self.pending.len() - hold].to_string(); + self.pending.drain(..self.pending.len() - hold); + StopScan { + emit, + stopped: false, + } + } + + /// Emit whatever is still held back (end of stream, no match). + pub fn flush(&mut self) -> String { + std::mem::take(&mut self.pending) + } + + /// Scan a complete text: returns the (possibly truncated) text and whether + /// a stop matched inside it. + pub fn scan_complete(&mut self, text: &str) -> (String, bool) { + let scan = self.push(text); + let mut out = scan.emit; + if !scan.stopped { + out.push_str(&self.flush()); + } + (out, scan.stopped) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn complete_scan_truncates_after_first_match() { + let mut scanner = TextStopScanner::new(vec![",".to_string()]); + let (out, stopped) = scanner.scan_complete("1, 2, 3"); + assert_eq!(out, "1,"); + assert!(stopped); + assert_eq!(scanner.matched(), Some(",")); + } + + #[test] + fn streaming_holds_back_partial_matches() { + let mut scanner = TextStopScanner::new(vec!["STOP".to_string()]); + let scan = scanner.push("abc ST"); + assert_eq!(scan.emit, "abc "); + assert!(!scan.stopped); + let scan = scanner.push("OP tail"); + assert_eq!(scan.emit, "STOP"); + assert!(scan.stopped); + // Post-match pushes are swallowed. + let scan = scanner.push("more"); + assert_eq!(scan.emit, ""); + assert!(scan.stopped); + } + + #[test] + fn flush_releases_unmatched_tail() { + let mut scanner = TextStopScanner::new(vec!["END".to_string()]); + let scan = scanner.push("value EN"); + assert_eq!(scan.emit, "value "); + assert_eq!(scanner.flush(), "EN"); + assert_eq!(scanner.matched(), None); + } + + #[test] + fn earliest_match_wins_across_stops() { + let mut scanner = TextStopScanner::new(vec!["zz".to_string(), "b".to_string()]); + let (out, stopped) = scanner.scan_complete("abzz"); + assert_eq!(out, "ab"); + assert!(stopped); + assert_eq!(scanner.matched(), Some("b")); + } +} diff --git a/model_gateway/src/routers/grpc/harmony/streaming.rs b/model_gateway/src/routers/grpc/harmony/streaming.rs index cef6b643a..66d231452 100644 --- a/model_gateway/src/routers/grpc/harmony/streaming.rs +++ b/model_gateway/src/routers/grpc/harmony/streaming.rs @@ -27,6 +27,7 @@ use tracing::{debug, error}; use super::{ builder::{convert_harmony_logprobs, try_harmony_encoding}, processor::ResponsesIterationResult, + stop::TextStopScanner, types::HarmonyChannelDelta, HarmonyParserAdapter, }; @@ -100,11 +101,15 @@ impl HarmonyStreamingProcessor { clippy::disallowed_methods, reason = "streaming tasks are fire-and-forget by design; client disconnect terminates them" )] + /// `router_stop_strings` is non-empty only when the router must enforce + /// string `stop` sequences itself (direct-ZMQ backends: the engine sees + /// token ids only). pub fn process_streaming_chat_response( self: Arc, execution_result: context::ExecutionResult, chat_request: Arc, dispatch: context::DispatchMetadata, + router_stop_strings: Vec, ) -> Response { // Create SSE channel let (tx, rx) = mpsc::unbounded_channel::>(); @@ -113,8 +118,14 @@ impl HarmonyStreamingProcessor { match execution_result { context::ExecutionResult::Single { stream } => { tokio::spawn(async move { - let result = - Self::process_single_stream(stream, dispatch, chat_request, &tx).await; + let result = Self::process_single_stream( + stream, + dispatch, + chat_request, + &tx, + router_stop_strings, + ) + .await; if let Err(e) = result { error!("Harmony streaming error: {}", e); @@ -137,6 +148,7 @@ impl HarmonyStreamingProcessor { dispatch, chat_request, &tx, + router_stop_strings, ) .await; @@ -179,6 +191,7 @@ impl HarmonyStreamingProcessor { dispatch: context::DispatchMetadata, original_request: Arc, tx: &mpsc::UnboundedSender>, + router_stop_strings: Vec, ) -> Result<(), String> { let mut prompt_tokens = HashMap::new(); let mut cached_tokens = HashMap::new(); @@ -189,6 +202,7 @@ impl HarmonyStreamingProcessor { tx, &mut prompt_tokens, &mut cached_tokens, + &router_stop_strings, ) .await } @@ -200,6 +214,7 @@ impl HarmonyStreamingProcessor { dispatch: context::DispatchMetadata, original_request: Arc, tx: &mpsc::UnboundedSender>, + router_stop_strings: Vec, ) -> Result<(), String> { // Phase 1: Process prefill stream (collect metadata) let mut prompt_tokens: HashMap = HashMap::new(); @@ -222,6 +237,7 @@ impl HarmonyStreamingProcessor { tx, &mut prompt_tokens, &mut cached_tokens, + &router_stop_strings, ) .await?; @@ -244,6 +260,7 @@ impl HarmonyStreamingProcessor { tx: &mpsc::UnboundedSender>, prompt_tokens: &mut HashMap, cached_tokens: &mut HashMap, + router_stop_strings: &[String], ) -> Result<(), String> { // Timing for metrics let start_time = Instant::now(); @@ -253,6 +270,12 @@ impl HarmonyStreamingProcessor { let mut parsers: HashMap = HashMap::new(); let mut is_firsts: HashMap = HashMap::new(); let mut matched_stops: HashMap> = HashMap::new(); + // Router-enforced string stops (direct-ZMQ): per-index, per-channel + // scanners. Once an index stops, its further deltas are swallowed and + // the engine's own Complete is not re-emitted. + let mut analysis_scanners: HashMap = HashMap::new(); + let mut final_scanners: HashMap = HashMap::new(); + let mut router_stopped: std::collections::HashSet = std::collections::HashSet::new(); let mut completion_tokens = CompletionTokenTracker::new(); // Reusable SSE encoder shared across every chunk emitted for this stream. let mut encoder = SseEncoder::new(); @@ -303,21 +326,76 @@ impl HarmonyStreamingProcessor { .map_err(|e| format!("Parse error: {e}"))?; // Emit SSE event if there's a delta - if let Some(delta) = delta_result { + if let Some(mut delta) = delta_result { + if router_stopped.contains(&index) { + continue; + } + let mut stop_matched: Option = None; + if !router_stop_strings.is_empty() { + if let Some(text) = delta.analysis_delta.take() { + let scanner = analysis_scanners.entry(index).or_insert_with(|| { + TextStopScanner::new(router_stop_strings.to_vec()) + }); + let scan = scanner.push(&text); + delta.analysis_delta = (!scan.emit.is_empty()).then_some(scan.emit); + if scan.stopped { + stop_matched = scanner.matched().map(str::to_string); + delta.final_delta = None; + delta.commentary_delta = None; + } + } + if stop_matched.is_none() { + if let Some(text) = delta.final_delta.take() { + let scanner = + final_scanners.entry(index).or_insert_with(|| { + TextStopScanner::new(router_stop_strings.to_vec()) + }); + let scan = scanner.push(&text); + delta.final_delta = + (!scan.emit.is_empty()).then_some(scan.emit); + if scan.stopped { + stop_matched = scanner.matched().map(str::to_string); + delta.commentary_delta = None; + } + } + } + } + + let has_payload = delta.analysis_delta.is_some() + || delta.final_delta.is_some() + || delta.commentary_delta.is_some(); let is_first = is_firsts.get(&index).copied().unwrap_or(false); - Self::emit_chunk_delta( - &delta, - index, - is_first, - dispatch, - original_request, - tx, - &mut encoder, - chunk_logprobs, - )?; + if has_payload || is_first { + Self::emit_chunk_delta( + &delta, + index, + is_first, + dispatch, + original_request, + tx, + &mut encoder, + chunk_logprobs, + )?; + + if is_first { + is_firsts.insert(index, false); + } + } - if is_first { - is_firsts.insert(index, false); + // A router-side stop fired: emit the final chunk now + // and swallow the rest of this index's stream (the + // engine keeps generating until its own limits). + if let Some(stop) = stop_matched { + Self::emit_final_chunk( + index, + "stop", + Some(&serde_json::Value::String(stop)), + dispatch, + original_request, + tx, + &mut encoder, + )?; + router_stopped.insert(index); } } } @@ -341,6 +419,37 @@ impl HarmonyStreamingProcessor { let final_output = parser.finalize(complete_wrapper.finish_reason().to_string()); + // A router-side stop already closed this choice. + if router_stopped.contains(&index) { + continue; + } + + // Release scanner-held text that never became a match. + let flushed = HarmonyChannelDelta { + analysis_delta: analysis_scanners + .get_mut(&index) + .map(TextStopScanner::flush) + .filter(|s| !s.is_empty()), + commentary_delta: None, + final_delta: final_scanners + .get_mut(&index) + .map(TextStopScanner::flush) + .filter(|s| !s.is_empty()), + is_final: false, + }; + if flushed.analysis_delta.is_some() || flushed.final_delta.is_some() { + Self::emit_chunk_delta( + &flushed, + index, + false, + dispatch, + original_request, + tx, + &mut encoder, + None, + )?; + } + Self::emit_final_chunk( index, &final_output.finish_reason, diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index 9a2c4b433..269c36d00 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -40,9 +40,16 @@ use engine_zmq_client::{ }; use futures::{stream::SelectAll, Stream, StreamExt}; use openai_protocol::worker::{SchedulerLoadSnapshot, WorkerLoadResponse}; -use smg_grpc_client::vllm_proto as vllm; +use smg_grpc_client::{tokenspeed_proto, vllm_proto as vllm}; -use crate::{routers::grpc::zmq_multimodal, worker::RuntimeType}; +use crate::{ + routers::grpc::{ + client::{ModelInfo, ServerInfo}, + proto_wrapper::ProtoGenerateRequest, + zmq_multimodal, + }, + worker::RuntimeType, +}; /// Loopback host for the same-host ZMQ transport (TCP handshake and local /// binds). Shared with the worker-side socket derivation. @@ -82,7 +89,7 @@ impl EosTokenIds { pub fn from_model_dir(dir: &Path) -> Self { let model_ids = eos_ids_from_file(&dir.join("config.json")); let gen_ids = eos_ids_from_file(&dir.join("generation_config.json")); - let primary = model_ids.first().copied().or(gen_ids.first().copied()); + let primary = (model_ids.first().or_else(|| gen_ids.first())).copied(); let mut extra = Vec::new(); for id in model_ids.into_iter().chain(gen_ids) { if Some(id) != primary && !extra.contains(&id) { @@ -132,6 +139,10 @@ impl ZmqEngineClient { /// engines connect to (chosen by SMG). `engine_count` is the number of DP /// ranks to await. `runtime` selects the wire protocol spoken over the shared /// transport (vLLM EngineCore vs TokenSpeed). + #[expect( + clippy::too_many_arguments, + reason = "transport constructor: endpoints, engine count, and runtime are all irreducible connection inputs" + )] pub async fn connect( handshake_address: &str, input_address: &str, @@ -211,7 +222,10 @@ impl ZmqEngineClient { } /// Submit a generate request and return a stream of vLLM-proto responses. - /// The request is translated into the backend's wire protocol. + /// The request is the engine's native proto (vLLM for a vLLM backend, + /// TokenSpeed for a TokenSpeed backend — the [`BackendClient`] builders emit + /// the matching variant per runtime); it is translated into the backend's + /// wire protocol here. /// /// Over gRPC the engine-side frontend (e.g. vLLM's AsyncLLM) fans `n` out /// itself and multiplexes the choices onto one stream. The raw ZMQ wire has @@ -219,15 +233,21 @@ impl ZmqEngineClient { /// single-sample engine requests (see [`fan_out_requests`]); their outputs /// are merged back into one stream with each sub tagged via the proto /// `index` field, exactly like the gRPC contract. + /// + /// [`BackendClient`]: crate::routers::grpc::backend_client::BackendClient pub async fn generate( &self, - req: vllm::GenerateRequest, + req: ProtoGenerateRequest, ) -> Result { - let subs = fan_out_requests(req); // Sub-streams submitted before a mid-loop failure are dropped with the // error, which auto-aborts their engine-side requests. match &self.backend { ZmqBackend::Vllm(client) => { + let ProtoGenerateRequest::Vllm(req) = req else { + return Err(tonic::Status::internal( + "vLLM ZMQ backend expects a vLLM generate request", + )); + }; // EngineCore needs a concrete `max_tokens`; vLLM's OpenAI frontend // (which the ZMQ path bypasses) defaults an unset value to // `max_model_len - prompt_len`. The context length comes from the @@ -238,7 +258,7 @@ impl ZmqEngineClient { .map(|e| (e.ready_response.max_model_len, e.ready_response.dtype)) .ok_or_else(|| tonic::Status::unavailable("no connected ZMQ engine"))?; let mut streams = SelectAll::new(); - for (index, sub) in subs.into_iter().enumerate() { + for (index, sub) in fan_out_requests(*req).into_iter().enumerate() { // The engine returns the sampled logprob plus up to // `logprobs` ranked candidates per position; carry the // requested count so the stream can shape `top_logprobs`. @@ -256,8 +276,13 @@ impl ZmqEngineClient { Ok(ZmqGenerateStream::Vllm(streams)) } ZmqBackend::TokenSpeed(client) => { + let ProtoGenerateRequest::TokenSpeed(req) = req else { + return Err(tonic::Status::internal( + "TokenSpeed ZMQ backend expects a TokenSpeed generate request", + )); + }; let mut streams = SelectAll::new(); - for (index, sub) in subs.into_iter().enumerate() { + for (index, sub) in fan_out_tokenspeed_requests(*req).into_iter().enumerate() { let request = translate_request_tokenspeed(sub) .map_err(tonic::Status::invalid_argument)?; let stream = client.submit(request).await.map_err(zmq_status)?; @@ -327,38 +352,55 @@ impl ZmqEngineClient { /// Model info derived from the handshake `EngineCoreReadyResponse` plus the /// configured model id (the engine does not report tokenizer/vocab metadata, - /// so those come from worker config). - pub fn get_model_info(&self) -> vllm::GetModelInfoResponse { + /// so those come from worker config). Returned as the runtime's native + /// metadata variant so the label mapping matches the gRPC path. + pub fn get_model_info(&self) -> ModelInfo { let max_context_length = self .engines() .first() .map(|e| e.ready_response.max_model_len) .unwrap_or(0); - vllm::GetModelInfoResponse { - model_path: self.model_id.clone(), - served_model_name: self.model_id.clone(), - tokenizer_path: self.model_id.clone(), - is_generation: true, - max_context_length: u32::try_from(max_context_length).unwrap_or(u32::MAX), - ..Default::default() + match &self.backend { + ZmqBackend::Vllm(_) => ModelInfo::Vllm(vllm::GetModelInfoResponse { + model_path: self.model_id.clone(), + served_model_name: self.model_id.clone(), + tokenizer_path: self.model_id.clone(), + is_generation: true, + max_context_length: u32::try_from(max_context_length).unwrap_or(u32::MAX), + ..Default::default() + }), + ZmqBackend::TokenSpeed(_) => { + ModelInfo::TokenSpeed(Box::new(tokenspeed_proto::GetModelInfoResponse { + model_path: self.model_id.clone(), + served_model_name: self.model_id.clone(), + tokenizer_path: self.model_id.clone(), + max_context_length: i32::try_from(max_context_length).unwrap_or(i32::MAX), + ..Default::default() + })) + } } } - /// Server info derived from the handshake response. - pub fn get_server_info(&self) -> vllm::GetServerInfoResponse { + /// Server info derived from the handshake response, as the runtime's native + /// metadata variant. + pub fn get_server_info(&self) -> ServerInfo { let data_parallel_size = self .engines() .first() .map(|e| e.ready_response.data_parallel_size) .unwrap_or(1); - let server_type = match &self.backend { - ZmqBackend::Vllm(_) => "vllm", - ZmqBackend::TokenSpeed(_) => "tokenspeed", - }; - vllm::GetServerInfoResponse { - data_parallel_size: i32::try_from(data_parallel_size).unwrap_or(i32::MAX), - server_type: server_type.to_string(), - ..Default::default() + match &self.backend { + ZmqBackend::Vllm(_) => ServerInfo::Vllm(vllm::GetServerInfoResponse { + data_parallel_size: i32::try_from(data_parallel_size).unwrap_or(i32::MAX), + server_type: "vllm".to_string(), + ..Default::default() + }), + // TokenSpeed's server-info proto carries no data-parallel size or + // server-type field; the ZMQ handshake supplies no `server_args` + // either, so only the fields it does expose are surfaced. + ZmqBackend::TokenSpeed(_) => { + ServerInfo::TokenSpeed(Box::::default()) + } } } } @@ -510,6 +552,14 @@ impl VllmGenerateStream { state.output_top_logprobs.extend(tick_top_logprobs); let response = match output.finish_reason { + // An engine-side request failure (e.g. grammar compilation) must + // surface as an error, not as a normal completion with empty + // output — that would produce a 200 with no content. + Some(EngineCoreFinishReason::Error) => { + return Err(tonic::Status::internal( + "engine finished the request with an error (see engine logs)", + )); + } Some(reason) => { let complete = vllm::GenerateResponse { response: Some(vllm::generate_response::Response::Complete( @@ -743,82 +793,63 @@ fn fan_out_requests(req: vllm::GenerateRequest) -> Vec { .collect() } -/// Translate a vLLM-proto generate request into a TokenSpeed +/// Split an `n > 1` TokenSpeed proto request into `n` single-sample +/// sub-requests, the TokenSpeed analogue of [`fan_out_requests`] (the wire has +/// no per-sample demux, so `generate` fans out here). An `n <= 1` request passes +/// through untouched. The TokenSpeed proto carries no seed, so the samples +/// differ by the engine's per-rid seeding alone — the suffixed request ids are +/// unique, so each rid seeds independently. +fn fan_out_tokenspeed_requests( + req: tokenspeed_proto::GenerateRequest, +) -> Vec { + let n = req.sampling_params.as_ref().map_or(1, |sp| sp.n.max(1)); + if n <= 1 { + return vec![req]; + } + (0..n) + .map(|i| { + let mut sub = req.clone(); + sub.request_id = format!("{}-{i}", req.request_id); + if let Some(sp) = sub.sampling_params.as_mut() { + sp.n = 1; + } + sub + }) + .collect() +} + +/// Translate a TokenSpeed proto `GenerateRequest` into the wire /// `TokenizedGenerateReqInput`. ZMQ mode requires pre-tokenized input (SMG /// tokenizes upstream). fn translate_request_tokenspeed( - req: vllm::GenerateRequest, + req: tokenspeed_proto::GenerateRequest, ) -> Result { - // The TokenSpeed wire has no multimodal slot yet; reject loudly rather + // The TokenSpeed ZMQ wire has no multimodal slot yet; reject loudly rather // than silently dropping pixels (assembly also refuses upstream). if req.mm_inputs.is_some() { return Err( "multimodal inputs are not supported over the TokenSpeed ZMQ backend".to_string(), ); } - let input_ids = match req.input { - Some(vllm::generate_request::Input::Tokenized(tokenized)) => tokenized.input_ids, - Some(vllm::generate_request::Input::Text(_)) => { - return Err("ZMQ mode requires pre-tokenized input (TokenizedInput)".to_string()); - } + let input_ids = match req.tokenized { + Some(tokenized) => tokenized.input_ids, None => { return Err("ZMQ mode requires pre-tokenized input; no input provided".to_string()); } }; - let stream = req.stream; - // Single-engine TokenSpeed: a pinned DP rank other than 0 cannot be honored. - if req.data_parallel_rank.is_some_and(|rank| rank != 0) { - return Err(format!( - "invalid data_parallel_rank {:?}: the TokenSpeed ZMQ backend is single-engine", - req.data_parallel_rank - )); + // Over the ZMQ wire TokenSpeed returns only the single sampled-token logprob + // per token: no top-k candidates (`top_logprobs_num > 1`) and no prompt + // logprobs (`token_ids_logprob`). Reject both rather than silently return + // fewer than asked. A bare `logprobs: true` (count 0/1) is the plain + // sampled-token logprob and is wired end-to-end via `return_logprob`. + if req.top_logprobs_num > 1 { + return Err("top_logprobs are not supported over the TokenSpeed ZMQ backend".to_string()); } - // TokenSpeed returns only the single sampled-token logprob per token - // (`top_logprobs_num = 0` on its wire) and no prompt logprobs. The vLLM - // sampling `logprobs` field is a count: the chat frontend maps a bare - // `logprobs: true` to `1` and `top_logprobs = k` to `k`, so a count above 1 - // (or `-1` = "all") is a top-k request that cannot be honored — reject it - // rather than silently return fewer. Counts of 0 or 1 are the plain - // sampled-token logprob and are wired end-to-end. Note the flip side: at - // this proto boundary a chat `top_logprobs: 1` is indistinguishable from a - // bare `logprobs: true` (both arrive as count 1), so it is accepted and its - // `top_logprobs` list simply stays empty. - if let Some(sp) = req.sampling_params.as_ref() { - if sp.logprobs.is_some_and(|n| !(0..=1).contains(&n)) { - return Err( - "top_logprobs are not supported over the TokenSpeed ZMQ backend".to_string(), - ); - } - if sp.prompt_logprobs.is_some() { - return Err( - "prompt logprobs are not supported over the TokenSpeed ZMQ backend".to_string(), - ); - } - // The response_format / forced-tool-choice constraint oneof is not - // translated onto the TokenSpeed structured-output fields yet; dropping - // it would return unconstrained text. - if sp.constraint.is_some() { - return Err( - "structured output constraints are not supported over the ZMQ backend yet" - .to_string(), - ); - } - // The TokenSpeed wire has no per-sample demux; n>1 is fanned out into - // single-sample sub-requests by `generate` before translation. - // String `stop` sequences are resolved upstream in request building - // (`resolve_string_stops`): the strings are dropped and single-token - // stops become `stop_token_ids`, while the router-side stop decoder - // trims the text. Any residual strings here are harmless — the decoder - // still enforces them — so this path forwards token ids only. - // logit_bias is not translated onto the TokenSpeed wire either. - if !sp.logit_bias.is_empty() { - return Err("logit_bias is not supported over the TokenSpeed ZMQ backend".to_string()); - } + if !req.token_ids_logprob.is_empty() { + return Err( + "prompt logprobs are not supported over the TokenSpeed ZMQ backend".to_string(), + ); } - let return_logprob = req - .sampling_params - .as_ref() - .is_some_and(|sp| sp.logprobs.is_some()); Ok(TokenizedGenerateReqInput { rid: req.request_id, input_ids, @@ -830,50 +861,68 @@ fn translate_request_tokenspeed( params.normalize(); params }), - return_logprob, - stream, + return_logprob: req.return_logprob, + stream: req.stream, // Every other field keeps its neutral default (the fields after // `stream` are not even emitted; the engine fills them from defaults). ..TokenizedGenerateReqInput::default() }) } -/// Map vLLM-proto sampling params onto TokenSpeed's native `SamplingParams`, -/// in the normalized form: the engine skips its decode-time re-derivation once +/// Map TokenSpeed proto sampling params onto the wire `SamplingParams`, in the +/// normalized form: the engine skips its decode-time re-derivation once /// `is_normalized` is set, so [`TokenSpeedSamplingParams::normalize`] resolves /// the derived fields (top_k sentinel, greedy collapse) before encoding. -fn translate_sampling_tokenspeed(sp: vllm::SamplingParams) -> TokenSpeedSamplingParams { +/// +/// String `stop` sequences are not forwarded — the token-only engine cannot +/// match them; the router-side stop decoder trims them from the text instead. +fn translate_sampling_tokenspeed( + sp: tokenspeed_proto::SamplingParams, +) -> TokenSpeedSamplingParams { let mut params = TokenSpeedSamplingParams { - max_new_tokens: sp.max_tokens, + max_new_tokens: sp.max_new_tokens, stop_token_ids: (!sp.stop_token_ids.is_empty()).then_some(sp.stop_token_ids), temperature: f64::from(sp.temperature.unwrap_or(1.0)), - top_p: f64::from(sp.top_p), - // vLLM proto uses `0` for "all tokens"; TokenSpeed's API form is `-1` - // (`normalize` resolves it to the engine's disabled sentinel). - top_k: if sp.top_k == 0 { - -1 - } else { - i32::try_from(sp.top_k).unwrap_or(-1) - }, - min_p: f64::from(sp.min_p), - frequency_penalty: f64::from(sp.frequency_penalty), - presence_penalty: f64::from(sp.presence_penalty), - repetition_penalty: f64::from(sp.repetition_penalty), - min_new_tokens: sp.min_tokens, + top_p: f64::from(sp.top_p.unwrap_or(1.0)), + // The proto keeps the API convention `-1` = "all tokens" (and unset); + // `normalize` resolves it to the engine's disabled sentinel. + top_k: sp.top_k.unwrap_or(-1), + min_p: f64::from(sp.min_p.unwrap_or(0.0)), + frequency_penalty: f64::from(sp.frequency_penalty.unwrap_or(0.0)), + presence_penalty: f64::from(sp.presence_penalty.unwrap_or(0.0)), + repetition_penalty: f64::from(sp.repetition_penalty.unwrap_or(1.0)), + min_new_tokens: sp.min_new_tokens, ignore_eos: sp.ignore_eos, - // A negative seed is a "no seed" sentinel; drop it rather than wrap - // (the engine then derives a per-rid seed). - seed: sp.seed.and_then(|seed| u64::try_from(seed).ok()), + skip_special_tokens: sp.skip_special_tokens, + spaces_between_special_tokens: sp.spaces_between_special_tokens, + no_stop_trim: sp.no_stop_trim, // Proto `0` means unspecified; TokenSpeed expects at least one sample. - // The engine stores n without acting on it — n>1 is fanned out before - // translation, so this is always 1 on the wire. + // n>1 is fanned out before translation, so this is always 1 on the wire. n: sp.n.max(1), ..TokenSpeedSamplingParams::default() }; + apply_tokenspeed_constraint(&mut params, sp.constraint); params.normalize(); params } +/// Map the proto structured-output `constraint` oneof onto the wire's dedicated +/// fields. The oneof is single-valued, so at most one field is set; the rest +/// stay `None`. +fn apply_tokenspeed_constraint( + params: &mut TokenSpeedSamplingParams, + constraint: Option, +) { + use tokenspeed_proto::sampling_params::Constraint; + match constraint { + Some(Constraint::JsonSchema(schema)) => params.json_schema = Some(schema), + Some(Constraint::Regex(regex)) => params.regex = Some(regex), + Some(Constraint::EbnfGrammar(grammar)) => params.ebnf = Some(grammar), + Some(Constraint::StructuralTag(tag)) => params.structural_tag = Some(tag), + None => {} + } +} + /// Translate a vLLM-proto generate request into an `EngineCoreRequest`. ZMQ mode /// requires pre-tokenized input (SMG tokenizes upstream). fn translate_request( @@ -1088,7 +1137,7 @@ mod tests { logprob: Option, finish: Option, ) -> EngineCoreOutputs { - let finished = finish.map(|_| std::collections::BTreeSet::from([request_id.to_string()])); + let finished = finish.map(|_| BTreeSet::from([request_id.to_string()])); let new_logprobs = logprob.map(|lp| { MaybeWireLogprobs::Direct(Logprobs { positions: vec![PositionLogprobs { @@ -1187,7 +1236,10 @@ mod tests { stream: true, ..Default::default() }; - let mut stream = client.generate(req).await.expect("generate"); + let mut stream = client + .generate(ProtoGenerateRequest::Vllm(Box::new(req))) + .await + .expect("generate"); let first = stream.next().await.expect("chunk item").expect("chunk ok"); match first.response { @@ -1293,7 +1345,7 @@ mod tests { finish_reason: Some(EngineCoreFinishReason::Length), ..Default::default() }], - finished_requests: Some(std::collections::BTreeSet::from(["r1".to_string()])), + finished_requests: Some(BTreeSet::from(["r1".to_string()])), ..Default::default() }); @@ -1328,7 +1380,10 @@ mod tests { stream: true, ..Default::default() }; - let mut stream = client.generate(req).await.expect("generate"); + let mut stream = client + .generate(ProtoGenerateRequest::Vllm(Box::new(req))) + .await + .expect("generate"); // The requested count is 2, so `top_logprobs` keeps the sampled entry // plus the leading candidate (the third entry is dropped). @@ -1456,24 +1511,25 @@ mod tests { .unwrap(); }); - let req = vllm::GenerateRequest { + let req = tokenspeed_proto::GenerateRequest { request_id: "r1".to_string(), - input: Some(vllm::generate_request::Input::Tokenized( - vllm::TokenizedInput { - original_text: String::new(), - input_ids: vec![1, 2, 3], - }, - )), - sampling_params: Some(vllm::SamplingParams { - max_tokens: Some(2), - // Plain sampled-token logprob (count 1); must be wired through. - logprobs: Some(1), + tokenized: Some(tokenspeed_proto::TokenizedInput { + input_ids: vec![1, 2, 3], + original_text: String::new(), + }), + sampling_params: Some(tokenspeed_proto::SamplingParams { + max_new_tokens: Some(2), ..Default::default() }), + // Plain sampled-token logprob; must be wired through. + return_logprob: true, stream: true, ..Default::default() }; - let mut stream = client.generate(req).await.expect("generate"); + let mut stream = client + .generate(ProtoGenerateRequest::TokenSpeed(Box::new(req))) + .await + .expect("generate"); let first = stream.next().await.expect("chunk item").expect("chunk ok"); match first.response { @@ -1533,35 +1589,33 @@ mod tests { } #[test] - fn tokenspeed_sampling_maps_top_k_sentinel_and_seed() { + fn tokenspeed_sampling_maps_top_k_sentinel_and_floors_n() { use engine_zmq_client::protocol::tokenspeed::sampling::TOP_K_DISABLED; - // Proto top_k=0 ("all tokens") normalizes to the engine's disabled - // sentinel; negative seed dropped; n floored to 1. - let mapped = translate_sampling_tokenspeed(vllm::SamplingParams { - top_k: 0, + // Unset top_k rides the API convention `-1` ("all tokens") and normalizes + // to the engine's disabled sentinel; n=0 floors to 1; max_new_tokens + // forwards. + let mapped = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams { + top_k: None, n: 0, - seed: Some(-1), - max_tokens: Some(8), + max_new_tokens: Some(8), ..Default::default() }); assert_eq!(mapped.top_k, TOP_K_DISABLED); assert_eq!(mapped.n, 1); - assert_eq!(mapped.seed, None); assert_eq!(mapped.max_new_tokens, Some(8)); // The wire form is always normalized (the engine skips re-derivation). assert!(mapped.is_normalized); - let mapped = translate_sampling_tokenspeed(vllm::SamplingParams { - top_k: 40, - seed: Some(7), + // An explicit top_k passes through unchanged. + let mapped = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams { + top_k: Some(40), ..Default::default() }); assert_eq!(mapped.top_k, 40); - assert_eq!(mapped.seed, Some(7)); // A near-zero temperature collapses to greedy on the wire. - let mapped = translate_sampling_tokenspeed(vllm::SamplingParams { + let mapped = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams { temperature: Some(0.0), ..Default::default() }); @@ -1569,7 +1623,7 @@ mod tests { assert_eq!(mapped.top_k, 1); // Empty stop_token_ids ride as None (the normalized encoding). - let mapped = translate_sampling_tokenspeed(vllm::SamplingParams::default()); + let mapped = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams::default()); assert_eq!(mapped.stop_token_ids, None); } @@ -1588,95 +1642,122 @@ mod tests { } } - #[test] - fn tokenspeed_plain_logprobs_set_return_logprob() { - // Counts 0 and 1 are the plain sampled-token case; both accepted. - for count in [0, 1] { - let req = translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { - logprobs: Some(count), - ..Default::default() - })) - .expect("plain logprobs accepted"); - assert!(req.return_logprob); + fn ts_tokenized_req( + sampling: tokenspeed_proto::SamplingParams, + ) -> tokenspeed_proto::GenerateRequest { + tokenspeed_proto::GenerateRequest { + request_id: "r1".to_string(), + tokenized: Some(tokenspeed_proto::TokenizedInput { + input_ids: vec![1, 2, 3], + original_text: String::new(), + }), + sampling_params: Some(sampling), + stream: true, + ..Default::default() } + } - // No logprobs -> the flag stays false. - let req = translate_request_tokenspeed(tokenized_req(vllm::SamplingParams::default())) - .expect("no logprobs accepted"); - assert!(!req.return_logprob); + #[test] + fn tokenspeed_return_logprob_flag_passes_through() { + // The request-level `return_logprob` drives the plain sampled-token + // logprob (count 0/1 in `top_logprobs_num` is the same case). + let mut req = ts_tokenized_req(tokenspeed_proto::SamplingParams::default()); + req.return_logprob = true; + let wire = translate_request_tokenspeed(req).expect("return_logprob accepted"); + assert!(wire.return_logprob); + + // Unset -> the flag stays false. + let wire = translate_request_tokenspeed(ts_tokenized_req( + tokenspeed_proto::SamplingParams::default(), + )) + .expect("no logprobs accepted"); + assert!(!wire.return_logprob); } #[test] - fn tokenspeed_rejects_top_k_and_prompt_logprobs() { - // Top-k (count > 1) cannot be honored. - assert!( - translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { - logprobs: Some(5), - ..Default::default() - })) - .is_err() - ); - // "all" (count -1) cannot be honored. - assert!( - translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { - logprobs: Some(-1), - ..Default::default() - })) - .is_err() - ); - // Prompt logprobs cannot be produced. - assert!( - translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { - prompt_logprobs: Some(1), - ..Default::default() - })) - .is_err() - ); + fn tokenspeed_rejects_top_logprobs_and_prompt_logprobs() { + // Top-k logprobs (count > 1) cannot be honored over the wire. + let mut req = ts_tokenized_req(tokenspeed_proto::SamplingParams::default()); + req.top_logprobs_num = 5; + assert!(translate_request_tokenspeed(req).is_err()); + + // Prompt (input) logprobs cannot be produced. + let mut req = ts_tokenized_req(tokenspeed_proto::SamplingParams::default()); + req.token_ids_logprob = vec![1, 2]; + assert!(translate_request_tokenspeed(req).is_err()); + + // A bare count of 0/1 is the plain sampled-token case: accepted. + for count in [0, 1] { + let mut req = ts_tokenized_req(tokenspeed_proto::SamplingParams::default()); + req.top_logprobs_num = count; + assert!(translate_request_tokenspeed(req).is_ok()); + } } #[test] - fn tokenspeed_rejects_unsupported_sampling_features() { - // Structured-output constraints have no TokenSpeed wire slot. - let err = translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { - constraint: Some(vllm::sampling_params::Constraint::JsonObject(true)), + fn tokenspeed_maps_structured_output_constraints() { + // The `constraint` oneof maps 1:1 onto the wire's dedicated fields; the + // oneof is single-valued, so the other three stay unset. + use tokenspeed_proto::sampling_params::Constraint; + + let json = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams { + constraint: Some(Constraint::JsonSchema("{\"type\":\"object\"}".into())), ..Default::default() - })) - .expect_err("constraint rejected"); - assert!(err.contains("structured output"), "{err}"); + }); + assert_eq!(json.json_schema.as_deref(), Some("{\"type\":\"object\"}")); + assert_eq!(json.regex, None); + assert_eq!(json.ebnf, None); + assert_eq!(json.structural_tag, None); - // logit_bias has no wire slot. - let err = translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { - logit_bias: HashMap::from([(7, 1.0)]), + let regex = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams { + constraint: Some(Constraint::Regex("[0-9]+".into())), ..Default::default() - })) - .expect_err("logit_bias rejected"); - assert!(err.contains("logit_bias"), "{err}"); + }); + assert_eq!(regex.regex.as_deref(), Some("[0-9]+")); + assert_eq!(regex.json_schema, None); + + let ebnf = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams { + constraint: Some(Constraint::EbnfGrammar("root ::= \"a\"".into())), + ..Default::default() + }); + assert_eq!(ebnf.ebnf.as_deref(), Some("root ::= \"a\"")); + + let tag = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams { + constraint: Some(Constraint::StructuralTag("".into())), + ..Default::default() + }); + assert_eq!(tag.structural_tag.as_deref(), Some("")); + + // No constraint leaves all four structured-output fields unset. + let none = translate_sampling_tokenspeed(tokenspeed_proto::SamplingParams::default()); + assert_eq!(none.json_schema, None); + assert_eq!(none.regex, None); + assert_eq!(none.ebnf, None); + assert_eq!(none.structural_tag, None); } #[test] - fn tokenspeed_ignores_residual_stop_strings() { - // String stops are resolved upstream (`resolve_string_stops`); any that - // reach here are harmless — the translator forwards token ids only and - // the router-side decoder trims the text. - let req = translate_request_tokenspeed(tokenized_req(vllm::SamplingParams { + fn tokenspeed_forwards_stop_token_ids_and_drops_stop_strings() { + // String stops are resolved upstream; any that reach here are dropped + // (the token-only engine cannot match them) while stop token ids ride + // through and the router-side decoder trims residual text. + let req = translate_request_tokenspeed(ts_tokenized_req(tokenspeed_proto::SamplingParams { stop: vec!["".to_string()], stop_token_ids: vec![13], ..Default::default() })) .expect("residual stop strings must not be rejected"); assert_eq!(req.sampling_params.stop_token_ids, Some(vec![13])); + assert_eq!(req.sampling_params.stop, None); } #[test] - fn tokenspeed_rejects_nonzero_dp_rank() { - // Single-engine backend: only rank 0 (or none) is valid. - let mut req = tokenized_req(vllm::SamplingParams::default()); - req.data_parallel_rank = Some(1); + fn tokenspeed_rejects_multimodal_inputs() { + // The TokenSpeed ZMQ wire has no multimodal slot yet; reject rather than + // silently drop pixels. + let mut req = ts_tokenized_req(tokenspeed_proto::SamplingParams::default()); + req.mm_inputs = Some(tokenspeed_proto::MultimodalInputs::default()); assert!(translate_request_tokenspeed(req).is_err()); - - let mut req = tokenized_req(vllm::SamplingParams::default()); - req.data_parallel_rank = Some(0); - assert!(translate_request_tokenspeed(req).is_ok()); } #[test] @@ -1960,7 +2041,10 @@ mod tests { ..Default::default() }); req.request_id = "r1".to_string(); - let mut stream = client.generate(req).await.expect("generate"); + let mut stream = client + .generate(ProtoGenerateRequest::Vllm(Box::new(req))) + .await + .expect("generate"); let mut completes = Vec::new(); while let Some(item) = stream.next().await { @@ -2035,12 +2119,15 @@ mod tests { let request: TokenizedGenerateReqInput = decode_msgpack(frames[1].as_ref()).unwrap(); assert_eq!(request.sampling_params.n, 1); - rids.push((request.rid.clone(), request.sampling_params.seed)); + // TokenSpeed has no seed on the wire; the engine derives one from + // the (unique) rid so all TP/DP ranks agree. + assert_eq!(request.sampling_params.seed, None); + rids.push(request.rid.clone()); } assert_eq!( rids, - vec![("r1-0".to_string(), Some(5)), ("r1-1".to_string(), Some(6))], - "sub-rids must be unique and seeds derived per sub" + vec!["r1-0".to_string(), "r1-1".to_string()], + "sub-rids must be unique per sub" ); // Both subs finish in one wire batch (the batch demux fans them // back out to their sub-streams). @@ -2060,13 +2147,15 @@ mod tests { .unwrap(); }); - let mut req = tokenized_req(vllm::SamplingParams { + let mut req = ts_tokenized_req(tokenspeed_proto::SamplingParams { n: 2, - seed: Some(5), ..Default::default() }); req.request_id = "r1".to_string(); - let mut stream = client.generate(req).await.expect("generate"); + let mut stream = client + .generate(ProtoGenerateRequest::TokenSpeed(Box::new(req))) + .await + .expect("generate"); let mut completes = Vec::new(); while let Some(item) = stream.next().await { @@ -2120,10 +2209,12 @@ mod tests { let (mut engine_input, _engine_output) = engine.split(); let stream = client - .generate(tokenized_req(vllm::SamplingParams { - n: 2, - ..Default::default() - })) + .generate(ProtoGenerateRequest::Vllm(Box::new(tokenized_req( + vllm::SamplingParams { + n: 2, + ..Default::default() + }, + )))) .await .expect("generate"); @@ -2138,7 +2229,7 @@ mod tests { drop(stream); // unfinished -> every sub auto-aborts - let mut aborted = std::collections::BTreeSet::new(); + let mut aborted = BTreeSet::new(); while aborted.len() < 2 { match engine_input.recv().await.unwrap() { EngineInbound::Abort(rids) => aborted.extend(rids), @@ -2147,7 +2238,7 @@ mod tests { } assert_eq!( aborted, - std::collections::BTreeSet::from(["r1-0".to_string(), "r1-1".to_string()]) + BTreeSet::from(["r1-0".to_string(), "r1-1".to_string()]) ); }