renderers 0.1.13.dev8__tar.gz → 0.1.13.dev13__tar.gz

{renderers-0.1.13.dev8 → renderers-0.1.13.dev13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: renderers
-Version: 0.1.13.dev8
+Version: 0.1.13.dev13
 Summary: Chat template renderers — deterministic message-to-token conversion for LLM training
 License-Expression: Apache-2.0
 License-File: LICENSE

{renderers-0.1.13.dev8 → renderers-0.1.13.dev13}/renderers/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.1.13.dev8'
-__version_tuple__ = version_tuple = (0, 1, 13, 'dev8')
+__version__ = version = '0.1.13.dev13'
+__version_tuple__ = version_tuple = (0, 1, 13, 'dev13')
 
 __commit_id__ = commit_id = None

{renderers-0.1.13.dev8 → renderers-0.1.13.dev13}/renderers/client.py

@@ -8,10 +8,12 @@ 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, pairwise
 from typing import Any, cast
 
 import httpx
@@ -93,7 +95,7 @@ async def _resolve_max_prompt_len(client: AsyncOpenAI, model: str) -> int | None
             return _max_prompt_len_cache[key]
         try:
             payload = await client.get("/models", cast_to=cast(Any, dict[str, Any]))
-        except Exception as exc:
+        except Exception as exc:  # noqa: BLE001
             _request_logger.debug("max_prompt_len lookup failed: %s", exc)
             _max_prompt_len_cache[key] = None
             return None
@@ -111,28 +113,45 @@ 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
+    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) -> dict[str, Any]:
-    stripped, routed_data = _strip_base64_field(raw, ROUTED_EXPERTS_DATA_PREFIX)
+    stripped, routed_data = _strip_base64_fields(raw, ROUTED_EXPERTS_DATA_PREFIX)
+    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
 
 
@@ -150,6 +169,80 @@ def _parse_packed_completion_logprobs(
     return values.astype(np.float64).tolist()
 
 
+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 pairwise(offsets))
+    ):
+        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="<u4", count=count)
+    values = np.frombuffer(raw, dtype="<f4", count=count, offset=count * 4)
+    if not np.isfinite(values).all() or (values > 1e-6).any():
+        raise MalformedGenerateResponseError("Invalid compact logprob values")
+    sampled = []
+    top_ids: list[list[int]] = []
+    top_logprobs: list[list[float]] = []
+    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)
+        if any(row[i] != p for i, p in items):
+            raise MalformedGenerateResponseError("Repeated compact logprobs disagree")
+        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,
+                supports[index] if supports 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 _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]:
@@ -180,7 +273,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}."
             )
