feat: parse kept_tokens sampling masks from /generate responses

renderers/base.py

@@ -1937,6 +1937,7 @@ def build_trajectory_step(
         "completion_mask": [True] * len(completion_ids),
         "completion_logprobs": [0.0] * len(completion_ids),
         "routed_experts": None,
+        "kept_tokens": None,
     }
     if (
         full_rendered.multi_modal_data is not None

renderers/client.py

@@ -32,6 +32,7 @@
 
 _request_logger = logging.getLogger("renderers.client")
 ROUTED_EXPERTS_DATA_PREFIX = b'"routed_experts":{"data":"'
+KEPT_TOKENS_IDS_PREFIX = b'"kept_tokens":{"ids":"'
 
 
 class OverlongPromptError(Exception):
@@ -121,23 +122,35 @@ async def _maybe_offload(renderer: Renderer | RendererPool, fn):
     return fn()
 
 
-def strip_routed_experts_data(raw: bytes) -> tuple[bytes, memoryview | None]:
-    data_start = raw.find(ROUTED_EXPERTS_DATA_PREFIX)
+def _strip_base64_field(raw: bytes, prefix: bytes) -> tuple[bytes, memoryview | None]:
+    """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(ROUTED_EXPERTS_DATA_PREFIX)
+    data_start += len(prefix)
     data_end = raw.index(b'"', data_start)
-    routed_data = memoryview(raw)[data_start:data_end]
+    data = memoryview(raw)[data_start:data_end]
     stripped = raw[:data_start] + raw[data_end:]
-    return stripped, routed_data
+    return stripped, data
+
+
+def strip_routed_experts_data(raw: bytes) -> tuple[bytes, memoryview | None]:
+    return _strip_base64_field(raw, ROUTED_EXPERTS_DATA_PREFIX)
 
 
 def parse_generate_response(raw: bytes) -> dict[str, Any]:
     stripped, routed_data = strip_routed_experts_data(raw)
+    stripped, kept_ids_data = _strip_base64_field(stripped, KEPT_TOKENS_IDS_PREFIX)
     payload: dict[str, Any] = json.loads(stripped)
     if routed_data is not None:
         payload["choices"][0]["routed_experts"]["data"] = routed_data
+    if kept_ids_data is not None:
+        payload["choices"][0]["kept_tokens"]["ids"] = kept_ids_data
     return payload
 
 
@@ -293,6 +306,7 @@ def _prepare():
     completion_logprobs = [float(c.get("logprob") or 0.0) for c in content_lp or []]
 
     routed_experts = choice.get("routed_experts")
+    kept_tokens = choice.get("kept_tokens")
 
     # /inference/v1/generate returns finish_reason in {"stop","length",...} —
     # never "tool_calls" (a chat-completions concept). Promote stop→tool_calls
@@ -319,6 +333,7 @@ def _prepare():
         "tool_calls": parsed.tool_calls,
         "finish_reason": finish_reason,
         "routed_experts": routed_experts,
+        "kept_tokens": kept_tokens,
         # The mm sidecar consumed on the request side, surfaced back so
         # callers can persist it on the trajectory step for downstream
         # multi-turn bridging and training-sample construction.