From 4e8eb1457bd54560fa88ba1443b6b9b55cd3fa16 Mon Sep 17 00:00:00 2001 From: Rockdu <89372739+Rockdu@users.noreply.github.com> Date: Mon, 20 Jul 2026 11:32:02 -0700 Subject: [PATCH] [Diffusion] msgpack raw-bytes transport (drop base64/JSON) (#31565) Co-authored-by: Claude Opus 4.8 --- .../entrypoints/post_training/rollout_api.py | 21 ++++++++-- .../entrypoints/post_training/utils.py | 16 +++----- .../test/unit/test_rollout_api.py | 38 +++++++++---------- 3 files changed, 42 insertions(+), 33 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py index 32e280896..ce2a77a9c 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/rollout_api.py @@ -4,9 +4,9 @@ from __future__ import annotations from typing import Any +import msgspec import torch -from fastapi import APIRouter, HTTPException -from fastapi.responses import ORJSONResponse +from fastapi import APIRouter, HTTPException, Response from sglang.multimodal_gen.configs.sample.sampling_params import generate_request_id from sglang.multimodal_gen.runtime.entrypoints.openai.utils import build_sampling_params @@ -298,7 +298,16 @@ def _build_sampling_kwargs(request: RolloutRequest) -> dict: return {k: v for k, v in sampling_kwargs.items() if v is not None} -@router.post("/generate", response_model=list[RolloutResponse]) +@router.post( + "/generate", + response_class=Response, + responses={ + 200: { + "model": list[RolloutResponse], + "content": {"application/msgpack": {}}, + } + }, +) async def rollout_generate(request: RolloutRequest): request_id = generate_request_id() server_args = get_global_server_args() @@ -326,4 +335,8 @@ async def rollout_generate(request: RolloutRequest): rollout_responses = _build_response( request_id, request.prompt, request.seed, request.rollout, output_batch ) - return ORJSONResponse(content=[r.model_dump() for r in rollout_responses]) + payload = [r.model_dump() for r in rollout_responses] + return Response( + content=msgspec.msgpack.encode(payload), + media_type="application/msgpack", + ) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py index f01f3857b..495fadf88 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/post_training/utils.py @@ -2,7 +2,6 @@ from __future__ import annotations -import base64 from typing import Any import numpy as np @@ -10,22 +9,19 @@ import torch from safetensors.torch import load, save -def tensor_to_base64(t: torch.Tensor) -> str: - t = t.detach().contiguous().cpu() - raw = save({"t": t}) - return base64.b64encode(raw).decode("ascii") +def tensor_to_bytes(t: torch.Tensor) -> bytes: + return save({"t": t.detach().contiguous().cpu()}) -def base64_to_tensor(s: str) -> torch.Tensor: - raw = base64.b64decode(s) - return load(raw)["t"] +def bytes_to_tensor(b: bytes) -> torch.Tensor: + return load(b)["t"] def _maybe_serialize(obj: Any) -> Any: if isinstance(obj, torch.Tensor): return { "__tensor__": True, - "data": tensor_to_base64(obj), + "data": tensor_to_bytes(obj), "shape": list(obj.shape), "dtype": str(obj.dtype), } @@ -41,7 +37,7 @@ def _maybe_serialize(obj: Any) -> Any: def _maybe_deserialize(obj: Any) -> Any: if isinstance(obj, dict): if obj.get("__tensor__"): - return base64_to_tensor(obj["data"]) + return bytes_to_tensor(obj["data"]) return {k: _maybe_deserialize(v) for k, v in obj.items()} if isinstance(obj, (list, tuple)): return [_maybe_deserialize(v) for v in obj] diff --git a/python/sglang/multimodal_gen/test/unit/test_rollout_api.py b/python/sglang/multimodal_gen/test/unit/test_rollout_api.py index ce2f5c54d..791e61ad5 100644 --- a/python/sglang/multimodal_gen/test/unit/test_rollout_api.py +++ b/python/sglang/multimodal_gen/test/unit/test_rollout_api.py @@ -8,8 +8,8 @@ import torch from sglang.multimodal_gen.runtime.entrypoints.post_training.utils import ( _maybe_deserialize, _maybe_serialize, - base64_to_tensor, - tensor_to_base64, + bytes_to_tensor, + tensor_to_bytes, ) from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( @@ -20,12 +20,12 @@ from sglang.multimodal_gen.runtime.post_training.rl_dataclasses import ( ) -class TestTensorToBase64Roundtrip(unittest.TestCase): +class TestTensorToBytesRoundtrip(unittest.TestCase): def _roundtrip(self, t: torch.Tensor): - encoded = tensor_to_base64(t) - self.assertIsInstance(encoded, str) - decoded = base64_to_tensor(encoded) + encoded = tensor_to_bytes(t) + self.assertIsInstance(encoded, bytes) + decoded = bytes_to_tensor(encoded) self.assertTrue( torch.equal(t, decoded), f"Mismatch for shape={t.shape} dtype={t.dtype}" ) @@ -55,21 +55,21 @@ class TestTensorToBase64Roundtrip(unittest.TestCase): if not torch.cuda.is_available(): self.skipTest("CUDA not available") t = torch.randn(4, device="cuda") - encoded = tensor_to_base64(t) - decoded = base64_to_tensor(encoded) + encoded = tensor_to_bytes(t) + decoded = bytes_to_tensor(encoded) self.assertTrue(torch.equal(t.cpu(), decoded)) def test_non_contiguous(self): t = torch.randn(4, 6)[:, ::2] self.assertFalse(t.is_contiguous()) self._roundtrip(t.contiguous()) - decoded = base64_to_tensor(tensor_to_base64(t)) + decoded = bytes_to_tensor(tensor_to_bytes(t)) self.assertTrue(torch.equal(t.contiguous(), decoded)) def test_grad_tensor_detaches(self): t = torch.randn(3, requires_grad=True) - encoded = tensor_to_base64(t) - decoded = base64_to_tensor(encoded) + encoded = tensor_to_bytes(t) + decoded = bytes_to_tensor(encoded) self.assertFalse(decoded.requires_grad) self.assertTrue(torch.equal(t.detach(), decoded)) @@ -82,7 +82,7 @@ class TestMaybeSerialize(unittest.TestCase): self.assertTrue(result["__tensor__"]) self.assertEqual(result["shape"], [2, 3]) self.assertEqual(result["dtype"], "torch.float32") - decoded = base64_to_tensor(result["data"]) + decoded = bytes_to_tensor(result["data"]) self.assertTrue(torch.equal(t, decoded)) def test_dict_with_tensors(self): @@ -243,7 +243,7 @@ class TestBuildResponse(unittest.TestCase): self.assertEqual(resp.seed, 42) self.assertIsNotNone(resp.generated_output) self.assertIsNotNone(resp.rollout_log_probs) - lp = base64_to_tensor(resp.rollout_log_probs["data"]) + lp = bytes_to_tensor(resp.rollout_log_probs["data"]) self.assertEqual(lp.shape, ()) self.assertAlmostEqual(resp.inference_time_s, 2.5) @@ -308,12 +308,12 @@ class TestBuildResponse(unittest.TestCase): batch.metrics = self._make_metrics(1.0) resps = _build_response("rb", "p", 0, True, batch) self.assertEqual(len(resps), B) - lp0 = base64_to_tensor(resps[0].rollout_log_probs["data"]) - lp1 = base64_to_tensor(resps[1].rollout_log_probs["data"]) + lp0 = bytes_to_tensor(resps[0].rollout_log_probs["data"]) + lp1 = bytes_to_tensor(resps[1].rollout_log_probs["data"]) self.assertEqual(lp0.shape, (T,)) self.assertEqual(lp1.shape, (T,)) - g0 = base64_to_tensor(resps[0].generated_output["data"]) - g1 = base64_to_tensor(resps[1].generated_output["data"]) + g0 = bytes_to_tensor(resps[0].generated_output["data"]) + g1 = bytes_to_tensor(resps[1].generated_output["data"]) self.assertEqual(g0.shape, (1, 8, 8)) self.assertEqual(g1.shape, (1, 8, 8)) self.assertFalse(torch.equal(g0, g1)) @@ -335,8 +335,8 @@ class TestBuildResponse(unittest.TestCase): self.assertEqual(len(resps), B) self.assertIsNotNone(resps[0].dit_trajectory) self.assertIsNotNone(resps[1].dit_trajectory) - ts0 = base64_to_tensor(resps[0].dit_trajectory["timesteps"]["data"]) - ts1 = base64_to_tensor(resps[1].dit_trajectory["timesteps"]["data"]) + ts0 = bytes_to_tensor(resps[0].dit_trajectory["timesteps"]["data"]) + ts1 = bytes_to_tensor(resps[1].dit_trajectory["timesteps"]["data"]) self.assertEqual(ts0.shape, (T,)) self.assertTrue(torch.equal(ts0, ts1)) self.assertEqual(