chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[model-gateway] reduce cpu overhead in grpc router (#14663)

Browse Source
This commit is contained in:
Simo Lin
2025-12-08 11:54:56 -08:00
committed by GitHub
parent d69ecc19b8
commit 39f9a9c2a5
7 changed files with 45 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
sgl-model-gateway/src/routers/grpc/common/response_collection.rs
View File

@@ -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());

8
sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs
View File

@@ -124,11 +124,9 @@ impl WorkerSelectionStage {
false, // get all workers, we'll filter by is_available() next
);
let available: Vec<Arc<dyn Worker>> = 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<Arc<dyn Worker>> =
workers.into_iter().filter(|w| w.is_available()).collect();
if available.is_empty() {
return None;

2
sgl-model-gateway/src/routers/grpc/harmony/processor.rs
View File

@@ -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(),
)
}

8
sgl-model-gateway/src/routers/grpc/harmony/streaming.rs
View File

@@ -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)

8
sgl-model-gateway/src/routers/grpc/regular/processor.rs
View File

@@ -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,
);

38
sgl-model-gateway/src/routers/grpc/regular/streaming.rs
View File

@@ -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(),
);
}

22
sgl-model-gateway/src/routers/grpc/utils.rs
View File

@@ -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<Box<dyn ReasoningParser>> {
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<Box<dyn ToolParser>> {
if let Some(parser_name) = configured_parser {