[Feat] RL: Add radix transfer plugins

.github/workflows/_pr-test-stage-cpu.yml

@@ -118,7 +118,7 @@ jobs:
         env:
           UV_SYSTEM_PYTHON: "1"
         run: |
-          uv pip install -e "python[dev]" --index-strategy unsafe-best-match --prerelease allow
+          uv pip install -e "python[dev]" -e radix_transfer_plugins --index-strategy unsafe-best-match --prerelease allow
 
       # Hosted runners are ephemeral, so models are re-fetched every run and the
       # Hub occasionally returns 429s. Persist the HF cache in GitHub's cache

docs/docs/advanced_features/server_arguments.mdx

@@ -2869,6 +2869,18 @@ Combining `--enable-response-store` with `--disaggregation-mode=prefill` or `dec
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`False`</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>bool flag (set to enable)</td>
     </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--output-store-backend`</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Storage backend (a radix_transfer_plugins backend name) for the routed experts, indexer top-k, and sampling masks of requests that set return_outputs_via_store; their responses carry meta_info.output_store_ref instead of the arrays. 'none' disables it.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`none`</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: str</td>
+    </tr>
+    <tr>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--output-store-backend-extra-config`</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>JSON object with the --output-store-backend config. For mooncake: master_server_address, local_hostname, local_buffer_size, key_prefix, and optionally protocol, metadata_server, device_name, replica_num. MOONCAKE_* env vars are not read.</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Type: str</td>
+    </tr>
     <tr>
       <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--disable-attn-tp-gather`</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Disable scheduler-side attn_tp_gather (the upstream SP path that pads num_tokens to attn_tp_size and pre-allocates a gathered buffer). Use for models that manage SP scatter/gather at the model level (e.g., perform their own all_gather/reduce_scatter inside attention) and do not consume the upstream gathered_buffer. Without this, the cuda graph runner pads num_tokens to attn_tp_size, which can cause kernel autotuners to select wrong-sized variants at small batches.</td>

python/sglang/srt/arg_groups/fields/exec_.py

@@ -148,6 +148,20 @@ class ExecFeatures(msgspec.Struct):
         "Use the same value on disaggregated prefill and decode nodes; clients "
         "should set top_k below the limit to leave headroom for cutoff ties.",
     ] = 4096
