From c57781b999d851a97bdabac6fcb5f883c584ac4a Mon Sep 17 00:00:00 2001 From: Phi Date: Thu, 8 Oct 2026 00:46:03 +0000 Subject: [PATCH 1/3] feat(client): capture score-centering heads with compact logprob transport --- renderers/client.py | 297 +++++++++++++++++++++++++++++++++++++++---- tests/test_client.py | 234 ++++++++++++++++++++++++++++++++-- 2 files changed, 499 insertions(+), 32 deletions(-) diff --git a/renderers/client.py b/renderers/client.py index fd13e6d6..b81d9289 100644 --- a/renderers/client.py +++ b/renderers/client.py @@ -7,15 +7,18 @@ from __future__ import annotations import asyncio +import base64 +import binascii import json import logging import math from collections.abc import Mapping +from itertools import islice from typing import Any, cast import httpx +import numpy as np from openai import AsyncOpenAI - from renderers.base import ( Message, MultiModalData, @@ -109,31 +112,115 @@ async def _resolve_max_prompt_len(client: AsyncOpenAI, model: str) -> int | None return value -def _strip_base64_field(raw: bytes, prefix: bytes) -> tuple[bytes, memoryview | None]: +def _strip_base64_fields(raw: bytes, prefix: bytes) -> tuple[bytes, list[memoryview]]: """Splice a large base64 string field out of raw JSON bytes. Avoids json-decoding megabytes of base64; the returned memoryview references ``raw`` and is re-inserted into the parsed payload. """ - data_start = raw.find(prefix) - if data_start < 0: - return raw, None - - data_start += len(prefix) - data_end = raw.index(b'"', data_start) - data = memoryview(raw)[data_start:data_end] - stripped = raw[:data_start] + raw[data_end:] - return stripped, data - - -def parse_generate_response(raw: bytes) -> dict[str, Any]: - stripped, routed_data = _strip_base64_field(raw, ROUTED_EXPERTS_DATA_PREFIX) + chunks = [] + fields = [] + cursor = 0 + while (start := raw.find(prefix, cursor)) >= 0: + start += len(prefix) + end = raw.index(b'"', start) + chunks.append(raw[cursor:start]) + fields.append(memoryview(raw)[start:end]) + cursor = end + if not fields: + return raw, fields + chunks.append(raw[cursor:]) + return b"".join(chunks), fields + + +def parse_generate_response( + raw: bytes, *, compact_logprobs: bool = False +) -> dict[str, Any]: + stripped, routed_data = _strip_base64_fields(raw, ROUTED_EXPERTS_DATA_PREFIX) + compact_data = [] + if compact_logprobs: + stripped, compact_data = _strip_base64_fields( + stripped, b'"compact_logprobs":{"data":"' + ) payload: dict[str, Any] = json.loads(stripped) - if routed_data is not None: - payload["choices"][0]["routed_experts"]["data"] = routed_data + for field, buffers in ( + ("routed_experts", routed_data), + ("compact_logprobs", compact_data), + ): + if buffers: + pending = iter(buffers) + for choice in payload["choices"]: + if ( + isinstance(choice.get(field), dict) + and choice[field].get("data") == "" + ): + choice[field]["data"] = next(pending) return payload +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]: + packed = choice.get("compact_logprobs") + if not isinstance(packed, Mapping) or packed.get("format") != "u32-f32-le-v1": + raise MalformedGenerateResponseError("Unsupported compact logprob format") + offsets = packed.get("offsets") + if ( + not isinstance(offsets, list) + or len(offsets) != len(completion_ids) + 1 + or any(type(x) is not int or x < 0 for x in offsets) + or offsets[0] != 0 + or any(end <= start for start, end in zip(offsets, offsets[1:])) + ): + raise MalformedGenerateResponseError("Invalid compact logprob offsets") + num_top = packed.get("num_top_logprobs") + if type(num_top) is not int or num_top < -1: + raise MalformedGenerateResponseError("Invalid compact logprob candidate count") + try: + raw = base64.b64decode(packed["data"], validate=True) + except (KeyError, TypeError, ValueError, binascii.Error) as exc: + raise MalformedGenerateResponseError("Invalid compact logprob base64") from exc + count = offsets[-1] + if len(raw) != count * 8: + raise MalformedGenerateResponseError("Invalid compact logprob buffer length") + ids = np.frombuffer(raw, dtype=" 1e-6).any(): + raise MalformedGenerateResponseError("Invalid compact logprob values") + sampled = [] + top_ids: list[list[int]] = [] + top_logprobs: list[list[float]] = [] + sampling_mask = ( + _parse_sampling_mask(choice, completion_ids) if k is not None else None + ) + for index, (token_id, start, end) in enumerate( + zip(completion_ids, offsets, offsets[1:]) + ): + # FlatLogprobs can repeat the sampled ID. Match vLLM's dict order/value semantics. + row = dict(zip(ids[start:end].tolist(), values[start:end].tolist())) + if token_id not in row or row[token_id] == VLLM_LOGPROB_SENTINEL: + raise MalformedGenerateResponseError( + "Missing compact sampled-token evidence" + ) + sampled.append(row[token_id]) + if k is not None: + head = dict(islice(row.items(), max(num_top, 1))) + selected_ids, selected_logprobs = _select_sampler_head( + head, + token_id, + row[token_id], + k, + sampling_mask[index] if sampling_mask is not None else None, + ) + top_ids.append(selected_ids) + top_logprobs.append(selected_logprobs) + return ( + sampled, + top_ids if k is not None else None, + top_logprobs if k is not None else None, + ) + + def _parse_completion_logprobs( choice: Mapping[str, Any], completion_ids: list[int] ) -> list[float]: @@ -188,6 +275,123 @@ def _parse_completion_logprobs( return completion_logprobs +def _parse_sampling_mask( + choice: Mapping[str, Any], completion_ids: list[int] +) -> list[list[int]] | None: + sampling_mask = choice.get("sampling_mask") + if sampling_mask is None: + return None + if not isinstance(sampling_mask, list) or len(sampling_mask) != len(completion_ids): + raise MalformedGenerateResponseError("Sampling mask token count mismatch") + for sampled_id, support in zip(completion_ids, sampling_mask, strict=True): + if ( + not isinstance(support, list) + or not support + or any(type(i) is not int or i < 0 for i in support) + or len(set(support)) != len(support) + or sampled_id not in support + ): + raise MalformedGenerateResponseError("Invalid recorded sampling support") + return sampling_mask + + +def _select_sampler_head( + head: dict[int, float], + sampled_id: int, + sampled_logprob: float, + k: int, + support: list[int] | None, +) -> tuple[list[int], list[float]]: + """Keep an action-independent head, or the complete recorded replay support.""" + if sampled_id in head and not math.isclose( + head[sampled_id], sampled_logprob, rel_tol=1e-6, abs_tol=1e-5 + ): + raise MalformedGenerateResponseError( + "Sampler head and sampled-token logprob disagree" + ) + # vLLM serializes filtered-out (-inf) candidates as -9999. + ranked = sorted( + ((i, p) for i, p in head.items() if p != VLLM_LOGPROB_SENTINEL), + key=lambda x: (-x[1], x[0]), + ) + if support is not None: + support_set = set(support) + if not support_set.issubset(i for i, _ in ranked): + raise MalformedGenerateResponseError( + "Missing sampling-support probabilities; increase sampling.logprobs for cutoff ties" + ) + ranked = [(i, p) for i, p in ranked if i in support_set] + if not math.isclose( + math.fsum(math.exp(p) for _, p in ranked), 1.0, rel_tol=0, abs_tol=1e-4 + ): + raise MalformedGenerateResponseError( + "Sampler probabilities must normalize over the recorded support" + ) + elif len(ranked) > k and ranked[k - 1][1] == ranked[k][1]: + boundary = ranked[k][1] + ranked = [(i, p) for i, p in ranked if p > boundary] + else: + ranked = ranked[:k] + if not ranked or math.fsum(math.exp(p) for _, p in ranked) > 1.0001: + raise MalformedGenerateResponseError("Invalid sampler head probability mass") + if support is None and len(head) < k + 2: + raise MalformedGenerateResponseError( + "Engine returned fewer top logprobs than requested" + ) + return [i for i, _ in ranked], [p for _, p in ranked] + + +def _parse_completion_top_logprobs( + choice: Mapping[str, Any], completion_ids: list[int], k: int +) -> tuple[list[list[int]], list[list[float]]]: + """Select the head from k+2 candidates, or keep every recorded support token.""" + content = choice["logprobs"]["content"] + sampling_mask = _parse_sampling_mask(choice, completion_ids) + top_ids: list[list[int]] = [] + top_logprobs: list[list[float]] = [] + for index, (sampled_id, entry) in enumerate( + zip(completion_ids, content, strict=True) + ): + candidates = entry.get("top_logprobs") + if not isinstance(candidates, list) or not candidates: + raise MalformedGenerateResponseError("Missing requested top logprobs") + head: dict[int, float] = {} + for candidate in candidates: + if not isinstance(candidate, Mapping): + raise MalformedGenerateResponseError("Top logprob must be an object") + token, logprob = candidate.get("token"), candidate.get("logprob") + if ( + not isinstance(token, str) + or not token.startswith("token_id:") + or not token[9:].isdigit() + ): + raise MalformedGenerateResponseError( + "Top logprob token must be token_id:" + ) + if ( + isinstance(logprob, bool) + or not isinstance(logprob, (int, float)) + or not math.isfinite(logprob) + ): + raise MalformedGenerateResponseError("Top logprob must be finite") + if logprob > 1e-6: + raise MalformedGenerateResponseError("Top logprob cannot be positive") + token_id = int(token[9:]) + if token_id in head: + raise MalformedGenerateResponseError("Duplicate top logprob token id") + head[token_id] = float(logprob) + selected_ids, selected_logprobs = _select_sampler_head( + head, + sampled_id, + entry["logprob"], + k, + sampling_mask[index] if sampling_mask is not None else None, + ) + top_ids.append(selected_ids) + top_logprobs.append(selected_logprobs) + return top_ids, top_logprobs + + async def generate( *, client: AsyncOpenAI, @@ -207,9 +411,17 @@ async def generate( ) -> dict[str, Any]: """Tokenize messages, call vLLM /inference/v1/generate, parse the response. - ``sampling_params`` is forwarded to vLLM verbatim. Two fields are always - set by us and override caller values: ``stop_token_ids`` (from the - renderer) and ``logprobs=1`` (we always emit completion_logprobs). Pass + ``sampling_params`` is forwarded to vLLM, with ``stop_token_ids`` taken + from the renderer and ``logprobs`` set to at least one for sampled-token + evidence. A positive integer requests a top-k head; we ask the engine + for two extra candidates to select that head independently of the sampled + action, excluding ties at its boundary. When the engine returns sampling + masks, capture every recorded support probability, including cutoff ties. + Head requests opt into compact numeric HTTP logprobs; set + ``extra_args.prl_compact_logprobs=False`` to use native JSON. Servers without + the extension can still return native JSON. Sampled-only requests keep the + native path by default. + Pass ``prompt_ids`` to skip rendering and use a prebuilt token sequence — pair it with ``multi_modal_data`` when the prebuilt prompt has image / video placeholders that need engine-side mm payload, and with @@ -242,7 +454,10 @@ async def generate( Returns a dict with: request_id, prompt_ids, renderer_prompt_ids, mm_placeholders, completion_ids, completion_logprobs, content, reasoning_content, tool_calls, finish_reason, routed_experts, - multi_modal_data, prompt_attribution. ``renderer_prompt_ids`` is the + multi_modal_data, prompt_attribution. When ``sampling_params["logprobs"]`` + is a positive int, also ``completion_top_ids`` / ``completion_top_logprobs`` + (original sampler probabilities on the head or complete replay support). + ``renderer_prompt_ids`` is the unexpanded logical prompt when ``process_multimodal=False`` and ``None`` otherwise. @@ -309,7 +524,22 @@ def _prepare(): sp: dict[str, Any] = dict(sampling_params or {}) sp["stop_token_ids"] = stop_token_ids - sp["logprobs"] = 1 + requested_logprobs = sp.get("logprobs") + head_size = ( + requested_logprobs + if type(requested_logprobs) is int and requested_logprobs > 0 + else None + ) + sp["logprobs"] = head_size + 2 if head_size is not None else 1 + if head_size is not None: + # Flat storage avoids millions of live Logprob objects for long completions. + sp.setdefault("flat_logprobs", True) + sp["extra_args"] = { + "prl_compact_logprobs": True, + **(sp.get("extra_args") or {}), + } + if not sp.get("stop"): + sp.setdefault("detokenize", False) sp.setdefault("skip_special_tokens", False) body: dict[str, Any] = { @@ -351,7 +581,10 @@ def _prepare(): if extra_headers: post_kwargs["options"] = cast(Any, {"headers": extra_headers}) raw_response = await client.post(endpoint, **post_kwargs) - data = parse_generate_response(raw_response.content) + data = parse_generate_response( + raw_response.content, + compact_logprobs=bool((sp.get("extra_args") or {}).get("prl_compact_logprobs")), + ) choice = (data.get("choices") or [{}])[0] completion_ids = choice.get("token_ids") or [] @@ -366,7 +599,21 @@ def _prepare(): "Engine response must include mm_placeholders when process_multimodal=False." ) - completion_logprobs = _parse_completion_logprobs(choice, completion_ids) + # Top-k sampler head (ids + logprobs), only parsed when the caller asked + # for more than the sampled token. Score centering consumes it to cancel + # trainer/sampler drift on off-policy rollouts. + completion_top_ids = None + completion_top_logprobs = None + if choice.get("compact_logprobs") is not None: + completion_logprobs, completion_top_ids, completion_top_logprobs = ( + _parse_compact_logprobs(choice, completion_ids, head_size) + ) + else: + completion_logprobs = _parse_completion_logprobs(choice, completion_ids) + if head_size is not None: + completion_top_ids, completion_top_logprobs = ( + _parse_completion_top_logprobs(choice, completion_ids, head_size) + ) parsed = renderer.parse_response( completion_ids, @@ -402,6 +649,8 @@ def _prepare(): "mm_placeholders": mm_placeholders, "completion_ids": list(completion_ids), "completion_logprobs": completion_logprobs, + "completion_top_ids": completion_top_ids, + "completion_top_logprobs": completion_top_logprobs, "content": parsed.content, "reasoning_content": parsed.reasoning_content, "tool_calls": parsed.tool_calls, diff --git a/tests/test_client.py b/tests/test_client.py index 16fdafef..8ce56271 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -12,7 +12,185 @@ RenderedTokens, ToolCallParseStatus, ) -from renderers.client import generate +from renderers.client import ( + _parse_compact_logprobs, + _parse_completion_top_logprobs, + generate, + parse_generate_response, +) + + +def _compact_choice(choice): + entries = choice["logprobs"]["content"] + ids, values, offsets = [], [], [0] + for entry in entries: + # Flat vLLM storage retains the sampled token even when it repeats in top-k. + for candidate in [entry, *entry.get("top_logprobs", [])]: + ids.append(int(candidate["token"].removeprefix("token_id:"))) + values.append(candidate["logprob"]) + offsets.append(len(ids)) + raw = ( + np.asarray(ids, dtype=" Date: Mon, 19 Oct 2026 13:19:14 +0200 Subject: [PATCH 2/3] Accept native numeric logprob token IDs from vLLM 0.32 --- renderers/client.py | 24 +++++++++++++++--------- tests/test_client.py | 30 ++++++++++++++++++++++++++---- 2 files changed, 41 insertions(+), 13 deletions(-) diff --git a/renderers/client.py b/renderers/client.py index 42fac65d..3eb76887 100644 --- a/renderers/client.py +++ b/renderers/client.py @@ -229,6 +229,16 @@ def _parse_compact_logprobs( ) +def _logprob_token_id(entry: Mapping[str, Any]) -> int | None: + if "token_id" in entry: + token_id = entry["token_id"] + return token_id if type(token_id) is int and token_id >= 0 else None + token = entry.get("token") + if isinstance(token, str) and token.startswith("token_id:") and token[9:].isdigit(): + return int(token[9:]) + return None + + def _parse_completion_logprobs( choice: Mapping[str, Any], completion_ids: list[int] ) -> list[float]: @@ -259,7 +269,7 @@ def _parse_completion_logprobs( f"Engine response choice.logprobs.content[{index}] must be an object." ) expected_token = f"token_id:{completion_ids[index]}" - if entry.get("token") != expected_token: + if _logprob_token_id(entry) != completion_ids[index]: raise MalformedGenerateResponseError( f"Engine response choice.logprobs.content[{index}].token must be {expected_token!r}." ) @@ -336,15 +346,12 @@ def _parse_completion_top_logprobs( for candidate in candidates: if not isinstance(candidate, Mapping): raise MalformedGenerateResponseError("Top logprob must be an object") - token, logprob = candidate.get("token"), candidate.get("logprob") - if ( - not isinstance(token, str) - or not token.startswith("token_id:") - or not token[9:].isdigit() - ): + token_id = _logprob_token_id(candidate) + if token_id is None: raise MalformedGenerateResponseError( - "Top logprob token must be token_id:" + "Top logprob must identify a nonnegative token id" ) + logprob = candidate.get("logprob") if ( isinstance(logprob, bool) or not isinstance(logprob, (int, float)) @@ -353,7 +360,6 @@ def _parse_completion_top_logprobs( raise MalformedGenerateResponseError("Top logprob must be finite") if logprob > 1e-6: raise MalformedGenerateResponseError("Top logprob cannot be positive") - token_id = int(token[9:]) if token_id in head: raise MalformedGenerateResponseError("Duplicate top logprob token id") head[token_id] = float(logprob) diff --git a/tests/test_client.py b/tests/test_client.py index 26b7614e..22ab77a7 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -27,7 +27,11 @@ def _compact_choice(choice): for entry in entries: # Flat vLLM storage retains the sampled token even when it repeats in top-k. for candidate in [entry, *entry.get("top_logprobs", [])]: - ids.append(int(candidate["token"].removeprefix("token_id:"))) + ids.append( + candidate["token_id"] + if "token_id" in candidate + else int(candidate["token"].removeprefix("token_id:")) + ) values.append(candidate["logprob"]) offsets.append(len(ids)) raw = ( @@ -293,7 +297,10 @@ def _run_generate(client, renderer=None): ) @pytest.mark.parametrize("head_size", [None, 2]) @pytest.mark.parametrize("transport", ["native", "compact", "opt_out"]) -def test_generate_builds_request_body_and_parses_response(usage, head_size, transport): +@pytest.mark.parametrize("numeric_ids", [False, True]) +def test_generate_builds_request_body_and_parses_response( + usage, head_size, transport, numeric_ids +): client = _FakeClient() client.usage = usage renderer = _FakeRenderer() @@ -307,6 +314,10 @@ def test_generate_builds_request_body_and_parses_response(usage, head_size, tran {"token": "token_id:21", "logprob": -5.0}, {"token": "token_id:22", "logprob": -6.0}, ] + if numeric_ids: + for entry in client.choice["logprobs"]["content"]: + for candidate in [entry, *entry.get("top_logprobs", [])]: + candidate["token_id"] = int(candidate.pop("token")[9:]) extra_args = ( {"prl_compact_logprobs": False, "custom_arg": 7} if transport == "opt_out" @@ -516,9 +527,20 @@ def test_generate_rejects_vllm_missing_logprob_sentinel(): _run_generate(client) -def test_generate_rejects_logprob_token_id_mismatch(): +@pytest.mark.parametrize( + "identity", + [ + {"token": "token_id:8"}, + {"token_id": 8}, + {"token_id": True}, + {"token_id": -1}, + {"token_id": "7"}, + {"token_id": None}, + ], +) +def test_generate_rejects_logprob_token_id_mismatch(identity): client = _FakeClient() - client.choice["logprobs"]["content"][0]["token"] = "token_id:8" + client.choice["logprobs"]["content"][0].update(identity) with pytest.raises( MalformedGenerateResponseError, From 5ffade9092d3ec9cdf438d542a9314114ce21537 Mon Sep 17 00:00:00 2001 From: Phi Date: Mon, 19 Oct 2026 17:56:24 +0000 Subject: [PATCH 3/3] fix(client): preserve complete sampling replay support --- renderers/client.py | 75 +++++++++++++++++++++++++++++++++++++++++--- tests/test_client.py | 39 +++++++++++++++++++++++ 2 files changed, 109 insertions(+), 5 deletions(-) diff --git a/renderers/client.py b/renderers/client.py index 3eb76887..88a151fc 100644 --- a/renderers/client.py +++ b/renderers/client.py @@ -201,7 +201,10 @@ def _parse_compact_logprobs( sampled = [] top_ids: list[list[int]] = [] top_logprobs: list[list[float]] = [] - for token_id, start, end in zip(completion_ids, offsets, offsets[1:]): + supports = _parse_sampling_mask(choice, completion_ids) if k is not None else None + for index, (token_id, start, end) in enumerate( + zip(completion_ids, offsets, offsets[1:]) + ): # FlatLogprobs can repeat the sampled ID. Match vLLM's dict order/value semantics. items = list(zip(ids[start:end].tolist(), values[start:end].tolist())) row = dict(items) @@ -219,6 +222,7 @@ def _parse_compact_logprobs( token_id, row[token_id], k, + supports[index] if supports is not None else None, ) top_ids.append(selected_ids) top_logprobs.append(selected_logprobs) @@ -296,13 +300,57 @@ def _parse_completion_logprobs( return completion_logprobs +def _parse_sampling_mask( + choice: Mapping[str, Any], completion_ids: list[int] +) -> list[list[int]] | None: + mask = choice.get("sampling_mask") + if mask is None: + return None + if isinstance(mask, Mapping): + arrays = {} + for name in ("ids", "counts"): + packed = mask.get(name) + if not isinstance(packed, Mapping) or packed.get("dtype") != "int32": + raise MalformedGenerateResponseError("Invalid packed sampling mask") + try: + values = np.frombuffer( + base64.b64decode(packed["data"], validate=True), dtype=np.int32 + ) + except (KeyError, TypeError, ValueError, binascii.Error) as exc: + raise MalformedGenerateResponseError( + "Invalid packed sampling mask" + ) from exc + if packed.get("shape") != [len(values)] or (values < 0).any(): + raise MalformedGenerateResponseError("Invalid packed sampling mask") + arrays[name] = values + counts, ids = arrays["counts"], arrays["ids"] + if len(counts) != len(completion_ids) or int(counts.sum()) != len(ids): + raise MalformedGenerateResponseError("Sampling mask token count mismatch") + mask = [row.tolist() for row in np.split(ids, np.cumsum(counts)[:-1])] + if not len(counts): + mask = [] + if not isinstance(mask, list) or len(mask) != len(completion_ids): + raise MalformedGenerateResponseError("Sampling mask token count mismatch") + for sampled, support in zip(completion_ids, mask, strict=True): + if ( + not isinstance(support, list) + or not support + or any(type(i) is not int or i < 0 for i in support) + or len(set(support)) != len(support) + or sampled not in support + ): + raise MalformedGenerateResponseError("Invalid recorded sampling support") + return mask + + def _select_sampler_head( head: dict[int, float], sampled_id: int, sampled_logprob: float, k: int, + support: list[int] | None = None, ) -> tuple[list[int], list[float]]: - """Keep a top-k head independent of the sampled action, excluding cutoff ties.""" + """Keep an action-independent head, or the complete recorded replay support.""" if sampled_id in head and not math.isclose( head[sampled_id], sampled_logprob, rel_tol=1e-6, abs_tol=1e-5 ): @@ -314,14 +362,27 @@ def _select_sampler_head( ((i, p) for i, p in head.items() if p != VLLM_LOGPROB_SENTINEL), key=lambda x: (-x[1], x[0]), ) - if len(ranked) > k and ranked[k - 1][1] == ranked[k][1]: + if support is not None: + support_set = set(support) + if not support_set.issubset(i for i, _ in ranked): + raise MalformedGenerateResponseError( + "Missing sampling-support probabilities" + ) + ranked = [(i, p) for i, p in ranked if i in support_set] + if not math.isclose( + math.fsum(math.exp(p) for _, p in ranked), 1.0, rel_tol=0, abs_tol=1e-4 + ): + raise MalformedGenerateResponseError( + "Sampler probabilities must normalize over the recorded support" + ) + elif len(ranked) > k and ranked[k - 1][1] == ranked[k][1]: boundary = ranked[k][1] ranked = [(i, p) for i, p in ranked if p > boundary] else: ranked = ranked[:k] if not ranked or math.fsum(math.exp(p) for _, p in ranked) > 1.0001: raise MalformedGenerateResponseError("Invalid sampler head probability mass") - if len(head) < k + 2: + if support is None and len(head) < k + 2: raise MalformedGenerateResponseError( "Engine returned fewer top logprobs than requested" ) @@ -338,7 +399,10 @@ def _parse_completion_top_logprobs( raise MalformedGenerateResponseError("Missing requested top logprobs") top_ids: list[list[int]] = [] top_logprobs: list[list[float]] = [] - for sampled_id, entry in zip(completion_ids, content, strict=True): + supports = _parse_sampling_mask(choice, completion_ids) + for index, (sampled_id, entry) in enumerate( + zip(completion_ids, content, strict=True) + ): candidates = entry.get("top_logprobs") if not isinstance(candidates, list) or not candidates: raise MalformedGenerateResponseError("Missing requested top logprobs") @@ -368,6 +432,7 @@ def _parse_completion_top_logprobs( sampled_id, entry["logprob"], k, + supports[index] if supports is not None else None, ) top_ids.append(selected_ids) top_logprobs.append(selected_logprobs) diff --git a/tests/test_client.py b/tests/test_client.py index 22ab77a7..6e1baac8 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -163,6 +163,45 @@ def test_top_logprobs_head_is_independent_of_sampled_action( _parse_test_head({"logprobs": {"content": [entry]}}, [sampled], 2, compact) +@pytest.mark.parametrize("compact", [False, True]) +@pytest.mark.parametrize("packed", [False, True]) +@pytest.mark.parametrize("error", [None, "missing", "mass", "mask"]) +def test_top_logprobs_preserve_complete_replay_support(compact, packed, error): + import math + + probabilities = {1: 0.6, 2: 0.2, 3: 0.2} + if error == "mass": + probabilities[1] = 0.5 + entry = { + "token": "token_id:3", + "logprob": math.log(probabilities[3]), + "top_logprobs": [ + {"token": f"token_id:{i}", "logprob": math.log(p)} + for i, p in probabilities.items() + if error != "missing" or i != 2 + ], + } + support = [1, 2] if error == "mask" else [3, 1, 2] + mask = [support] + if packed: + mask = { + name: { + "data": base64.b64encode(np.asarray(values, dtype=np.int32)).decode(), + "shape": [len(values)], + "dtype": "int32", + } + for name, values in (("ids", support), ("counts", [len(support)])) + } + choice = {"logprobs": {"content": [entry]}, "sampling_mask": mask} + if error is not None: + with pytest.raises(MalformedGenerateResponseError): + _parse_test_head(choice, [3], 2, compact) + else: + ids, logprobs = _parse_test_head(choice, [3], 2, compact) + assert ids == [[1, 2, 3]] + np.testing.assert_allclose(np.exp(logprobs), [[0.6, 0.2, 0.2]]) + + @pytest.mark.parametrize("compact", [False, True]) def test_top_logprobs_preserve_filtered_candidates(compact): entry = {