[slime-router] support pd disaggregation and remove radix tree middleware

docs/en/advanced/slime-router.md

@@ -8,11 +8,10 @@ slime includes an optional slime router used during rollout / data generation. I
 
 slime router is a small FastAPI service that:
 
-- Registers workers (SGLang HTTP servers) into a local pool
-- Routes requests to a selected worker (simple least-inflight load balancing)
-- Proxies arbitrary paths to the selected worker (e.g. `/generate`)
+- Registers workers (SGLang HTTP servers) into a local pool, with support for **prefill / decode / regular** worker types
+- Routes requests to a selected worker via least-inflight load balancing or **PD dual-dispatch routing**
+- Streams proxied responses (e.g. `/generate`) without buffering the full body, improving throughput under high concurrency
 - Runs periodic health checks and quarantines unhealthy workers
-- Supports middleware plugins (via `--slime-router-middleware-paths`) to implement rollout-specific processing (e.g. caching, request/response transforms)
 
 In slime's architecture, the router is part of the rollout system ("SGLang + router") that generates samples and pushes them into the data buffer.
 
@@ -27,26 +26,9 @@ In distributed training, slime will start a router automatically when `--sglang-
 
 ## 2. Why we need slime router
 
-Unlike production inference, RL rollout needs to capture additional metadata for training: token-level logprobs, loss masks, and (for MoE models) expert routing decisions. slime router provides these capabilities through its middleware system and passthrough proxy design.
+Unlike production inference, RL rollout needs to capture additional metadata for training: token-level logprobs, loss masks, and (for MoE models) expert routing decisions. slime router provides these capabilities through its passthrough proxy design.
 
-### 2.1 Radix-tree cache (transparent token management)
-
-> Use this when your rollout pipeline is text-in/text-out and you cannot reliably persist token IDs; if you already control token-in/token-out (e.g. search r1, multiturn VLM examples), you likely don't need the radix-tree cache.
-
-Text-in text-out interfaces can cause token retokenization mismatches - re-tokenizing text at training time may produce different token sequences than rollout, breaking per-token alignment needed for PPO/GRPO losses.
-
-The radix-tree cache solves this transparently: it intercepts text-based requests, tokenizes them, and stores trajectories (text, token IDs, logprobs, loss masks) keyed by the text prefix. After rollout finishes, calling `/retrieve_from_text` returns the exact token sequence with aligned metadata, without requiring any changes to your rollout code.
-
-Implementation-wise, the radix-tree cache:
-
-- Accepts text plus tokens/metadata and stores them in a radix tree
-- Uses longest-prefix matching to reuse cached token sequences (enabling token-in/token-out downstream)
-- Allows insertion of new text continuations as rollout proceeds (multiple trajectories per prompt, e.g. GRPO)
-- Periodically cleans up stale nodes to control memory usage
-
-Use the radix cache when you have text-based rollout code and want token-level precision without rewriting, or when running GRPO with multiple trajectories sharing the same prompt prefix.
-
-### 2.2 Rollout routing replay (R3) for MoE
+### 2.1 Rollout routing replay (R3) for MoE
 
 For MoE models, slime supports rollout routing replay (R3): record expert routing decisions during rollout and replay them during training to improve stability.
 
@@ -72,6 +54,17 @@ slime consumes the routing data and replays it during training:
 
 We need slime router because the SGLang worker returns routed experts in the response (`meta_info.routed_experts`) when the request sets `return_routed_experts=true`, and slime router preserves this field end-to-end. SGLang Model Gateway may drop this extra metadata when it reconstructs responses with a fixed schema (see section 3).
 
+### 2.2 PD disaggregation
+
+slime router supports **Prefill-Decode (PD) disaggregation**. When prefill and decode workers are registered, the router automatically enables PD mode:
+
+- Workers register themselves with a `worker_type` (`prefill`, `decode`, or `regular`) via the `POST /workers` endpoint.
+- For each request, the router picks a (prefill, decode) worker pair via least-inflight load balancing, injects bootstrap information (`bootstrap_host`, `bootstrap_port`, `bootstrap_room`) into the request body, and sends the same modified request to **both** workers concurrently.
+- The decode worker's response is returned to the caller. The actual KV-cache transfer between workers is coordinated internally via the bootstrap connection.
+- If no prefill/decode workers exist, the router falls back to standard single-worker routing.
+
+This mirrors the dual-dispatch approach used by SGLang Model Gateway's PD router.
+
 ---
 
 ## 3. Differences vs SGLang Model Gateway
@@ -80,13 +73,13 @@ slime router and SGLang Model Gateway can both route requests to workers, but th
 
 ### Key differences
 
-slime router is a lightweight Python/FastAPI proxy that acts as a passthrough to SGLang workers. This passthrough design enables RL-specific features like radix-tree trajectory caching and R3 (which require preserving raw response metadata like `routed_experts`).
+slime router is a lightweight Python/FastAPI proxy that acts as a passthrough to SGLang workers. This passthrough design enables RL-specific features like R3 (which require preserving raw response metadata like `routed_experts`).
 
 SGLang Model Gateway is a high-performance Rust-based router optimized for large-scale inference: async non-blocking routing, advanced fault tolerance (retries, circuit breakers), multiple load balancing policies (including cache-aware routing), and PD disaggregation support. However, it reconstructs responses with a fixed schema, so it does not preserve the metadata needed for slime's R3 flow.
 
 For more details on SGLang Model Gateway, see the [official documentation](https://docs.sglang.io/advanced_features/sgl_model_gateway.html).
 
 ### When to use which
 
-- Use slime router when you need R3 or radix-tree caching
+- Use slime router when you need R3 or PD disaggregation with metadata preservation
 - Use SGLang Model Gateway for everything else (recommended default)

docs/en/get_started/customization.md

@@ -28,7 +28,6 @@ Below is a summary of all available customization interfaces and their purposes.
 | [`--custom-megatron-init-path`](#17-megatron-hooks) | Custom initialization after Megatron setup. |
 | [`--custom-megatron-before-log-prob-hook-path`](#17-megatron-hooks) | Custom logic before log probability computation. |
 | [`--custom-megatron-before-train-step-hook-path`](#17-megatron-hooks) | Custom logic before each training step. |
-| [`--slime-router-middleware-paths`](#18-slime-router-middleware---slime-router-middleware-paths) | Add custom middleware to slime router. |
 
 ## Detailed Interface Reference
 
@@ -400,18 +399,7 @@ def custom_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler
 
 ---
 
-### 18. slime Router Middleware (`--slime-router-middleware-paths`)
-
-**Purpose**: Add custom middleware to the slime router for request processing.
-
-**Use Cases**:
-- Request/response transformation
-- Custom routing logic
-- Caching and optimization
-
----
-
-### 19. MoE Routing Replay
+### 18. MoE Routing Replay
 
 Stabilize MoE RL training by recording and replaying expert routing decisions to ensure consistency.
 

docs/zh/advanced/slime-router.md

@@ -8,11 +8,10 @@ slime 提供一个可选的 slime router，用于 rollout / data generation 阶
 
 slime router 是一个小型 FastAPI 服务，主要能力包括：
 
-- 注册 worker（SGLang HTTP server）到本地池
-- 路由请求到选定的 worker（简单的 least-inflight load balancing）
-- 代理任意路径到选定的 worker（例如 `/generate`）
+- 注册 worker（SGLang HTTP server）到本地池，支持 **prefill / decode / regular** worker 类型
+- 路由请求到选定的 worker——支持 least-inflight 负载均衡和 **PD 双发路由**
+- 流式代理请求到选定的 worker（例如 `/generate`），不缓冲完整 response body，提高高并发下的吞吐
 - 定期 health checks，并隔离不健康的 worker
-- 支持 middleware plugins（通过 `--slime-router-middleware-paths`）实现 rollout 特定处理（例如 caching、request/response transform）
 
 在 slime 架构中，router 是 rollout 系统（"SGLang + router"）的一部分：负责生成样本并将其推入数据缓冲区。
 
@@ -27,26 +26,9 @@ slime router 是一个小型 FastAPI 服务，主要能力包括：
 
 ## 2. 为什么需要 slime router
 
-与 production inference 不同，RL rollout 往往需要捕获用于训练的额外 metadata：token-level logprobs、loss masks，以及（对 MoE 模型）expert routing decisions。slime router 通过 middleware system 和 passthrough proxy 设计提供这些能力。
+与 production inference 不同，RL rollout 往往需要捕获用于训练的额外 metadata：token-level logprobs、loss masks，以及（对 MoE 模型）expert routing decisions。slime router 通过 passthrough proxy 设计提供这些能力。
 
-### 2.1 Radix-tree cache（透明的 token 管理）
-
-> 当你的 rollout 流程是 text-in/text-out、并且很难可靠地保存 token IDs 时，适合用 radix-tree cache；如果你已经能自己控制 token-in/token-out（例如 search r1、multiturn VLM 这些 example），通常不需要 radix-tree cache。
-
-text-in text-out 接口可能导致 token retokenization mismatches：训练阶段重新 tokenize 文本，得到的 token 序列可能与 rollout 阶段不同，从而破坏 PPO/GRPO 这类方法所需的 per-token alignment。
-
-radix-tree cache 可以透明地解决这个问题：它拦截 text-based request，对其进行 tokenize，并将 trajectory（text、token IDs、logprobs、loss masks）按文本前缀作为 key 存储。rollout 结束后，调用 `/retrieve_from_text` 就能取回与 rollout 完全一致的 token 序列及其对齐的 metadata，无需修改现有 rollout 代码。
-
-实现上，radix-tree cache 会做几件事：
-
-- 接收 text 以及 tokens/metadata，并写入 radix tree
-- 通过 longest-prefix matching 复用已缓存的 token 序列（使后续流程可以走 token-in/token-out）
-- rollout 过程中持续插入新的 text continuation（同一 prompt 下可有多条 trajectory，例如 GRPO）
-- 定期清理 stale nodes，控制内存占用
-
-当你有 text-based rollout 代码、想获得 token-level 精度但又不想重写，或者在 GRPO 场景中多个 trajectory 共享相同 prompt 前缀时，建议使用 radix-tree cache。
-
-### 2.2 Rollout routing replay (R3) for MoE
+### 2.1 Rollout routing replay (R3) for MoE
 
 对 MoE 模型，slime 支持 rollout routing replay (R3)：在 rollout 期间记录 expert routing decisions，并在训练期间 replay，以提升训练稳定性。
 
@@ -72,6 +54,17 @@ slime 侧消费路由数据，并在训练中完成 replay：
 
 我们需要 slime router，是因为当请求设置 `return_routed_experts=true` 时，SGLang worker 会在响应里返回路由信息（`meta_info.routed_experts`），而 slime router 会端到端保留这个字段。SGLang Model Gateway 会用固定 schema 重建响应，可能会丢掉这类额外 metadata（细节见第 3 节）。
 
+### 2.2 PD 分离（Prefill-Decode 分离）
+
+slime router 支持 **Prefill-Decode (PD) 分离**。当 prefill 和 decode worker 注册后，router 会自动启用 PD 模式：
+
+- Worker 注册时携带 `worker_type`（`prefill`、`decode` 或 `regular`），通过 `POST /workers` 端点。
+- 对每个请求，router 通过 least-inflight 负载均衡选择一对 (prefill, decode) worker，向请求体注入 bootstrap 信息（`bootstrap_host`、`bootstrap_port`、`bootstrap_room`），然后将同一个修改后的请求**并发发送**给两个 worker。
+- Decode worker 的响应返回给调用方。实际的 KV-cache 传输由 worker 通过 bootstrap 连接内部协调完成。
+- 如果没有 prefill/decode worker 存在，router 回退到标准的单 worker 路由。
+
+这与 SGLang Model Gateway 的 PD router 所使用的双发方式一致。
+
 ---
 
 ## 3. 与 SGLang Model Gateway 的区别
@@ -80,13 +73,13 @@ slime router 与 SGLang Model Gateway 都能将请求路由到 worker，但它
 
 ### 主要区别
 
-slime router 是一个轻量级的 Python/FastAPI proxy，作为 SGLang worker 的 passthrough proxy。这种 passthrough 设计使得 RL 特定功能成为可能，例如 radix-tree trajectory caching 和 R3（需要保留原始 response metadata，如 `routed_experts`）。
+slime router 是一个轻量级的 Python/FastAPI proxy，作为 SGLang worker 的 passthrough proxy。这种 passthrough 设计使得 RL 特定功能成为可能，例如 R3（需要保留原始 response metadata，如 `routed_experts`）。
 
 SGLang Model Gateway 是一个高性能 Rust router，面向大规模 inference 优化：async non-blocking routing、高级 fault tolerance（retries、circuit breakers）、多种 load balancing policy（包括 cache-aware routing），以及 PD disaggregation 支持。但它会用固定 schema 重建响应，因此不保留 slime 的 R3 流程所需 metadata。
 
 更多关于 SGLang Model Gateway 的信息，请参阅[官方文档](https://docs.sglang.io/advanced_features/sgl_model_gateway.html)。
 
 ### 如何选择
 
-- 当你需要 R3 或 radix-tree cache 时，使用 SlimeRouter
+- 当你需要 R3 或需要保留 metadata 的 PD 分离时，使用 SlimeRouter
 - 其他情况使用 SGLang Model Gateway（推荐默认选项）

docs/zh/get_started/customization.md

@@ -28,7 +28,6 @@ slime 通过函数路径参数提供了广泛的自定义能力。这些参数
 | [`--custom-megatron-init-path`](#17-megatron-hook) | Megatron 设置后的自定义初始化。 |
 | [`--custom-megatron-before-log-prob-hook-path`](#17-megatron-hook) | log probability 计算前的自定义逻辑。 |
 | [`--custom-megatron-before-train-step-hook-path`](#17-megatron-hook) | 每个训练步骤前的自定义逻辑。 |
-| [`--slime-router-middleware-paths`](#18-slime-router-中间件---slime-router-middleware-paths) | 向 slime router 添加自定义中间件。 |
 
 ## 详细接口参考
 
@@ -402,18 +401,7 @@ def custom_hook(args, rollout_id, step_id, model, optimizer, opt_param_scheduler
 
 ---
 
-### 18. slime Router 中间件 (`--slime-router-middleware-paths`)
-
-**用途**: 向 slime router 添加自定义中间件用于请求处理。
-
-**使用场景**:
-- 请求/响应转换
-- 自定义路由逻辑
-- 缓存和优化
-
----
-
-### 19. MoE 路由重放
+### 18. MoE 路由重放
 
 通过记录和重放专家路由决策来稳定 MoE RL 训练。
 

slime/backends/sglang_utils/sglang_engine.py

@@ -195,10 +195,8 @@ def _init_normal(self, server_args_dict):
             return
 
         if self.node_rank == 0 and self.router_ip and self.router_port:
-            if parse(sglang_router.__version__) <= parse("0.2.1") or self.args.use_slime_router:
-                assert (
-                    self.worker_type == "regular"
-                ), "pd disaggregation is not supported in old router or slime router."
+            if not self.args.use_slime_router and parse(sglang_router.__version__) <= parse("0.2.1"):
+                assert self.worker_type == "regular", "pd disaggregation is not supported in old router."
                 response = requests.post(
                     f"http://{self.router_ip}:{self.router_port}/add_worker?url=http://{self.server_host}:{self.server_port}"
                 )
@@ -317,7 +315,7 @@ def shutdown(self):
         if self.worker_type != "encoder" and self.node_rank == 0:
             worker_url = f"http://{self.server_host}:{self.server_port}"
             response = None
-            if parse(sglang_router.__version__) <= parse("0.2.1") or self.args.use_slime_router:
+            if self.args.use_slime_router or parse(sglang_router.__version__) <= parse("0.2.1"):
                 response = requests.post(
                     f"http://{self.router_ip}:{self.router_port}/remove_worker?url=http://{self.server_host}:{self.server_port}"
                 )

slime/ray/rollout.py

@@ -912,9 +912,6 @@ def _start_router(args, *, has_pd_disaggregation: bool = False, force_new: bool
         router_args = copy.copy(args)
         router_args.sglang_router_ip = router_ip
         router_args.sglang_router_port = router_port
-        if has_pd_disaggregation:
-            router_args.slime_router_pd_disaggregation = True
-
     else:
         from sglang_router.launch_router import RouterArgs
 

slime/rollout/sglang_rollout.py

@@ -177,37 +177,33 @@ async def generate(args: Namespace, sample: Sample, sampling_params: dict[str, A
         if not sample.tokens:  # Initialize sample.tokens for the first turn
             sample.tokens = prompt_ids
 
-    # Use session_id for consistent hashing routing if router uses consistent_hashing policy
+    # Use session_id for consistent hashing routing (SGLang Model Gateway)
     headers = None
-    if getattr(args, "router_policy", None) == "consistent_hashing" and sample.session_id:
-        headers = {"X-SMG-Routing-Key": sample.session_id}
+    if sample.session_id:
+        if getattr(args, "router_policy", None) == "consistent_hashing":
+            headers = {"X-SMG-Routing-Key": sample.session_id}
 
     output = await post(url, payload, headers=headers)
 
-    if args.use_slime_router and "RadixTreeMiddleware" in args.slime_router_middleware_paths:
-        from slime.router.middleware_hub.radix_tree_middleware import postprocess_sample_with_radix_tree
-
-        sample = await postprocess_sample_with_radix_tree(args, sample, output)
+    if "output_token_logprobs" in output["meta_info"]:
+        new_response_tokens = [item[1] for item in output["meta_info"]["output_token_logprobs"]]
+        new_response_log_probs = [item[0] for item in output["meta_info"]["output_token_logprobs"]]
     else:
-        if "output_token_logprobs" in output["meta_info"]:
-            new_response_tokens = [item[1] for item in output["meta_info"]["output_token_logprobs"]]
-            new_response_log_probs = [item[0] for item in output["meta_info"]["output_token_logprobs"]]
-        else:
-            new_response_tokens, new_response_log_probs = [], []
+        new_response_tokens, new_response_log_probs = [], []
 
-        # Update sample with tokens directly - avoiding re-tokenization
-        sample.tokens = sample.tokens + new_response_tokens
-        sample.response_length += len(new_response_tokens)
-        sample.response += output["text"]
+    # Update sample with tokens directly - avoiding re-tokenization
+    sample.tokens = sample.tokens + new_response_tokens
+    sample.response_length += len(new_response_tokens)
+    sample.response += output["text"]
 
-        # When partial rollout and masking off policy is enabled, update the loss mask
-        if sample.loss_mask is not None:
-            assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout
-            sample.loss_mask += [1] * len(new_response_tokens)
+    # When partial rollout and masking off policy is enabled, update the loss mask
+    if sample.loss_mask is not None:
+        assert args.partial_rollout and args.mask_offpolicy_in_partial_rollout
+        sample.loss_mask += [1] * len(new_response_tokens)
 
-        if sample.rollout_log_probs is None:
-            sample.rollout_log_probs = []
-        sample.rollout_log_probs += new_response_log_probs
+    if sample.rollout_log_probs is None:
+        sample.rollout_log_probs = []
+    sample.rollout_log_probs += new_response_log_probs
 
     if "routed_experts" in output["meta_info"]:
         sample.rollout_routed_experts = np.frombuffer(