diff --git a/renderers/client.py b/renderers/client.py index 88a151fc..656f424a 100644 --- a/renderers/client.py +++ b/renderers/client.py @@ -157,7 +157,7 @@ def parse_generate_response(raw: bytes) -> dict[str, Any]: def _parse_packed_completion_logprobs( packed: Mapping[str, Any], completion_ids: list[int] -) -> list[float]: +) -> np.ndarray: """prime-rl's packed form: base64 float32 ``{data, shape, dtype}``, one per token.""" values = np.frombuffer(base64.b64decode(packed["data"]), dtype=packed["dtype"]) valid = np.isfinite(values) & (values > VLLM_LOGPROB_SENTINEL) @@ -166,12 +166,12 @@ def _parse_packed_completion_logprobs( f"Engine response choice.completion_logprobs: {len(values)} values for " f"{len(completion_ids)} completion tokens, or a non-finite or missing (-9999) logprob." ) - return values.astype(np.float64).tolist() + return values.astype(np.float32, copy=False) def _parse_compact_logprobs( choice: Mapping[str, Any], completion_ids: list[int], k: int | None -) -> tuple[list[float], list[list[int]] | None, list[list[float]] | None]: +) -> tuple[np.ndarray, list[list[int]] | None, list[list[float]] | None]: packed = choice.get("compact_logprobs") if not isinstance(packed, Mapping) or packed.get("format") != "u32-f32-le-v1": raise MalformedGenerateResponseError("Unsupported compact logprob format") @@ -227,7 +227,7 @@ def _parse_compact_logprobs( top_ids.append(selected_ids) top_logprobs.append(selected_logprobs) return ( - sampled, + np.asarray(sampled, dtype=np.float32), top_ids if k is not None else None, top_logprobs if k is not None else None, ) @@ -245,7 +245,7 @@ def _logprob_token_id(entry: Mapping[str, Any]) -> int | None: def _parse_completion_logprobs( choice: Mapping[str, Any], completion_ids: list[int] -) -> list[float]: +) -> np.ndarray: packed = choice.get("completion_logprobs") if packed is not None: return _parse_packed_completion_logprobs(packed, completion_ids) @@ -297,7 +297,7 @@ def _parse_completion_logprobs( f"Engine response choice.logprobs.content[{index}].logprob does not contain sampling evidence." ) completion_logprobs.append(logprob) - return completion_logprobs + return np.asarray(completion_logprobs, dtype=np.float32) def _parse_sampling_mask( @@ -501,6 +501,9 @@ async def generate( mm_placeholders, completion_ids, completion_logprobs, content, reasoning_content, tool_calls, finish_reason, routed_experts, weight_versions, multi_modal_data, prompt_attribution. + ``completion_logprobs`` is a float32 numpy array, one value per completion + token: vLLM samples in float32, so the packed form decodes without a copy + and the JSON forms narrow losslessly. ``renderer_prompt_ids`` is the unexpanded logical prompt when ``process_multimodal=False`` and ``None`` otherwise. ``weight_versions`` is vLLM's (>= 0.32) list of ``{version, start, end}`` spans over diff --git a/tests/replay.py b/tests/replay.py index b4aac4c7..67774852 100644 --- a/tests/replay.py +++ b/tests/replay.py @@ -44,5 +44,5 @@ def respond(request): def assert_evidence(result, ids): assert result["completion_ids"] == ids - assert result["completion_logprobs"] == [-0.25] * len(ids) + assert result["completion_logprobs"].tolist() == [-0.25] * len(ids) assert result["sampling_mask"] == [[tid] for tid in ids] diff --git a/tests/test_client.py b/tests/test_client.py index 6e1baac8..3b53fe25 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -122,9 +122,14 @@ def test_parse_generate_response_preserves_multiple_compact_choices(separators): json.dumps({"choices": choices}, separators=separators).encode(), ) for i, choice in zip((7, 8), payload["choices"]): - assert _parse_compact_logprobs(choice, [i], None) == ([-0.5], None, None) + sampled, top_ids, top_logprobs = _parse_compact_logprobs(choice, [i], None) + assert sampled.dtype == np.float32 and sampled.tolist() == [-0.5] + assert top_ids is None and top_logprobs is None assert base64.b64decode(choice["routed_experts"]["data"]) == bytes([i]) - assert _parse_compact_logprobs(payload["choices"][2], [], 2) == ([], [], []) + sampled, top_ids, top_logprobs = _parse_compact_logprobs( + payload["choices"][2], [], 2 + ) + assert sampled.tolist() == [] and top_ids == [] and top_logprobs == [] @pytest.mark.parametrize("sampled", [1, 3, 5]) @@ -594,7 +599,7 @@ def test_generate_preserves_zero_completion_logprob(): result = _run_generate(client) - assert result["completion_logprobs"] == [0.0, -0.2] + assert result["completion_logprobs"] == pytest.approx([0.0, -0.2]) def test_generate_reads_packed_completion_logprobs(): @@ -609,7 +614,8 @@ def test_generate_reads_packed_completion_logprobs(): result = _run_generate(client) - assert result["completion_logprobs"] == values.tolist() + assert result["completion_logprobs"].dtype == np.float32 + assert result["completion_logprobs"].tolist() == values.tolist() client.choice["completion_logprobs"]["data"] = base64.b64encode( np.array([-0.1, -9999.0], dtype=np.float32).tobytes() @@ -1141,5 +1147,5 @@ def parse_response(self, completion_ids, *, tools=None, prompt_ids=None): assert result["reasoning_complete"] is reasoning_complete assert result["finish_reason"] == finish assert result["completion_ids"] == [7, 8] - assert result["completion_logprobs"] == [-0.1, -0.2] + assert result["completion_logprobs"] == pytest.approx([-0.1, -0.2]) assert result["content"] == ""