sglang 0.4.10.post2__py3-none-any.whl → 0.5.0rc0__py3-none-any.whl

sglang/srt/entrypoints/openai/tool_server.py

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import logging
+from abc import ABC, abstractmethod
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+from typing import Any
+
+logger = logging.getLogger(__name__)
+try:
+    from mcp import ClientSession
+    from mcp.client.sse import sse_client
+    from mcp.types import ListToolsResult
+except ImportError:
+    logger.warning("Ignoring mcp import error")
+
+from openai_harmony import ToolDescription, ToolNamespaceConfig
+
+
+async def list_server_and_tools(server_url: str):
+
+    async with sse_client(url=server_url) as streams, ClientSession(
+        *streams
+    ) as session:
+        initialize_response = await session.initialize()
+        list_tools_response = await session.list_tools()
+        return initialize_response, list_tools_response
+
+
+def trim_schema(schema: dict) -> dict:
+    # Turn JSON Schema from MCP generated into Harmony's variant.
+    if "title" in schema:
+        del schema["title"]
+    if "default" in schema and schema["default"] is None:
+        del schema["default"]
+    if "anyOf" in schema:
+        # Turn "anyOf": [{"type": "type-1"}, {"type": "type-2"}]
+        # into "type": ["type-1", "type-2"]
+        # if there's more than 1 types, also remove "null" type as Harmony will
+        # just ignore it
+        types = [
+            type_dict["type"]
+            for type_dict in schema["anyOf"]
+            if type_dict["type"] != "null"
+        ]
+        schema["type"] = types
+        del schema["anyOf"]
+    if "properties" in schema:
+        schema["properties"] = {
+            k: trim_schema(v) for k, v in schema["properties"].items()
+        }
+    return schema
+
+
+def post_process_tools_description(
+    list_tools_result: "ListToolsResult",
+) -> "ListToolsResult":
+    # Adapt the MCP tool result for Harmony
+    for tool in list_tools_result.tools:
+        tool.inputSchema = trim_schema(tool.inputSchema)
+
+    # Some tools schema don't need to be part of the prompt (e.g. simple text
+    # in text out for Python)
+    list_tools_result.tools = [
+        tool
+        for tool in list_tools_result.tools
+        if getattr(tool.annotations, "include_in_prompt", True)
+    ]
+
+    return list_tools_result
+
+
+class ToolServer(ABC):
+
+    @abstractmethod
+    def has_tool(self, tool_name: str):
+        pass
+
+    @abstractmethod
+    def get_tool_description(self, tool_name: str):
+        pass
+
+    @abstractmethod
+    def get_tool_session(self, tool_name: str) -> AbstractAsyncContextManager[Any]: ...
+
+
+class MCPToolServer(ToolServer):
+
+    def __init__(self):
+        self.harmony_tool_descriptions = {}
+
+    async def add_tool_server(self, server_url: str):
+        tool_urls = server_url.split(",")
+        self.harmony_tool_descriptions = {}
+        self.urls: dict[str, str] = {}
+        for url in tool_urls:
+            url = f"http://{url}/sse"
+            initialize_response, list_tools_response = await list_server_and_tools(url)
+
+            list_tools_response = post_process_tools_description(list_tools_response)
+
+            tool_from_mcp = ToolNamespaceConfig(
+                name=initialize_response.serverInfo.name,
+                description=initialize_response.instructions,
+                tools=[
+                    ToolDescription.new(
+                        name=tool.name,
+                        description=tool.description,
+                        parameters=tool.inputSchema,
+                    )
+                    for tool in list_tools_response.tools
+                ],
+            )
+            self.harmony_tool_descriptions[tool_from_mcp.name] = tool_from_mcp
+            if tool_from_mcp.name not in self.urls:
+                self.urls[tool_from_mcp.name] = url
+            else:
+                logger.warning(
+                    "Tool %s already exists. Ignoring duplicate tool server %s",
+                    tool_from_mcp.name,
+                    url,
+                )
+
+    def has_tool(self, tool_name: str):
+        return tool_name in self.harmony_tool_descriptions
+
+    def get_tool_description(self, tool_name: str):
+        return self.harmony_tool_descriptions.get(tool_name)
+
+    @asynccontextmanager
+    async def get_tool_session(self, tool_name: str):
+        url = self.urls.get(tool_name)
+        if url:
+            async with sse_client(url=url) as streams, ClientSession(
+                *streams
+            ) as session:
+                await session.initialize()
+                yield session
+        else:
+            logger.warning("Tool %s not found", tool_name)
+
+
+class DemoToolServer(ToolServer):
+
+    def __init__(self):
+        from sglang.srt.entrypoints.tool import (
+            HarmonyBrowserTool,
+            HarmonyPythonTool,
+            Tool,
+        )
+
+        self.tools: dict[str, Tool] = {}
+        browser_tool = HarmonyBrowserTool()
+        if browser_tool.enabled:
+            self.tools["browser"] = browser_tool
+        python_tool = HarmonyPythonTool()
+        if python_tool.enabled:
+            self.tools["python"] = python_tool
+
+    def has_tool(self, tool_name: str):
+        return tool_name in self.tools
+
+    def get_tool_description(self, tool_name: str):
+        if tool_name not in self.tools:
+            return None
+        if tool_name == "browser":
+            return ToolNamespaceConfig.browser()
+        elif tool_name == "python":
+            return ToolNamespaceConfig.python()
+        else:
+            raise ValueError(f"Unknown tool {tool_name}")
+
+    @asynccontextmanager
+    async def get_tool_session(self, tool_name: str):
+        yield self.tools[tool_name]

