From dc743fe4ba850926b1943037d9d49306ed628390 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 17 Jan 2026 15:47:49 +0800 Subject: [PATCH] [diffusion] chore: clean srt imports (#17252) --- .../runtime/entrypoints/http_server.py | 4 +- .../runtime/entrypoints/openai/common_api.py | 15 +- .../runtime/entrypoints/openai/protocol.py | 25 +- .../multimodal_gen/runtime/launch_server.py | 43 +++- .../layers/attention/backends/flash_attn.py | 2 +- .../multimodal_gen/runtime/layers/utils.py | 223 ++++++++++++++++++ .../runtime/loader/component_loader.py | 2 +- 7 files changed, 308 insertions(+), 6 deletions(-) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py index 49205c86e..b803373c1 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/http_server.py @@ -12,13 +12,15 @@ from fastapi.responses import ORJSONResponse from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams from sglang.multimodal_gen.runtime.entrypoints.openai import image_api, video_api +from sglang.multimodal_gen.runtime.entrypoints.openai.protocol import ( + VertexGenerateReqInput, +) from sglang.multimodal_gen.runtime.entrypoints.utils import ( post_process_sample, prepare_request, ) from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client from sglang.multimodal_gen.runtime.server_args import ServerArgs, get_global_server_args -from sglang.srt.managers.io_struct import VertexGenerateReqInput DEFAULT_SEED = 1024 VERTEX_ROUTE = os.environ.get("AIP_PREDICT_ROUTE", "/vertex_generate") diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py index 63f7444de..15b702514 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/common_api.py @@ -1,7 +1,9 @@ +import time from typing import Any, List, Optional, Union from fastapi import APIRouter, Body, HTTPException from fastapi.responses import ORJSONResponse +from pydantic import BaseModel, Field from sglang.multimodal_gen.registry import get_model_info from sglang.multimodal_gen.runtime.entrypoints.openai.utils import ( @@ -15,12 +17,23 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBa from sglang.multimodal_gen.runtime.scheduler_client import async_scheduler_client from sglang.multimodal_gen.runtime.server_args import get_global_server_args from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger -from sglang.srt.entrypoints.openai.protocol import ModelCard router = APIRouter(prefix="/v1") logger = init_logger(__name__) +class ModelCard(BaseModel): + """Model cards.""" + + id: str + object: str = "model" + created: int = Field(default_factory=lambda: int(time.time())) + owned_by: str = "sglang" + root: Optional[str] = None + parent: Optional[str] = None + max_model_len: Optional[int] = None + + class DiffusionModelCard(ModelCard): """Extended ModelCard with diffusion-specific fields.""" diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py index 358daba16..67c16c727 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/openai/protocol.py @@ -1,5 +1,8 @@ import time -from typing import Any, Dict, List, Optional +import uuid +from abc import ABC +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional, Union from pydantic import BaseModel, Field @@ -95,3 +98,23 @@ class VideoListResponse(BaseModel): class VideoRemixRequest(BaseModel): prompt: str + + +@dataclass +class BaseReq(ABC): + rid: Optional[Union[str, List[str]]] = field(default=None, kw_only=True) + http_worker_ipc: Optional[str] = field(default=None, kw_only=True) + + def regenerate_rid(self): + """Generate a new request ID and return it.""" + if isinstance(self.rid, list): + self.rid = [uuid.uuid4().hex for _ in range(len(self.rid))] + else: + self.rid = uuid.uuid4().hex + return self.rid + + +@dataclass +class VertexGenerateReqInput(BaseReq): + instances: List[dict] + parameters: Optional[dict] = None diff --git a/python/sglang/multimodal_gen/runtime/launch_server.py b/python/sglang/multimodal_gen/runtime/launch_server.py index de0d805ec..272c4ae16 100644 --- a/python/sglang/multimodal_gen/runtime/launch_server.py +++ b/python/sglang/multimodal_gen/runtime/launch_server.py @@ -2,8 +2,11 @@ import multiprocessing as mp import os +import signal import sys +import threading +import psutil import uvicorn from sglang.multimodal_gen.runtime.entrypoints.http_server import create_app @@ -14,7 +17,45 @@ from sglang.multimodal_gen.runtime.server_args import ( set_global_server_args, ) from sglang.multimodal_gen.runtime.utils.logging_utils import configure_logger, logger -from sglang.srt.utils import kill_process_tree + + +def kill_process_tree(parent_pid, include_parent: bool = True, skip_pid: int = None): + """Kill the process and all its child processes.""" + # Remove sigchld handler to avoid spammy logs. + if threading.current_thread() is threading.main_thread(): + signal.signal(signal.SIGCHLD, signal.SIG_DFL) + + if parent_pid is None: + parent_pid = os.getpid() + include_parent = False + + try: + itself = psutil.Process(parent_pid) + except psutil.NoSuchProcess: + return + + children = itself.children(recursive=True) + for child in children: + if child.pid == skip_pid: + continue + try: + child.kill() + except psutil.NoSuchProcess: + pass + + if include_parent: + try: + if parent_pid == os.getpid(): + itself.kill() + sys.exit(0) + + itself.kill() + + # Sometime processes cannot be killed with SIGKILL (e.g, PID=1 launched by kubernetes), + # so we send an additional signal to kill them. + itself.send_signal(signal.SIGQUIT) + except psutil.NoSuchProcess: + pass def launch_server(server_args: ServerArgs, launch_http_server: bool = True): diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py index 6d8d5bdc0..af84009e1 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/flash_attn.py @@ -6,12 +6,12 @@ from typing import Any, List, Optional, Tuple import torch +from sglang.multimodal_gen.runtime.layers.utils import register_custom_op from sglang.multimodal_gen.runtime.managers.forward_context import get_forward_context from sglang.multimodal_gen.runtime.platforms import ( AttentionBackendEnum, current_platform, ) -from sglang.srt.utils.custom_op import register_custom_op try: from sgl_kernel.flash_attn import flash_attn_varlen_func diff --git a/python/sglang/multimodal_gen/runtime/layers/utils.py b/python/sglang/multimodal_gen/runtime/layers/utils.py index 615ebc385..363129124 100644 --- a/python/sglang/multimodal_gen/runtime/layers/utils.py +++ b/python/sglang/multimodal_gen/runtime/layers/utils.py @@ -3,8 +3,13 @@ # SPDX-License-Identifier: Apache-2.0 # Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/utils.py """Utility methods for model layers.""" +import inspect +from typing import Any, Callable, List, Optional import torch +from torch.library import Library + +from sglang.multimodal_gen.runtime.platforms import current_platform def get_token_bin_counts_and_mask( @@ -22,3 +27,221 @@ def get_token_bin_counts_and_mask( mask = bin_counts > 0 return bin_counts, mask + + +sglang_lib = Library("sglang", "FRAGMENT") # noqa + + +def direct_register_custom_op( + op_name: str, + op_func: Callable, + mutates_args: List[str], + fake_impl: Optional[Callable] = None, + target_lib: Optional[Library] = None, +): + """ + `torch.library.custom_op` can have significant overhead because it + needs to consider complicated dispatching logic. This function + directly registers a custom op and dispatches it to the CUDA backend. + See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5 + for more details. + + By default, the custom op is registered to the vLLM library. If you + want to register it to a different library, you can pass the library + object to the `target_lib` argument. + + IMPORTANT: the lifetime of the operator is tied to the lifetime of the + library object. If you want to bind the operator to a different library, + make sure the library object is alive when the operator is used. + + Note: This function will silently skip registration if the operator + with the same name is already registered to avoid RuntimeError in + multi-engine scenarios (e.g., VERL framework). + """ + import torch.library + + my_lib = target_lib or sglang_lib + + # Check if operator is already registered to avoid duplicate registration + # This is important for scenarios where multiple SGLang engines run in the same process + try: + # Try to access the operator to see if it's already registered + lib_name = my_lib.m.name if hasattr(my_lib.m, "name") else "sglang" + if hasattr(torch.ops, lib_name) and hasattr( + getattr(torch.ops, lib_name), op_name + ): + # Operator already exists, skip registration + return + except (AttributeError, RuntimeError): + # Operator doesn't exist, proceed with registration + pass + + if hasattr(torch.library, "infer_schema"): + schema_str = torch.library.infer_schema(op_func, mutates_args=mutates_args) + else: + # for pytorch 2.4 + import torch._custom_op.impl + + schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args) + + try: + my_lib.define(op_name + schema_str) + my_lib.impl( + op_name, op_func, "CUDA" if not current_platform.is_npu() else "PrivateUse1" + ) + if fake_impl is not None: + my_lib._register_fake(op_name, fake_impl) + except RuntimeError as error: + if "Tried to register an operator" in str(error) and "multiple times" in str( + error + ): + # Silently ignore duplicate registration errors + # This can happen in multi-engine scenarios + pass + else: + # Re-raise other RuntimeErrors + raise error + except AttributeError as error: + # Always re-raise AttributeError as it indicates missing dependencies + raise error + + +class CustomOpWrapper: + def __init__( + self, + op_name: str, + op_func: Callable, + mutates_args: List[str], + **extra_kwargs, + ): + self.op_name = op_name + self.op_func = op_func + self.mutates_args = mutates_args + self.extra_kwargs = extra_kwargs + self._impl: Optional[Callable] = None + + def __call__(self, *args, **kwargs): + return self.real_impl(*args, **kwargs) + + @property + def real_impl(self) -> Callable: + if self._impl is None: + if not hasattr(torch.ops.sglang, self.op_name): + + # NOTE(dark): if torch compile fail here, mark the decorator as eager + # lazy registration does not work with torch compile + direct_register_custom_op( + op_name=self.op_name, + op_func=self.op_func, + mutates_args=self.mutates_args, + fake_impl=self.fake_impl, + ) + self._impl = getattr(torch.ops.sglang, self.op_name) + assert self._impl is not None + return self._impl + + @property + def fake_impl(self) -> Callable: + if "fake_impl" in self.extra_kwargs: + return self.extra_kwargs["fake_impl"] + assert "out_shape" in self.extra_kwargs + signature = inspect.signature(self.op_func) + out_shape = self.extra_kwargs["out_shape"] + + # check out_shape in signature + + def fake_impl(*args, **kwargs): + if out_shape is None: + return None + bound = signature.bind(*args, **kwargs) + bound.apply_defaults() + try: + return torch.empty_like( + bound.args[out_shape] + if isinstance(out_shape, int) + else bound.arguments[out_shape] + ) + except (IndexError, KeyError): + raise RuntimeError( + f"Cannot find output argument at position `{out_shape}` for " + f"custom operator `{self.op_name}` with signature `{signature}`." + ) + + return fake_impl + + +# Real implementation +def register_custom_op( + fn: Optional[Callable] = None, + *, + op_name: Optional[str] = None, + mutates_args: Optional[List[str]] = None, + eager: bool = True, + **extra_kwargs, +) -> Any: + """ + A decorator to register a custom operator. + + Example usage: + ```python + # inplace operator, out_shape is None by default + @register_custom_op(mutates_args=["x"]) + def add_1_(x: torch.Tensor) -> None: + x.add_(1) + + # operator with output, out_shape indicates the position of output + @register_custom_op(mutates_args=["x"], out_shape=0) + def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x.add_(y) + ``` + + :param fn: The function to be registered as a custom operator. + If None, return a decorator. + :type fn: Callable + :param op_name: The name of the operator. If None, use the function name + :type op_name: Optional[str] + :param mutates_args: A list of argument names that are mutated in-place. + :type mutates_args: List[str] + :param out_shape: The position (int for positional, str for keyword) of the output-shape tensor. + It is used to generate a fake implementation for torch.compile compatibility. + If the operator is inplace and has no output, set to None. + :type out_shape: Optional[List[Union[int, str]]] + :param fake_impl: A fake implementation for the operator. + Only one of `out_shape` or `fake_impl` should be provided. + :type fake_impl: Optional[Callable] + :param eager: Whether to register the operator eagerly. + If False, the registration will be deferred until the first call. + If you met any issue with torch.compile, try to set eager=True. + Currently, to avoid misuse, we set eager=True by default. + :type eager: bool + :return: The registered JIT custom operator, or a decorator. + NOTE: the real register will occur at the first call of the function. + :rtype: Callable + """ + extra_kwarg_keys = set(extra_kwargs.keys()) + expected_kwarg_keys = set({"out_shape", "fake_impl"}) + assert ( + expected_kwarg_keys >= extra_kwarg_keys + ), f"Unexpected extra kwargs: {extra_kwarg_keys - expected_kwarg_keys}" + + has_out_shape = "out_shape" in extra_kwargs + has_fake_impl = "fake_impl" in extra_kwargs + assert not ( + has_out_shape and has_fake_impl + ), "Only one of `out_shape` or `fake_impl` should be provided." + # Assume inplace if neither out_shape nor fake_impl is provided + if not (has_out_shape or has_fake_impl): + extra_kwargs["out_shape"] = None + + def decorator(op_func: Callable) -> Callable: + wrapper = CustomOpWrapper( + op_name=op_name or op_func.__name__, + op_func=op_func, + mutates_args=mutates_args or [], + **extra_kwargs, + ) + return wrapper.real_impl if eager else wrapper + + if fn is not None: + return decorator(fn) + return decorator diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loader.py index d29c03bf3..fde23b2a0 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loader.py @@ -149,7 +149,7 @@ class ComponentLoader(ABC): """ gpu_mem_before_loading = current_platform.get_available_gpu_memory() logger.info( - "Loading %s from %s. avail mem: %.2f GB", + "Loading %s. avail mem: %.2f GB", module_name, component_model_path, gpu_mem_before_loading,