perf: return completion logprobs as a float32 array

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

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]

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"] == ""