From 36c4d142792b430b954385e909a16ffa0a6abdcd Mon Sep 17 00:00:00 2001 From: Jiajun Li Date: Mon, 5 Oct 2026 12:04:58 -0700 Subject: [PATCH] feat(protocol): forward return_outputs_via_store on generate and chat SGLang's Mooncake output store (sgl-project/sglang#42352) is opted into per request with return_outputs_via_store. The router re-serializes typed requests, so the flag was dropped: a Miles R3 training run with --sglang-output-store-backend mooncake silently got inline replay outputs through this router, while MilesRouter forwarded the flag. Add the field to GenerateRequest and ChatCompletionRequest as an optional bool that is serialized only when the client sets it, so requests without it are forwarded unchanged. Non-streaming responses already pass through as raw bytes, so meta_info.output_store_ref reaches the client. Co-Authored-By: Claude Opus 5.5 --- src/protocols/chat.rs | 5 +++++ src/protocols/generate.rs | 5 +++++ tests/routing/test_openai_routing.rs | 1 + tests/spec/chat_completion.rs | 5 +++++ tests/spec/generate.rs | 5 +++++ 5 files changed, 21 insertions(+) diff --git a/src/protocols/chat.rs b/src/protocols/chat.rs index 521f808..4e96c9d 100644 --- a/src/protocols/chat.rs +++ b/src/protocols/chat.rs @@ -369,6 +369,11 @@ pub struct ChatCompletionRequest { #[serde(skip_serializing_if = "Option::is_none")] pub sampling_logprobs_mode: Option, + /// Return routed experts, indexer top-k and sampling masks through the + /// server's output store instead of inline in meta_info (SGLang extension) + #[serde(skip_serializing_if = "Option::is_none")] + pub return_outputs_via_store: Option, + /// Absolute start position for returned routed experts (SGLang extension) #[serde(default)] pub routed_experts_start_len: i32, diff --git a/src/protocols/generate.rs b/src/protocols/generate.rs index 0e92ff9..a272f6c 100644 --- a/src/protocols/generate.rs +++ b/src/protocols/generate.rs @@ -102,6 +102,11 @@ pub struct GenerateRequest { #[serde(skip_serializing_if = "Option::is_none")] pub sampling_logprobs_mode: Option, + /// Return routed experts, indexer top-k and sampling masks through the + /// server's output store instead of inline in meta_info + #[serde(skip_serializing_if = "Option::is_none")] + pub return_outputs_via_store: Option, + /// Absolute start position for returned routed experts #[serde(default)] pub routed_experts_start_len: i32, diff --git a/tests/routing/test_openai_routing.rs b/tests/routing/test_openai_routing.rs index 27edbbd..3a66ab4 100644 --- a/tests/routing/test_openai_routing.rs +++ b/tests/routing/test_openai_routing.rs @@ -619,6 +619,7 @@ async fn test_unsupported_endpoints() { return_routed_experts: false, return_sampling_mask: None, sampling_logprobs_mode: None, + return_outputs_via_store: None, return_indexer_topk: false, routed_experts_start_len: 0, return_prompt_token_ids: false, diff --git a/tests/spec/chat_completion.rs b/tests/spec/chat_completion.rs index 1a1ac18..5c71741 100644 --- a/tests/spec/chat_completion.rs +++ b/tests/spec/chat_completion.rs @@ -580,6 +580,7 @@ fn test_sglang_extension_fields_roundtrip() { "return_routed_experts": true, "return_sampling_mask": true, "sampling_logprobs_mode": "support", + "return_outputs_via_store": true, "routed_experts_start_len": 10, "return_cached_tokens_details": true, "return_prompt_token_ids": true, @@ -608,6 +609,7 @@ fn test_sglang_extension_fields_roundtrip() { assert!(req.return_routed_experts); assert!(req.return_sampling_mask); assert_eq!(req.sampling_logprobs_mode.as_deref(), Some("support")); + assert_eq!(req.return_outputs_via_store, Some(true)); assert_eq!(req.routed_experts_start_len, 10); assert!(req.return_cached_tokens_details); assert!(req.return_prompt_token_ids); @@ -625,6 +627,7 @@ fn test_sglang_extension_fields_roundtrip() { assert_eq!(serialized["return_routed_experts"], true); assert_eq!(serialized["return_sampling_mask"], true); assert_eq!(serialized["sampling_logprobs_mode"], "support"); + assert_eq!(serialized["return_outputs_via_store"], true); assert_eq!(serialized["routed_experts_start_len"], 10); assert_eq!(serialized["return_cached_tokens_details"], true); assert_eq!(serialized["return_prompt_token_ids"], true); @@ -660,6 +663,7 @@ fn test_sglang_extension_fields_default_values() { assert!(!req.return_routed_experts); assert!(!req.return_sampling_mask); assert!(req.sampling_logprobs_mode.is_none()); + assert!(req.return_outputs_via_store.is_none()); assert_eq!(req.routed_experts_start_len, 0); assert!(!req.return_cached_tokens_details); assert!(!req.return_prompt_token_ids); @@ -700,6 +704,7 @@ fn test_sglang_optional_fields_omitted_when_none() { "priority", "bootstrap_host", "bootstrap_port", "bootstrap_room", "routed_dp_rank", "disagg_prefill_dp_rank", "data_parallel_rank", "sampling_logprobs_mode", + "return_outputs_via_store", ]; for field in omitted { assert!( diff --git a/tests/spec/generate.rs b/tests/spec/generate.rs index 90a5d5c..5072f36 100644 --- a/tests/spec/generate.rs +++ b/tests/spec/generate.rs @@ -9,6 +9,7 @@ fn test_generate_extension_fields_roundtrip() { "return_hidden_states": true, "return_routed_experts": true, "sampling_logprobs_mode": "support", + "return_outputs_via_store": true, "routed_experts_start_len": 10, "return_prompt_token_ids": true, "require_reasoning": true, @@ -26,6 +27,7 @@ fn test_generate_extension_fields_roundtrip() { assert!(req.return_hidden_states); assert!(req.return_routed_experts); assert_eq!(req.sampling_logprobs_mode.as_deref(), Some("support")); + assert_eq!(req.return_outputs_via_store, Some(true)); assert_eq!(req.routed_experts_start_len, 10); assert!(req.return_prompt_token_ids); assert!(req.require_reasoning); @@ -44,6 +46,7 @@ fn test_generate_extension_fields_roundtrip() { assert_eq!(serialized["return_hidden_states"], true); assert_eq!(serialized["return_routed_experts"], true); assert_eq!(serialized["sampling_logprobs_mode"], "support"); + assert_eq!(serialized["return_outputs_via_store"], true); assert_eq!(serialized["routed_experts_start_len"], 10); assert_eq!(serialized["return_prompt_token_ids"], true); assert_eq!(serialized["require_reasoning"], true); @@ -66,6 +69,7 @@ fn test_generate_extension_fields_defaults() { assert!(!req.return_hidden_states); assert!(!req.return_routed_experts); assert!(req.sampling_logprobs_mode.is_none()); + assert!(req.return_outputs_via_store.is_none()); assert_eq!(req.routed_experts_start_len, 0); assert!(!req.return_prompt_token_ids); assert!(!req.require_reasoning); @@ -96,6 +100,7 @@ fn test_generate_optional_fields_omitted_when_none() { "min_dynamic_patch", "lora_backfill_paths", "sampling_logprobs_mode", + "return_outputs_via_store", ]; for field in omitted { assert!(