From c244237c0df2413787631933ea176376bf7e427a Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Tue, 6 Oct 2026 11:50:59 +0000 Subject: [PATCH 1/3] perf(client): read prime-rl's packed completion_logprobs --- renderers/client.py | 36 +++++++++++++++++++++++++++++++++++- tests/test_client.py | 21 +++++++++++++++++++++ 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/renderers/client.py b/renderers/client.py index fd13e6d61..6972923cb 100644 --- a/renderers/client.py +++ b/renderers/client.py @@ -7,6 +7,7 @@ from __future__ import annotations import asyncio +import base64 import json import logging import math @@ -14,6 +15,7 @@ from typing import Any, cast import httpx +import numpy as np from openai import AsyncOpenAI from renderers.base import ( @@ -134,9 +136,40 @@ def parse_generate_response(raw: bytes) -> dict[str, Any]: return payload +def _parse_packed_completion_logprobs( + packed: Mapping[str, Any], completion_ids: list[int] +) -> list[float]: + """prime-rl's packed form: base64 float32 ``{data, shape, dtype}``, one per token.""" + try: + values = np.frombuffer( + base64.b64decode(packed["data"]), dtype=packed["dtype"] + ).astype(np.float64) + except (KeyError, TypeError, ValueError) as exc: + raise MalformedGenerateResponseError( + "Engine response choice.completion_logprobs must be a packed float array." + ) from exc + if len(values) != len(completion_ids): + raise MalformedGenerateResponseError( + "Engine response completion token count " + f"({len(completion_ids)}) does not match logprob count ({len(values)})." + ) + if not np.isfinite(values).all(): + raise MalformedGenerateResponseError( + "Engine response choice.completion_logprobs must be finite." + ) + if (values == VLLM_LOGPROB_SENTINEL).any(): + raise MalformedGenerateResponseError( + "Engine response choice.completion_logprobs does not contain sampling evidence." + ) + return values.tolist() + + def _parse_completion_logprobs( choice: Mapping[str, Any], completion_ids: list[int] ) -> list[float]: + packed = choice.get("completion_logprobs") + if packed is not None: + return _parse_packed_completion_logprobs(packed, completion_ids) raw_logprobs = choice.get("logprobs") if not isinstance(raw_logprobs, Mapping): raise MalformedGenerateResponseError( @@ -376,7 +409,8 @@ def _prepare(): routed_experts = choice.get("routed_experts") # vLLM's native kept-set sampling masks (``--return-sampling-mask``): - # one list of surviving vocab ids per completion token. + # one list of surviving vocab ids per completion token, or (prime-rl's + # server) packed CSR ``{"ids", "counts"}`` base64 int32 arrays. sampling_mask = choice.get("sampling_mask") # /inference/v1/generate returns finish_reason in {"stop","length",...} — diff --git a/tests/test_client.py b/tests/test_client.py index 16fdafef6..daca8c476 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -336,6 +336,27 @@ def test_generate_preserves_zero_completion_logprob(): assert result["completion_logprobs"] == [0.0, -0.2] +def test_generate_reads_packed_completion_logprobs(): + client = _FakeClient() + client.choice.pop("logprobs") + values = np.array([-0.1, 0.0], dtype=np.float32) + client.choice["completion_logprobs"] = { + "data": base64.b64encode(values.tobytes()).decode(), + "shape": [2], + "dtype": "float32", + } + + result = _run_generate(client) + + assert result["completion_logprobs"] == values.tolist() + + client.choice["completion_logprobs"]["data"] = base64.b64encode( + np.array([-0.1, -9999.0], dtype=np.float32).tobytes() + ).decode() + with pytest.raises(MalformedGenerateResponseError, match="sampling evidence"): + _run_generate(client) + + class _MalformedToolRenderer(_FakeRenderer): """Returns only a malformed tool-call attempt — finish_reason must stay "stop".""" From 582005230e2178054c60d5fda835db738ecb8f76 Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Wed, 7 Oct 2026 14:30:28 +0000 Subject: [PATCH 2/3] simplify packed completion_logprobs parsing --- renderers/client.py | 25 +++++-------------------- tests/test_client.py | 2 +- 2 files changed, 6 insertions(+), 21 deletions(-) diff --git a/renderers/client.py b/renderers/client.py index 6972923cb..cdd625573 100644 --- a/renderers/client.py +++ b/renderers/client.py @@ -140,28 +140,13 @@ def _parse_packed_completion_logprobs( packed: Mapping[str, Any], completion_ids: list[int] ) -> list[float]: """prime-rl's packed form: base64 float32 ``{data, shape, dtype}``, one per token.""" - try: - values = np.frombuffer( - base64.b64decode(packed["data"]), dtype=packed["dtype"] - ).astype(np.float64) - except (KeyError, TypeError, ValueError) as exc: - raise MalformedGenerateResponseError( - "Engine response choice.completion_logprobs must be a packed float array." - ) from exc - if len(values) != len(completion_ids): - raise MalformedGenerateResponseError( - "Engine response completion token count " - f"({len(completion_ids)}) does not match logprob count ({len(values)})." - ) - if not np.isfinite(values).all(): - raise MalformedGenerateResponseError( - "Engine response choice.completion_logprobs must be finite." - ) - if (values == VLLM_LOGPROB_SENTINEL).any(): + values = np.frombuffer(base64.b64decode(packed["data"]), dtype=packed["dtype"]) + if len(values) != len(completion_ids) or not (values > VLLM_LOGPROB_SENTINEL).all(): raise MalformedGenerateResponseError( - "Engine response choice.completion_logprobs does not contain sampling evidence." + f"Engine response choice.completion_logprobs: {len(values)} values for " + f"{len(completion_ids)} completion tokens, or a missing (-9999) logprob." ) - return values.tolist() + return values.astype(np.float64).tolist() def _parse_completion_logprobs( diff --git a/tests/test_client.py b/tests/test_client.py index daca8c476..f8cdb44d7 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -353,7 +353,7 @@ def test_generate_reads_packed_completion_logprobs(): client.choice["completion_logprobs"]["data"] = base64.b64encode( np.array([-0.1, -9999.0], dtype=np.float32).tobytes() ).decode() - with pytest.raises(MalformedGenerateResponseError, match="sampling evidence"): + with pytest.raises(MalformedGenerateResponseError, match="missing"): _run_generate(client) From 19c558f1db2434f4f6f226f58931f11b1d28e9aa Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Wed, 7 Oct 2026 14:49:17 +0000 Subject: [PATCH 3/3] reject non-finite packed completion logprobs --- renderers/client.py | 5 +++-- tests/test_client.py | 6 ++++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/renderers/client.py b/renderers/client.py index cdd625573..d663f82cf 100644 --- a/renderers/client.py +++ b/renderers/client.py @@ -141,10 +141,11 @@ def _parse_packed_completion_logprobs( ) -> list[float]: """prime-rl's packed form: base64 float32 ``{data, shape, dtype}``, one per token.""" values = np.frombuffer(base64.b64decode(packed["data"]), dtype=packed["dtype"]) - if len(values) != len(completion_ids) or not (values > VLLM_LOGPROB_SENTINEL).all(): + valid = np.isfinite(values) & (values > VLLM_LOGPROB_SENTINEL) + if len(values) != len(completion_ids) or not valid.all(): raise MalformedGenerateResponseError( f"Engine response choice.completion_logprobs: {len(values)} values for " - f"{len(completion_ids)} completion tokens, or a missing (-9999) logprob." + f"{len(completion_ids)} completion tokens, or a non-finite or missing (-9999) logprob." ) return values.astype(np.float64).tolist() diff --git a/tests/test_client.py b/tests/test_client.py index f8cdb44d7..d9fc53724 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -356,6 +356,12 @@ def test_generate_reads_packed_completion_logprobs(): with pytest.raises(MalformedGenerateResponseError, match="missing"): _run_generate(client) + client.choice["completion_logprobs"]["data"] = base64.b64encode( + np.array([-0.1, np.inf], dtype=np.float32).tobytes() + ).decode() + with pytest.raises(MalformedGenerateResponseError, match="non-finite"): + _run_generate(client) + class _MalformedToolRenderer(_FakeRenderer): """Returns only a malformed tool-call attempt — finish_reason must stay "stop"."""