perf(r3): accelerate routed-expert transport with packed arrays

pyproject.toml

@@ -92,6 +92,8 @@ skyrl-train = [
     "polars",
     "s3fs",
     "fastapi",
+    "orjson>=3.11.9",
+    "pybase64>=1.4.2",
     "uvicorn",
     "vllm-router; sys_platform == 'linux'",
     "pybind11",

skyrl/backends/skyrl_train/inference_servers/base.py

@@ -1,6 +1,8 @@
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Hashable, List, Optional, Tuple, TypedDict
 
+from skyrl.utils.routed_experts import RoutedExpertIndices
+
 if TYPE_CHECKING:
     from skyrl.backends.skyrl_train.weight_sync import WeightUpdateRequest
     from skyrl.backends.skyrl_train.weight_sync.transfer_strategy import (
@@ -47,7 +49,7 @@ class InferenceEngineOutput(TypedDict):
     stop_reasons: List[str]
     response_logprobs: Optional[List[List[float]]]
     prompt_logprobs: Optional[List[List[float]]]  # per-prompt-token logprobs under the current model
-    rollout_expert_indices: Optional[List[List[List[int]]]]  # [seq_len, layer_num, topk]
+    rollout_expert_indices: Optional[List[RoutedExpertIndices]]
 
 
 class InferenceEngineInterface(ABC):

skyrl/backends/skyrl_train/inference_servers/remote_inference_client.py

@@ -64,6 +64,7 @@
 )
 
 import aiohttp
+import orjson
 
 from skyrl.backends.skyrl_train.inference_servers.base import (
     InferenceEngineInput,
@@ -72,6 +73,9 @@
     MMPlaceholderRangeInfo,
     MultiModalFeatures,
 )
+from skyrl.backends.skyrl_train.inference_servers.routed_experts_wire import (
+    decode_packed_routed_experts,
+)
 from skyrl.env_vars import (
     SKYRL_GENERATE_CONCURRENCY_PER_ENGINE,
     SKYRL_HTTP_CONNECTION_LIMIT,
@@ -312,8 +316,8 @@ async def _post(self, url: str, json: Dict[str, Any], headers: Optional[Dict[str
             try:
                 async with session.post(url, json=json, headers=headers) as resp:
                     try:
-                        body = await resp.json(content_type=None)
-                    except Exception as e:
+                        body = orjson.loads(await resp.read())
+                    except orjson.JSONDecodeError as e:
                         if 400 <= resp.status < 500:
                             # Non-JSON client error (e.g. plain text 422 from vllm-router).
                             # Raise immediately — client errors won't succeed on retry.
@@ -446,15 +450,16 @@ async def _throttled_detokenize(token_ids: List[int]) -> str:
         raw_results = await asyncio.gather(*[_throttled_generate(idx) for idx in range(batch_size)])
         responses = await asyncio.gather(*[_throttled_detokenize(r["response_ids"]) for r in raw_results])
 
-        rollout_expert_indices = [r.get("routed_experts") for r in raw_results]
-        has_routed_experts = any(x is not None for x in rollout_expert_indices)
+        rollout_expert_indices = (
+            [result["routed_experts"] for result in raw_results] if self.enable_return_routed_experts else None
+        )
 
         return InferenceEngineOutput(
             responses=responses,
             stop_reasons=[r["stop_reason"] for r in raw_results],
             response_ids=[r["response_ids"] for r in raw_results],
             response_logprobs=[r["response_logprobs"] for r in raw_results] if get_logprobs else None,
-            rollout_expert_indices=rollout_expert_indices if has_routed_experts else None,
+            rollout_expert_indices=rollout_expert_indices,
         )
 
     async def _generate_single(
@@ -511,7 +516,12 @@ async def _generate_single(
             if logprobs_content:
                 response_logprobs = [logprob_info["logprob"] for logprob_info in logprobs_content]
 
-        routed_experts = choice.get("routed_experts")
+        routed_experts = None
+        if self.enable_return_routed_experts:
+            packed_routed_experts = choice.get("routed_experts")
+            if not isinstance(packed_routed_experts, dict):
+                raise ValueError("/skyrl/v1/generate must return packed routed_experts")
+            routed_experts = decode_packed_routed_experts(packed_routed_experts)
 
         return {
             "stop_reason": stop_reason,

skyrl/backends/skyrl_train/inference_servers/routed_experts_wire.py

@@ -0,0 +1,45 @@
+"""Compact routed-expert HTTP payloads."""
+
+import math
+from typing import Any
+
+import numpy as np
+import pybase64
+
+from skyrl.utils.routed_experts import (
+    ROUTED_EXPERT_DTYPES,
+    RoutedExpertIndices,
+    compact_routed_expert_indices,
+)
+
+_DTYPES = {dtype.name: dtype for dtype in ROUTED_EXPERT_DTYPES}
+
+
+def pack_routed_experts(routed_experts: RoutedExpertIndices) -> dict[str, Any]:
+    compact = compact_routed_expert_indices(routed_experts)
+    return {
+        "data": pybase64.b64encode(memoryview(compact)).decode("ascii"),
+        "shape": list(compact.shape),
+        "dtype": compact.dtype.name,
+    }
+
+
+def decode_packed_routed_experts(payload: dict[str, Any]) -> RoutedExpertIndices:
+    if not isinstance(payload, dict):
+        raise TypeError("packed routed expert indices must be an object")
+    try:
+        dtype = _DTYPES[payload["dtype"]]
+        shape = tuple(payload["shape"])
+        data = pybase64.b64decode_as_bytearray(payload["data"], validate=True)
+    except (KeyError, TypeError, ValueError) as exc:
+        raise ValueError("invalid packed routed_experts payload") from exc
+    if len(shape) != 3 or any(type(dim) is not int or dim < 0 for dim in shape):
+        raise ValueError(f"invalid packed routed_experts shape: {shape}")
+    expected_size = math.prod(shape) * dtype.itemsize
+    if len(data) != expected_size:
+        raise ValueError(f"packed routed_experts has {len(data)} bytes, expected {expected_size}")
+    decoded = np.frombuffer(data, dtype=dtype).reshape(shape)
+    compact = compact_routed_expert_indices(decoded)
+    if compact.dtype != dtype:
+        raise ValueError(f"packed routed_experts uses non-canonical dtype {dtype.name}; expected {compact.dtype.name}")
+    return compact

skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py

@@ -4,15 +4,18 @@
 
 import asyncio
 import logging
+import math
 import os
 import time
 from argparse import Namespace
 from typing import List, Optional, Tuple
 
 import httpx
+import numpy as np
+import orjson
 import uvicorn
 import vllm.envs as envs
-from fastapi import HTTPException, Request
+from fastapi import HTTPException, Request, Response
 from ray.util.placement_group import PlacementGroup
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.engine.async_llm_engine import AsyncLLMEngine
@@ -34,6 +37,9 @@
     get_node_ip,
 )
 from skyrl.backends.skyrl_train.inference_servers.protocols import ServerActorProtocol
+from skyrl.backends.skyrl_train.inference_servers.routed_experts_wire import (
+    pack_routed_experts,
+)
 from skyrl.env_vars import (
     SKYRL_HTTP_CONNECTION_LIMIT,
     SKYRL_VLLM_DP_PORT_OFFSET,
@@ -424,20 +430,20 @@ async def _skyrl_generate(request: Request):
                 content = []
                 for tid, lp_dict in zip(token_ids_out, resp.logprobs):
                     if lp_dict and tid in lp_dict:
-                        content.append({"logprob": lp_dict[tid].logprob})
+                        logprob = lp_dict[tid].logprob
+                        if not math.isfinite(logprob):
+                            raise ValueError("Out of range float values are not JSON compliant")
+                        content.append({"logprob": logprob})
                     else:
                         # -9999.0 is the default in vLLM's ChatCompletionLogProb
                         content.append({"logprob": -9999.0})
                 logprobs = {"content": content}
 
             routed_experts = None
             if resp.routed_experts is not None:
-                if hasattr(resp.routed_experts, "tolist"):
-                    routed_experts = resp.routed_experts.tolist()
-                else:
-                    routed_experts = resp.routed_experts
+                routed_experts = pack_routed_experts(np.asarray(resp.routed_experts))
 
-            return {
+            payload = {
                 "choices": [
                     {
                         "token_ids": token_ids_out,
@@ -447,6 +453,7 @@ async def _skyrl_generate(request: Request):
                     }
                 ]
             }
+            return Response(content=orjson.dumps(payload), media_type="application/json")
 
     async def shutdown(self) -> None:
         """Gracefully shutdown the server."""

skyrl/backends/skyrl_train/training_batch.py

@@ -7,7 +7,9 @@
 
 import numpy as np
 import torch
-from jaxtyping import Float, Integer
+from jaxtyping import Bool, Float, Integer
+
+from skyrl.utils.routed_experts import make_replay_padding_indices
 
 DictType = TypeVar("DictType")
 
@@ -476,6 +478,7 @@ class TrainingInput(TypedDict, total=False):
     rewards: Optional[Float[torch.Tensor, "batch_size seq_len"]]
     rollout_logprobs: Optional[Float[torch.Tensor, "batch_size seq_len"]]
     rollout_expert_indices: Optional[Integer[torch.Tensor, "batch_size seq_len layer_num topk"]]
+    router_padding_mask: Optional[Bool[torch.Tensor, "batch_size seq_len"]]
     pixel_values: Optional[TensorList]  # list of `batch_size` [num_patches_i, dim] tensors
     image_grid_thw: Optional[TensorList]  # list of `batch_size` [num_images_i, 3] tensors
 
@@ -524,6 +527,18 @@ def pad_training_input_batch(unpadded_batch: TrainingInputBatch, pad_size: int)
             additional_dims = tensor.shape[1:]
             padding_tensor = torch.zeros(pad_size, *additional_dims, dtype=tensor.dtype, device=tensor.device)
             new_tensors[key] = torch.cat([tensor, padding_tensor], dim=0)
+        elif key == "rollout_expert_indices":
+            additional_dims = tensor.shape[1:]
+            padding_tensor = make_replay_padding_indices(
+                (pad_size, *additional_dims),
+                dtype=tensor.dtype,
+                device=tensor.device,
+            )
+            new_tensors[key] = torch.cat([tensor, padding_tensor], dim=0)
+        elif key == "router_padding_mask":
+            additional_dims = tensor.shape[1:]
+            padding_tensor = torch.ones(pad_size, *additional_dims, dtype=torch.bool, device=tensor.device)
+            new_tensors[key] = torch.cat([tensor, padding_tensor], dim=0)
         else:
             # Copy row 0 `pad_size` times. Loss masked so values don't affect the loss. Just need valid shape/dtype.
             assert tensor.shape[0] > 0, f"Cannot pad empty tensor field {key!r}"