sglang/srt/entrypoints/tool.py

@@ -0,0 +1,87 @@
+# SPDX-License-Identifier: Apache-2.0
+import logging
+import os
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    # Avoid circular import.
+    from sglang.srt.entrypoints.context import ConversationContext
+
+logger = logging.getLogger(__name__)
+
+
+class Tool(ABC):
+
+    @abstractmethod
+    async def get_result(self, context: "ConversationContext") -> Any:
+        pass
+
+
+class HarmonyBrowserTool(Tool):
+
+    def __init__(self):
+        self.enabled = True
+        exa_api_key = os.getenv("EXA_API_KEY")
+        if not exa_api_key:
+            self.enabled = False
+            logger.warning_once("EXA_API_KEY is not set, browsing is disabled")
+            return
+
+        try:
+            from gpt_oss.tools.simple_browser import SimpleBrowserTool
+            from gpt_oss.tools.simple_browser.backend import ExaBackend
+        except ImportError:
+            self.enabled = False
+            logger.warning_once("gpt_oss is not installed, browsing is disabled")
+            return
+
+        browser_backend = ExaBackend(source="web", api_key=exa_api_key)
+        self.browser_tool = SimpleBrowserTool(backend=browser_backend)
+        logger.info_once("Browser tool initialized")
+
+    async def get_result(self, context: "ConversationContext") -> Any:
+        from sglang.srt.entrypoints.context import HarmonyContext
+
+        assert isinstance(context, HarmonyContext)
+        last_msg = context.messages[-1]
+        tool_output_msgs = []
+        async for msg in self.browser_tool.process(last_msg):
+            tool_output_msgs.append(msg)
+        return tool_output_msgs
+
+    @property
+    def tool_config(self) -> Any:
+        return self.browser_tool.tool_config
+
+
+class HarmonyPythonTool(Tool):
+
+    def __init__(self):
+        self.enabled = True
+
+        try:
+            from gpt_oss.tools.python_docker.docker_tool import PythonTool
+        except ImportError:
+            self.enabled = False
+            logger.warning_once(
+                "gpt_oss is not installed, code interpreter is disabled"
+            )
+            return
+
+        self.python_tool = PythonTool()
+        logger.info_once("Code interpreter tool initialized")
+
+    async def get_result(self, context: "ConversationContext") -> Any:
+        from sglang.srt.entrypoints.context import HarmonyContext
+
+        assert isinstance(context, HarmonyContext)
+        last_msg = context.messages[-1]
+        tool_output_msgs = []
+        async for msg in self.python_tool.process(last_msg):
+            tool_output_msgs.append(msg)
+        return tool_output_msgs
+
+    @property
+    def tool_config(self) -> Any:
+        return self.python_tool.tool_config

sglang/srt/eplb/expert_location.py

@@ -35,6 +35,7 @@ class ExpertLocationMetadata:
     physical_to_logical_map: torch.Tensor  # (layers, num_physical_experts)
     physical_to_logical_map_cpu: torch.Tensor
     logical_to_all_physical_map: torch.Tensor  # (layers, num_logical_experts, X)
+    logical_to_all_physical_map_cpu: torch.Tensor  # CPU copy for performance
     logical_to_all_physical_map_num_valid: torch.Tensor  # (layers, num_logical_experts)
     # (layers, num_logical_experts)
     logical_to_rank_dispatch_physical_map: Optional[torch.Tensor]
@@ -221,6 +222,7 @@ class ExpertLocationMetadata:
             physical_to_logical_map=physical_to_logical_map,
             physical_to_logical_map_cpu=physical_to_logical_map.cpu(),
             logical_to_all_physical_map=logical_to_all_physical_map_padded,
