diff --git a/src/aiperf/endpoints/vllm_generate.py b/src/aiperf/endpoints/vllm_generate.py new file mode 100644 index 0000000000..4fa01d2592 --- /dev/null +++ b/src/aiperf/endpoints/vllm_generate.py @@ -0,0 +1,142 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +"""vLLM token-in/token-out ``/inference/v1/generate`` endpoint.""" + +from __future__ import annotations + +from typing import Any + +from aiperf.common.models import ( + BaseResponseData, + ExtractedPayload, + InferenceServerResponse, + ParsedResponse, + RequestInfo, + RequestRecord, +) +from aiperf.endpoints.base_endpoint import BaseEndpoint + + +class VllmGenerateEndpoint(BaseEndpoint): + """Send and measure vLLM ``GenerateRequest`` token arrays. + + The endpoint is intentionally non-streaming so the same payload works with + vLLM and Dynamo's vLLM-compatible engine API. Use ``--endpoint-path`` only + when the server mounts the API at a custom path. + """ + + def format_payload(self, request_info: RequestInfo) -> dict[str, Any]: + if len(request_info.turns) != 1: + raise ValueError( + "vLLM generate endpoint requires one token payload per turn" + ) + + turn = request_info.turns[0] + extra_body = dict(turn.extra_body or {}) + token_ids = extra_body.pop("token_ids", None) + if not self._valid_token_ids(token_ids): + raise ValueError( + "turn.extra_body.token_ids must be a non-empty list of integers" + ) + + sampling_params = dict(extra_body.pop("sampling_params", {}) or {}) + if turn.max_tokens is not None: + sampling_params.setdefault("max_tokens", turn.max_tokens) + if extra_body.pop("stream", False): + raise ValueError("vLLM generate endpoint does not support streaming") + + payload: dict[str, Any] = { + "model": turn.model or request_info.model_endpoint.primary_model_name, + "token_ids": token_ids, + "sampling_params": sampling_params, + "stream": False, + } + if request_info.x_request_id: + payload["request_id"] = request_info.x_request_id + payload.update(dict(self.model_endpoint.endpoint.extra or [])) + payload.update(extra_body) + if payload.get("stream") is not False: + raise ValueError("vLLM generate endpoint requires stream=false") + return payload + + def extract_payload_inputs(self, payload: dict[str, Any]) -> ExtractedPayload: + result = ExtractedPayload() + token_ids = payload.get("token_ids") + if self._valid_token_ids(token_ids): + result.pretokenised_token_count = len(token_ids) + return result + + def parse_response( + self, response: InferenceServerResponse + ) -> ParsedResponse | None: + return self._parse_response(response, prompt_tokens=None) + + def extract_response_data(self, record: RequestRecord) -> list[ParsedResponse]: + prompt_tokens = self._prompt_token_count(record) + return [ + parsed + for response in record.responses or [] + if (parsed := self._parse_response(response, prompt_tokens)) is not None + ] + + def _parse_response( + self, + response: InferenceServerResponse, + prompt_tokens: int | None, + ) -> ParsedResponse | None: + payload = response.get_json() + if not isinstance(payload, dict): + return None + choices = payload.get("choices") + if ( + not isinstance(choices, list) + or not choices + or not isinstance(choices[0], dict) + ): + return None + completion_ids = choices[0].get("token_ids") + if not isinstance(completion_ids, list) or not all( + isinstance(token_id, int) and not isinstance(token_id, bool) + for token_id in completion_ids + ): + return None + + completion_tokens = len(completion_ids) + usage = { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": ( + prompt_tokens + completion_tokens if prompt_tokens is not None else None + ), + } + return ParsedResponse( + perf_ns=response.perf_ns, + data=BaseResponseData(), + usage=usage, + metadata={ + "request_id": payload.get("request_id"), + "finish_reason": choices[0].get("finish_reason"), + "completion_token_ids": completion_ids, + }, + ) + + @classmethod + def _prompt_token_count(cls, record: RequestRecord) -> int | None: + if not record.turns: + return None + turn = record.turns[-1] + payload = turn.raw_payload or turn.extra_body or {} + token_ids = payload.get("token_ids") if isinstance(payload, dict) else None + return len(token_ids) if cls._valid_token_ids(token_ids) else None + + @staticmethod + def _valid_token_ids(value: Any) -> bool: + return ( + isinstance(value, list) + and bool(value) + and all( + isinstance(token_id, int) and not isinstance(token_id, bool) + for token_id in value + ) + ) diff --git a/src/aiperf/plugin/plugins.yaml b/src/aiperf/plugin/plugins.yaml index d72121a9d0..5fe3c5782f 100644 --- a/src/aiperf/plugin/plugins.yaml +++ b/src/aiperf/plugin/plugins.yaml @@ -198,6 +198,19 @@ endpoint: supports_videos: true metrics_title: LLM Metrics + vllm_generate: + class: aiperf.endpoints.vllm_generate:VllmGenerateEndpoint + description: | + Non-streaming vLLM token-in/token-out generate endpoint. Accepts exact + token_ids and reconstructs usage from request and response token arrays. + Override the endpoint path for Dynamo's vLLM-compatible engine API. + metadata: + endpoint_path: /inference/v1/generate + supports_streaming: false + produces_tokens: true + tokenizes_input: false + metrics_title: LLM Metrics + cohere_rankings: class: aiperf.endpoints.cohere_rankings:CohereRankingsEndpoint description: | diff --git a/tests/unit/common/enums/test_endpoints_enums.py b/tests/unit/common/enums/test_endpoints_enums.py index de01bd6541..8becefada5 100644 --- a/tests/unit/common/enums/test_endpoints_enums.py +++ b/tests/unit/common/enums/test_endpoints_enums.py @@ -69,6 +69,14 @@ class TestEndpointType: "/generate", "LLM Metrics", ), + ( + EndpointType.VLLM_GENERATE, + "vllm_generate", + False, + True, + "/inference/v1/generate", + "LLM Metrics", + ), ], ) def test_endpoint_type_metadata( @@ -106,6 +114,7 @@ def test_endpoint_type_metadata( "hf_tei_rankings", "cohere_rankings", "huggingface_generate", + "vllm_generate", ], ) def test_enum_string_comparison(self, tag_value): diff --git a/tests/unit/endpoints/test_vllm_generate.py b/tests/unit/endpoints/test_vllm_generate.py new file mode 100644 index 0000000000..905e5d42f7 --- /dev/null +++ b/tests/unit/endpoints/test_vllm_generate.py @@ -0,0 +1,83 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +# SPDX-License-Identifier: Apache-2.0 + +import orjson +import pytest + +from aiperf.common.models import RequestRecord, TextResponse, Turn +from aiperf.endpoints.vllm_generate import VllmGenerateEndpoint +from aiperf.plugin import plugins +from aiperf.plugin.enums import EndpointType +from tests.unit.endpoints.conftest import create_model_endpoint, create_request_info + + +@pytest.fixture +def endpoint(): + return VllmGenerateEndpoint( + create_model_endpoint(EndpointType.VLLM_GENERATE, model_name="test-model") + ) + + +def test_metadata(): + metadata = plugins.get_endpoint_metadata(EndpointType.VLLM_GENERATE) + assert metadata.endpoint_path == "/inference/v1/generate" + assert metadata.supports_streaming is False + assert metadata.produces_tokens is True + assert metadata.tokenizes_input is False + + +def test_format_payload(endpoint): + request = create_request_info( + model_endpoint=endpoint.model_endpoint, + max_tokens=17, + extra_body={"token_ids": [1, 2, 3], "sampling_params": {"temperature": 0}}, + ) + + payload = endpoint.format_payload(request) + + assert payload == { + "model": "test-model", + "token_ids": [1, 2, 3], + "sampling_params": {"temperature": 0, "max_tokens": 17}, + "stream": False, + "request_id": "test-request-id", + } + + +def test_format_payload_rejects_missing_tokens(endpoint): + request = create_request_info(model_endpoint=endpoint.model_endpoint) + with pytest.raises(ValueError, match="token_ids"): + endpoint.format_payload(request) + + +def test_extract_payload_inputs_counts_exact_ids(endpoint): + extracted = endpoint.extract_payload_inputs({"token_ids": [10, 11, 12, 13]}) + assert extracted.pretokenised_token_count == 4 + + +def test_extract_response_data_reconstructs_usage(endpoint): + response = TextResponse( + perf_ns=123, + text=orjson.dumps( + { + "request_id": "req-1", + "choices": [ + {"index": 0, "token_ids": [20, 21], "finish_reason": "stop"} + ], + } + ).decode(), + content_type="application/json", + ) + record = RequestRecord( + model_name="test-model", + responses=[response], + turns=[Turn(role="user", raw_payload={"token_ids": [1, 2, 3]})], + ) + + parsed = endpoint.extract_response_data(record) + + assert len(parsed) == 1 + assert parsed[0].usage.prompt_tokens == 3 + assert parsed[0].usage.completion_tokens == 2 + assert parsed[0].usage.total_tokens == 5 + assert parsed[0].metadata["completion_token_ids"] == [20, 21]