fix: eliminate JSON serialization bottleneck in r3 + NeMo-Gym rollouts

nemo_rl/environments/nemo_gym.py

@@ -27,6 +27,7 @@
     _get_node_ip_local,
 )
 from nemo_rl.environments.interfaces import EnvironmentInterface
+from nemo_rl.utils.routed_experts_codec import decode_routed_experts
 from nemo_rl.utils.timer import Timer
 
 # Kept local (not imported from models.generation) so the gym actor stays free of
@@ -366,7 +367,7 @@ def _postprocess_nemo_gym_to_nemo_rl_result(
                 routed_experts_dtype = _ROUTED_EXPERTS_DTYPES[
                     self.cfg.get("routed_experts_dtype", "int16")
                 ]
-                routed_experts = torch.as_tensor(
+                routed_experts = decode_routed_experts(
                     routed_experts_raw, dtype=routed_experts_dtype
                 )
                 if routed_experts.dim() != 3:

nemo_rl/models/generation/vllm/utils.py

@@ -22,6 +22,7 @@
     ROUTED_EXPERTS_FALLBACK_DTYPE,
     GenerationDatumSpec,
 )
+from nemo_rl.utils.routed_experts_codec import encode_routed_experts
 
 R3_MISSING_ROUTE_SENTINEL = -1
 
@@ -292,9 +293,13 @@ def attach_routed_experts_to_chat_response_choices(
                 r3_stats["actual_routes"],
                 r3_stats["expected_routes"],
             )
-        choice.message.routed_experts = routed_experts.to(
-            dtype=routed_experts_dtype
-        ).tolist()
+        # Base64 envelope instead of .tolist(): nested JSON int lists cost
+        # ~1s of CPU per serialize/parse hop at long context lengths and get
+        # re-validated at every gym HTTP hop; a single string passes through
+        # the gym chain opaquely.
+        choice.message.routed_experts = encode_routed_experts(
+            routed_experts.to(dtype=routed_experts_dtype)
+        )
 
     if len(attached_choice_indices) != len(choices):
         missing_choice_indices = sorted(

nemo_rl/utils/routed_experts_codec.py

@@ -0,0 +1,106 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Compact wire codec for router-replay routed-expert indices.
+
+Routed experts for a long-context sample are millions of ints (shape
+[tokens, num_moe_layers, topk]). Serialized as nested JSON lists they cost
+~1s of single-threaded CPU per serialize/parse hop, and every HTTP hop on
+the NeMo Gym path (model server, agent, resources server) pays that again
+for pydantic validation and re-serialization. Encoded as a single base64
+string the payload stays one opaque Python object end to end, so
+intermediate hops only pay a string copy.
+
+Envelope format (version 1):
+    "nrlre1:<dtype>:<S>x<L>x<K>:<base64 of C-contiguous array bytes>"
+
+This module must stay importable inside the NeMo Gym actor, so it may only
+depend on torch.
+"""
+
+import base64
+from typing import Any, Union
+
+import torch
+
+_MAGIC = "nrlre1"
+_WIRE_TORCH_DTYPES = {
+    "int8": torch.int8,
+    "int16": torch.int16,
+    "int32": torch.int32,
+}
+_TORCH_DTYPE_NAMES = {v: k for k, v in _WIRE_TORCH_DTYPES.items()}
+
+
+def encode_routed_experts(routed_experts: torch.Tensor) -> str:
+    """Encode a [tokens, num_moe_layers, topk] tensor as a base64 envelope.
+
+    The tensor's own dtype (int8/int16/int32, as resolved by
+    ``resolve_routed_experts_dtype``) is preserved on the wire.
+    """
+    if routed_experts.dim() != 3:
+        raise ValueError(
+            "routed_experts must have shape [tokens, num_moe_layers, topk], "
+            f"got {tuple(routed_experts.shape)}."
+        )
+    dtype_name = _TORCH_DTYPE_NAMES.get(routed_experts.dtype)
+    if dtype_name is None:
+        raise ValueError(
+            f"Unsupported routed_experts dtype {routed_experts.dtype}; "
+            f"expected one of {sorted(_WIRE_TORCH_DTYPES)}."
+        )
+    arr = routed_experts.detach().cpu().contiguous().numpy()
+    tokens, num_layers, topk = arr.shape
+    # memoryview feeds the buffer to b64encode zero-copy; tobytes() would
+    # materialize a full-size intermediate copy of the payload.
+    data = base64.b64encode(memoryview(arr)).decode("ascii")
+    return f"{_MAGIC}:{dtype_name}:{tokens}x{num_layers}x{topk}:{data}"
+
+
+def decode_routed_experts(payload: Union[str, Any], dtype: torch.dtype) -> torch.Tensor:
+    """Decode routed experts into a tensor of the requested dtype.
+
+    Accepts the base64 envelope produced by ``encode_routed_experts`` or the
+    legacy nested-list format.
+    """
+    if not isinstance(payload, str):
+        return torch.as_tensor(payload, dtype=dtype)
+    parts = payload.split(":", 3)
+    if len(parts) != 4 or parts[0] != _MAGIC:
+        raise ValueError(
+            "routed_experts string payload is not a valid "
+            f"'{_MAGIC}:<dtype>:<SxLxK>:<base64>' envelope."
+        )
+    _, dtype_name, shape_str, data = parts
+    wire_dtype = _WIRE_TORCH_DTYPES.get(dtype_name)
+    if wire_dtype is None:
+        raise ValueError(f"Unsupported routed_experts dtype '{dtype_name}'.")
+    shape = tuple(int(dim) for dim in shape_str.split("x"))
+    if len(shape) != 3:
+        raise ValueError(
+            f"routed_experts envelope shape '{shape_str}' is not 3-dimensional."
+        )
+    # Decode into a writable bytearray so torch.frombuffer can share its
+    # memory directly (frombuffer writability follows the underlying buffer;
+    # the bytearray construction is the single unavoidable copy).
+    raw = bytearray(base64.b64decode(data, validate=True))
+    expected = shape[0] * shape[1] * shape[2]
+    if len(raw) != expected * wire_dtype.itemsize:
+        raise ValueError(
+            f"routed_experts envelope has {len(raw) // wire_dtype.itemsize} "
+            f"elements, expected {expected} for shape {shape}."
+        )
+    if expected == 0:
+        return torch.empty(shape, dtype=dtype)
+    tensor = torch.frombuffer(raw, dtype=wire_dtype).reshape(shape)
+    return tensor if tensor.dtype == dtype else tensor.to(dtype)

pyrefly.toml

@@ -195,6 +195,7 @@ project-includes = [
   "nemo_rl/utils/packed_tensor.py",
   "nemo_rl/utils/prefetch_venvs.py",
   "nemo_rl/utils/r3_trace.py",
+  "nemo_rl/utils/routed_experts_codec.py",
   "nemo_rl/utils/timer.py",
   "nemo_rl/utils/venvs.py",
   "nemo_rl/utils/weight_transfer_http.py",

tests/unit/models/generation/test_vllm_utils.py

@@ -35,6 +35,11 @@
     model_dump_chat_response_with_routed_experts,
     pad_and_align_routed_expert_indices,
 )
+from nemo_rl.utils.routed_experts_codec import decode_routed_experts
+
+
+def _decoded_routes(payload: str) -> list:
+    return decode_routed_experts(payload, dtype=torch.int32).tolist()
 
 
 def _mk_inputs(batch_size: int = 2, seq_len: int = 5):
@@ -449,13 +454,15 @@ def test_attach_routed_experts_to_chat_response_choices_reassociates_by_choice_i
         device=torch.device("cpu"),
     )
 
-    assert response.choices[0].message.routed_experts == [
+    # Routes travel as a base64 string envelope, one opaque object per choice.
+    assert isinstance(response.choices[0].message.routed_experts, str)
+    assert _decoded_routes(response.choices[0].message.routed_experts) == [
         [[10]],
         [[11]],
         [[30]],
         [[0]],
     ]
-    assert response.choices[1].message.routed_experts == [
+    assert _decoded_routes(response.choices[1].message.routed_experts) == [
         [[10]],
         [[11]],
         [[31]],
@@ -517,7 +524,7 @@ def test_attach_routed_experts_to_chat_response_choices_warns_on_missing_routes(
         2,
         4,
     )
-    assert response.choices[0].message.routed_experts == [
+    assert _decoded_routes(response.choices[0].message.routed_experts) == [
         [[10]],
         [[11]],
         [[R3_MISSING_ROUTE_SENTINEL]],

tests/unit/utils/test_routed_experts_codec.py

@@ -0,0 +1,62 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+
+from nemo_rl.utils.routed_experts_codec import (
+    decode_routed_experts,
+    encode_routed_experts,
+)
+
+
+@pytest.mark.parametrize("dtype", [torch.int8, torch.int16, torch.int32])
+def test_round_trip_preserves_values_and_wire_dtype(dtype):
+    routes = torch.randint(0, 100, (7, 3, 2), dtype=torch.int32)
+    routes[2, 1, 0] = -1  # missing-route sentinel survives signed dtypes
+    routes = routes.to(dtype)
+
+    payload = encode_routed_experts(routes)
+    assert isinstance(payload, str)
+    assert payload.startswith(f"nrlre1:{str(dtype).removeprefix('torch.')}:7x3x2:")
+
+    decoded = decode_routed_experts(payload, dtype=torch.int32)
+    assert decoded.dtype == torch.int32
+    assert torch.equal(decoded, routes.to(torch.int32))
+
+
+def test_decode_accepts_legacy_nested_lists():
+    decoded = decode_routed_experts([[[1, 2]], [[3, 4]]], dtype=torch.int16)
+    assert decoded.dtype == torch.int16
+    assert decoded.tolist() == [[[1, 2]], [[3, 4]]]
+
+
+def test_decode_rejects_malformed_payloads():
+    with pytest.raises(ValueError, match="envelope"):
+        decode_routed_experts("not-an-envelope", dtype=torch.int32)
+
+    good = encode_routed_experts(torch.zeros(2, 3, 2, dtype=torch.int16))
+    magic, dtype_name, _, data = good.split(":", 3)
+    with pytest.raises(ValueError, match="expected"):
+        # Element count does not match the declared shape.
+        decode_routed_experts(f"{magic}:{dtype_name}:5x3x2:{data}", dtype=torch.int32)
+    with pytest.raises(ValueError, match="dtype"):
+        decode_routed_experts(f"{magic}:int64:2x3x2:{data}", dtype=torch.int32)
+
+
+def test_encode_rejects_bad_inputs():
+    with pytest.raises(ValueError, match="shape"):
+        encode_routed_experts(torch.zeros(3, 2, dtype=torch.int16))
+    with pytest.raises(ValueError, match="dtype"):
+        encode_routed_experts(torch.zeros(2, 3, 2, dtype=torch.int64))