+            logical_to_all_physical_map_cpu=logical_to_all_physical_map_padded.cpu(),
             logical_to_all_physical_map_num_valid=logical_to_all_physical_map_num_valid,
             logical_to_rank_dispatch_physical_map=(
                 compute_logical_to_rank_dispatch_physical_map(
@@ -251,6 +253,7 @@ class ExpertLocationMetadata:
             "physical_to_logical_map",
             "physical_to_logical_map_cpu",
             "logical_to_all_physical_map",
+            "logical_to_all_physical_map_cpu",
             "logical_to_all_physical_map_num_valid",
             "logical_to_rank_dispatch_physical_map",
         ]:
@@ -270,9 +273,10 @@ class ExpertLocationMetadata:
     def logical_to_all_physical(
         self, layer_id: int, logical_expert_id: int
     ) -> List[int]:
+        # Use CPU copy to avoid GPU→CPU sync on every call, which is expensive in update weights scenario
         return [
             physical_expert_id
-            for physical_expert_id in self.logical_to_all_physical_map[
+            for physical_expert_id in self.logical_to_all_physical_map_cpu[
                 layer_id, logical_expert_id
             ].tolist()
             if physical_expert_id != -1

sglang/srt/function_call/harmony_tool_parser.py

@@ -0,0 +1,130 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Harmony tool call parser for processing tool calls in harmony models."""
+
+import uuid
+from typing import List, Optional, Tuple
+
+from sglang.srt.entrypoints.openai.protocol import (
+    ChatMessage,
+    FunctionResponse,
+    ToolCall,
+)
+
+
+class HarmonyToolCallParser:
+    """Parser for extracting tool calls from harmony model outputs."""
+
+    def extract_tool_calls_from_message(self, msg) -> Optional[ToolCall]:
+        """
+        Extract tool call from a single message if it's a tool call.
+
+        Args:
+            msg: The harmony message
+
+        Returns:
+            ToolCall if the message is a tool call, None otherwise
+        """
+        if (
+            msg.channel == "commentary"
+            and msg.recipient
+            and msg.recipient.startswith("functions.")
+        ):
+            function_name = msg.recipient.split(".")[-1]
+            arguments = msg.content[0].text if msg.content else "{}"
+
+            return ToolCall(
+                id=f"call_{uuid.uuid4().hex[:24]}",
+                function=FunctionResponse(
+                    name=function_name,
+                    arguments=arguments,
+                ),
+            )
+        return None
+
+    def process_streaming_chunk(
+        self,
+        harmony_parser,
+        index: int,
+        tool_call_trackers: dict,
+        stream_buffers: dict,
+    ) -> Tuple[Optional[dict], bool, Optional[str]]:
+        """
+        Process a streaming chunk for tool calls.
+
+        Args:
+            harmony_parser: The harmony parser instance
+            index: The choice index
+            tool_call_trackers: Dict tracking tool calls per choice
+            stream_buffers: Dict for buffering content
+
+        Returns:
+            Tuple of (tool_call_data, is_tool_call, delta)
+        """
+        # Check if we're in a tool call
+        is_tool_call = (
+            harmony_parser.current_channel == "commentary"
+            and harmony_parser.current_recipient
+            and harmony_parser.current_recipient.startswith("functions.")
+        )
+
+        delta = harmony_parser.last_content_delta or ""
+        tool_call_data = None
+
+        if is_tool_call:
+            # Handle tool call streaming
+            function_name = harmony_parser.current_recipient.split(".")[-1]
+
+            # Track tool call indices per choice
+            if index not in tool_call_trackers:
+                tool_call_trackers[index] = {"count": 0, "current_function": None}
+
+            # Check if we just started a new tool call
+            tool_call_tracker = tool_call_trackers[index]
+            if tool_call_tracker["current_function"] != function_name:
+                # New tool call started
+                tool_call_tracker["current_function"] = function_name
+                tool_call_index = tool_call_tracker["count"]
+                tool_call_tracker["count"] += 1
+
+                # Store the tool call index for this function
+                tool_call_key = f"{index}_{function_name}"
+                stream_buffers[tool_call_key] = {
+                    "index": tool_call_index,
+                    "content": "",
+                }
+
+                tool_call_data = {
+                    "id": f"call_{uuid.uuid4().hex[:24]}",
+                    "index": tool_call_index,
+                    "function_name": function_name,
+                    "arguments": delta,
+                    "is_first_chunk": True,
+                }
+            else:
+                # Subsequent chunks for the same tool call
+                tool_call_key = f"{index}_{function_name}"
+                tool_call_index = stream_buffers[tool_call_key]["index"]
+
+                tool_call_data = {
+                    "id": None,
+                    "index": tool_call_index,
+                    "function_name": None,
+                    "arguments": delta,
+                    "is_first_chunk": False,
+                }
+
+            stream_buffers[tool_call_key]["content"] += delta
+
+        return tool_call_data, is_tool_call, delta

sglang/srt/hf_transformers_utils.py

@@ -14,10 +14,11 @@
 """Utilities for Huggingface Transformers."""
 
 import contextlib
+import json
 import os
 import warnings
 from pathlib import Path
-from typing import Dict, Optional, Type, Union
+from typing import Any, Dict, Optional, Type, Union
 
 import torch
 from huggingface_hub import snapshot_download
@@ -62,11 +63,17 @@ for name, cls in _CONFIG_REGISTRY.items():
         AutoConfig.register(name, cls)
 
 
-def download_from_hf(model_path: str):
+def download_from_hf(
+    model_path: str,
+    allow_patterns: Optional[Union[str, list]] = None,
+):
     if os.path.exists(model_path):
         return model_path
 
-    return snapshot_download(model_path, allow_patterns=["*.json", "*.bin", "*.model"])
+    if not allow_patterns:
+        allow_patterns = ["*.json", "*.bin", "*.model"]
+
+    return snapshot_download(model_path, allow_patterns=allow_patterns)
 
 
 def get_hf_text_config(config: PretrainedConfig):
@@ -171,6 +178,26 @@ def get_generation_config(
         return None
 
 
+# Qwen-1M related
+def get_sparse_attention_config(
+    model: str,
+    sparse_attention_config_filename: str = "sparse_attention_config.json",
+) -> Dict[str, Any]:
+    is_local = os.path.isdir(model)
+    if not is_local:
+        # Download the config files.
+        model = download_from_hf(model, allow_patterns=["*.json"])
+
+    config_file = os.path.join(model, sparse_attention_config_filename)
+    if not os.path.exists(config_file):
+        return {}
+
+    # Load the sparse attention config.
+    with open(config_file) as f:
+        config = json.load(f)
+    return config
+
+
 # Models don't use the same configuration key for determining the maximum
 # context length.  Store them here so we can sanely check them.
 # NOTE: The ordering here is important. Some models have two of these and we

sglang/srt/jinja_template_utils.py

@@ -9,6 +9,8 @@ import logging
 import jinja2
 import transformers.utils.chat_template_utils as hf_chat_utils
 
+from sglang.srt.utils import ImageData
+
 logger = logging.getLogger(__name__)
 
 # ============================================================================
@@ -140,7 +142,12 @@ def process_content_for_template_format(
                 chunk_type = chunk.get("type")
 
                 if chunk_type == "image_url":
-                    image_data.append(chunk["image_url"]["url"])
+                    image_data.append(
+                        ImageData(
+                            url=chunk["image_url"]["url"],
+                            detail=chunk["image_url"].get("detail", "auto"),
+                        )
+                    )
                     if chunk.get("modalities"):
                         modalities.append(chunk.get("modalities"))
                     # Normalize to simple 'image' type for template compatibility

sglang/srt/layers/attention/aiter_backend.py

@@ -720,11 +720,6 @@ class AiterIndicesUpdaterPrefill:
         self.req_to_token = model_runner.req_to_token_pool.req_to_token
         self.update = self.update_single_wrapper
 
-        # get the last index of the pool
-        self.pool_size = (
-            model_runner.token_to_kv_pool.size + model_runner.token_to_kv_pool.page_size
-        ) - 1
-
         self.kv_indices = None
         self.max_q_len = 0
         self.max_kv_len = 0
@@ -769,9 +764,8 @@ class AiterIndicesUpdaterPrefill:
             # but the 0 location will be made nan (noqa) in cuda graph capture mode
             # this will cause the output tensor value becomes nan
             # WA is to assure that last index of pool not changed
-            kv_indices = torch.full(
-                (paged_kernel_lens_sum + 128,),
-                self.pool_size,
+            kv_indices = torch.empty(
+                paged_kernel_lens_sum + 256,
                 dtype=torch.int32,
                 device=req_pool_indices.device,
             )
@@ -785,6 +779,9 @@ class AiterIndicesUpdaterPrefill:
                 self.req_to_token.shape[1],
             )
 
+            token_num = kv_indptr[-1]
+            kv_indices[token_num:] = kv_indices[0]
+
             self.max_kv_len = torch.max(paged_kernel_lens).item()
 
             extend_lens = seq_lens - prefix_lens