[diffusion] chore: clean srt imports (#17252)

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

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."""
 

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

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

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

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

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,