From 39f9a9c2a5ce337dfdfc48b181316303148d6bfa Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Mon, 8 Dec 2025 11:54:56 -0800 Subject: [PATCH] [model-gateway] reduce cpu overhead in grpc router (#14663) --- .../grpc/common/response_collection.rs | 4 +- .../grpc/common/stages/worker_selection.rs | 8 ++-- .../src/routers/grpc/harmony/processor.rs | 2 +- .../src/routers/grpc/harmony/streaming.rs | 8 ++-- .../src/routers/grpc/regular/processor.rs | 8 ++-- .../src/routers/grpc/regular/streaming.rs | 38 ++++++++----------- sgl-model-gateway/src/routers/grpc/utils.rs | 22 +++++++---- 7 files changed, 45 insertions(+), 45 deletions(-) diff --git a/sgl-model-gateway/src/routers/grpc/common/response_collection.rs b/sgl-model-gateway/src/routers/grpc/common/response_collection.rs index ff415871d..68daa138f 100644 --- a/sgl-model-gateway/src/routers/grpc/common/response_collection.rs +++ b/sgl-model-gateway/src/routers/grpc/common/response_collection.rs @@ -74,7 +74,9 @@ fn merge_prefill_logprobs( ) { // Only SGLang supports PD mode and has input_logprobs if let Some(ProtoGenerateComplete::Sglang(prefill_first)) = prefill_responses.first() { - if let Some(prefill_input_logprobs) = prefill_first.input_logprobs.clone() { + // Use ref to borrow input_logprobs instead of cloning upfront + // This avoids one allocation when the Option is Some + if let Some(ref prefill_input_logprobs) = prefill_first.input_logprobs { for response in decode_responses.iter_mut() { if let ProtoGenerateComplete::Sglang(decode_resp) = response { decode_resp.input_logprobs = Some(prefill_input_logprobs.clone()); diff --git a/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs b/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs index 969a45294..f043f57f8 100644 --- a/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs +++ b/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs @@ -124,11 +124,9 @@ impl WorkerSelectionStage { false, // get all workers, we'll filter by is_available() next ); - let available: Vec> = workers - .iter() - .filter(|w| w.is_available()) - .cloned() - .collect(); + // Use into_iter() to take ownership of Arcs without cloning (avoids atomic inc/dec) + let available: Vec> = + workers.into_iter().filter(|w| w.is_available()).collect(); if available.is_empty() { return None; diff --git a/sgl-model-gateway/src/routers/grpc/harmony/processor.rs b/sgl-model-gateway/src/routers/grpc/harmony/processor.rs index deabab1a2..ad7556574 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/processor.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/processor.rs @@ -128,7 +128,7 @@ impl HarmonyResponseProcessor { .created(dispatch.created) .choices(choices) .usage(usage) - .maybe_system_fingerprint(dispatch.weight_version.clone()) + .maybe_system_fingerprint(dispatch.weight_version.as_deref()) .build(), ) } diff --git a/sgl-model-gateway/src/routers/grpc/harmony/streaming.rs b/sgl-model-gateway/src/routers/grpc/harmony/streaming.rs index b7b3db36d..afcf7dd54 100644 --- a/sgl-model-gateway/src/routers/grpc/harmony/streaming.rs +++ b/sgl-model-gateway/src/routers/grpc/harmony/streaming.rs @@ -459,7 +459,7 @@ impl HarmonyStreamingProcessor { ) .created(dispatch.created) .add_choice_role(index, "assistant") - .maybe_system_fingerprint(dispatch.weight_version.clone()) + .maybe_system_fingerprint(dispatch.weight_version.as_deref()) .build(); let chunk_json = serde_json::to_string(&role_chunk) @@ -499,7 +499,7 @@ impl HarmonyStreamingProcessor { finish_reason: None, matched_stop: None, }) - .maybe_system_fingerprint(dispatch.weight_version.clone()) + .maybe_system_fingerprint(dispatch.weight_version.as_deref()) .build(); let chunk_json = serde_json::to_string(&chunk) @@ -525,7 +525,7 @@ impl HarmonyStreamingProcessor { ChatCompletionStreamResponse::builder(&dispatch.request_id, &original_request.model) .created(dispatch.created) .add_choice_finish_reason(index, finish_reason, matched_stop.cloned()) - .maybe_system_fingerprint(dispatch.weight_version.clone()) + .maybe_system_fingerprint(dispatch.weight_version.as_deref()) .build(); let chunk_json = serde_json::to_string(&chunk) @@ -555,7 +555,7 @@ impl HarmonyStreamingProcessor { total_tokens: prompt_tokens + completion_tokens, completion_tokens_details: None, }) - .maybe_system_fingerprint(dispatch.weight_version.clone()) + .maybe_system_fingerprint(dispatch.weight_version.as_deref()) .build(); let chunk_json = serde_json::to_string(&usage_chunk) diff --git a/sgl-model-gateway/src/routers/grpc/regular/processor.rs b/sgl-model-gateway/src/routers/grpc/regular/processor.rs index 8a7fd0f49..49dabd1e6 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/processor.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/processor.rs @@ -101,7 +101,7 @@ impl ResponseProcessor { if original_request.separate_reasoning && reasoning_parser_available { let pooled_parser = utils::get_reasoning_parser( &self.reasoning_parser_factory, - self.configured_reasoning_parser.as_ref(), + self.configured_reasoning_parser.as_deref(), &original_request.model, ); @@ -227,7 +227,7 @@ impl ResponseProcessor { let reasoning_parser_available = chat_request.separate_reasoning && utils::check_reasoning_parser_availability( &self.reasoning_parser_factory, - self.configured_reasoning_parser.as_ref(), + self.configured_reasoning_parser.as_deref(), &chat_request.model, ); @@ -240,7 +240,7 @@ impl ResponseProcessor { && chat_request.tools.is_some() && utils::check_tool_parser_availability( &self.tool_parser_factory, - self.configured_tool_parser.as_ref(), + self.configured_tool_parser.as_deref(), &chat_request.model, ); @@ -308,7 +308,7 @@ impl ResponseProcessor { // Get pooled parser for this model let pooled_parser = utils::get_tool_parser( &self.tool_parser_factory, - self.configured_tool_parser.as_ref(), + self.configured_tool_parser.as_deref(), model, ); diff --git a/sgl-model-gateway/src/routers/grpc/regular/streaming.rs b/sgl-model-gateway/src/routers/grpc/regular/streaming.rs index 5ab0c0bec..afc42211b 100644 --- a/sgl-model-gateway/src/routers/grpc/regular/streaming.rs +++ b/sgl-model-gateway/src/routers/grpc/regular/streaming.rs @@ -204,7 +204,7 @@ impl StreamingProcessor { let reasoning_parser_available = separate_reasoning && utils::check_reasoning_parser_availability( &self.reasoning_parser_factory, - self.configured_reasoning_parser.as_ref(), + self.configured_reasoning_parser.as_deref(), model, ); @@ -222,7 +222,7 @@ impl StreamingProcessor { let tool_parser_available = tools.is_some() && utils::check_tool_parser_availability( &self.tool_parser_factory, - self.configured_tool_parser.as_ref(), + self.configured_tool_parser.as_deref(), model, ); @@ -300,7 +300,7 @@ impl StreamingProcessor { let first_chunk = ChatCompletionStreamResponse::builder(request_id, model) .created(created) .add_choice_role(index, "assistant") - .maybe_system_fingerprint(system_fingerprint.map(|s| s.to_string())) + .maybe_system_fingerprint(system_fingerprint) .build(); Self::format_sse_chunk_into(&mut sse_buffer, &first_chunk); tx.send(Ok(Bytes::from(sse_buffer.clone()))) @@ -419,9 +419,7 @@ impl StreamingProcessor { ChatCompletionStreamResponse::builder(request_id, model) .created(created) .add_choice_content(index, "assistant", text) - .maybe_system_fingerprint( - system_fingerprint.map(|s| s.to_string()), - ) + .maybe_system_fingerprint(system_fingerprint) .build(); let sse_chunk = @@ -489,7 +487,7 @@ impl StreamingProcessor { let tool_chunk = ChatCompletionStreamResponse::builder(request_id, model) .created(created) .add_choice_tool_call_delta(*index, tool_call_delta) - .maybe_system_fingerprint(system_fingerprint.map(|s| s.to_string())) + .maybe_system_fingerprint(system_fingerprint) .build(); let sse_chunk = serde_json::to_string(&tool_chunk) @@ -514,7 +512,7 @@ impl StreamingProcessor { let finish_chunk = ChatCompletionStreamResponse::builder(request_id, model) .created(created) .add_choice_finish_reason(*index, final_finish_reason, matched_stop_value) - .maybe_system_fingerprint(system_fingerprint.map(|s| s.to_string())) + .maybe_system_fingerprint(system_fingerprint) .build(); let sse_chunk = serde_json::to_string(&finish_chunk) @@ -537,7 +535,7 @@ impl StreamingProcessor { total_tokens: total_prompt + total_completion, completion_tokens_details: None, }) - .maybe_system_fingerprint(system_fingerprint.map(|s| s.to_string())) + .maybe_system_fingerprint(system_fingerprint) .build(); let sse_chunk = serde_json::to_string(&usage_chunk) @@ -1023,7 +1021,7 @@ impl StreamingProcessor { reasoning_parsers.entry(index).or_insert_with(|| { let parser = utils::create_reasoning_parser( &self.reasoning_parser_factory, - self.configured_reasoning_parser.as_ref(), + self.configured_reasoning_parser.as_deref(), model, ) .expect("Parser should be available - checked upfront"); @@ -1048,7 +1046,7 @@ impl StreamingProcessor { ChatCompletionStreamResponse::builder(request_id, model) .created(created) .add_choice_reasoning(index, reasoning_text) - .maybe_system_fingerprint(system_fingerprint.map(|s| s.to_string())) + .maybe_system_fingerprint(system_fingerprint) .build(), ) } else { @@ -1098,7 +1096,7 @@ impl StreamingProcessor { ChatCompletionStreamResponse::builder(request_id, model) .created(created) .add_choice_tool_name(index, tool_call_id, function.name.clone()) - .maybe_system_fingerprint(system_fingerprint.map(|s| s.to_string())) + .maybe_system_fingerprint(system_fingerprint) .build(), ); } @@ -1109,7 +1107,7 @@ impl StreamingProcessor { ChatCompletionStreamResponse::builder(request_id, model) .created(created) .add_choice_tool_args(index, delta.to_string()) - .maybe_system_fingerprint(system_fingerprint.map(|s| s.to_string())) + .maybe_system_fingerprint(system_fingerprint) .build(), ); } @@ -1139,16 +1137,12 @@ impl StreamingProcessor { // Create fresh parser for this index (not pooled, to avoid state pollution) tool_parsers.entry(index).or_insert_with(|| { let parser = if use_json_parser { - utils::create_tool_parser( - &self.tool_parser_factory, - Some(&"json".to_string()), - model, - ) - .expect("JSON parser should be available") + utils::create_tool_parser(&self.tool_parser_factory, Some("json"), model) + .expect("JSON parser should be available") } else { utils::create_tool_parser( &self.tool_parser_factory, - self.configured_tool_parser.as_ref(), + self.configured_tool_parser.as_deref(), model, ) .expect("Parser should be available - checked upfront") @@ -1167,7 +1161,7 @@ impl StreamingProcessor { ChatCompletionStreamResponse::builder(request_id, model) .created(created) .add_choice_content(index, "assistant", normal_text) - .maybe_system_fingerprint(system_fingerprint.map(|s| s.to_string())) + .maybe_system_fingerprint(system_fingerprint) .build(), ); } @@ -1209,7 +1203,7 @@ impl StreamingProcessor { ChatCompletionStreamResponse::builder(request_id, model) .created(created) .add_choice_tool_call_delta(index, tool_call_delta) - .maybe_system_fingerprint(system_fingerprint.map(|s| s.to_string())) + .maybe_system_fingerprint(system_fingerprint) .build(), ); } diff --git a/sgl-model-gateway/src/routers/grpc/utils.rs b/sgl-model-gateway/src/routers/grpc/utils.rs index 64bde9295..e470aec01 100644 --- a/sgl-model-gateway/src/routers/grpc/utils.rs +++ b/sgl-model-gateway/src/routers/grpc/utils.rs @@ -197,7 +197,7 @@ pub fn generate_tool_constraints( // Return the tool's parameters schema directly (not wrapped in array) let params_schema = serde_json::to_string(&tool.function.parameters) .map_err(|e| format!("Failed to serialize tool parameters: {}", e))?; - Ok(Some(("json_schema".to_string(), params_schema))) + Ok(Some((String::from("json_schema"), params_schema))) } // Required: Array of tool calls with minItems: 1 @@ -673,7 +673,13 @@ pub fn generate_tool_call_id( tool_index: usize, history_count: usize, ) -> String { - if model.to_lowercase().contains("kimi") { + // Case-insensitive check without allocation (search for "kimi" substring) + let is_kimi = model + .as_bytes() + .windows(4) // "kimi".len() + .any(|window| window.eq_ignore_ascii_case(b"kimi")); + + if is_kimi { // KimiK2 format: functions.{name}:{global_index} format!("functions.{}:{}", tool_name, history_count + tool_index) } else { @@ -685,7 +691,7 @@ pub fn generate_tool_call_id( /// Check if a reasoning parser is available for the given model pub fn check_reasoning_parser_availability( reasoning_parser_factory: &ReasoningParserFactory, - configured_parser: Option<&String>, + configured_parser: Option<&str>, model: &str, ) -> bool { if let Some(parser_name) = configured_parser { @@ -700,7 +706,7 @@ pub fn check_reasoning_parser_availability( /// Check if a tool parser is available for the given model pub fn check_tool_parser_availability( tool_parser_factory: &ToolParserFactory, - configured_parser: Option<&String>, + configured_parser: Option<&str>, model: &str, ) -> bool { if let Some(parser_name) = configured_parser { @@ -717,7 +723,7 @@ pub fn check_tool_parser_availability( /// Get a pooled reasoning parser (for non-streaming where state doesn't matter) pub fn get_reasoning_parser( reasoning_parser_factory: &ReasoningParserFactory, - configured_parser: Option<&String>, + configured_parser: Option<&str>, model: &str, ) -> ReasoningPooledParser { if let Some(parser_name) = configured_parser { @@ -741,7 +747,7 @@ pub fn get_reasoning_parser( /// Create a fresh reasoning parser instance (for streaming where state isolation is needed) pub fn create_reasoning_parser( reasoning_parser_factory: &ReasoningParserFactory, - configured_parser: Option<&String>, + configured_parser: Option<&str>, model: &str, ) -> Option> { if let Some(parser_name) = configured_parser { @@ -769,7 +775,7 @@ pub fn create_reasoning_parser( /// Get a pooled tool parser (for non-streaming where state doesn't matter) pub fn get_tool_parser( tool_parser_factory: &ToolParserFactory, - configured_parser: Option<&String>, + configured_parser: Option<&str>, model: &str, ) -> ToolPooledParser { if let Some(parser_name) = configured_parser { @@ -793,7 +799,7 @@ pub fn get_tool_parser( /// Create a fresh tool parser instance (for streaming where state isolation is needed) pub fn create_tool_parser( tool_parser_factory: &ToolParserFactory, - configured_parser: Option<&String>, + configured_parser: Option<&str>, model: &str, ) -> Option> { if let Some(parser_name) = configured_parser {