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

[model-gateway] optimize Vec and HashMap allocations in responses api (#16406)

Browse Source
This commit is contained in:
Simo Lin
2026-01-04 10:07:55 -08:00
committed by GitHub
parent f16606d6f6
commit cf6800f6cc
2 changed files with 16 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -279,8 +279,7 @@ pub fn generate_tool_constraints(
/// Build JSON schema for required tool calls (array with minItems: 1)
/// Includes $defs consolidation from all tools (matching Python's behavior)
fn build_required_array_schema(tools: &[Tool]) -> Result<String, String> {
// Build anyOf schemas for each tool
let mut any_of_schemas = Vec::new();
let mut any_of_schemas = Vec::with_capacity(tools.len());
for tool in tools {
let tool_schema = json!({
"properties": {
@@ -296,7 +295,7 @@ fn build_required_array_schema(tools: &[Tool]) -> Result<String, String> {
}
// Consolidate $defs from all tools (matching Python's _get_tool_schema_defs)
let mut all_defs: HashMap<String, Value> = HashMap::new();
let mut all_defs: Map<String, Value> = Map::new();
for tool in tools {
if let Value::Object(params) = &tool.function.parameters {
if let Some(Value::Object(defs)) = params.get("$defs") {
@@ -332,8 +331,7 @@ fn build_required_array_schema(tools: &[Tool]) -> Result<String, String> {
// Add $defs if any were found (matching Python's behavior)
if !all_defs.is_empty() {
if let Value::Object(ref mut schema_obj) = array_schema {
let defs_value = Value::Object(all_defs.into_iter().collect::<Map<String, Value>>());
schema_obj.insert("$defs".to_string(), defs_value);
schema_obj.insert("$defs".to_string(), Value::Object(all_defs));
}
}
@@ -439,8 +437,8 @@ pub fn process_chat_messages(
.transpose()
.map_err(|e| format!("Failed to serialize tools: {}", e))?;
// Build template kwargs, merging reasoning_effort if present
let mut combined_template_kwargs = HashMap::new();
let kwargs_capacity = 1 + request.chat_template_kwargs.as_ref().map_or(0, |k| k.len());
let mut combined_template_kwargs = HashMap::with_capacity(kwargs_capacity);
// Add reasoning_effort if present (like Python does)
if let Some(reasoning_effort) = &request.reasoning_effort {
@@ -878,7 +876,7 @@ pub fn convert_proto_to_openai_logprobs(
proto_logprobs: &OutputLogProbs,
tokenizer: &Arc<dyn Tokenizer>,
) -> Result<ChatLogProbs, String> {
let mut content_items = Vec::new();
let mut content_items = Vec::with_capacity(proto_logprobs.token_logprobs.len());
// Decode token IDs to text (always with skip_special_tokens=false for logprobs)
let token_texts: Vec<String> = proto_logprobs
@@ -901,8 +899,9 @@ pub fn convert_proto_to_openai_logprobs(
let bytes = Some(token_text.as_bytes().to_vec());
// Build top_logprobs for this position
let mut top_logprobs = Vec::new();
if let Some(top_logprobs_entry) = proto_logprobs.top_logprobs.get(i) {
let top_logprobs = if let Some(top_logprobs_entry) = proto_logprobs.top_logprobs.get(i) {
let mut top_logprobs = Vec::with_capacity(top_logprobs_entry.values.len());
// Decode top token IDs (always with skip_special_tokens=false)
let top_token_texts: Vec<String> = top_logprobs_entry
.token_ids
@@ -928,7 +927,10 @@ pub fn convert_proto_to_openai_logprobs(
});
}
}
}
top_logprobs
} else {
Vec::new()
};
content_items.push(ChatLogProbsContent {
token: token_text,

5
sgl-model-gateway/src/routers/openai/mcp.rs
View File

@@ -217,8 +217,8 @@ pub(super) fn prepare_mcp_payload_for_streaming(
}
// Build function tools for all discovered MCP tools
let mut tools_json = Vec::new();
let tools = active_mcp.list_tools();
let mut tools_json = Vec::with_capacity(tools.len());
for t in tools {
let parameters = Value::Object((*t.input_schema).clone());
let tool = serde_json::json!({
@@ -252,7 +252,8 @@ pub(super) fn build_resume_payload(
.ok_or_else(|| "payload not an object".to_string())?;
// Build input array: start with original user input
let mut input_array = Vec::new();
// Pre-allocate: 1 for user message + conversation history
let mut input_array = Vec::with_capacity(1 + conversation_history.len());
// Add original user message
// For structured input, serialize the original input items