@@ -207,6 +300,145 @@ 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 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"
+            )
+        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 an action-independent head from k+2 candidates."""
+    logprobs = choice.get("logprobs")
+    content = logprobs.get("content") if isinstance(logprobs, Mapping) else None
+    if not isinstance(content, list) or len(content) != len(completion_ids):
+        raise MalformedGenerateResponseError("Missing requested top logprobs")
+    top_ids: list[list[int]] = []
+    top_logprobs: list[list[float]] = []
+    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")
+        head: dict[int, float] = {}
+        for candidate in candidates:
+            if not isinstance(candidate, Mapping):
+                raise MalformedGenerateResponseError("Top logprob must be an object")
+            token_id = _logprob_token_id(candidate)
+            if token_id is None:
+                raise MalformedGenerateResponseError(
+                    "Top logprob must identify a nonnegative token id"
+                )
+            logprob = candidate.get("logprob")
+            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")
+            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,
+            supports[index] if supports 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,
@@ -226,9 +458,16 @@ 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.
+    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
@@ -267,6 +506,10 @@ async def generate(
     is vLLM's (>= 0.32) list of ``{version, start, end}`` spans over
     completion-token indexes, or ``None`` when the engine labels none.
 
+    When ``sampling_params["logprobs"]`` is a positive int, also
+    ``completion_top_ids`` / ``completion_top_logprobs``
+    (original sampler probabilities, retaining full-vocabulary normalization).
+
     ``prompt_attribution`` is the renderer's :class:`RenderedTokens` for
     the prompt — either the one this call computed via
     ``renderer.render(...)`` or the one the caller threaded in alongside
@@ -330,7 +573,18 @@ async def generate(
 
     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:
+        sp["extra_args"] = {
+            "prl_compact_logprobs": True,
+            **(sp.get("extra_args") or {}),
+        }
     sp.setdefault("skip_special_tokens", False)
 
     body: dict[str, Any] = {
@@ -389,7 +643,18 @@ async def generate(
             "Engine response must include mm_placeholders when process_multimodal=False."
         )
 
-    completion_logprobs = _parse_completion_logprobs(choice, completion_ids)
+    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,
@@ -426,6 +691,8 @@ async def generate(
         "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,

{renderers-0.1.13.dev8 → renderers-0.1.13.dev13}/tests/test_client.py

@@ -5,6 +5,7 @@ import json
 import httpx
 import numpy as np
 import pytest
+
 from renderers import MalformedGenerateResponseError
 from renderers.base import (
     ParsedResponse,
@@ -12,7 +13,210 @@ from renderers.base import (
     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(
+                candidate["token_id"]
+                if "token_id" in candidate
+                else int(candidate["token"].removeprefix("token_id:"))
+            )
+            values.append(candidate["logprob"])
+        offsets.append(len(ids))
+    raw = (
+        np.asarray(ids, dtype="<u4").tobytes()
+        + np.asarray(values, dtype="<f4").tobytes()
+    )
+    return {
+        **choice,
+        "logprobs": None,
+        "compact_logprobs": {
+            "data": base64.b64encode(raw).decode("ascii"),
+            "format": "u32-f32-le-v1",
+            "offsets": offsets,
+            "num_top_logprobs": max(
+                (len(e.get("top_logprobs", [])) for e in entries), default=0
+            ),
+        },
+    }
+
+
+def _parse_test_head(choice, sampled, k, compact):
+    if compact:
+        _, ids, logprobs = _parse_compact_logprobs(_compact_choice(choice), sampled, k)
+        return ids, logprobs
+    return _parse_completion_top_logprobs(choice, sampled, k)
+
+
+@pytest.mark.parametrize(
+    "update",
+    [
+        {"format": "unknown"},
+        {"offsets": None},
+        {"offsets": []},
+        {"offsets": [1, 2]},
+        {"offsets": [0, 0]},
+        {"offsets": [0, -1]},
+        {"offsets": [False, 1]},
+        {"offsets": [0, 1.5]},
+        {"offsets": [0, 100]},
+        {"num_top_logprobs": True},
+        {"num_top_logprobs": -2},
+        {"data": "?"},
+        {"data": None},
+        {"data": ""},
+    ],
+)
+def test_compact_logprobs_reject_malformed_payload(update):
+    choice = _compact_choice(
+        {"logprobs": {"content": [{"token": "token_id:7", "logprob": -0.1}]}}
+    )
+    choice["compact_logprobs"].update(update)
+    with pytest.raises(MalformedGenerateResponseError):
+        _parse_compact_logprobs(choice, [7], None)
+
+
+@pytest.mark.parametrize(
+    "logprob", [float("nan"), float("inf"), float("-inf"), 0.1, -9999.0]
+)
+def test_compact_logprobs_require_valid_sampled_evidence(logprob):
+    choice = _compact_choice(
+        {"logprobs": {"content": [{"token": "token_id:7", "logprob": logprob}]}}
+    )
+    with pytest.raises(MalformedGenerateResponseError):
+        _parse_compact_logprobs(choice, [7], None)
+    choice = _compact_choice(
+        {"logprobs": {"content": [{"token": "token_id:7", "logprob": -0.1}]}}
+    )
+    with pytest.raises(MalformedGenerateResponseError, match="Missing compact sampled"):
+        _parse_compact_logprobs(choice, [8], None)
+
+
+@pytest.mark.parametrize("separators", [(",", ":"), (", ", ": ")])
+def test_parse_generate_response_preserves_multiple_compact_choices(separators):
+    choices = [
+        _compact_choice(
+            {
+                "logprobs": {"content": [{"token": f"token_id:{i}", "logprob": -0.5}]},
+                "routed_experts": {
+                    "data": base64.b64encode(bytes([i])).decode("ascii"),
+                    "shape": [1, 1, 1],
+                },
+            }
+        )
+        for i in (7, 8)
+    ]
+    choices.append(_compact_choice({"logprobs": {"content": []}}))
+    payload = parse_generate_response(
+        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)
+        assert base64.b64decode(choice["routed_experts"]["data"]) == bytes([i])
+    assert _parse_compact_logprobs(payload["choices"][2], [], 2) == ([], [], [])
+
+
+@pytest.mark.parametrize("sampled", [1, 3, 5])
+@pytest.mark.parametrize("compact", [False, True])
+@pytest.mark.parametrize(
+    "probabilities, expected_ids",
+    [
+        ({1: 0.6, 2: 0.25, 3: 0.1, 4: 0.04, 5: 0.01}, [1, 2]),
+        ({1: 0.4, 2: 0.2, 3: 0.2, 4: 0.1, 5: 0.1}, [1]),
+    ],
+)
+def test_top_logprobs_head_is_independent_of_sampled_action(
+    sampled, probabilities, expected_ids, compact
+):
+    import math
+
+    order = [sampled, *[i for i in probabilities if i != sampled]][:4]
+    entry = {
+        "token": f"token_id:{sampled}",
+        "logprob": math.log(probabilities[sampled]),
+        "top_logprobs": [
+            {"token": f"token_id:{i}", "logprob": math.log(probabilities[i])}
+            for i in order
+        ],
+    }
+    ids, logprobs = _parse_test_head(
+        {"logprobs": {"content": [entry]}}, [sampled], 2, compact
+    )
+    assert ids == [expected_ids]
+    np.testing.assert_allclose(
+        np.exp(logprobs), [[probabilities[i] for i in expected_ids]]
+    )
+
+    entry["top_logprobs"] = entry["top_logprobs"][:2]
+    with pytest.raises(MalformedGenerateResponseError, match="fewer top logprobs"):
+        _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 = {
+        "token": "token_id:1",
+        "logprob": -0.5,
+        "top_logprobs": [
+            {"token": f"token_id:{i}", "logprob": p}
+            for i, p in [(1, -0.5), (2, -2.0), (3, -9999.0)]
+        ],
+    }
+    choice = {"logprobs": {"content": [entry]}}
+    assert _parse_test_head(choice, [1], 1, compact) == ([[1]], [[-0.5]])
+    entry["top_logprobs"][0]["logprob"] = -0.6
+    with pytest.raises(MalformedGenerateResponseError, match="disagree"):
+        _parse_test_head(choice, [1], 1, compact)
 
 
 class _FakeRenderer:
@@ -130,10 +334,38 @@ def _run_generate(client, renderer=None):
         },
     ],
 )
-def test_generate_builds_request_body_and_parses_response(usage):
+@pytest.mark.parametrize("head_size", [None, 2])
+@pytest.mark.parametrize("transport", ["native", "compact", "opt_out"])
+@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()
+    sampling_params = {"temperature": 0.3, "max_tokens": 7, "min_tokens": 2}
+    if head_size is not None:
+        sampling_params["logprobs"] = head_size
+        for entry in client.choice["logprobs"]["content"]:
+            entry["top_logprobs"] = [
+                {"token": entry["token"], "logprob": entry["logprob"]},
+                {"token": "token_id:20", "logprob": -4.0},
+                {"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"
+        else None
+    )
+    if extra_args is not None:
+        sampling_params["extra_args"] = extra_args
+    if head_size is not None and transport == "compact":
+        client.choice = _compact_choice(client.choice)
 
     result = asyncio.run(
         generate(
@@ -142,7 +374,7 @@ def test_generate_builds_request_body_and_parses_response(usage):
             messages=[{"role": "user", "content": "hi"}],
             model="test-model",
             tools=[{"type": "function", "function": {"name": "echo"}}],
-            sampling_params={"temperature": 0.3, "max_tokens": 7, "min_tokens": 2},
+            sampling_params=sampling_params,
             cache_salt="ckpt-42",
         )
     )
@@ -167,7 +399,15 @@ def test_generate_builds_request_body_and_parses_response(usage):
             "max_tokens": 7,
             "min_tokens": 2,
             "stop_token_ids": [99],
-            "logprobs": 1,
+            "logprobs": head_size + 2 if head_size is not None else 1,
+            **(
+                {
+                    "extra_args": {"prl_compact_logprobs": True},
+                }
+                if head_size
+                else {}
+            ),
+            **({"extra_args": extra_args} if extra_args is not None else {}),
             "skip_special_tokens": False,
         },
     }
@@ -178,7 +418,17 @@ def test_generate_builds_request_body_and_parses_response(usage):
     assert result["reasoning_content"] == "think"
     assert result["prompt_ids"] == [1, 2, 3]
     assert result["completion_ids"] == [7, 8]
-    assert result["completion_logprobs"] == [-0.1, -0.2]
+    np.testing.assert_allclose(result["completion_logprobs"], [-0.1, -0.2])
+    if head_size is not None:
+        assert result["completion_top_ids"] == [[7, 20], [8, 20]]
+        np.testing.assert_allclose(
+            result["completion_top_logprobs"], [[-0.1, -4.0], [-0.2, -4.0]]
+        )
+    else:
+        assert result["completion_top_ids"] is None
+        assert result["completion_top_logprobs"] is None
+    if extra_args is not None:
+        assert extra_args == {"prl_compact_logprobs": False, "custom_arg": 7}
     assert result["routed_experts"]["shape"] == [2, 1, 1]
     assert isinstance(result["routed_experts"]["data"], memoryview)
     assert result["routed_experts"]["data"].tobytes() == base64.b64encode(b"\x01\x02")
@@ -316,9 +566,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,
@@ -494,6 +755,7 @@ def test_generate_serializes_multimodal_features_for_qwen_vl_family(
     pytest.importorskip("vllm", reason="vllm needed for features serialization")
 
     import torch as _torch
+
     from renderers.base import (
         MultiModalData,
         PlaceholderRange,
@@ -570,6 +832,7 @@ def test_generate_serializes_multimodal_features_for_gemma4():
     pytest.importorskip("vllm", reason="vllm needed for features serialization")
 
     import torch as _torch
+
     from renderers.base import MultiModalData, PlaceholderRange, load_tokenizer
     from renderers.gemma4 import Gemma4Renderer