+    output_store_backend: A[
+        str,
+        "Storage backend (a radix_transfer_plugins backend name) for the routed "
+        "experts, indexer top-k, and sampling masks of requests that set "
+        "return_outputs_via_store; their responses carry meta_info.output_store_ref "
+        "instead of the arrays. 'none' disables it.",
+    ] = "none"
+    output_store_backend_extra_config: A[
+        Optional[str],
+        "JSON object with the --output-store-backend config. For mooncake: "
+        "master_server_address, local_hostname, local_buffer_size, key_prefix, and "
+        "optionally protocol, metadata_server, device_name, replica_num. "
+        "MOONCAKE_* env vars are not read.",
+    ] = None
     disable_outlines_disk_cache: A[
         bool,
         "Disable disk cache of outlines to avoid possible crashes related to file system or high concurrency.",

python/sglang/srt/entrypoints/engine.py

@@ -1358,6 +1358,8 @@ def sigterm_handler(signum, frame):
         finally:
             if not startup_complete and isinstance(tokenizer_manager, TokenizerManager):
                 tokenizer_manager.cuda_vmm_feature_transport.shutdown()
+                if tokenizer_manager.output_store is not None:
+                    tokenizer_manager.output_store.shutdown()
 
         return (
             tokenizer_manager,
@@ -1445,6 +1447,8 @@ def shutdown(self):
                 if mm_processor is not None:
                     mm_processor.shutdown()
                 self.tokenizer_manager.cuda_vmm_feature_transport.shutdown()
+                if self.tokenizer_manager.output_store is not None:
+                    self.tokenizer_manager.output_store.shutdown()
             return
         shm = getattr(self, "_multi_tokenizer_shm", None)
         if shm is not None:
@@ -1483,6 +1487,8 @@ def shutdown(self):
                 if mm_processor is not None:
                     mm_processor.shutdown()
                 self.tokenizer_manager.cuda_vmm_feature_transport.shutdown()
+                if self.tokenizer_manager.output_store is not None:
+                    self.tokenizer_manager.output_store.shutdown()
 
     def __enter__(self):
         return self

python/sglang/srt/entrypoints/openai/protocol.py

@@ -903,6 +903,7 @@ class ChatCompletionRequest(PDRoutingFields):
     return_routed_experts: bool = False
     routed_experts_start_len: int = 0
     return_indexer_topk: bool = False
+    return_outputs_via_store: bool = False
     return_cached_tokens_details: bool = False
     return_spec_tokens_details: bool = False
     return_prompt_token_ids: bool = False

python/sglang/srt/entrypoints/openai/serving_chat.py

@@ -100,6 +100,7 @@ class ThinkingMode(str, Enum):
     strip_structural_tag_excludes,
 )
 from sglang.srt.managers.io_struct import GenerateReqInput
+from sglang.srt.managers.output_store import OUTPUT_STORE_REF_KEY
 from sglang.srt.parser.conversation import generate_chat_conv
 from sglang.srt.parser.hunyuan_reasoning import (
     normalize_hunyuan_reasoning_effort,
@@ -1050,6 +1051,9 @@ def _validate_request(self, request: ChatCompletionRequest) -> str | None:
         if request.return_indexer_topk and not request.return_meta_info:
             return "return_indexer_topk requires return_meta_info=true."
 
+        if request.return_outputs_via_store and not request.return_meta_info:
+            return "return_outputs_via_store requires return_meta_info=true."
+
         if request.return_flat_raw_output_top_logprobs and not request.return_meta_info:
             return "return_flat_raw_output_top_logprobs requires return_meta_info=true."
 
@@ -1321,6 +1325,7 @@ def _convert_to_internal_request(
             return_routed_experts=request.return_routed_experts,
             routed_experts_start_len=request.routed_experts_start_len,
             return_indexer_topk=request.return_indexer_topk,
+            return_outputs_via_store=request.return_outputs_via_store,
             rid=request.rid,
             session_id=request.session_id,
             extra_key=request.extra_key,
@@ -2324,14 +2329,29 @@ async def _handle_non_streaming_request(
         if not isinstance(ret, list):
             ret = [ret]
 
-        response = self._build_chat_response(
-            request,
-            ret,
-            int(time.time()),
-        )
-
+        try:
+            response = self._build_chat_response(
+                request,
+                ret,
+                int(time.time()),
+            )
+        except Exception:
+            self._cleanup_stored_outputs(ret)
+            raise
+        if not isinstance(response, ChatCompletionResponse):
+            self._cleanup_stored_outputs(ret)
         return response
 
+    def _cleanup_stored_outputs(self, ret: list[dict[str, Any]]) -> None:
+        """The response is an error, so nobody can read the outputs it stored."""
+        refs = [
+            item["meta_info"][OUTPUT_STORE_REF_KEY]
+            for item in ret
+            if OUTPUT_STORE_REF_KEY in item["meta_info"]
+        ]
+        if refs:
+            self.tokenizer_manager.output_store.cleanup_refs(refs)
+
     def _build_chat_response(
         self,
         request: ChatCompletionRequest,

python/sglang/srt/managers/detokenizer_manager.py

@@ -46,6 +46,7 @@
 from sglang.srt.plugins import load_plugins
 from sglang.srt.runtime_context import (
     get_device,
+    get_exec,
     get_model,
     get_observability,
     get_serving,
@@ -156,6 +157,7 @@ def init_running_status(self, server_args: ServerArgs):
             get_serving().disable_tokenizer_batch_decode
         )
         self.is_tool_call_parser_gpt_oss = get_serving().tool_call_parser == "gpt-oss"
+        self.output_store_enabled = get_exec().features.output_store_backend != "none"
 
         self.soft_watchdog = Watchdog.create(
             debug_name="DetokenizerManager",
@@ -457,8 +459,16 @@ def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput):
             if len(recv_obj.rids) > 0
             else []
         )
-        routed_experts = self._b64_encode_per_request(recv_obj.routed_experts)
-        indexer_topk = self._b64_encode_per_request(recv_obj.indexer_topk)
+        # Requests that put these into the output store need the tensors, so the
+        # tokenizer encodes the remaining ones itself.
+        if self.output_store_enabled:
+            routed_experts = indexer_topk = None
+            routed_experts_raw = recv_obj.routed_experts
+            indexer_topk_raw = recv_obj.indexer_topk
+        else:
+            routed_experts = self._b64_encode_per_request(recv_obj.routed_experts)
+            indexer_topk = self._b64_encode_per_request(recv_obj.indexer_topk)
+            routed_experts_raw = indexer_topk_raw = None
         return BatchStrOutput(
             rids=recv_obj.rids,
             http_worker_ipcs=recv_obj.http_worker_ipcs,
@@ -500,6 +510,8 @@ def handle_batch_token_id_out(self, recv_obj: BatchTokenIDOutput):
             routed_experts=routed_experts,
             indexer_topk=indexer_topk,
             indexer_topk_num_layers=recv_obj.indexer_topk_num_layers,
+            routed_experts_raw=routed_experts_raw,
+            indexer_topk_raw=indexer_topk_raw,
             customized_info=recv_obj.customized_info,
             placeholder_tokens_idx=None,
             placeholder_tokens_val=None,

python/sglang/srt/managers/io_struct.py

@@ -238,7 +238,8 @@ class GenerateReqInput:
     # Whether to return each output token's sampling support and behavior logprob.
     return_sampling_mask: Optional[Union[List[bool], bool]] = None
     # Return either the selected token's behavior logprob or behavior logprobs for
-    # the full sampling support. When omitted, selected mode is used if
+    # the full sampling support; support mode also returns each sampled token's logprob
+    # in output_token_sampling_selected_logprobs. When omitted, selected mode is used if
     # return_sampling_mask is enabled.
     sampling_logprobs_mode: Optional[
         Union[List[Optional[SamplingLogprobsMode]], SamplingLogprobsMode]
@@ -267,6 +268,9 @@ class GenerateReqInput:
     # 0 = full sequence.
     routed_experts_start_len: int = 0
     return_indexer_topk: bool = False
+    # Return routed experts, indexer top-k, and sampling mask via --output-store-backend
+    # as meta_info.output_store_ref; one prompt, n=1, no streaming or beam search.
+    return_outputs_via_store: bool = False
 
     # The modalities of the image data [image, multi-images, video]
     modalities: Optional[List[str]] = None
@@ -434,6 +438,15 @@ def normalize_batch_and_arguments(self):
         if self.session_id is not None and self.session_params is not None:
             raise ValueError("session_id and session_params cannot both be set.")
         self._handle_parallel_sampling()
+        # TODO(guapisolo): support batch, n>1 (deleting the others' stored objects
+        # when a sub-request fails), and beam search; RL sends one prompt per request.
+        if self.return_outputs_via_store and (
+            not self.is_single or self._sampling_params_beam_width() > 1
+        ):
+            raise ValueError(
+                "return_outputs_via_store supports one prompt with n=1 "
+                "and no beam search."
+            )
 
         if self.is_single:
             self._normalize_single_inputs()
@@ -1002,6 +1015,7 @@ def __getitem__(self, i):
             return_routed_experts=self.return_routed_experts,
             routed_experts_start_len=self.routed_experts_start_len,
             return_indexer_topk=self.return_indexer_topk,
+            return_outputs_via_store=self.return_outputs_via_store,
             modalities=self.modalities[i] if self.modalities else None,
             session_params=self.session_params,
             lora_path=self.lora_path[i] if self.lora_path is not None else None,
@@ -1663,8 +1677,9 @@ class BatchTokenIDOutput(BaseBatchReq, kw_only=True):
 
     # Per-request routed experts (input + output tokens), shape
     # (token, layer, top_k). DetokenizerManager encodes to base64 into
-    # BatchStrOutput; on the skip_tokenizer_init path the scheduler sends this
-    # straight to TokenizerManager, which encodes on demand.
+    # BatchStrOutput unless --output-store-backend is set (see routed_experts_raw);
+    # on the skip_tokenizer_init path the scheduler sends this straight to
+    # TokenizerManager, which encodes on demand.
     routed_experts: Optional[List[Optional[torch.Tensor]]]
 
     indexer_topk: Optional[List[Optional[torch.Tensor]]]
@@ -1762,7 +1777,8 @@ class BatchStrOutput(BaseBatchReq, kw_only=True):
     output_hidden_states: OutputHiddenStates
 
     # Per-request routed experts, base64-encoded by DetokenizerManager off the
-    # tokenizer hot path. Underlying tensor shape is (token, layer, top_k);
+    # tokenizer hot path; None with --output-store-backend (see routed_experts_raw).
+    # Underlying tensor shape is (token, layer, top_k);
     # see BatchTokenIDOutput.routed_experts.
     routed_experts: Optional[List[Optional[str]]]
 
@@ -1822,6 +1838,11 @@ class BatchStrOutput(BaseBatchReq, kw_only=True):
     input_top_logprobs_idx_flat: Optional[List[Optional[np.ndarray]]] = None
     input_top_logprobs_flat_null_prefix: Optional[List[Optional[int]]] = None
 
+    # Unencoded routed_experts / indexer_topk, sent instead of the base64 fields when
+    # --output-store-backend is set; TokenizerManager encodes them per request.
+    routed_experts_raw: Optional[List[Optional[torch.Tensor]]] = None
+    indexer_topk_raw: Optional[List[Optional[torch.Tensor]]] = None
+
 
 class BatchEmbeddingOutput(BaseBatchReq, kw_only=True):
     # The finish reason

python/sglang/srt/managers/multi_tokenizer_mixin.py

@@ -380,6 +380,12 @@ def _handle_output_by_index(output, i):
                 output, "indexer_topk", i, check_length=False
             ),
             indexer_topk_num_layers=getattr(output, "indexer_topk_num_layers", None),
+            routed_experts_raw=_extract_field_by_index(
+                output, "routed_experts_raw", i, check_length=False
+            ),
+            indexer_topk_raw=_extract_field_by_index(
+                output, "indexer_topk_raw", i, check_length=False
+            ),
             customized_info=_extract_field_by_index(
                 output, "customized_info", i, check_length=False
             ),
@@ -469,6 +475,8 @@ def _extract_field_by_indices(
         "output_hidden_states",
         "routed_experts",
         "indexer_topk",
+        "routed_experts_raw",
+        "indexer_topk_raw",
         "token_steps",
         "customized_info",
         "dp_ranks",