chore: import upstream snapshot with attribution

python/sglang/srt/utils/__init__.py

@@ -0,0 +1,2 @@
+# Temporarily do this to avoid changing all imports in the repo
+from sglang.srt.utils.common import *

python/sglang/srt/utils/aio_rwlock.py

@@ -0,0 +1,104 @@
+import asyncio
+
+
+class RWLock:
+    def __init__(self):
+        # Protects internal state
+        self._lock = asyncio.Lock()
+
+        # Condition variable used to wait for state changes
+        self._cond = asyncio.Condition(self._lock)
+
+        # Number of readers currently holding the lock
+        self._readers = 0
+
+        # Whether a writer is currently holding the lock
+        self._writer_active = False
+
+        # How many writers are queued waiting for a turn
+        self._waiting_writers = 0
+
+    @property
+    def reader_lock(self):
+        """
+        A context manager for acquiring a shared (reader) lock.
+
+        Example:
+            async with rwlock.reader_lock:
+                # read-only access
+        """
+        return _ReaderLock(self)
+
+    @property
+    def writer_lock(self):
+        """
+        A context manager for acquiring an exclusive (writer) lock.
+
+        Example:
+            async with rwlock.writer_lock:
+                # exclusive access
+        """
+        return _WriterLock(self)
+
+    async def acquire_reader(self):
+        async with self._lock:
+            # Wait until there is no active writer or waiting writer
+            # to ensure fairness.
+            while self._writer_active or self._waiting_writers > 0:
+                await self._cond.wait()
+            self._readers += 1
+
+    async def release_reader(self):
+        async with self._lock:
+            self._readers -= 1
+            # If this was the last reader, wake up anyone waiting
+            # (potentially a writer or new readers).
+            if self._readers == 0:
+                self._cond.notify_all()
+
+    async def acquire_writer(self):
+        async with self._lock:
+            # Increment the count of writers waiting
+            self._waiting_writers += 1
+            try:
+                # Wait while either a writer is active or readers are present
+                while self._writer_active or self._readers > 0:
+                    await self._cond.wait()
+                self._writer_active = True
+            finally:
+                # Decrement waiting writers only after we've acquired the writer lock
+                self._waiting_writers -= 1
+
+    async def release_writer(self):
+        async with self._lock:
+            self._writer_active = False
+            # Wake up anyone waiting (readers or writers)
+            self._cond.notify_all()
+
+    async def is_locked(self):
+        async with self._lock:
+            return self._writer_active or self._readers > 0
+
+
+class _ReaderLock:
+    def __init__(self, rwlock: RWLock):
+        self._rwlock = rwlock
+
+    async def __aenter__(self):
+        await self._rwlock.acquire_reader()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self._rwlock.release_reader()
+
+
+class _WriterLock:
+    def __init__(self, rwlock: RWLock):
+        self._rwlock = rwlock
+
+    async def __aenter__(self):
+        await self._rwlock.acquire_writer()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self._rwlock.release_writer()

python/sglang/srt/utils/async_probe.py

@@ -0,0 +1,149 @@
+"""Async invariant probes — fire torch._assert_async without CPU sync.
+
+All probes are gated on SGLANG_ENABLE_ASYNC_ASSERT (default off in prod).
+When the gate is on, a violation surfaces as an assertion at the next CUDA
+sync point instead of as a silent NaN cascade or illegal-address crash.
+"""
+
+import logging
+from typing import Optional
+
+import torch
+
+from sglang.srt.environ import envs
+
+logger = logging.getLogger(__name__)
+
+
+class _AsyncNanWarner:
+    """One-shot NaN monitor: device-side detection lands in pinned host
+    memory without any stream sync; the host reads the (slightly stale) flag
+    on a later call, warns once, and stops detecting."""
+
+    def __init__(self):
+        self._dev = None
+        self._host = None
+        self._warned = False
+
+    def check(self, tensor: torch.Tensor, msg: str):
+        if self._warned or not tensor.is_cuda:
+            return
+        if self._dev is None:
+            self._dev = torch.zeros(1, dtype=torch.int32, device=tensor.device)
+            self._host = torch.zeros(1, dtype=torch.int32, pin_memory=True)
+
+        # Report a hit enqueued on an earlier step (pinned read, no sync).
+        if int(self._host[0]):
+            logger.warning(
+                "NaN detected in %s; values were sanitized before sampling. "
+                "This usually indicates numerical overflow (e.g. fp16 "
+                "activations) or an upstream bug producing NaN. "
+                "Logged once; further occurrences are silent.",
+                msg,
+            )
+            self._warned = True
+            return
+
+        # Enqueue this step's detection (async, no sync).
+        self._dev.add_(torch.isnan(tensor).any().to(torch.int32))
+        self._host.copy_(self._dev, non_blocking=True)
+
+
+_nan_warner = _AsyncNanWarner()
+
+
+def maybe_warn_nan(tensor: Optional[torch.Tensor], msg: str = ""):
+    """Non-fatal counterpart of maybe_detect_nan: throttled sync-free warning
+    instead of crashing. Callers sanitize the tensor themselves."""
+    if envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
+        # The hard assert path already covers detection.
+        return
+    if tensor is None:
+        return
+    _nan_warner.check(tensor, msg)
+
+
+def sanitize_nan_logits(logits: torch.Tensor, msg: str = ""):
+    """Detect NaN (assert in CI, throttled warning in prod), then sanitize in
+    place: NaN logits (e.g. fp16 activation overflow) are undefined behavior
+    in sampling kernels and can come back as out-of-vocab token ids. +-1e30
+    rather than dtype min/max because callers divide logits by temperature,
+    which would overflow dtype min/max to +-Inf and softmax back to NaN."""
+    maybe_detect_nan(logits, msg)
+    if not envs.SGLANG_SANITIZE_NAN_LOGITS.get():
+        return
+    maybe_warn_nan(logits, msg)
+    torch.nan_to_num_(logits, nan=-1e30, posinf=1e30, neginf=-1e30)
+
+
+def maybe_assert_async(cond: torch.Tensor, msg: str = ""):
+    if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
+        return
+    torch._assert_async(cond, msg)
+
+
+def maybe_detect_nan(tensor: Optional[torch.Tensor], msg: str = ""):
+    """Async NaN check — no GPU-CPU sync, error surfaces at next sync point."""
+    if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
+        return
+    # A None tensor means there is nothing to probe, e.g. hidden_states on
+    # capture_hidden_mode=NULL paths (STANDALONE speculative decoding).
+    if tensor is None:
+        return
+    torch._assert_async(~torch.any(torch.isnan(tensor)), f"NaN detected! {msg}")
+
+
+def maybe_detect_inf(tensor: Optional[torch.Tensor], msg: str = ""):
+    """Async Inf check — fp16 overflow surfaces as Inf before NaN."""
+    if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
+        return
+    if tensor is None:
+        return
+    torch._assert_async(~torch.any(torch.isinf(tensor)), f"Inf detected! {msg}")
+
+
+def maybe_detect_in_closed_range(
+    tensor: Optional[torch.Tensor], low: float, high: float, msg: str = ""
+):
+    if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
+        return
+    if tensor is None or tensor.numel() == 0:
+        return
+    torch._assert_async(
+        ((tensor >= low) & (tensor <= high)).all(),
+        f"value outside [{low}, {high}]: {msg}",
+    )
+
+
+def maybe_detect_oob(indices: Optional[torch.Tensor], low: int, high: int, msg: str):
+    """Async OOB check — no GPU-CPU sync, error surfaces at next sync point.
+
+    Low/high asserted separately so the message names which failed (low =
+    negative/sentinel, high = out of range).
+    """
+    if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
+        return
+    if indices is None or indices.numel() == 0:
+        return
+    torch._assert_async(
+        indices.min() >= low,
+        f"index < {low} (negative / unmasked sentinel?): {msg}",
+    )
+    torch._assert_async(
+        indices.max() < high,
+        f"index >= {high} (out of range): {msg}",
+    )
+
+
+def maybe_detect_page_aligned(
+    indices: Optional[torch.Tensor], page_size: int, msg: str
+):
+    """Async page-alignment check on slot ids."""
+    if not envs.SGLANG_ENABLE_ASYNC_ASSERT.get():
+        return
+    if indices is None or indices.numel() == 0 or page_size <= 1:
+        return
+    torch._assert_async(
+        (indices % page_size == 0).all(),
+        f"page-misaligned indices (page_size={page_size}): {msg}",
+    )

python/sglang/srt/utils/auth.py

@@ -0,0 +1,208 @@
+"""Auth utilities for HTTP servers.
+
+This module is intentionally lightweight (no torch import) so it can be used in unit tests.
+"""
+
+from __future__ import annotations
+
+import secrets
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Optional
+
+
+@dataclass(frozen=True)
+class AuthDecision:
+    allowed: bool
+    error_status_code: int = 401  # Only meaningful when allowed=False
+
+
+class AuthLevel(str, Enum):
+    """Per-endpoint auth level (attached to endpoint function via `@auth_level`)."""
+
+    NORMAL = "normal"
+    ADMIN_OPTIONAL = "admin_optional"
+    ADMIN_FORCE = "admin_force"
+
+
+def auth_level(level: AuthLevel):
+    """Mark endpoint with auth level (stored in endpoint metadata)."""
+
+    def decorator(func):
+        func._auth_level = level
+        return func
+
+    return decorator
+
+
+def _get_auth_level_from_app_and_scope(app: Any, scope: dict) -> AuthLevel:
+    """Best-effort resolve auth level by matching the request to a route."""
+    # Import lazily to keep this module unit-test friendly (FastAPI/Starlette are not
+    # required unless you actually use the middleware / route matching).
+    from starlette.routing import Match
+
+    # Prefer app.router.routes when available; fall back to app.routes.
+    routes = getattr(getattr(app, "router", None), "routes", None) or getattr(
+        app, "routes", []
+    )
+
+    for route in routes:
+        try:
+            match, child_scope = route.matches(scope)
+        except Exception:
+            continue
+        if match == Match.FULL:
+            endpoint = child_scope.get("endpoint") or getattr(route, "endpoint", None)
+            level = getattr(endpoint, "_auth_level", None)
+            return level if isinstance(level, AuthLevel) else AuthLevel.NORMAL
+
+    return AuthLevel.NORMAL
+
+
+def app_has_admin_force_endpoints(app: Any) -> bool:
+    """Return True if any route endpoint is marked as ADMIN_FORCE."""
+    routes = getattr(getattr(app, "router", None), "routes", None) or getattr(
+        app, "routes", []
+    )
+    for route in routes:
+        endpoint = getattr(route, "endpoint", None)
+        if getattr(endpoint, "_auth_level", None) == AuthLevel.ADMIN_FORCE:
+            return True
+    return False
+
+
+def decide_request_auth(
+    *,
+    method: str,
+    path: str,
+    authorization_header: Optional[str],
+    api_key: Optional[str],
+    admin_api_key: Optional[str],
+    auth_level: AuthLevel,
+) -> AuthDecision:
+    """Pure auth decision function (easy to unit test).
+
+    Auth levels:
+    - NORMAL: legacy behavior (api_key protects all endpoints when configured)
+    - ADMIN_OPTIONAL: can be accessed without any key (if no keys configured),
+      or with api_key/admin_api_key depending on server config.
+    - ADMIN_FORCE: requires admin_api_key; if admin_api_key is NOT configured,
+      it must be rejected (403) even if api_key is provided.
+
+    NOTE :
+    - Health/metrics endpoints are always allowed (even when api_key/admin_api_key is set),
+      to support k8s/liveness/readiness and Prometheus scraping without embedding secrets.
+    - We match them by prefix to cover common variants like /health_generate.
+    """
+    if method == "OPTIONS":
+        return AuthDecision(allowed=True)
+
+    if path.startswith("/health") or path.startswith("/metrics"):
+        return AuthDecision(allowed=True)
+
+    def _check_bearer_token(
+        authorization_header: Optional[str], expected_token: str
+    ) -> bool:
+        """Check bearer token with constant-time comparison."""
+        if not authorization_header:
+            return False
+        parts = authorization_header.split(" ", 1)
+        if len(parts) != 2 or parts[0].lower() != "bearer":
+            return False
+        return secrets.compare_digest(parts[1], expected_token)
+
+    # Force-auth endpoints: only admin_api_key can unlock them; if admin_api_key is unset,
+    # reject them unconditionally (explicitly "not allowed").
+    if auth_level == AuthLevel.ADMIN_FORCE:
+        if not admin_api_key:
+            return AuthDecision(allowed=False, error_status_code=403)
+        if not _check_bearer_token(authorization_header, admin_api_key):
+            return AuthDecision(allowed=False)
+        return AuthDecision(allowed=True)
+
+    # Optional-auth endpoints:
+    # - no keys configured: allow
+    # - only api_key: require api_key
+    # - only admin_api_key: require admin_api_key
+    # - both: require admin_api_key (api_key is NOT accepted)
+    if auth_level == AuthLevel.ADMIN_OPTIONAL:
+        if admin_api_key:
+            return AuthDecision(
+                allowed=_check_bearer_token(authorization_header, admin_api_key)
+            )
+        elif api_key:
+            return AuthDecision(
+                allowed=_check_bearer_token(authorization_header, api_key)
+            )
+        else:
+            return AuthDecision(allowed=True)
+
+    # Normal endpoints:
+    # - if api_key is configured, require api_key (even if admin_api_key is also configured)
+    # - otherwise allow (including the "admin_api_key only" case)
+    if api_key:
+        return AuthDecision(allowed=_check_bearer_token(authorization_header, api_key))
+
+    return AuthDecision(allowed=True)
+
+
+def add_api_key_middleware(
+    app,
+    *,
+    api_key: Optional[str],
+    admin_api_key: Optional[str],
+):
+    """Add middleware for three endpoint auth levels: normal/admin_optional/admin_force."""
+    # Import lazily so `decide_request_auth()` can be unit-tested without FastAPI installed.
+    from fastapi.responses import ORJSONResponse
+    from starlette.requests import Request
+
+    class _ApiKeyASGIMiddleware:
+        """ASGI-native middleware to preserve client disconnect events."""
+
+        def __init__(self, app, *, api_key, admin_api_key, fastapi_app):
+            self.app = app
+            self.api_key = api_key
+            self.admin_api_key = admin_api_key
+            self.fastapi_app = fastapi_app
+
+        async def __call__(self, scope, receive, send):
+            if scope["type"] != "http":
+                await self.app(scope, receive, send)
+                return
+
+            request = Request(scope, receive=receive)
+            path = request.url.path
+            authz = request.headers.get("Authorization")
+            level = _get_auth_level_from_app_and_scope(self.fastapi_app, scope)
+            decision = decide_request_auth(
+                method=request.method,
+                path=path,
+                authorization_header=authz,
+                api_key=self.api_key,
+                admin_api_key=self.admin_api_key,
+                auth_level=level,
+            )
+
+            if not decision.allowed:
+                response = ORJSONResponse(
+                    content={
+                        "error": (
+                            "Unauthorized"
+                            if decision.error_status_code == 401
+                            else "Forbidden"
+                        )
+                    },
+                    status_code=decision.error_status_code,
+                )
+                await response(scope, receive, send)
+                return
+
+            await self.app(scope, receive, send)
+
+    app.add_middleware(
+        _ApiKeyASGIMiddleware,
+        api_key=api_key,
+        admin_api_key=admin_api_key,
+        fastapi_app=app,
+    )

python/sglang/srt/utils/bench_utils.py

@@ -0,0 +1,157 @@
+import os
+import re
+import sys
+from contextlib import nullcontext
+
+import torch
+
+
+# NOTE copied and modified from DeepGEMM
+class suppress_stdout_stderr:
+    def __enter__(self):
+        self.outnull_file = open(os.devnull, "w")
+        self.errnull_file = open(os.devnull, "w")
+
+        self.old_stdout_fileno_undup = sys.stdout.fileno()
+        self.old_stderr_fileno_undup = sys.stderr.fileno()
+
+        self.old_stdout_fileno = os.dup(sys.stdout.fileno())
+        self.old_stderr_fileno = os.dup(sys.stderr.fileno())
+
+        self.old_stdout = sys.stdout
+        self.old_stderr = sys.stderr
+
+        os.dup2(self.outnull_file.fileno(), self.old_stdout_fileno_undup)
+        os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup)
+
+        sys.stdout = self.outnull_file
+        sys.stderr = self.errnull_file
+        return self
+
+    def __exit__(self, *_):
+        sys.stdout = self.old_stdout
+        sys.stderr = self.old_stderr
+
+        os.dup2(self.old_stdout_fileno, self.old_stdout_fileno_undup)
+        os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup)
+
+        os.close(self.old_stdout_fileno)
+        os.close(self.old_stderr_fileno)
+
+        self.outnull_file.close()
+        self.errnull_file.close()
+
+
+# NOTE copied and modified from DeepGEMM
+def bench_kineto(
+    fn,
+    kernel_names,
+    num_tests: int = 30,
+    suppress_kineto_output: bool = False,
+    trace_path: str = None,
+    flush_l2: bool = True,
+    with_multiple_kernels: bool = False,
+):
+    # Conflict with Nsight Systems
+    using_nsys = int(os.environ.get("SGLANG_NSYS_PROFILING", 0))
+
+    # By default, flush L2 with an excessive 8GB memset to give the GPU some (literal) chill time without full idle
+    flush_l2_size = int(8e9 // 4)
+
+    # For some auto-tuning kernels with prints
+    fn()
+
+    # Profile
+    suppress = (
+        suppress_stdout_stderr
+        if suppress_kineto_output and not using_nsys
+        else nullcontext
+    )
+    with suppress():
+        schedule = (
+            torch.profiler.schedule(wait=0, warmup=1, active=1, repeat=1)
+            if not using_nsys
+            else None
+        )
+        profiler = (
+            torch.profiler.profile(
+                activities=[torch.profiler.ProfilerActivity.CUDA],
+                schedule=schedule,
+                acc_events=True,
+            )
+            if not using_nsys
+            else nullcontext()
+        )
+        with profiler:
+            for i in range(2):
+                for _ in range(num_tests):
+                    if flush_l2:
+                        torch.empty(
+                            flush_l2_size, dtype=torch.int, device="cuda"
+                        ).zero_()
+                    fn()
+                if not using_nsys:
+                    torch.cuda.synchronize()
+                    profiler.step()
+
+    # Return 1 if using Nsight Systems
+    if using_nsys:
+        return 1
+
+    # Parse the profiling table
+    assert isinstance(kernel_names, str) or isinstance(kernel_names, tuple)
+    is_tuple = isinstance(kernel_names, tuple)
+    prof_lines = (
+        profiler.key_averages()
+        .table(sort_by="cuda_time_total", max_name_column_width=100)
+        .split("\n")
+    )
+    kernel_names = (kernel_names,) if isinstance(kernel_names, str) else kernel_names
+    assert all([isinstance(name, str) for name in kernel_names])
+    # Check if profiler captured any events (can be empty with some CUDA versions)
+    non_empty_lines = [l for l in prof_lines if l.strip() and not l.startswith("-")]
+    if len(non_empty_lines) <= 1:
+        print(
+            "WARNING: Profiler returned empty table — falling back to wall-clock timing"
+        )
+        import time
+
+        torch.cuda.synchronize()
+        start = time.perf_counter()
+        for _ in range(num_tests):
+            fn()
+        torch.cuda.synchronize()
+        elapsed = (time.perf_counter() - start) / num_tests
+        return tuple([elapsed] * len(kernel_names)) if is_tuple else elapsed
+
+    if not with_multiple_kernels:
+        for name in kernel_names:
+            assert (
+                sum([int(re.search(name, line) is not None) for line in prof_lines])
+                == 1
+            ), f"Errors of the kernel {name} in the profiling table (table: {prof_lines})"
+
+    # Save chrome traces
+    if trace_path is not None:
+        profiler.export_chrome_trace(trace_path)
+
+    # Return average kernel times
+    units = {"ms": 1e3, "us": 1e6}
+    kernel_times = []
+    for name in kernel_names:
+        total_time = 0
+        total_num = 0
+        for line in prof_lines:
+            if re.search(name, line) is not None:
+                time_str = line.split()[-2]
+                num_str = line.split()[-1]
+                for unit, scale in units.items():
+                    if unit in time_str:
+                        total_time += (
+                            float(time_str.replace(unit, "")) / scale * int(num_str)
+                        )
+                        total_num += int(num_str)
+                        break
+        kernel_times.append(total_time / total_num)
+
+    return tuple(kernel_times) if is_tuple else kernel_times[0]

python/sglang/srt/utils/cuda_ipc_transport_utils.py

@@ -0,0 +1,514 @@
+import fcntl
+import logging
+import threading
+import time
+from multiprocessing import shared_memory
+from typing import Any, Tuple
+
+import numpy as np
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.runtime_context import get_server_args
+from sglang.srt.utils.stale_shm_cleanup import make_shm_name
+
+logger = logging.getLogger(__name__)
+
+MM_FEATURE_CACHE_SIZE = envs.SGLANG_MM_FEATURE_CACHE_MB.get() * 1024 * 1024
+
+MM_ITEM_MEMORY_POOL_RECYCLE_INTERVAL = (
+    envs.SGLANG_MM_ITEM_MEM_POOL_RECYCLE_INTERVAL_SEC.get()
+)
+
+SHM_LOCK_FILE = "/tmp/shm_wr_lock.lock"
+
+
+# Cache for pool-level IPC handles on the consumer side.
+# Key: the pool CUDA IPC handle tuple. Value: opened UntypedStorage.
+_pool_storage_cache: dict = {}
+_pool_cache_lock = threading.Lock()
+
+
+def _normalize_pool_cache_key(pool_handle, pool_device_index: int) -> tuple[Any, ...]:
+    normalized_handle = (
+        pool_handle if isinstance(pool_handle, tuple) else tuple(pool_handle)
+    )
+    return (pool_device_index, normalized_handle)
+
+
+def _open_pooled_storage_uncached(pool_handle):
+    return torch.UntypedStorage._new_shared_cuda(*pool_handle)
+
+
+def _pool_handle_cache_get_or_open(cache_key, pool_handle):
+    storage = _pool_storage_cache.get(cache_key)
+    if storage is None:
+        with _pool_cache_lock:
+            storage = _pool_storage_cache.get(cache_key)
+            if storage is None:
+                storage = _open_pooled_storage_uncached(pool_handle)
+                _pool_storage_cache[cache_key] = storage
+    return storage
+
+
+def _pool_handle_cache_set(cache_key, storage):
+    with _pool_cache_lock:
+        _pool_storage_cache[cache_key] = storage
+
+
+def _pool_handle_cache_invalidate(cache_key):
+    with _pool_cache_lock:
+        _pool_storage_cache.pop(cache_key, None)
+
+
+def _pool_handle_cache_clear():
+    with _pool_cache_lock:
+        _pool_storage_cache.clear()
+
+
+class ShmSyncBuffer:
+    def __init__(self, byte_size: int = 4):
+        self.buffer = shared_memory.SharedMemory(
+            create=True, size=byte_size, name=make_shm_name("sync")
+        )
+        self.buffer_wrapper = np.ndarray(1, dtype=np.float32, buffer=self.buffer.buf)
+        self.buffer_wrapper *= 0
+        self.meta_data = {
+            "handle": self.buffer.name,
+            "shape": self.buffer_wrapper.shape,
+            "dtype": str(self.buffer_wrapper.dtype),
+        }
+
+    def __del__(self):
+        if isinstance(self.buffer, shared_memory.SharedMemory):
+            self.buffer.close()
+            self.buffer.unlink()
+
+
+class MmItemMemoryChunk:
+    def __init__(self, area: Tuple, sync_buffer: ShmSyncBuffer):
+        self.area = area
+        self.sync_flag = sync_buffer
+
+    @property
+    def mem_size(self):
+        return self.area[1] - self.area[0]
+
+    @property
+    def start(self):
+        return self.area[0]
+
+    @property
+    def end(self):
+        return self.area[1]
+
+    def try_to_recycle(self) -> bool:
+        try:
+            tp_num = get_server_args().tp_size
+        except Exception:
+            logger.info(
+                "server_args has not been published yet, skip this turn's recycle"
+            )
+            return False
+
+        val = float(self.sync_flag.buffer_wrapper.item())
+        logger.debug(f"[try_to_recycle] area={self.area}, flag={val}, tp_size={tp_num}")
+
+        if val == float(tp_num):
+            self.sync_flag.buffer_wrapper *= 0.0
+            return True
+
+        return False
+
+
+class MmItemMemoryPool:
+    def __init__(self, memory_size, recycle_interval, base_gpu_id):
+        self.memory_pool = torch.empty(
+            memory_size, dtype=torch.int8, device=f"cuda:{base_gpu_id}"
+        ).contiguous()
+        storage = self.memory_pool.untyped_storage()
+        self._pool_ipc_handle = storage._share_cuda_()
+        self._pool_device_index = self.memory_pool.device.index
+
+        self.sync_flag_list = []
+
+        init_chunk = MmItemMemoryChunk((0, memory_size), self.pop_sync_buffer())
+        self.available_chunks = [init_chunk]
+        self.occupied_chunks = []
+
+        self._lock = threading.Lock()
+        self._pool_full_warned = False
+
+        self._recycle_interval = recycle_interval
+        self._stop_recycler = False
+        self._recycle_thread = threading.Thread(
+            target=self._recycle_loop, name="MmItemMemoryPoolRecycler", daemon=True
+        )
+        self._recycle_thread.start()
+
+        logger.debug(
+            f"[MmItemMemoryPool] init: memory_size={memory_size}, "
+            f"recycle_interval={recycle_interval}s"
+        )
+
+    def shutdown(self):
+        self._stop_recycler = True
+        if self._recycle_thread.is_alive():
+            self._recycle_thread.join(timeout=1.0)
+
+    def _recycle_loop(self):
+        while not self._stop_recycler:
+            try:
+                with self._lock:
+                    self.recycle_chunks()
+                    self.merge_chunks()
+            except Exception as e:
+                logger.warning(
+                    f"[MmItemMemoryPool] recycle loop error: {e}", exc_info=True
+                )
+
+            time.sleep(self._recycle_interval)
+
+    def clear_sync_flag_list(self):
+        # call each chunk's __del__
+        self.sync_flag_list.clear()
+
+    def pop_sync_buffer(self):
+        if len(self.sync_flag_list) == 0:
+            try:
+                new_sync_buffer = ShmSyncBuffer()
+                return new_sync_buffer
+            except:
+                logger.info("allocate shm buffer failed")
+                raise RuntimeError
+        else:
+            return self.sync_flag_list.pop()
+
+    def push_sync_buffer(self, sync_buffer):
+        self.sync_flag_list.append(sync_buffer)
+
+    def get_available_chunk(self, src_tensor: torch.Tensor) -> MmItemMemoryChunk:
+        # find currently available_chunks contain a available chunk or not
+        # if not, return None
+        src_tensor_size = src_tensor.numel() * src_tensor.element_size()
+        min_size = self.memory_pool.numel() * self.memory_pool.element_size() + 1
+        selected_chunk = None
+        for chunk in self.available_chunks:
+            if chunk.mem_size >= src_tensor_size:
+                if chunk.mem_size < min_size:
+                    min_size = chunk.mem_size
+                    selected_chunk = chunk
+
+        if selected_chunk:
+            occupied_chunk_area = (
+                selected_chunk.start,
+                selected_chunk.start + src_tensor_size,
+            )
+            occupied_chunk_sync_flag = selected_chunk.sync_flag
+            new_occupied_chunk = MmItemMemoryChunk(
+                occupied_chunk_area, occupied_chunk_sync_flag
+            )
+
+            self.occupied_chunks.append(new_occupied_chunk)
+            self.available_chunks.remove(selected_chunk)
+
+            available_split_chunk_area = (new_occupied_chunk.end, selected_chunk.end)
+            # add a new chunk
+            if available_split_chunk_area[0] != available_split_chunk_area[1]:
+                split_available_chunk = MmItemMemoryChunk(
+                    available_split_chunk_area, self.pop_sync_buffer()
+                )
+                self.available_chunks.append(split_available_chunk)
+
+            return new_occupied_chunk
+
+        return None
+
+    def return_a_slice_tensor_with_flag(self, src_tensor: torch.Tensor):
+        with self._lock:
+            available_chunk = self.get_available_chunk(src_tensor)
+            if available_chunk is not None:
+                return (
+                    available_chunk.sync_flag.meta_data,
+                    self.memory_pool[available_chunk.start : available_chunk.end],
+                    available_chunk.start,
+                )
+        self._warn_pool_full_once(src_tensor)
+        return None, None, None
+
+    def _warn_pool_full_once(self, src_tensor: torch.Tensor):
+        if self._pool_full_warned:
+            return
+        self._pool_full_warned = True
+        pool_mb = (
+            self.memory_pool.numel() * self.memory_pool.element_size() / (1024 * 1024)
+        )
+        need_mb = src_tensor.numel() * src_tensor.element_size() / (1024 * 1024)
+        logger.warning(
+            "MmItemMemoryPool has no free chunk large enough for a %.2f MiB tensor "
+            "(pool size: %.2f MiB); falling back to non-IPC transport. "
+            "Consider increasing SGLANG_MM_FEATURE_CACHE_MB.",
+            need_mb,
+            pool_mb,
+        )
+
+    def recycle_chunks(self):
+
+        new_occupied_chunks = []
+        for chunk in self.occupied_chunks:
+            if chunk.try_to_recycle():
+                self.available_chunks.append(chunk)
+            else:
+                new_occupied_chunks.append(chunk)
+        self.occupied_chunks = new_occupied_chunks
+
+    def merge_chunks(self):
+        # merge_all_available_chunks
+        merged_chunks = []
+        for chunk in sorted(self.available_chunks, key=lambda x: x.start):
+            if len(merged_chunks) == 0:
+                merged_chunks.append(chunk)
+            else:
+                if chunk.start == merged_chunks[-1].end:
+                    to_merge_chunk = merged_chunks.pop()
+                    to_merge_chunk_sync = to_merge_chunk.sync_flag
+                    merged_chunk_area = (to_merge_chunk.start, chunk.end)
+                    merged_chunks.append(
+                        MmItemMemoryChunk(merged_chunk_area, to_merge_chunk_sync)
+                    )
+                    self.push_sync_buffer(chunk.sync_flag)
+                else:
+                    merged_chunks.append(chunk)
+
+        self.available_chunks = merged_chunks
+
+
+class CudaIpcTensorTransportProxy:
+    """
+    A torch.tensor's proxy used to do inter-process data-sharing
+    including:
+
+    torch.tensor(on gpu)'s cuda-ipc-hande infos
+    a shm sync buffer's meta data which is used to sync between different process
+    """
+
+    def __init__(
+        self,
+        data: torch.Tensor,
+        info_data: torch.Tensor,
+        sync_buffer_meta,
+        pool_ipc_handle=None,
+        pool_byte_offset: int = 0,
+        pool_device_index: int = 0,
+    ):
+
+        if (not isinstance(data, torch.Tensor)) or (
+            not isinstance(info_data, torch.Tensor)
+        ):
+            raise TypeError(
+                f"Input 'data' must be a torch.Tensor, but got {type(data)}"
+            )
+
+        if pool_ipc_handle is not None:
+            self.proxy_state = {
+                "ipc_extra": {
+                    "pool_handle": pool_ipc_handle,
+                    "pool_byte_offset": pool_byte_offset,
+                    "pool_device_index": pool_device_index,
+                    "shape": data.shape,
+                    "dtype": data.dtype,
+                    "stride": data.stride(),
+                    "storage_offset": 0,
+                    "nbytes": data.numel() * data.element_size(),
+                    "recons_shape": info_data.shape,
+                    "recons_dtype": info_data.dtype,
+                },
+                "tensor_data": None,
+            }
+        else:
+            self.proxy_state = self.get_proxy_state(data, info_data)
+        self.reconstruct_tensor = None
+        self.sync_data_meta = sync_buffer_meta
+        self.sync_buffer = None
+
+    @property
+    def get_sync_flag(self):
+        if not self.sync_buffer:
+            shm_name = self.sync_data_meta["handle"]
+            self.sync_buffer = shared_memory.SharedMemory(name=shm_name)
+
+        shape = self.sync_data_meta["shape"]
+        dtype = self.sync_data_meta["dtype"]
+        return np.ndarray(shape, dtype=dtype, buffer=self.sync_buffer.buf)
+
+    def close_shm(self):
+        self.sync_buffer.close()
+        self.sync_buffer = None
+
+    def get_proxy_state(self, data, info_data):
+        # acquire all serialize metadata from _metadata
+        state = {}
+
+        try:
+            storage = data.untyped_storage()
+            handle = storage._share_cuda_()
+
+            state["ipc_extra"] = {
+                "handle": handle,
+                "shape": data.shape,
+                "dtype": data.dtype,
+                "stride": data.stride(),
+                "device_index": data.device.index,
+                "storage_offset": data.storage_offset(),
+                "recons_shape": info_data.shape,
+                "recons_dtype": info_data.dtype,
+            }
+            state["tensor_data"] = None
+        except Exception:
+            # Failed to get CUDA IPC handle (possibly tp). Falling back to default transport.
+            state["ipc_extra"] = None
+            state["tensor_data"] = data
+
+        return state
+
+    def _reconstruct_from_ipc_extra(
+        self, ipc_extra, *, use_cache: bool, rebuild_device_idx: int
+    ):
+        shape = ipc_extra["shape"]
+        dtype = ipc_extra["dtype"]
+        stride = ipc_extra["stride"]
+        # Redirect handle[0] to the consumer's device so _new_shared_cuda's
+        # CUDAGuard stays there; peer access handles the cross-GPU open.
+        pool_handle = ipc_extra["pool_handle"]
+        redirected_handle = (rebuild_device_idx,) + tuple(pool_handle)[1:]
+        target_device = torch.device(f"cuda:{rebuild_device_idx}")
+        cache_key = _normalize_pool_cache_key(pool_handle, rebuild_device_idx)
+
+        with torch.cuda.device(target_device):
+            if use_cache:
+                storage = _pool_handle_cache_get_or_open(cache_key, redirected_handle)
+                storage_to_cache = None
+            else:
+                storage = _open_pooled_storage_uncached(redirected_handle)
+                storage_to_cache = storage
+            slice_storage = storage[
+                ipc_extra["pool_byte_offset"] : ipc_extra["pool_byte_offset"]
+                + ipc_extra["nbytes"]
+            ]
+            slice_tensor = torch.empty(0, dtype=dtype, device=target_device).set_(
+                slice_storage,
+                storage_offset=ipc_extra["storage_offset"],
+                size=shape,
+                stride=stride,
+            )
+
+        return slice_tensor, target_device, cache_key, storage_to_cache
+
+    def _copy_slice_tensor_to_target(
+        self,
+        slice_tensor: torch.Tensor,
+        rebuild_device: torch.device,
+        recons_shape,
+        recons_dtype,
+    ):
+        with torch.cuda.device(rebuild_device):
+            reconstructed_tensor = torch.empty(
+                recons_shape, dtype=recons_dtype, device=rebuild_device
+            ).contiguous()
+            reconstructed_tensor.view(torch.int8).view(-1).copy_(slice_tensor)
+
+            open(SHM_LOCK_FILE, "a").close()
+            # write the shm_sync_buffer with a file lock
+            with open(SHM_LOCK_FILE, "w+") as f:
+                fcntl.flock(f, fcntl.LOCK_EX)
+                sync_flag = self.get_sync_flag
+                sync_flag += 1
+                fcntl.flock(f, fcntl.LOCK_UN)
+
+            self.close_shm()
+
+        return reconstructed_tensor
+
+    def reconstruct_on_target_device(self, rebuild_device_idx):
+        rebuild_device = torch.device(f"cuda:{rebuild_device_idx}")
+        if (
+            isinstance(self.reconstruct_tensor, torch.Tensor)
+            and self.reconstruct_tensor.device == rebuild_device
+        ):
+            return self.reconstruct_tensor
+
+        if self.proxy_state["ipc_extra"]:
+            ipc_extra = self.proxy_state["ipc_extra"]
+            recons_shape = ipc_extra["recons_shape"]
+            recons_dtype = ipc_extra["recons_dtype"]
+
+            if "pool_handle" in ipc_extra:
+                try:
+                    (
+                        slice_tensor,
+                        _target_device,
+                        cache_key,
+                        storage_to_cache,
+                    ) = self._reconstruct_from_ipc_extra(
+                        ipc_extra,
+                        use_cache=True,
+                        rebuild_device_idx=rebuild_device_idx,
+                    )
+                except Exception as e:
+                    cache_key = _normalize_pool_cache_key(
+                        ipc_extra["pool_handle"], rebuild_device_idx
+                    )
+                    logger.info(
+                        "Failed to deserialize from cached pooled CUDA IPC handle (%s). "
+                        "Invalidating cache entry and retrying uncached.",
+                        e,
+                    )
+                    _pool_handle_cache_invalidate(cache_key)
+                    (
+                        slice_tensor,
+                        _target_device,
+                        _cache_key,
+                        storage_to_cache,
+                    ) = self._reconstruct_from_ipc_extra(
+                        ipc_extra,
+                        use_cache=False,
+                        rebuild_device_idx=rebuild_device_idx,
+                    )
+                    if storage_to_cache is not None:
+                        _pool_handle_cache_set(cache_key, storage_to_cache)
+            else:
+                # Non-pooled path: redirect handle[0] the same way as the pooled path.
+                try:
+                    original_handle = ipc_extra["handle"]
+                    redirected_handle = (rebuild_device_idx,) + tuple(original_handle)[
+                        1:
+                    ]
+                    target_device = torch.device(f"cuda:{rebuild_device_idx}")
+                    with torch.cuda.device(target_device):
+                        storage = torch.UntypedStorage._new_shared_cuda(
+                            *redirected_handle
+                        )
+                        slice_tensor = torch.empty(
+                            0, dtype=ipc_extra["dtype"], device=target_device
+                        ).set_(
+                            storage,
+                            storage_offset=ipc_extra["storage_offset"],
+                            size=ipc_extra["shape"],
+                            stride=ipc_extra["stride"],
+                        )
+                except Exception as e:
+                    logger.info("Failed to deserialize from CUDA IPC handle (%s).", e)
+                    raise
+
+            reconstructed_tensor = self._copy_slice_tensor_to_target(
+                slice_tensor, rebuild_device, recons_shape, recons_dtype
+            )
+        elif isinstance(self.proxy_state["tensor_data"], torch.Tensor):
+            reconstructed_tensor = self.proxy_state["tensor_data"].to(
+                rebuild_device, non_blocking=True
+            )
+        else:
+            raise TypeError("invalid proxy_state")
+
+        self.reconstruct_tensor = reconstructed_tensor
+        return self.reconstruct_tensor

python/sglang/srt/utils/cudacore_pyspy_dump_utils.py

@@ -0,0 +1,144 @@
+# 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.
+# ==============================================================================
+"""CUDA core dump and py-spy dump utilities."""
+
+from __future__ import annotations
+
+import logging
+import os
+import platform
+import subprocess
+import time
+from errno import ENXIO
+from pathlib import Path
+from typing import List
+
+import psutil
+
+logger = logging.getLogger(__name__)
+
+
+def _resolve_cuda_coredump_pipe_path(proc: psutil.Process) -> Path:
+    pipe_template = os.environ.get("CUDA_COREDUMP_PIPE")
+    if pipe_template is None:
+        pipe_path = f"corepipe.cuda.{platform.node()}.{proc.pid}"
+    else:
+        pipe_path = (
+            pipe_template.replace("%h", platform.node())
+            .replace("%p", str(proc.pid))
+            .replace("%t", str(int(time.time())))
+        )
+
+    path = Path(pipe_path)
+    if path.is_absolute():
+        return path
+
+    try:
+        return Path(proc.cwd()) / path
+    except (psutil.Error, OSError):
+        return Path.cwd() / path
+
+
+def _is_sglang_scheduler_process(proc: psutil.Process) -> bool:
+    try:
+        proc_title = " ".join(proc.cmdline())
+    except (psutil.Error, OSError):
+        return False
+    return proc_title.startswith("sglang::scheduler")
+
+
+def collect_scheduler_processes() -> List[psutil.Process]:
+    current = psutil.Process()
+    return [
+        proc
+        for proc in current.children(recursive=True)
+        if _is_sglang_scheduler_process(proc)
+    ]
+
+
+def pyspy_dump_schedulers(scheduler_only=False):
+    """py-spy dump on all scheduler in a local node."""
+    if scheduler_only:
+        procs = collect_scheduler_processes()
+        if not procs:
+            logger.error("No sglang scheduler processes found for py-spy dump.")
+            return
+        pids = [proc.pid for proc in procs]
+    else:
+        pids = [psutil.Process().pid]
+    for pid in pids:
+        for attempt, native_flag in enumerate(["--native", ""]):
+            try:
+                cmd = f"py-spy dump {native_flag} --pid {pid}".strip()
+                result = subprocess.run(
+                    cmd, shell=True, capture_output=True, text=True, check=True
+                )
+                logger.error(f"Pyspy dump for PID {pid} ({cmd}):\n{result.stdout}")
+                break
+            except subprocess.CalledProcessError as e:
+                logger.error(f"Pyspy failed ({cmd}). Error: {e.stderr}")
+                if attempt == 1:
+                    logger.error(f"All pyspy dump attempts failed for PID {pid}.")
+
+
+def trigger_cuda_user_coredump(scheduler_only=False):
+    """Trigger CUDA user-induced GPU core dumps by writing to coredump pipes."""
+    if os.environ.get("CUDA_ENABLE_USER_TRIGGERED_COREDUMP") != "1":
+        logger.error(
+            "CUDA user-triggered coredump is not enabled. Set "
+            "CUDA_ENABLE_USER_TRIGGERED_COREDUMP=1 before CUDA initialization."
+        )
+
+    if scheduler_only:
+        procs = collect_scheduler_processes()
+        if not procs:
+            logger.error("No sglang scheduler processes found for CUDA coredump.")
+            return
+    else:
+        procs = [psutil.Process()]
+
+    for proc in procs:
+        pipe_path = _resolve_cuda_coredump_pipe_path(proc)
+        try:
+            fd = os.open(pipe_path, os.O_WRONLY | os.O_NONBLOCK)
+            try:
+                os.write(fd, b"1")
+            finally:
+                os.close(fd)
+            logger.error(
+                "Triggered CUDA user coredump for PID %s via %s",
+                proc.pid,
+                pipe_path,
+            )
+        except FileNotFoundError:
+            logger.error(
+                "CUDA coredump pipe not found for PID %s: %s. Ensure "
+                "CUDA_ENABLE_USER_TRIGGERED_COREDUMP=1 was set before this "
+                "process initialized CUDA.",
+                proc.pid,
+                pipe_path,
+            )
+        except OSError as e:
+            if e.errno == ENXIO:
+                logger.error(
+                    "CUDA coredump pipe has no reader for PID %s: %s",
+                    proc.pid,
+                    pipe_path,
+                )
+            else:
+                logger.exception(
+                    "Failed to trigger CUDA user coredump for PID %s via %s",
+                    proc.pid,
+                    pipe_path,
+                )

python/sglang/srt/utils/custom_op.py

@@ -0,0 +1,337 @@
+from __future__ import annotations
+
+import inspect
+from typing import Any, Callable, List, Optional, TypeVar, Union, overload
+
+import torch
+import torch.library
+
+from sglang.kernel_api_logging import debug_torch_op
+
+F = TypeVar("F", bound=Callable)
+
+
+@overload
+def register_custom_op(
+    fn: F,
+    *,
+    op_name: Optional[str] = None,
+    mutates_args: Optional[List[str]] = None,
+    out_shape: Optional[Union[int, str]] = None,
+    eager: bool = True,
+) -> F: ...
+
+
+@overload
+def register_custom_op(
+    fn: F,
+    *,
+    op_name: Optional[str] = None,
+    mutates_args: Optional[List[str]] = None,
+    fake_impl: Optional[Callable],
+    eager: bool = True,
+) -> F: ...
+
+
+@overload
+def register_custom_op(
+    *,
+    op_name: Optional[str] = None,
+    mutates_args: Optional[List[str]] = None,
+    out_shape: Optional[Union[int, str]] = None,
+    eager: bool = True,
+) -> Callable[[F], F]: ...
+
+
+@overload
+def register_custom_op(
+    *,
+    op_name: Optional[str] = None,
+    mutates_args: Optional[List[str]] = None,
+    fake_impl: Optional[Callable],
+    eager: bool = True,
+) -> Callable[[F], F]: ...
+
+
+# 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
+
+
+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):
+                from sglang.srt.utils.common import direct_register_custom_op
+
+                # 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 = debug_torch_op(self.op_func, 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
+
+
+def register_custom_op_from_extern(
+    fn: Callable,
+    *,
+    op_name: Optional[str] = None,
+    mutates_args: Optional[List[str]] = None,
+    out_shape: Optional[Union[int, str]] = None,
+    out_dtype: Optional[torch.dtype] = None,
+    fake_impl: Optional[Callable] = None,
+    computed_args: Optional[dict] = None,
+) -> Callable:
+    """Wrap an external library function as a custom op for torch.compile compatibility.
+
+    Use this to wrap functions from external libraries (e.g. flashinfer kernels) that
+    perform operations incompatible with torch.compile/dynamo tracing, such as JIT
+    compilation, file I/O, or dynamic module loading.
+
+    The wrapped function becomes an opaque node in the compiled graph. Dynamo will
+    not trace inside it, avoiding tracing failures. A fake implementation is used
+    for shape/dtype propagation during compilation.
+
+    The external function must have type annotations compatible with
+    ``torch.library.infer_schema`` (``torch.Tensor``, ``int``, ``float``, ``bool``,
+    ``Optional[torch.Tensor]``, etc.).
+
+    This function is idempotent: calling it multiple times with the same ``op_name``
+    (or ``fn.__name__``) safely skips re-registration.
+
+    Example usage::
+
+        from flashinfer.fused_moe import trtllm_fp8_block_scale_moe
+
+        trtllm_fp8_block_scale_moe = register_custom_op_from_extern(
+            trtllm_fp8_block_scale_moe,
+            out_shape="hidden_states",
+            out_dtype=torch.bfloat16,
+            computed_args={
+                "tune_max_num_tokens": lambda hidden_states, **kw: next_power_of_2(
+                    hidden_states.shape[0]
+                ),
+            },
+        )
+
+    :param fn: The external function to wrap.
+    :param op_name: The name of the custom operator.
+                    Defaults to ``fn.__name__``.
+    :param mutates_args: A list of argument names that are mutated in-place.
+                         Defaults to ``[]``.
+    :param out_shape: The position (int) or name (str) of the argument whose shape
+                      matches the output tensor. Used to auto-generate a fake
+                      implementation. Set to ``None`` for inplace-only operators.
+    :param out_dtype: Override the output dtype in the fake implementation.
+                      If ``None``, ``torch.empty_like`` is used (same dtype as the
+                      reference tensor). Useful when the output dtype differs from
+                      the input (e.g. fp8 input -> bf16 output).
+    :param fake_impl: A custom fake implementation for shape/dtype propagation.
+                      Only one of ``out_shape`` or ``fake_impl`` should be provided.
+    :param computed_args: A dict mapping argument names to callables. These arguments
+                          are excluded from the custom op schema and computed inside
+                          the op body at runtime. Each callable receives the other
+                          arguments as keyword args and returns the computed value.
+                          Use this for arguments that vary dynamically (e.g.
+                          ``tune_max_num_tokens``) to avoid torch.compile recompilation.
+    :return: The registered custom op callable (``torch.ops.sglang.<op_name>``).
+    """
+    name = op_name or fn.__name__
+    computed_args = computed_args or {}
+
+    assert not (
+        out_shape is not None and fake_impl is not None
+    ), "Only one of `out_shape` or `fake_impl` should be provided."
+
+    # If computed_args specified, create a wrapper with a reduced signature
+    # that computes the excluded args inside the op body.
+    if computed_args:
+        original_fn = fn
+        original_sig = inspect.signature(fn)
+
+        # Build new signature excluding computed args
+        new_params = [
+            p
+            for param_name, p in original_sig.parameters.items()
+            if param_name not in computed_args
+        ]
+        new_sig = original_sig.replace(parameters=new_params)
+
+        def wrapper(*args, **kwargs):
+            bound = new_sig.bind(*args, **kwargs)
+            bound.apply_defaults()
+            # Compute excluded args from the bound arguments
+            for arg_name, compute_fn in computed_args.items():
+                bound.arguments[arg_name] = compute_fn(**bound.arguments)
+            return original_fn(**bound.arguments)
+
+        wrapper.__name__ = fn.__name__
+        wrapper.__qualname__ = fn.__qualname__
+        wrapper.__module__ = fn.__module__
+        wrapper.__signature__ = new_sig  # type: ignore[attr-defined]
+        # Build annotations without computed args, preserving return type
+        wrapper.__annotations__ = {
+            k: v
+            for k, v in getattr(fn, "__annotations__", {}).items()
+            if k not in computed_args
+        }
+        fn = wrapper
+
+    # Generate fake_impl from out_shape if needed
+    fake_sig = inspect.signature(fn)
+    if fake_impl is None and out_shape is not None:
+
+        def _fake_impl(*args, **kwargs):
+            bound = fake_sig.bind(*args, **kwargs)
+            bound.apply_defaults()
+            try:
+                ref = (
+                    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"external function `{name}` with signature `{fake_sig}`."
+                )
+            if out_dtype is not None:
+                return torch.empty(ref.shape, dtype=out_dtype, device=ref.device)
+            return torch.empty_like(ref)
+
+        fake_impl = _fake_impl
+    elif fake_impl is None:
+        fake_impl = lambda *args, **kwargs: None
+
+    from sglang.srt.utils.common import direct_register_custom_op
+
+    direct_register_custom_op(
+        op_name=name,
+        op_func=fn,
+        mutates_args=mutates_args or [],
+        fake_impl=fake_impl,
+    )
+
+    return debug_torch_op(fn, name)

python/sglang/srt/utils/device_timer.py

@@ -0,0 +1,88 @@
+from collections import deque
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import Callable, Deque, Dict, List, Optional
+
+import torch
+
+
+class DeviceTimer:
+    def __init__(self, reporter: Callable):
+        self._intervals: Deque[_TimingInterval] = deque()
+        self._reporters: List[Callable] = [reporter]
+
+    def add_reporter(self, reporter: Callable):
+        self._reporters.append(reporter)
+
+    @contextmanager
+    def wrap(self, metadata: Dict):
+        self._intervals.append(_TimingInterval.create())
+        try:
+            yield
+        finally:
+            self._intervals[-1].end(metadata=metadata)
+            self._report()
+
+    def _report(self):
+        while len(self._intervals) > 0:
+            interval = self._intervals[0]
+            if not interval.end_event.query():
+                break
+
+            self._intervals.popleft()
+            elapsed = interval.elapsed_time() / 1000.0
+            for reporter in self._reporters:
+                reporter(t=elapsed, **interval.metadata)
+
+
+class GapTimer(DeviceTimer):
+    """Measures GPU idle gaps between consecutive uses of a stream.
+
+    Where DeviceTimer.wrap() measures the duration *inside* a block,
+    GapTimer.wrap() measures the time *between* consecutive blocks
+    (gap = next_block_start - last_block_end).
+    """
+
+    def __init__(self, reporter: Callable):
+        super().__init__(reporter)
+        self._pending: Optional[_TimingInterval] = None
+
+    @contextmanager
+    def wrap(self, metadata: Dict):
+        if self._pending is not None:
+            self._pending.end(metadata=metadata)
+            self._intervals.append(self._pending)
+            self._pending = None
+            self._report()
+        try:
+            yield
+        finally:
+            self._pending = _TimingInterval.create()
+
+    def cancel(self):
+        """Discard a pending gap (e.g. server went idle)."""
+        self._pending = None
+
+
+@dataclass
+class _TimingInterval:
+    start_event: torch.cuda.Event
+    end_event: Optional[torch.cuda.Event] = None
+    metadata: Optional[Dict] = None
+
+    @staticmethod
+    def create():
+        start_event = torch.cuda.Event(enable_timing=True)
+        start_event.record()
+        return _TimingInterval(start_event=start_event)
+
+    def end(self, metadata: Dict):
+        end_event = torch.cuda.Event(enable_timing=True)
+        end_event.record()
+
+        assert self.end_event is None
+        self.end_event = end_event
+        self.metadata = metadata
+
+    def elapsed_time(self) -> float:
+        return self.start_event.elapsed_time(self.end_event)

python/sglang/srt/utils/field_validators.py

@@ -0,0 +1,81 @@
+# 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.
+# ==============================================================================
+"""Lightweight, reusable validators for hot-path API fields.
+
+These are intended to be paired with ``pydantic.PlainValidator`` on
+dataclass fields whose JSON shape is large or homogeneously typed, where
+pydantic's default per-element walk has been measured to dominate
+request latency.
+
+Usage::
+
+    from typing import Annotated, List, Optional, Union
+    from pydantic import PlainValidator
+    from sglang.srt.utils.field_validators import validate_optional_list_i64_1d_2d
+
+    @dataclass
+    class MyReq:
+        input_ids: Annotated[
+            Optional[Union[List[List[int]], List[int]]],
+            PlainValidator(validate_optional_list_i64_1d_2d),
+        ] = None
+"""
+
+from __future__ import annotations
+
+from array import array
+from typing import Any
+
+
+def validate_list_i64_1d(v: Any) -> list[int]:
+    """Validates type: list[int]"""
+    if v is None:
+        raise ValueError("must not be None")
+    if not isinstance(v, list):
+        raise ValueError(f"must be list; got {type(v).__name__}")
+    if not v:
+        return v
+    if not isinstance(v[0], int):
+        raise ValueError(f"elements must be int; got {type(v[0]).__name__}")
+    try:
+        array("q", v)
+    except (TypeError, OverflowError) as e:
+        raise ValueError(f"contains non-int64 element: {e}") from None
+    return v
+
+
+def validate_optional_list_i64_1d_2d(
+    v: Any,
+) -> list[int] | list[list[int]] | None:
+    """Validates type: list[int] | list[list[int]] | None"""
+    if v is None:
+        # Accept None
+        return v
+    if not isinstance(v, list):
+        raise ValueError(f"must be list or null; got {type(v).__name__}")
+    if not v:
+        # Accept empty list
+        return v
+    if isinstance(v[0], int):
+        # Accept list[int]
+        return validate_list_i64_1d(v)
+    if isinstance(v[0], list):
+        # Accept list[list[int]]
+        for i, row in enumerate(v):
+            try:
+                validate_list_i64_1d(row)
+            except ValueError as e:
+                raise ValueError(f"row {i}: {e}") from None
+        return v
+    raise ValueError(f"elements must be int or list; got {type(v[0]).__name__}")

python/sglang/srt/utils/gauge_histogram.py

@@ -0,0 +1,76 @@
+"""Gauge with gt/le bucket labels for Grafana heatmap visualization.
+
+Unlike Prometheus Histogram which uses cumulative buckets, this uses
+non-cumulative buckets (gt < value <= le) suitable for heatmap display.
+
+Note: Keep in sync with Rust implementation in
+sgl-model-gateway/src/observability/gauge_histogram.rs
+"""
+
+import bisect
+from typing import Dict, Iterator, List, Tuple, Union
+
+
+class BucketLabels:
+    """Bucket label pairs and count computation for a GaugeHistogram."""
+
+    def __init__(self, upper_bounds: List[Union[int, float]]):
+        self._upper_bounds = upper_bounds
+        self._labels: List[Tuple[str, str]] = []
+        for i, upper in enumerate(upper_bounds):
+            lower = upper_bounds[i - 1] if i > 0 else 0
+            self._labels.append((str(lower), str(upper)))
+        self._labels.append((str(upper_bounds[-1]), "+Inf"))
+
+    def __len__(self) -> int:
+        return len(self._labels)
+
+    def __iter__(self) -> Iterator[Tuple[str, str]]:
+        return iter(self._labels)
+
+    def compute_bucket_counts(self, observations: List[Union[int, float]]) -> List[int]:
+        """Compute how many observations fall into each bucket. O(n) complexity."""
+        counts = [0] * len(self)
+        for v in observations:
+            # bisect_left finds insertion point; values at boundary go to current bucket
+            idx = bisect.bisect_left(self._upper_bounds, v)
+            counts[idx] += 1
+        return counts
+
+
+class GaugeHistogram:
+    """Gauge with gt/le bucket labels for Grafana heatmap visualization."""
+
+    def __init__(
+        self,
+        name: str,
+        documentation: str,
+        labelnames: List[str],
+        bucket_bounds: List[Union[int, float]],
+        multiprocess_mode: str = "mostrecent",
+    ):
+        from prometheus_client import Gauge
+
+        self._buckets = BucketLabels(bucket_bounds)
+
+        self._gauge = Gauge(
+            name=name,
+            documentation=documentation,
+            labelnames=list(labelnames) + ["gt", "le"],
+            multiprocess_mode=multiprocess_mode,
+        )
+
+    def set_raw(self, labels: Dict[str, str], values: List[int]):
+        """Set bucket counts directly."""
+        for (gt, le), count in zip(self._buckets, values):
+            self._gauge.labels(**labels, gt=gt, le=le).set(count)
+
+    def set_by_current_observations(
+        self, labels: Dict[str, str], observations: List[Union[int, float]]
+    ):
+        """Compute bucket counts from observations and set them."""
+        counts = self._buckets.compute_bucket_counts(observations)
+        self.set_raw(labels, counts)
+
+    def buckets(self) -> BucketLabels:
+        return self._buckets

python/sglang/srt/utils/hf_transformers/__init__.py

@@ -0,0 +1,65 @@
+# 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.
+# ==============================================================================
+"""Hugging Face Transformers utilities.
+
+This package provides HF Transformers helpers, split into submodules
+(common, config, tokenizer, processor, mistral_utils).  Compatibility
+monkey-patches live in the sibling ``sglang.srt.utils.hf_transformers_patches``
+module and are applied at sglang import time.
+All public symbols are re-exported here for convenience.  The old import
+path ``sglang.srt.utils.hf_transformers_utils`` is preserved by a
+separate shim module.
+"""
+
+from ..hf_transformers_patches import normalize_rope_scaling_compat
+from .common import (
+    CONTEXT_LENGTH_KEYS,
+    AutoConfig,
+    attach_additional_stop_token_ids,
+    check_gguf_file,
+    download_from_hf,
+    get_context_length,
+    get_generation_config,
+    get_hf_text_config,
+    get_rope_config,
+    get_sparse_attention_config,
+    get_tokenizer_from_processor,
+)
+from .config import get_config
+from .processor import get_processor
+from .tokenizer import (
+    _fix_added_tokens_encoding,
+    _fix_v5_add_bos_eos_token,
+    get_tokenizer,
+)
+
+__all__ = [
+    "AutoConfig",
+    "CONTEXT_LENGTH_KEYS",
+    "_fix_added_tokens_encoding",
+    "_fix_v5_add_bos_eos_token",
+    "attach_additional_stop_token_ids",
+    "check_gguf_file",
+    "download_from_hf",
+    "get_config",
+    "get_context_length",
+    "get_generation_config",
+    "get_hf_text_config",
+    "get_processor",
+    "get_rope_config",
+    "get_sparse_attention_config",
+    "get_tokenizer",
+    "get_tokenizer_from_processor",
+    "normalize_rope_scaling_compat",
+]

python/sglang/srt/utils/hf_transformers/common.py

@@ -0,0 +1,499 @@
+# 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.
+# ==============================================================================
+"""Shared helpers used by config, tokenizer, and processor modules."""
+
+import json
+import os
+from pathlib import Path
+from typing import Any, Dict, Optional, Type, Union
+
+import torch
+from huggingface_hub import snapshot_download
+
+from sglang.srt.configs import (
+    AfmoeConfig,
+    BailingHybridConfig,
+    ChatGLMConfig,
+    DbrxConfig,
+    DeepseekVL2Config,
+    DotsOCRConfig,
+    DotsVLMConfig,
+    ExaoneConfig,
+    FalconH1Config,
+    GraniteMoeHybridConfig,
+    InternS2PreviewConfig,
+    JetNemotronConfig,
+    JetVLMConfig,
+    KimiK25Config,
+    KimiLinearConfig,
+    KimiVLConfig,
+    LagunaConfig,
+    LocateAnythingConfig,
+    LongcatFlashConfig,
+    MiniCPMV4_6Config,
+    MiniCPMV4_6VisionConfig,
+    MiniMaxM3VLConfig,
+    MultiModalityConfig,
+    NemotronH_Nano_Omni_Reasoning_V3_Config,
+    NemotronH_Nano_VL_V2_Config,
+    NemotronHConfig,
+    NemotronHPuzzleConfig,
+    Olmo3Config,
+    Qwen3_5Config,
+    Qwen3_5MoeConfig,
+    Qwen3NextConfig,
+    Step3p5Config,
+    Step3p7Config,
+    Step3VLConfig,
+)
+from sglang.srt.configs.deepseek_ocr import DeepseekVLV2Config
+from sglang.srt.configs.internvl import InternVLChatConfig
+from sglang.srt.utils import get_bool_env_var, logger, lru_cache_frozenset
+from sglang.srt.utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
+
+from ..hf_transformers_patches import normalize_rope_scaling_compat
+
+if get_bool_env_var("SGLANG_USE_MODELSCOPE"):
+    from modelscope import AutoConfig, GenerationConfig
+else:
+    from transformers import AutoConfig, GenerationConfig
+
+from transformers import PretrainedConfig
+
+# ---------------------------------------------------------------------------
+# Config registry
+# ---------------------------------------------------------------------------
+
+_CONFIG_REGISTRY: Dict[str, Type[PretrainedConfig]] = {
+    cls.model_type: cls
+    for cls in [
+        AfmoeConfig,
+        BailingHybridConfig,
+        ChatGLMConfig,
+        DbrxConfig,
+        ExaoneConfig,
+        DeepseekVL2Config,
+        MultiModalityConfig,
+        KimiVLConfig,
+        LocateAnythingConfig,
+        InternVLChatConfig,
+        LagunaConfig,
+        Step3VLConfig,
+        LongcatFlashConfig,
+        Olmo3Config,
+        KimiLinearConfig,
+        Qwen3NextConfig,
+        FalconH1Config,
+        GraniteMoeHybridConfig,
+        DotsVLMConfig,
+        DotsOCRConfig,
+        NemotronH_Nano_VL_V2_Config,
+        NemotronH_Nano_Omni_Reasoning_V3_Config,
+        NemotronHConfig,
+        NemotronHPuzzleConfig,
+        DeepseekVLV2Config,
+        Qwen3_5Config,
+        Qwen3_5MoeConfig,
+        InternS2PreviewConfig,
+        JetNemotronConfig,
+        JetVLMConfig,
+        KimiK25Config,
+        Step3p5Config,
+        Step3p7Config,
+        MiniCPMV4_6Config,
+        MiniCPMV4_6VisionConfig,
+        MiniMaxM3VLConfig,
+    ]
+}
+
+# DeepSeek V3.2 / V4 reuse the V3 config schema. Subclass the upstream
+# transformers class with each model_type so AutoConfig.register passes its
+# consistency check (which requires class.model_type == registered key).
+# Default-value divergences (e.g. V4's topk_group) are handled in
+# model_config.py post-load.
+try:
+    from transformers import DeepseekV3Config as _HFDeepseekV3Config
+
+    class _DeepseekV32ConfigAlias(_HFDeepseekV3Config):
+        model_type = "deepseek_v32"
+
+    class _DeepseekV4ConfigAlias(_HFDeepseekV3Config):
+        model_type = "deepseek_v4"
+
+    _CONFIG_REGISTRY["deepseek_v32"] = _DeepseekV32ConfigAlias
+    _CONFIG_REGISTRY["deepseek_v4"] = _DeepseekV4ConfigAlias
+
+    # For kimi_k25_eagle3
+    class _KimiK2ConfigAlias(_HFDeepseekV3Config):
+        model_type = "kimi_k2"
+
+    _CONFIG_REGISTRY["kimi_k2"] = _KimiK2ConfigAlias
+except ImportError:
+    pass
+
+try:
+    from transformers import Gemma4Config as _HFGemma4Config
+
+    class _Gemma4UnifiedConfigAlias(_HFGemma4Config):
+        model_type = "gemma4_unified"
+
+    _CONFIG_REGISTRY["gemma4_unified"] = _Gemma4UnifiedConfigAlias
+except ImportError:
+    pass
+
+for name, cls in _CONFIG_REGISTRY.items():
+    try:
+        AutoConfig.register(name, cls)
+    except ValueError as e:
+        err = str(e).lower()
+        if "already registered" not in err and "already used" not in err:
+            logger.warning("Failed to register config %s: %s", name, e)
+
+
+# ---------------------------------------------------------------------------
+# Download / path helpers
+# ---------------------------------------------------------------------------
+
+
+def download_from_hf(
+    model_path: str,
+    allow_patterns: Optional[Union[str, list]] = None,
+):
+    if os.path.exists(model_path):
+        return model_path
+
+    if not allow_patterns:
+        allow_patterns = ["*.json", "*.bin", "*.model"]
+
+    return snapshot_download(model_path, allow_patterns=allow_patterns)
+
+
+def resolve_runai_obj_uri(model_name_or_path: str) -> str:
+    if is_runai_obj_uri(model_name_or_path):
+        return ObjectStorageModel.get_path(model_name_or_path)
+    return model_name_or_path
+
+
+def _resolve_local_or_cached_file(model_name_or_path, filename, revision=None):
+    """Resolve a file from a local directory or HF hub cache (no network)."""
+    local_path = Path(model_name_or_path) / filename
+    if local_path.is_file():
+        return str(local_path)
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(
+        model_name_or_path, filename, revision=revision, local_files_only=True
+    )
+
+
+def _cached_file_exists(model_name_or_path, filename, revision=None) -> bool:
+    """Whether *filename* is available locally or in the HF cache (no network)."""
+    try:
+        _resolve_local_or_cached_file(model_name_or_path, filename, revision)
+        return True
+    except Exception:
+        return False
+
+
+def _remote_file_exists(repo_id, filename, revision=None) -> bool:
+    """Whether *filename* exists on the HF hub (HEAD request only, no download).
+
+    Returns False on any error (offline, gated, network, invalid id) so callers
+    fall back to their default path instead of crashing.
+    """
+    from huggingface_hub.constants import HF_HUB_OFFLINE
+
+    if HF_HUB_OFFLINE:
+        return False
+    try:
+        from huggingface_hub import HfApi
+
+        return HfApi().file_exists(repo_id, filename, revision=revision)
+    except Exception:
+        return False
+
+
+def check_gguf_file(model: Union[str, os.PathLike]) -> bool:
+    model = Path(model)
+    if not model.is_file():
+        return False
+    elif model.suffix == ".gguf":
+        return True
+
+    with open(model, "rb") as f:
+        header = f.read(4)
+    return header == b"GGUF"
+
+
+# ---------------------------------------------------------------------------
+# Rope / text config helpers
+# ---------------------------------------------------------------------------
+
+
+def get_rope_config(config):
+    """Get (rope_theta, rope_params) from config, supporting both v4 and v5.
+
+    Trust-remote-code configs or parent configs passed to sub-models may not
+    have the v5 ``rope_parameters`` property, so we fall back to the v4-style
+    ``config.rope_theta`` / ``config.rope_scaling`` attributes.
+
+    Returns:
+        (rope_theta, rope_params): In v5, rope_params is the full
+        rope_parameters dict (which subsumes rope_scaling and includes
+        rope_theta). In v4, rope_params is the rope_scaling dict or None.
+    """
+    rope_params = getattr(config, "rope_parameters", None)
+    if rope_params is not None:
+        return rope_params["rope_theta"], rope_params
+    return getattr(config, "rope_theta", 10000), getattr(config, "rope_scaling", None)
+
+
+def _patch_text_config(parent_config: PretrainedConfig, text_config):
+    """Synchronize standard attributes between parent config and text sub-config.
+
+    In transformers v5, the "untangle config" refactor removed automatic
+    inheritance of top-level PretrainedConfig attributes (pad_token_id,
+    tie_word_embeddings, etc.) from sub-configs. Downstream code expects
+    these attributes to be present on both configs (some models pass the
+    parent directly to the language model, others pass the text sub-config),
+    so we propagate in both directions when an attribute is missing.
+    (See https://github.com/huggingface/transformers/pull/41541)
+    """
+    _ATTRS_TO_PROPAGATE = [
+        "pad_token_id",
+        "bos_token_id",
+        "eos_token_id",
+        "tie_word_embeddings",
+    ]
+    for attr in _ATTRS_TO_PROPAGATE:
+        parent_has = hasattr(parent_config, attr)
+        text_has = hasattr(text_config, attr)
+        if parent_has and not text_has:
+            setattr(text_config, attr, getattr(parent_config, attr))
+        elif text_has and not parent_has:
+            setattr(parent_config, attr, getattr(text_config, attr))
+    return text_config
+
+
+def get_hf_text_config(config: PretrainedConfig):
+    """Get the "sub" config relevant to llm for multi modal models.
+    No op for pure text models.
+    """
+    if config.architectures is not None:
+        class_name = config.architectures[0]
+        if class_name.startswith("Llava") and class_name.endswith("ForCausalLM"):
+            # We support non-hf version of llava models, so we do not want to
+            # read the wrong values from the unused default text_config.
+            # NOTE(HandH1998): We set `torch_dtype` of config to `torch.float16` for the weights, as
+            # `torch.float16` is default used for image features in `python/sglang/srt/models/llava.py`.
+            setattr(config, "dtype", torch.float16)
+            return config
+
+    text_config = None
+
+    # Some models (e.g. DeepSeek-OCR) store sub-configs as plain dicts.
+    # Convert to PretrainedConfig early so hasattr() checks and asserts work.
+    parent_dtype = getattr(config, "dtype", None)
+    for _attr in ("text_config", "llm_config", "language_config", "thinker_config"):
+        _sub = getattr(config, _attr, None)
+        if isinstance(_sub, dict):
+            _converted = PretrainedConfig(**_sub)
+            if getattr(_converted, "dtype", None) is None and parent_dtype is not None:
+                _converted.dtype = parent_dtype
+            setattr(config, _attr, _converted)
+        elif _sub is not None and parent_dtype is not None:
+            # transformers v5 multimodal configs (e.g. Mistral3Config) carry
+            # `dtype` only on the top-level config, leaving the sub-configs at
+            # None. Without this, _get_and_verify_dtype falls back to float32
+            # and then "auto" downcasts to float16, which overflows the Pixtral
+            # vision tower on real images and produces NaN features.
+            if getattr(_sub, "dtype", None) is None:
+                _sub.dtype = parent_dtype
+
+    # Priority: thinker_config > llm_config > language_config > text_config
+    if hasattr(config, "thinker_config"):
+        # qwen2.5 omni
+        thinker_config = config.thinker_config
+        if hasattr(thinker_config, "text_config"):
+            setattr(
+                thinker_config.text_config,
+                "dtype",
+                getattr(thinker_config, "dtype", None),
+            )
+            text_config = thinker_config.text_config
+        else:
+            text_config = thinker_config
+    elif hasattr(config, "llm_config"):
+        # PointsV1.5 Chat Model
+        assert hasattr(config.llm_config, "num_attention_heads")
+        text_config = config.llm_config
+    elif hasattr(config, "language_config"):
+        text_config = config.language_config
+    elif hasattr(config, "text_config"):
+        # The code operates under the assumption that text_config should have
+        # `num_attention_heads` (among others). Assert here to fail early
+        # if transformers config doesn't align with this assumption.
+        assert hasattr(config.text_config, "num_attention_heads")
+        text_config = config.text_config
+
+    # Ensure rope_scaling dicts have "type" for remote-code compat (v5).
+    normalize_rope_scaling_compat(config)
+
+    if text_config is not None:
+        return _patch_text_config(config, text_config)
+    return config
+
+
+# ---------------------------------------------------------------------------
+# Model-specific helpers
+# ---------------------------------------------------------------------------
+
+
+def _ensure_sub_configs(config: PretrainedConfig, *attr_names: str) -> None:
+    """Convert dict-valued sub-configs to proper AutoConfig objects in-place."""
+    for attr in attr_names:
+        sub = getattr(config, attr, None)
+        if sub is not None and isinstance(sub, dict):
+            setattr(config, attr, AutoConfig.for_model(**sub))
+
+
+def _is_deepseek_ocr_model(config: PretrainedConfig) -> bool:
+    # TODO: Remove this workaround once AutoConfig correctly identifies deepseek-ocr.
+    # Hugging Face's AutoConfig currently misidentifies it as deepseekvl2.
+    auto_map = getattr(config, "auto_map", None) or {}
+    return auto_map.get("AutoModel") == "modeling_deepseekocr.DeepseekOCRForCausalLM"
+
+
+def _is_deepseek_ocr2_model(config: PretrainedConfig) -> bool:
+    auto_map = getattr(config, "auto_map", None) or {}
+    return auto_map.get("AutoModel") == "modeling_deepseekocr2.DeepseekOCR2ForCausalLM"
+
+
+def _override_v_head_dim_if_zero(config: PretrainedConfig, patch: int = 128) -> None:
+    patched = False
+    for attr in ("text_config", "language_config"):
+        sub = getattr(config, attr, None)
+        if sub is None:
+            continue
+        if isinstance(sub, dict):
+            if sub.get("v_head_dim") == 0:
+                sub["v_head_dim"] = patch
+                patched = True
+        elif getattr(sub, "v_head_dim", None) == 0:
+            sub.v_head_dim = patch
+            patched = True
+    if patched:
+        logger.warning(
+            f"Overriding v_head_dim from 0 to {patch} to avoid potential issues."
+        )
+
+
+# ---------------------------------------------------------------------------
+# Context length / generation config / sparse attention
+# ---------------------------------------------------------------------------
+
+# 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
+# have a preference for which value gets used.
+CONTEXT_LENGTH_KEYS = [
+    "max_sequence_length",
+    "seq_length",
+    "max_seq_len",
+    "model_max_length",
+    "max_position_embeddings",
+]
+
+
+def get_context_length(config):
+    """Get the context length of a model from a huggingface model configs."""
+    text_config = config
+    rope_scaling = getattr(text_config, "rope_scaling", None)
+    if rope_scaling:
+        rope_scaling_factor = rope_scaling.get("factor", 1)
+        if "original_max_position_embeddings" in rope_scaling:
+            rope_scaling_factor = 1
+        if rope_scaling.get("rope_type", None) == "llama3":
+            rope_scaling_factor = 1
+    else:
+        rope_scaling_factor = 1
+
+    for key in CONTEXT_LENGTH_KEYS:
+        val = getattr(text_config, key, None)
+        if val is not None:
+            return int(rope_scaling_factor * val)
+    return 2048
+
+
+@lru_cache_frozenset(maxsize=32)
+def get_generation_config(
+    model: str,
+    trust_remote_code: bool,
+    revision: Optional[str] = None,
+    **kwargs,
+):
+    try:
+        return GenerationConfig.from_pretrained(
+            model, trust_remote_code=trust_remote_code, revision=revision, **kwargs
+        )
+    except FileNotFoundError:
+        return None
+    except OSError as e:
+        logger.warning(
+            "Failed to load generation config for %s: %s. "
+            "Proceeding without generation config.",
+            model,
+            e,
+        )
+        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:
+        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 {}
+
+    with open(config_file) as f:
+        config = json.load(f)
+    return config
+
+
+# ---------------------------------------------------------------------------
+# Tokenizer / processor helpers
+# ---------------------------------------------------------------------------
+
+
+# Some models don't have an available processor, e.g.: InternVL
+def get_tokenizer_from_processor(processor):
+    from transformers import PreTrainedTokenizerBase
+
+    if isinstance(processor, PreTrainedTokenizerBase):
+        return processor
+    return processor.tokenizer
+
+
+def attach_additional_stop_token_ids(tokenizer):
+    added = tokenizer.get_added_vocab()
+    if "<|eom_id|>" in added:
+        tokenizer.additional_stop_token_ids = {added["<|eom_id|>"]}
+    else:
+        tokenizer.additional_stop_token_ids = None

python/sglang/srt/utils/hf_transformers/config.py

@@ -0,0 +1,264 @@
+# 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.
+# ==============================================================================
+"""Config loading utilities."""
+
+from pathlib import Path
+from typing import Optional
+
+from transformers import PretrainedConfig
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+
+from sglang.srt.configs.model_config_parser_registry import (
+    ModelConfigParserBase,
+    get_model_config_parser,
+    register_model_config_parser,
+)
+from sglang.srt.connector import create_remote_connector
+from sglang.srt.utils import is_remote_url, lru_cache_frozenset
+
+from ..hf_transformers_patches import _ensure_gguf_version
+from .common import (
+    _CONFIG_REGISTRY,
+    AutoConfig,
+    DeepseekVLV2Config,
+    _is_deepseek_ocr2_model,
+    _is_deepseek_ocr_model,
+    _override_v_head_dim_if_zero,
+    check_gguf_file,
+    get_hf_text_config,
+    resolve_runai_obj_uri,
+)
+from .mistral_utils import is_mistral_model, load_mistral_config
+
+
+def _set_architectures(config, arch_name):
+    config.update({"architectures": [arch_name]})
+
+
+def _apply_deepseek_ocr_overrides(config, model):
+    _override_v_head_dim_if_zero(config)
+    _set_architectures(config, "DeepseekOCRForCausalLM")
+    config._name_or_path = model
+
+
+_LONGCAT_ARCHS = {
+    "LongcatCausalLM",
+    "LongcatFlashForCausalLM",
+    "LongcatFlashNgramForCausalLM",
+}
+
+
+def _try_load_longcat_config(model, revision: Optional[str], **kwargs):
+    config_dict, _ = PretrainedConfig.get_config_dict(
+        model, revision=revision, **kwargs
+    )
+    architectures = config_dict.get("architectures") or []
+    if not any(arch in _LONGCAT_ARCHS for arch in architectures):
+        return None
+
+    return _CONFIG_REGISTRY["longcat_flash"].from_pretrained(
+        model, revision=revision, **kwargs
+    )
+
+
+@register_model_config_parser("hf")
+class HfModelConfigParser(ModelConfigParserBase):
+    def parse(
+        self,
+        model,
+        trust_remote_code: bool,
+        revision: Optional[str] = None,
+        **kwargs,
+    ):
+        config = _try_load_longcat_config(model, revision, **kwargs)
+        if config is None:
+            config = AutoConfig.from_pretrained(
+                model,
+                trust_remote_code=trust_remote_code,
+                revision=revision,
+                **kwargs,
+            )
+
+        if (
+            config.architectures is not None
+            and config.architectures[0] == "Phi4MMForCausalLM"
+        ):
+            from transformers import SiglipVisionConfig
+
+            config.vision_config = SiglipVisionConfig(
+                hidden_size=1152,
+                image_size=448,
+                intermediate_size=4304,
+                model_type="siglip_vision_model",
+                num_attention_heads=16,
+                num_hidden_layers=26,
+                patch_size=14,
+            )
+
+        if config.architectures in [
+            ["LongcatCausalLM"],
+            ["LongcatFlashForCausalLM"],
+            ["LongcatFlashNgramForCausalLM"],
+        ]:
+            config.model_type = "longcat_flash"
+
+        text_config = get_hf_text_config(config=config)
+
+        if isinstance(model, str) and text_config is not None:
+            items = (
+                text_config.items()
+                if hasattr(text_config, "items")
+                else vars(text_config).items()
+            )
+            for key, val in items:
+                if not hasattr(config, key) and val is not None:
+                    setattr(config, key, val)
+
+        is_ocr = _is_deepseek_ocr_model(config)
+        is_ocr2 = _is_deepseek_ocr2_model(config)
+
+        if is_ocr2:
+            _override_v_head_dim_if_zero(config)
+            config.model_type = "deepseek-ocr"
+            _set_architectures(config, "DeepseekOCRForCausalLM")
+            config = DeepseekVLV2Config.from_pretrained(model, revision=revision)
+            _apply_deepseek_ocr_overrides(config, model)
+        elif config.model_type in _CONFIG_REGISTRY:
+            model_type = config.model_type
+            if model_type == "deepseek_vl_v2" and is_ocr:
+                model_type = "deepseek-ocr"
+            config = _CONFIG_REGISTRY[model_type].from_pretrained(
+                model, revision=revision
+            )
+
+            # Re-check after reloading config from registry
+            if _is_deepseek_ocr_model(config) or _is_deepseek_ocr2_model(config):
+                _apply_deepseek_ocr_overrides(config, model)
+            else:
+                config._name_or_path = model
+
+        if isinstance(model, str) and config.model_type == "internvl_chat":
+            for key, val in config.llm_config.__dict__.items():
+                if not hasattr(config, key):
+                    setattr(config, key, val)
+
+        if config.model_type == "multi_modality":
+            _set_architectures(config, "MultiModalityCausalLM")
+
+        if config.model_type in (
+            "gemma4",
+            "gemma4_assistant",
+            "gemma4_unified",
+            "gemma4_unified_assistant",
+        ):
+            # Gemma4 configs use base attributes for SWA layers and `global_*`
+            # variants for full-attention layers.  SGLang expects the opposite:
+            # base = full-attention, `swa_*` = sliding-window overrides.
+            text_config = config.text_config
+            global_head_dim = getattr(text_config, "global_head_dim", None)
+            global_kv_heads = getattr(text_config, "num_global_key_value_heads", None)
+
+            swa_head_dim = text_config.head_dim
+            swa_kv_heads = text_config.num_key_value_heads
+
+            text_config.swa_head_dim = swa_head_dim
+            text_config.swa_v_head_dim = swa_head_dim
+            text_config.swa_num_key_value_heads = swa_kv_heads
+
+            if global_head_dim is not None:
+                text_config.head_dim = global_head_dim
+            if global_kv_heads is not None:
+                text_config.num_key_value_heads = global_kv_heads
+
+            if not hasattr(text_config, "v_head_dim"):
+                text_config.v_head_dim = text_config.head_dim
+            if not hasattr(text_config, "swa_v_head_dim"):
+                text_config.swa_v_head_dim = text_config.swa_head_dim
+
+            # Unified Gemma4 names the end-of-audio token `eoa_token_index`,
+            # but the multimodal processor expects `eoa_token_id`.
+            if not hasattr(config, "eoa_token_id") and hasattr(
+                config, "eoa_token_index"
+            ):
+                config.eoa_token_id = config.eoa_token_index
+
+        if config.model_type == "longcat_flash":
+            _set_architectures(config, "LongcatFlashForCausalLM")
+
+        return config
+
+
+@register_model_config_parser("mistral")
+class MistralModelConfigParser(ModelConfigParserBase):
+    def parse(
+        self,
+        model,
+        trust_remote_code: bool,
+        revision: Optional[str] = None,
+        **kwargs,
+    ):
+        del kwargs
+        return load_mistral_config(
+            model, trust_remote_code=trust_remote_code, revision=revision
+        )
+
+
+@lru_cache_frozenset(maxsize=32)
+def get_config(
+    model: str,
+    trust_remote_code: bool,
+    revision: Optional[str] = None,
+    model_override_args: Optional[dict] = None,
+    model_config_parser: str = "auto",
+    **kwargs,
+):
+    is_gguf = check_gguf_file(model)
+    if is_gguf:
+        if model_config_parser not in ("auto", "hf"):
+            raise ValueError(
+                f"model_config_parser={model_config_parser!r} is incompatible "
+                "with GGUF inputs; only 'hf' (or 'auto') is supported."
+            )
+        _ensure_gguf_version()
+        kwargs["gguf_file"] = model
+        model = Path(model).parent
+        # Skip auto-resolution for GGUF: the name-based Mistral heuristic
+        # would misfire on the rewritten parent dir.
+        model_config_parser = "hf"
+
+    model = resolve_runai_obj_uri(model)
+
+    if is_remote_url(model):
+        client = create_remote_connector(model)
+        client.pull_files(ignore_pattern=["*.pt", "*.safetensors", "*.bin"])
+        model = client.get_local_dir()
+
+    if model_config_parser == "auto":
+        # `model` is post-rewrite (gguf parent / runai uri / remote pull).
+        model_config_parser = "mistral" if is_mistral_model(model) else "hf"
+
+    parser = get_model_config_parser(model_config_parser)
+    config = parser.parse(
+        model, trust_remote_code=trust_remote_code, revision=revision, **kwargs
+    )
+
+    if model_override_args:
+        config.update(model_override_args)
+
+    if is_gguf:
+        if config.model_type not in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES:
+            raise RuntimeError(f"Can't get gguf config for {config.model_type}.")
+        _set_architectures(config, MODEL_FOR_CAUSAL_LM_MAPPING_NAMES[config.model_type])
+
+    return config

python/sglang/srt/utils/hf_transformers/mistral_utils.py

@@ -0,0 +1,637 @@
+# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/transformers_utils/configs/mistral.py
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
+import tempfile
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, Optional
+
+from transformers import AutoConfig, PretrainedConfig, WhisperConfig
+
+from sglang.srt.utils import logger
+
+from .common import (
+    _cached_file_exists,
+    _ensure_sub_configs,
+    _remote_file_exists,
+    download_from_hf,
+)
+
+
+def adapt_config_dict(
+    config_dict: dict[str, Any], model: str, **kwargs
+) -> tuple[dict, PretrainedConfig]:
+    config_dict.update(kwargs)
+    config_dict = _remap_general_mistral_args(config_dict)
+
+    if bool(config_dict.get("quantization")):
+        config_dict = _remap_mistral_quantization_args(config_dict)
+
+    is_moe = bool(config_dict.get("moe"))
+    is_mistral_large_3 = (
+        is_moe and (config_dict["moe"].get("num_shared_experts") or 0) > 0
+    )
+    is_eagle = "eagle" in model.lower()
+    is_mla_eagle = is_eagle and any(
+        config_dict.get(k) is not None
+        for k in ("kv_lora_rank", "q_lora_rank", "v_head_dim")
+    )
+    if is_eagle and not is_moe and is_mla_eagle:
+        # Dense MLA EAGLE draft model (e.g. Mistral Small 4 EAGLE).
+        # Uses MLA attention like MistralLarge3 but has no MoE layers.
+        # Set model_type to deepseek_v3 for MLA support, and override
+        # MoE fields so all layers are dense.
+        config_dict["model_type"] = "deepseek_v3"
+        config_dict["architectures"] = ["MistralLarge3ForCausalLMEagle"]
+        num_layers = config_dict.get("num_hidden_layers", 0)
+        config_dict["n_routed_experts"] = 1
+        config_dict["first_k_dense_replace"] = num_layers
+        config_dict["moe_layer_freq"] = 1
+        config_dict["n_shared_experts"] = 0
+        config_dict["n_group"] = 1
+        config_dict["topk_group"] = 1
+        config_dict["num_experts_per_tok"] = 1
+        config_dict["moe_intermediate_size"] = 1
+        config_dict["routed_scaling_factor"] = 1.0
+        config_dict["topk_method"] = None
+        config_dict["scoring_func"] = "softmax"
+        config_dict["routing_method_type"] = 1
+    elif is_eagle and not is_moe:
+        # Dense GQA EAGLE draft model (e.g. Mistral Medium 3.5 EAGLE).
+        # Routes to a Llama-backbone draft body — no MoE shimming required.
+        config_dict["architectures"] = ["MistralForCausalLMEagle"]
+        config_dict["model_type"] = "mistral"
+        config_dict["rope_is_neox_style"] = False
+        for mla_key in (
+            "q_lora_rank",
+            "qk_rope_head_dim",
+            "qk_nope_head_dim",
+            "kv_lora_rank",
+            "v_head_dim",
+        ):
+            if config_dict.get(mla_key) is None:
+                config_dict.pop(mla_key, None)
+    elif is_moe:
+        if is_mistral_large_3:
+            config_dict = _remap_moe_args(config_dict)
+            config_dict["model_type"] = "deepseek_v3"
+            if is_eagle:
+                config_dict["architectures"] = ["MistralLarge3ForCausalLMEagle"]
+            else:
+                config_dict["architectures"] = ["MistralLarge3ForCausalLM"]
+
+            assert (
+                "llama_4_scaling" in config_dict
+            ), "MistralLarge3 expect llama4 scaling config."
+            llama_4_scaling_config_keys = ["original_max_position_embeddings", "beta"]
+            assert all(
+                [
+                    key in config_dict["llama_4_scaling"]
+                    for key in llama_4_scaling_config_keys
+                ]
+            ), (
+                "llama_4_scaling config should define the keys: "
+                f"{','.join(llama_4_scaling_config_keys)}"
+            )
+        else:
+            config_dict["architectures"] = ["MixtralForCausalLM"]
+    else:
+        config_dict["architectures"] = ["MistralForCausalLM"]
+        config_dict["model_type"] = "mistral"
+        # Mistral models use non-interleaved RoPE (is_neox_style=False),
+        # unlike Llama which defaults to True.
+        config_dict["rope_is_neox_style"] = False
+        # Remove None-valued MLA fields that would shadow defaults in
+        # model_config._derive_model_shapes (getattr returns None instead
+        # of the fallback when the attribute exists but is None).
+        for mla_key in (
+            "q_lora_rank",
+            "qk_rope_head_dim",
+            "qk_nope_head_dim",
+            "kv_lora_rank",
+            "v_head_dim",
+        ):
+            if config_dict.get(mla_key) is None:
+                config_dict.pop(mla_key, None)
+
+    if bool(config_dict.get("yarn")):
+        config_dict = _remap_mistral_yarn_args(config_dict)
+
+    is_vision = bool(
+        (config_dict.get("multimodal") or {}).get("vision_encoder_args")
+        or config_dict.get("vision_encoder")
+    )
+    is_audio = bool(
+        ((config_dict.get("multimodal") or {}).get("whisper_model_args") or {}).get(
+            "encoder_args"
+        )
+    )
+
+    assert not (is_vision and is_audio), "Vision and audio are mutually exclusive"
+
+    if is_vision:
+        config_dict = _remap_mistral_vision_args(config_dict)
+    if is_audio:
+        config_dict = _remap_mistral_audio_args(config_dict)
+
+    config = PretrainedConfig.from_dict(config_dict)
+
+    logger.debug("Initialized config %s", config)
+
+    return config_dict, config
+
+
+def _remap_mistral_vision_args(config: dict) -> dict:
+    if config.get("multimodal"):
+        vision_config = config.pop("multimodal")
+    else:
+        vision_config = config.pop("vision_encoder")
+
+    quant_config = config.get("quantization_config")
+
+    config = {
+        "model_type": "pixtral",
+        "architectures": ["PixtralForConditionalGeneration"],
+        "text_config": config,
+        "vision_config": {"model_type": "pixtral", **vision_config},
+    }
+    if quant_config:
+        config["quantization_config"] = quant_config
+    return config
+
+
+def _remap_mistral_yarn_args(config: dict) -> dict:
+    yarn_config_map = {
+        "factor": "factor",
+        "original_max_position_embeddings": "original_max_position_embeddings",
+        "beta": "beta_fast",
+        "alpha": "beta_slow",
+        "apply_scale": "apply_yarn_scaling",
+    }
+    yarn_config = config.get("yarn") or {}
+    config["rope_scaling"] = {
+        "rope_type": "deepseek_yarn",
+        "mscale_all_dim": 1,
+    }
+    # Include rope_theta in rope_scaling if present at the top level,
+    # as transformers yarn validation requires it.
+    if "rope_theta" in config:
+        config["rope_scaling"]["rope_theta"] = config["rope_theta"]
+    for old_name, new_name in yarn_config_map.items():
+        if old_name in yarn_config:
+            value = yarn_config.pop(old_name)
+            if new_name is not None:
+                config["rope_scaling"][new_name] = value
+
+    assert len(yarn_config) == 0, f"Unparsed yarn config: {yarn_config}"
+
+    return config
+
+
+def _remap_general_mistral_args(config: dict) -> dict:
+    # Mistral key -> HF key
+    config_mapping = {
+        "dim": "hidden_size",
+        "norm_eps": "rms_norm_eps",
+        "n_kv_heads": "num_key_value_heads",
+        "n_layers": "num_hidden_layers",
+        "n_heads": "num_attention_heads",
+        "hidden_dim": "intermediate_size",
+    }
+    # HF key -> (Mistral key, default value)
+    top_level_mapping_with_default = {
+        "model_type": ("model_type", "transformer"),
+        "hidden_act": ("activation", "silu"),
+        "tie_word_embeddings": ("tied_embeddings", False),
+        "max_seq_len": ("max_seq_len", 128_000),
+        "max_position_embeddings": ("max_position_embeddings", 128_000),
+    }
+
+    for key, new_key in config_mapping.items():
+        if key in config:
+            config[new_key] = config.pop(key)
+
+    for new_key, (key, default_value) in top_level_mapping_with_default.items():
+        config[new_key] = config.pop(key, default_value)
+
+    return config
+
+
+def _remap_mistral_quantization_args(config: dict) -> dict:
+    if config.get("quantization"):
+        quantization = config.pop("quantization", {})
+        if quantization.get("qformat_weight") == "fp8_e4m3":
+            qscheme_act = quantization.get("qscheme_act")
+            assert qscheme_act in (
+                "NO_SCALES",
+                "TENSOR",
+                None,
+            ), "Only NO_SCALES and TENSOR (default) are supported for qscheme_act"
+            is_dynamic = qscheme_act == "NO_SCALES"
+            config["quantization_config"] = {
+                "quant_method": "fp8",
+                "activation_scheme": "dynamic" if is_dynamic else "static",
+            }
+        else:
+            raise ValueError(f"Found unknown quantization='{quantization}' in config")
+
+    return config
+
+
+def _remap_mistral_audio_args(config: dict) -> dict:
+    whisper_args = config["multimodal"].pop("whisper_model_args")
+    encoder_args = whisper_args["encoder_args"]
+    downsample_args = whisper_args["downsample_args"]
+
+    quant_config = config.get("quantization_config")
+    config = {
+        "model_type": "whixtral",
+        "architectures": ["VoxtralForConditionalGeneration"],
+        "text_config": PretrainedConfig.from_dict(config),
+        "audio_config": WhisperConfig(
+            num_mel_bins=encoder_args["audio_encoding_args"]["num_mel_bins"],
+            window_size=encoder_args["audio_encoding_args"]["window_size"],
+            sampling_rate=encoder_args["audio_encoding_args"]["sampling_rate"],
+            hop_length=encoder_args["audio_encoding_args"]["hop_length"],
+            downsample_factor=downsample_args["downsample_factor"],
+            d_model=encoder_args["dim"],
+            encoder_layers=encoder_args["n_layers"],
+            encoder_ffn_dim=encoder_args["hidden_dim"],
+            encoder_attention_heads=encoder_args["n_heads"],
+            vocab_size=encoder_args["vocab_size"],
+            max_source_positions=encoder_args["max_source_positions"],
+            is_encoder_decoder=False,  # Override WhisperConfig default
+        ),
+    }
+    if quant_config:
+        config["quantization_config"] = quant_config
+    return config
+
+
+def _remap_moe_args(config: dict) -> dict:
+    moe_config_map = {
+        "route_every_n": "moe_layer_freq",
+        "first_k_dense_replace": "first_k_dense_replace",
+        "num_experts_per_tok": "num_experts_per_tok",
+        "num_experts": "n_routed_experts",
+        "expert_hidden_dim": "moe_intermediate_size",
+        "routed_scale": "routed_scaling_factor",
+        "num_shared_experts": "n_shared_experts",
+        "num_expert_groups": "n_group",
+        "num_expert_groups_per_tok": "topk_group",
+    }
+    moe_config = config.get("moe", {})
+    for old_name, new_name in moe_config_map.items():
+        if old_name in moe_config:
+            value = moe_config.pop(old_name)
+            config[new_name] = value
+
+    config["topk_method"] = None
+    config["scoring_func"] = "softmax"
+    config["routing_method_type"] = 1  # RoutingMethodType.Renormalize
+
+    return config
+
+
+class MistralConfigParser:
+    def get_hf_file_to_dict(
+        self, file_name: str, model: str | Path, revision: str | None = "main"
+    ):
+        file_path = Path(model) / file_name
+        if not file_path.is_file():
+            raise FileNotFoundError(f"File not found {model}, {file_name}")
+
+        with open(file_path) as file:
+            return json.load(file)
+
+    def _download_mistral_config_file(self, model, revision) -> dict:
+        config_file_name = "params.json"
+        config_dict = self.get_hf_file_to_dict(config_file_name, model, revision)
+        if config_dict is None:
+            raise ValueError(
+                f"Failed to load mistral '{config_file_name}' config for model "
+                f"{model}. Please check if the model is a mistral-format model "
+                f"and if the config file exists."
+            )
+        assert isinstance(config_dict, dict)
+        return config_dict
+
+    def parse(
+        self,
+        model: str | Path,
+        revision: str | None = None,
+        **kwargs,
+    ) -> tuple[dict, PretrainedConfig]:
+        config_dict = self._download_mistral_config_file(model, revision)
+        if config_dict.get("max_position_embeddings") is None:
+            logger.warning(
+                "The params.json file is missing 'max_position_embeddings'"
+                " and could not get a value from the HF config."
+                " Defaulting to 128000"
+            )
+            config_dict["max_position_embeddings"] = 128_000
+
+        config_dict, config = adapt_config_dict(config_dict, model)
+
+        # Mistral configs may define sliding_window as list[int]. Convert it
+        # to int and add the layer_types list[str] to make it HF compatible
+        if (sliding_window := getattr(config, "sliding_window", None)) and isinstance(
+            sliding_window, list
+        ):
+            pattern_repeats = config.num_hidden_layers // len(sliding_window)
+            layer_types = sliding_window * pattern_repeats
+            config.layer_types = [
+                "full_attention" if layer_type is None else "sliding_attention"
+                for layer_type in layer_types
+            ]
+            config.sliding_window = next(filter(None, sliding_window), None)
+
+        return config_dict, config
+
+
+def is_mistral_model(name) -> bool:
+    """Return True if *name* refers to a Mistral model needing the custom parser."""
+    lower = str(name).lower()
+    if "mistral-large-3" in lower or "mistral-small-4" in lower or "leanstral" in lower:
+        return True
+    # EAGLE drafts for Mistral targets ship native-format only (params.json +
+    # consolidated.safetensors, no config.json), so route them through the
+    # custom parser regardless of the base model name.
+    if "eagle" in lower and "mistral" in lower:
+        return True
+    return False
+
+
+@lru_cache(maxsize=2)
+def load_mistral_config(
+    model_path: str,
+    trust_remote_code: bool = False,
+    revision: Optional[str] = None,
+):
+    """Load and parse a Mistral model config via the custom params.json format.
+
+    Returns a ``PretrainedConfig`` with dict sub-configs (text_config,
+    vision_config) converted to proper AutoConfig objects.
+    """
+    local_path = download_from_hf(model_path)
+    parser = MistralConfigParser()
+    config_dict, _ = parser.parse(local_path)
+
+    with tempfile.NamedTemporaryFile(mode="w+", suffix=".json") as f:
+        json.dump(config_dict, f)
+        f.flush()
+        loaded_config = AutoConfig.from_pretrained(
+            f.name, trust_remote_code=trust_remote_code, revision=revision
+        )
+    _ensure_sub_configs(loaded_config, "text_config", "vision_config")
+
+    return loaded_config
+
+
+def wrap_as_pixtral(processor, config):
+    """Wrap a tokenizer as a PixtralProcessor for Mistral vision models."""
+    from transformers.models.pixtral.image_processing_pixtral import (
+        PixtralImageProcessor,
+    )
+    from transformers.models.pixtral.processing_pixtral import (
+        PixtralProcessor as HFPixtralProcessor,
+    )
+
+    vision_config = config.vision_config
+    patch_size = vision_config.patch_size
+    image_size = vision_config.image_size
+    spatial_merge_size = getattr(vision_config, "spatial_merge_size", 1)
+
+    effective_patch = patch_size * spatial_merge_size
+    image_processor = PixtralImageProcessor(
+        do_resize=True,
+        size={"longest_edge": image_size},
+        patch_size={"height": effective_patch, "width": effective_patch},
+    )
+    return HFPixtralProcessor(
+        image_processor=image_processor,
+        tokenizer=processor,
+        patch_size=patch_size,
+        spatial_merge_size=spatial_merge_size,
+    )
+
+
+# kwargs that MistralCommon tokenizers reject.
+_MISTRAL_COMMON_REJECTED_KWARGS = frozenset(
+    {
+        "trust_remote_code",
+        "tokenizer_revision",
+        "use_fast",
+        "_from_auto",
+        "clean_up_tokenization_spaces",
+    }
+)
+
+# Models whose tokenizer should be loaded from a different checkpoint.
+_MISTRAL_TOKENIZER_REDIRECTS = {
+    # TODO(Xinyuan): Remove this once we have a proper tokenizer for Devstral
+    "mistralai/Devstral-Small-2505": "mistralai/Mistral-Small-3.1-24B-Instruct-2503",
+}
+
+
+def is_bare_tekken_checkpoint(tokenizer_name, revision=None) -> bool:
+    """True iff the checkpoint ships tekken.json but no tokenizer.json.
+
+    AutoTokenizer converts tekken.json on the fly, but the converter assigns
+    BPE ids from rank 0, dropping the 1000 special-token slots that precede
+    the BPE vocab in tekken's id space — every encoded id is shifted and
+    generation produces garbage. Such checkpoints must load through the
+    mistral-common backed tokenizer instead.
+    """
+
+    local_dir = Path(tokenizer_name)
+    if local_dir.is_dir():
+        return (local_dir / "tekken.json").is_file() and not (
+            local_dir / "tokenizer.json"
+        ).is_file()
+
+    if _cached_file_exists(tokenizer_name, "tokenizer.json", revision):
+        return False
+    if _cached_file_exists(tokenizer_name, "tekken.json", revision):
+        return True
+
+    # Cold cache: the tokenizer loads before weights, so tekken.json isn't
+    # cached yet on a first launch — HEAD-probe the hub to still detect it.
+    if not _remote_file_exists(tokenizer_name, "tekken.json", revision):
+        return False
+    return not _remote_file_exists(tokenizer_name, "tokenizer.json", revision)
+
+
+def retry_without_mistral_common_kwargs(tokenizer_name, *args, **common_kwargs):
+    """Retry ``AutoTokenizer.from_pretrained`` without kwargs that MistralCommon rejects.
+
+    Returns the loaded tokenizer, or *None* if the error is not a
+    MistralCommon kwargs rejection.
+    """
+    from transformers import AutoTokenizer
+
+    stripped = {
+        k: v
+        for k, v in common_kwargs.items()
+        if k not in _MISTRAL_COMMON_REJECTED_KWARGS
+    }
+    return AutoTokenizer.from_pretrained(tokenizer_name, *args, **stripped)
+
+
+def patch_mistral_common_tokenizer(tokenizer):
+    """Patch MistralCommonTokenizer/Backend to be compatible with HF tokenizer API.
+
+    MistralCommon tokenizers (used by Voxtral, Pixtral, etc.) reject several
+    standard kwargs and lack some attributes that sglang expects.  We wrap the
+    offending methods once at load time so that the rest of the codebase does
+    not need any special-casing.
+    """
+    cls_name = type(tokenizer).__name__
+    if "MistralCommon" not in cls_name:
+        return tokenizer
+    if getattr(tokenizer, "_mistral_common_patched", False):
+        return tokenizer
+    tokenizer._mistral_common_patched = True
+
+    if not hasattr(tokenizer, "get_added_vocab"):
+        tokenizer.get_added_vocab = lambda: {}
+
+    # Keep the old no-op pad add working on transformers 5.12 MistralCommon.
+    _orig_add_special_tokens = tokenizer.add_special_tokens
+
+    def _safe_add_special_tokens(special_tokens_dict, *args, **kwargs):
+        if set(special_tokens_dict) == {"pad_token"}:
+            tokenizer.pad_token = special_tokens_dict["pad_token"]
+            return 0
+        return _orig_add_special_tokens(special_tokens_dict, *args, **kwargs)
+
+    tokenizer.add_special_tokens = _safe_add_special_tokens
+
+    # Set a chat_template containing "audio" so that sglang's content format
+    # detector returns "openai" (which preserves audio_url extraction).
+    if not hasattr(tokenizer, "chat_template") or tokenizer.chat_template is None:
+        tokenizer.chat_template = "<!-- audio/image multimodal -->"
+
+    _orig_convert = tokenizer.convert_tokens_to_ids
+
+    def _safe_convert(val):
+        try:
+            return _orig_convert(val)
+        except AssertionError:
+            logger.debug(
+                "convert_tokens_to_ids failed for %r, returning unk_token_id", val
+            )
+            return getattr(tokenizer, "unk_token_id", None)
+
+    tokenizer.convert_tokens_to_ids = _safe_convert
+
+    def _drop_kwargs(fn, keys):
+        def wrapper(*args, **kwargs):
+            for k in keys:
+                kwargs.pop(k, None)
+            return fn(*args, **kwargs)
+
+        return wrapper
+
+    tokenizer.decode = _drop_kwargs(tokenizer.decode, ["spaces_between_special_tokens"])
+    tokenizer.batch_decode = _drop_kwargs(
+        tokenizer.batch_decode, ["spaces_between_special_tokens"]
+    )
+
+    if hasattr(tokenizer, "_text_to_ids"):
+        _orig_text_to_ids = tokenizer._text_to_ids
+        marker_to_id = {
+            "[IMG]": tokenizer.convert_tokens_to_ids("[IMG]"),
+            "[IMG_BREAK]": tokenizer.convert_tokens_to_ids("[IMG_BREAK]"),
+            "[IMG_END]": tokenizer.convert_tokens_to_ids("[IMG_END]"),
+        }
+
+        def _text_to_ids_with_pixtral_markers(text, add_special_tokens):
+            if not isinstance(text, str) or not any(
+                marker in text for marker in marker_to_id
+            ):
+                return _orig_text_to_ids(text, add_special_tokens)
+
+            ids = []
+            pos = 0
+            while pos < len(text):
+                next_marker = None
+                next_idx = len(text)
+                for marker in marker_to_id:
+                    marker_idx = text.find(marker, pos)
+                    if marker_idx != -1 and marker_idx < next_idx:
+                        next_marker = marker
+                        next_idx = marker_idx
+
+                if next_marker is None:
+                    ids.extend(_orig_text_to_ids(text[pos:], False))
+                    break
+                if next_idx > pos:
+                    ids.extend(_orig_text_to_ids(text[pos:next_idx], False))
+                ids.append(marker_to_id[next_marker])
+                pos = next_idx + len(next_marker)
+
+            if add_special_tokens:
+                return tokenizer.build_inputs_with_special_tokens(ids)
+            return ids
+
+        tokenizer._text_to_ids = _text_to_ids_with_pixtral_markers
+
+    tokenizer._orig_apply_chat_template = tokenizer.apply_chat_template
+
+    def _adapt_placeholder_content_for_mistral_common(content):
+        if not isinstance(content, list):
+            return content
+
+        rendered_parts = []
+        has_placeholder = False
+        for part in content:
+            if not isinstance(part, dict):
+                return content
+            part_type = part.get("type")
+            if part_type in ("text", "input_text"):
+                rendered_parts.append(part.get("text", ""))
+            elif part_type == "image" and not any(
+                key in part for key in ("url", "path", "base64")
+            ):
+                has_placeholder = True
+                rendered_parts.append("[IMG]")
+            elif part_type in ("audio", "video") and not any(
+                key in part for key in ("url", "path", "base64")
+            ):
+                has_placeholder = True
+                continue
+            else:
+                return content
+
+        return "".join(rendered_parts) if has_placeholder else content
+
+    def _adapt_placeholder_messages_for_mistral_common(messages):
+        if not isinstance(messages, (list, tuple)):
+            return messages
+
+        adapted = []
+        for msg in messages:
+            if isinstance(msg, (list, tuple)):
+                adapted.append(_adapt_placeholder_messages_for_mistral_common(msg))
+            elif isinstance(msg, dict):
+                adapted.append(
+                    {
+                        **msg,
+                        "content": _adapt_placeholder_content_for_mistral_common(
+                            msg.get("content", "")
+                        ),
+                    }
+                )
+            else:
+                adapted.append(msg)
+        return adapted
+
+    def _safe_apply_chat_template(messages, **kwargs):
+        kwargs.pop("add_generation_prompt", None)
+        messages = _adapt_placeholder_messages_for_mistral_common(messages)
+        return tokenizer._orig_apply_chat_template(messages, **kwargs)
+
+    tokenizer.apply_chat_template = _safe_apply_chat_template
+    return tokenizer

python/sglang/srt/utils/hf_transformers/processor.py

@@ -0,0 +1,306 @@
+# 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.
+# ==============================================================================
+"""Processor loading utilities."""
+
+import json
+from pathlib import Path
+from typing import Optional
+
+from transformers import (
+    AutoProcessor,
+    AutoTokenizer,
+    PreTrainedTokenizerBase,
+)
+
+from sglang.srt.multimodal.customized_mm_processor_utils import _CUSTOMIZED_MM_PROCESSOR
+from sglang.srt.utils import logger
+
+from .common import (
+    AutoConfig,
+    _is_deepseek_ocr2_model,
+    _is_deepseek_ocr_model,
+    _override_v_head_dim_if_zero,
+    _resolve_local_or_cached_file,
+    attach_additional_stop_token_ids,
+    download_from_hf,
+    get_tokenizer_from_processor,
+    resolve_runai_obj_uri,
+)
+from .mistral_utils import (
+    is_mistral_model,
+    load_mistral_config,
+    patch_mistral_common_tokenizer,
+    wrap_as_pixtral,
+)
+from .tokenizer import (
+    _TOKENIZERS_BACKEND,
+    _fix_added_tokens_encoding,
+    _fix_special_tokens_pattern,
+)
+
+
+def _build_processor_manually(
+    model_path, config, trust_remote_code, revision, **kwargs
+):
+    """Build processor when AutoProcessor fails to resolve feature_extractor_type.
+
+    In transformers v5, AutoProcessor.from_pretrained calls
+    AutoFeatureExtractor.from_pretrained which fails if
+    preprocessor_config.json lacks 'feature_extractor_type'. This resolves
+    the processor class via dynamic module resolution and constructs it with
+    individually-loaded components.
+    """
+    import transformers
+    from transformers import AutoImageProcessor, AutoTokenizer
+    from transformers.dynamic_module_utils import get_class_from_dynamic_module
+
+    # Resolve processor class from auto_map -- check both the model config
+    # and the preprocessor_config.json (some models like MiniCPM-o only
+    # declare AutoProcessor in the latter).
+    auto_map = getattr(config, "auto_map", None) or {}
+    proc_ref = auto_map.get("AutoProcessor")
+    if not proc_ref:
+        try:
+            pp_file = _resolve_local_or_cached_file(
+                model_path, "preprocessor_config.json", revision
+            )
+            with open(pp_file) as f:
+                pp_auto_map = json.load(f).get("auto_map", {})
+            proc_ref = pp_auto_map.get("AutoProcessor")
+        except (OSError, json.JSONDecodeError, ValueError) as e:
+            logger.warning(
+                "_build_processor_manually: could not read preprocessor_config.json "
+                "for %s: %s",
+                model_path,
+                e,
+            )
+    if not proc_ref:
+        raise ValueError(f"Cannot determine processor class for {model_path}")
+
+    proc_cls = get_class_from_dynamic_module(
+        proc_ref, model_path, code_revision=revision
+    )
+
+    # Load sub-components individually (these succeed)
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_path, trust_remote_code=trust_remote_code, revision=revision
+    )
+    init_kwargs = {"tokenizer": tokenizer}
+
+    if "image_processor" in getattr(proc_cls, "attributes", []):
+        try:
+            init_kwargs["image_processor"] = AutoImageProcessor.from_pretrained(
+                model_path, trust_remote_code=trust_remote_code, revision=revision
+            )
+        except (ImportError, OSError, ValueError) as e:
+            raise RuntimeError(
+                f"Failed to load image_processor for {model_path}: {e}. "
+                f"This model requires an image processor for multimodal features. "
+                f"Check that the model files are complete and accessible."
+            ) from e
+
+    # Instantiate feature extractor from its declared class
+    fe_class_name = getattr(proc_cls, "feature_extractor_class", None)
+    if fe_class_name:
+        fe_class = getattr(transformers, fe_class_name, None)
+        if fe_class is not None:
+            try:
+                init_kwargs["feature_extractor"] = fe_class()
+            except TypeError as e:
+                logger.warning(
+                    "Cannot instantiate feature extractor %s with no arguments "
+                    "for %s: %s",
+                    fe_class_name,
+                    model_path,
+                    e,
+                )
+        else:
+            logger.warning(
+                "Feature extractor class %s not found in transformers for %s",
+                fe_class_name,
+                model_path,
+            )
+
+    return proc_cls(**init_kwargs)
+
+
+def get_processor(
+    tokenizer_name: str,
+    *args,
+    tokenizer_mode: str = "auto",
+    trust_remote_code: bool = False,
+    tokenizer_revision: Optional[str] = None,
+    use_fast: Optional[bool] = True,
+    tokenizer_backend: str = "huggingface",
+    model_name: Optional[str] = None,
+    **kwargs,
+):
+    if tokenizer_backend == "fastokens":
+        from .tokenizer import _ensure_fastokens_patched
+
+        _ensure_fastokens_patched()
+
+    revision = kwargs.pop("revision", tokenizer_revision)
+    tokenizer_name = resolve_runai_obj_uri(tokenizer_name)
+
+    if is_mistral_model(tokenizer_name):
+        config = load_mistral_config(
+            tokenizer_name,
+            trust_remote_code=trust_remote_code,
+            revision=revision,
+        )
+    elif model_name is not None:
+        config = AutoConfig.from_pretrained(
+            model_name,
+            trust_remote_code=trust_remote_code,
+            revision=revision,
+            **kwargs,
+        )
+    else:
+        config = AutoConfig.from_pretrained(
+            tokenizer_name,
+            trust_remote_code=trust_remote_code,
+            revision=revision,
+            **kwargs,
+        )
+    is_ocr2 = _is_deepseek_ocr2_model(config)
+    if _is_deepseek_ocr_model(config) or is_ocr2:
+        config.model_type = "deepseek-ocr"
+        config.update({"architectures": ["DeepseekOCRForCausalLM"]})
+        if is_ocr2:
+            _override_v_head_dim_if_zero(config)
+
+    if config.model_type in {"qwen2_vl", "sarashina2_vision"}:
+        if "size" not in kwargs:
+            kwargs["size"] = {"shortest_edge": 3136, "longest_edge": 1003520}
+
+    if config.model_type not in {"llava", "clip"}:
+        kwargs["use_fast"] = use_fast
+    try:
+        if "InternVL3_5" in tokenizer_name:
+            processor = AutoTokenizer.from_pretrained(
+                tokenizer_name,
+                *args,
+                trust_remote_code=trust_remote_code,
+                revision=revision,
+                **kwargs,
+            )
+        else:
+            if config.model_type in _CUSTOMIZED_MM_PROCESSOR:
+                processor = _CUSTOMIZED_MM_PROCESSOR[config.model_type].from_pretrained(
+                    tokenizer_name,
+                    *args,
+                    trust_remote_code=trust_remote_code,
+                    revision=revision,
+                    **kwargs,
+                )
+            else:
+                processor = AutoProcessor.from_pretrained(
+                    tokenizer_name,
+                    *args,
+                    trust_remote_code=trust_remote_code,
+                    revision=revision,
+                    **kwargs,
+                )
+
+    except ValueError as e:
+        error_message = str(e)
+        if "does not have a slow version" in error_message:
+            logger.info(
+                "Processor %s does not have a slow version. Automatically use fast version",
+                tokenizer_name,
+            )
+            kwargs["use_fast"] = True
+            processor = AutoProcessor.from_pretrained(
+                tokenizer_name,
+                *args,
+                trust_remote_code=trust_remote_code,
+                revision=revision,
+                **kwargs,
+            )
+        elif "Unrecognized feature extractor" in error_message:
+            logger.info(
+                "AutoProcessor failed on feature extractor for %s, "
+                "constructing processor manually",
+                tokenizer_name,
+            )
+            processor = _build_processor_manually(
+                tokenizer_name,
+                config,
+                trust_remote_code,
+                revision,
+                **kwargs,
+            )
+        elif (
+            "are not supported by" in error_message and "MistralCommon" in error_message
+        ):
+            logger.info(
+                "AutoProcessor for %s rejected standard kwargs, "
+                "retrying without trust_remote_code/use_fast",
+                tokenizer_name,
+            )
+            kwargs.pop("use_fast", None)
+            kwargs.pop("_from_auto", None)
+            processor = AutoProcessor.from_pretrained(
+                tokenizer_name,
+                *args,
+                revision=revision,
+                **kwargs,
+            )
+        else:
+            raise
+    if (
+        isinstance(processor, PreTrainedTokenizerBase)
+        and getattr(config, "model_type", None) == "pixtral"
+    ):
+        processor = wrap_as_pixtral(processor, config)
+
+    tokenizer = get_tokenizer_from_processor(processor)
+
+    # AutoProcessor may internally create a TokenizersBackend tokenizer
+    # (same issue as get_tokenizer). Replace it with a properly loaded one.
+    if type(tokenizer).__name__ == _TOKENIZERS_BACKEND:
+        from .tokenizer import get_tokenizer
+
+        logger.warning(
+            "Processor tokenizer for %s is TokenizersBackend, "
+            "reloading via get_tokenizer",
+            tokenizer_name,
+        )
+        tokenizer = get_tokenizer(
+            tokenizer_name,
+            tokenizer_mode=tokenizer_mode,
+            trust_remote_code=trust_remote_code,
+            tokenizer_revision=revision,
+            tokenizer_backend=tokenizer_backend,
+        )
+        if isinstance(processor, PreTrainedTokenizerBase):
+            processor = tokenizer
+        else:
+            processor.tokenizer = tokenizer
+
+    if tokenizer.chat_template is None:
+        local_path = download_from_hf(
+            tokenizer_name, allow_patterns=["*.json", "*.jinja", "*.model"]
+        )
+        jinja_path = Path(local_path) / "chat_template.jinja"
+        if jinja_path.is_file():
+            tokenizer.chat_template = jinja_path.read_text()
+            logger.info("Loaded chat_template from %s", jinja_path)
+
+    patch_mistral_common_tokenizer(tokenizer)
+    _fix_special_tokens_pattern(tokenizer)
+    _fix_added_tokens_encoding(tokenizer)
+    attach_additional_stop_token_ids(tokenizer)
+    return processor

python/sglang/srt/utils/hf_transformers/tokenizer.py

@@ -0,0 +1,613 @@
+# 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.
+# ==============================================================================
+"""Tokenizer loading utilities."""
+
+import json
+import logging
+import warnings
+from pathlib import Path
+from typing import Optional, Union
+
+from transformers import (
+    AutoTokenizer,
+    PreTrainedTokenizer,
+    PreTrainedTokenizerFast,
+)
+
+from sglang.srt.connector import create_remote_connector
+from sglang.srt.utils import is_remote_url, logger
+from sglang.srt.utils.patch_tokenizer import patch_tokenizer
+
+from ..hf_transformers_patches import _ensure_gguf_version
+from .common import (
+    _resolve_local_or_cached_file,
+    attach_additional_stop_token_ids,
+    check_gguf_file,
+    resolve_runai_obj_uri,
+)
+from .mistral_utils import (
+    _MISTRAL_TOKENIZER_REDIRECTS,
+    is_bare_tekken_checkpoint,
+    patch_mistral_common_tokenizer,
+    retry_without_mistral_common_kwargs,
+)
+
+# A fast LLaMA tokenizer with the pre-processed `tokenizer.json` file.
+_FAST_LLAMA_TOKENIZER = "hf-internal-testing/llama-tokenizer"
+
+# Class name used by transformers v5 when no tokenizer mapping exists for a model_type.
+_TOKENIZERS_BACKEND = "TokenizersBackend"
+
+
+def _load_tokenizer_by_declared_class(tokenizer_name, *args, **kwargs):
+    """Load tokenizer by the class declared in tokenizer_config.json.
+
+    AutoTokenizer resolves to TokenizersBackend when the model's config
+    model_type has no tokenizer class mapping (e.g. deepseek_vl_v2), even
+    though tokenizer_config.json declares a standard class like
+    LlamaTokenizerFast.  Returns None if it cannot improve on AutoTokenizer.
+    """
+    import transformers
+
+    try:
+        revision = kwargs.get("revision") or kwargs.get("tokenizer_revision")
+        config_file = _resolve_local_or_cached_file(
+            tokenizer_name, "tokenizer_config.json", revision
+        )
+        with open(config_file) as f:
+            tok_config = json.load(f)
+        tok_class_name = tok_config.get("tokenizer_class")
+    except FileNotFoundError:
+        return None
+    except (OSError, json.JSONDecodeError) as e:
+        logger.debug(
+            "Failed to read tokenizer_config.json for %s: %s", tokenizer_name, e
+        )
+        return None
+
+    if not tok_class_name:
+        return None
+
+    # Skip base classes that don't implement required methods (e.g. get_vocab)
+    if tok_class_name in ("PreTrainedTokenizer", "PreTrainedTokenizerBase"):
+        return None
+
+    tok_cls = getattr(transformers, tok_class_name, None)
+    if tok_cls is None and kwargs.get("trust_remote_code"):
+        # Class not in transformers — try loading via auto_map.
+        try:
+            auto_map = tok_config.get("auto_map", {})
+            auto_tok_ref = auto_map.get("AutoTokenizer")
+            if isinstance(auto_tok_ref, (list, tuple)):
+                auto_tok_ref = auto_tok_ref[0]
+            if auto_tok_ref:
+                from transformers.dynamic_module_utils import (
+                    get_class_from_dynamic_module,
+                )
+
+                tok_cls = get_class_from_dynamic_module(
+                    auto_tok_ref,
+                    tokenizer_name,
+                    code_revision=revision,
+                )
+        except (OSError, ImportError, ValueError, RuntimeError) as e:
+            logger.debug("Dynamic module lookup for %s failed: %s", tok_class_name, e)
+    if tok_cls is None:
+        return None
+
+    logger.debug(
+        "Loading tokenizer for %s directly as %s (bypassing AutoTokenizer)",
+        tokenizer_name,
+        tok_class_name,
+    )
+    try:
+        return tok_cls.from_pretrained(tokenizer_name, *args, **kwargs)
+    except (OSError, ValueError, TypeError, ImportError) as e:
+        logger.warning(
+            "Direct load as %s failed for %s: %s. "
+            "Falling back to AutoTokenizer result.",
+            tok_class_name,
+            tokenizer_name,
+            e,
+        )
+        return None
+
+
+# Filter warnings like: https://github.com/sgl-project/sglang/issues/8082
+class TokenizerWarningsFilter(logging.Filter):
+    def filter(self, record: logging.LogRecord) -> bool:
+        return "Calling super().encode with" not in record.getMessage()
+
+
+# ---------------------------------------------------------------------------
+# Helpers for get_tokenizer
+# ---------------------------------------------------------------------------
+
+
+def _resolve_tokenizer_name(tokenizer_name, kwargs):
+    """Resolve special name formats (GGUF, remote URLs, etc.) to a local path.
+
+    May mutate *kwargs* (e.g. to add ``gguf_file``).
+    """
+    tokenizer_name = _MISTRAL_TOKENIZER_REDIRECTS.get(tokenizer_name, tokenizer_name)
+
+    if check_gguf_file(tokenizer_name):
+        _ensure_gguf_version()
+        kwargs["gguf_file"] = tokenizer_name
+        tokenizer_name = Path(tokenizer_name).parent
+
+    tokenizer_name = resolve_runai_obj_uri(tokenizer_name)
+
+    if is_remote_url(tokenizer_name):
+        # BaseConnector implements __del__() to clean up the local dir.
+        # Since config files need to exist all the time, so we DO NOT use
+        # with statement to avoid closing the client.
+        client = create_remote_connector(tokenizer_name)
+        client.pull_files(ignore_pattern=["*.pt", "*.safetensors", "*.bin"])
+        tokenizer_name = client.get_local_dir()
+
+    return tokenizer_name
+
+
+def _auto_tokenizer_from_pretrained(tokenizer_name, *args, **common_kwargs):
+    """Call ``AutoTokenizer.from_pretrained`` with error handling."""
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_name, *args, **common_kwargs
+        )
+        logging.getLogger(tokenizer.__class__.__module__).addFilter(
+            TokenizerWarningsFilter()
+        )
+        return tokenizer
+    except TypeError as e:
+        err_msg = (
+            "Failed to load the tokenizer. If you are using a LLaMA V1 model "
+            f"consider using '{_FAST_LLAMA_TOKENIZER}' instead of the "
+            "original tokenizer."
+        )
+        raise RuntimeError(err_msg) from e
+    except ValueError as e:
+        # MistralCommon tokenizers reject standard HF kwargs like
+        # trust_remote_code, use_fast etc. Retry without them.
+        if "are not supported by" in str(e) and "MistralCommon" in str(e):
+            return retry_without_mistral_common_kwargs(
+                tokenizer_name, *args, **common_kwargs
+            )
+        # If the error pertains to the tokenizer class not existing or not
+        # currently being imported, suggest using the --trust-remote-code flag.
+        if not common_kwargs.get("trust_remote_code") and (
+            "does not exist or is not currently imported." in str(e)
+            or "requires you to execute the tokenizer file" in str(e)
+        ):
+            err_msg = (
+                "Failed to load the tokenizer. If the tokenizer is a custom "
+                "tokenizer not yet available in the HuggingFace transformers "
+                "library, consider setting `trust_remote_code=True` in LLM "
+                "or using the `--trust-remote-code` flag in the CLI."
+            )
+            raise RuntimeError(err_msg) from e
+        raise
+
+
+def _resolve_tokenizers_backend(tokenizer_name, *args, **common_kwargs):
+    """Resolve generic ``TokenizersBackend`` to a proper tokenizer class.
+
+    In transformers v5, ``AutoTokenizer`` falls back to ``TokenizersBackend``
+    when the model_type has no tokenizer mapping.  This retries with
+    ``use_fast=False``, then attempts loading by the class declared in
+    ``tokenizer_config.json``.  May still return a ``TokenizersBackend``
+    if all retries fail (with a warning).
+    """
+    logger.debug(
+        "Tokenizer loaded as generic TokenizersBackend for %s, "
+        "retrying with use_fast=False",
+        tokenizer_name,
+    )
+    common_kwargs = {**common_kwargs, "use_fast": False}
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_name, *args, **common_kwargs
+        )
+    except (ValueError, TypeError, OSError, ImportError, RuntimeError) as e:
+        raise RuntimeError(
+            f"Retry with use_fast=False for {tokenizer_name} also failed "
+            f"(initial load returned TokenizersBackend): {e}"
+        ) from e
+
+    if type(tokenizer).__name__ == _TOKENIZERS_BACKEND:
+        tokenizer = (
+            _load_tokenizer_by_declared_class(tokenizer_name, *args, **common_kwargs)
+            or tokenizer
+        )
+
+    if type(tokenizer).__name__ == _TOKENIZERS_BACKEND:
+        if common_kwargs.get("trust_remote_code"):
+            logger.warning(
+                "Tokenizer for %s is still TokenizersBackend after retries "
+                "with --trust-remote-code. Model-specific tokenizer attributes "
+                "may be missing.",
+                tokenizer_name,
+            )
+        else:
+            logger.debug(
+                "Tokenizer for %s loaded as generic TokenizersBackend. "
+                "Set --trust-remote-code to load the model-specific tokenizer.",
+                tokenizer_name,
+            )
+
+    return tokenizer
+
+
+# ---------------------------------------------------------------------------
+# Post-load fixups
+# ---------------------------------------------------------------------------
+
+
+def _fix_v5_tokenizer_components(tokenizer, model_name_or_path, revision=None):
+    """Fix pre_tokenizer/decoder when a v5 tokenizer class overwrites them.
+
+    In transformers v5, some tokenizer classes (e.g. LlamaTokenizer) have a
+    custom __init__ that rebuilds the pre_tokenizer and decoder from scratch
+    with class-specific components, discarding the originals from tokenizer.json.
+    This breaks models that specify LlamaTokenizerFast but actually use a
+    different tokenizer architecture (e.g. DeepSeek-V3.2 uses ByteLevel).
+
+    Detects the mismatch by comparing against the raw tokenizer.json and
+    restores the original components when they differ.
+    """
+    backend = getattr(tokenizer, "_tokenizer", None)
+    if backend is None:
+        return
+
+    try:
+        from tokenizers import Tokenizer as RawTokenizer
+
+        tok_file = _resolve_local_or_cached_file(
+            model_name_or_path, "tokenizer.json", revision
+        )
+        raw = RawTokenizer.from_file(tok_file)
+    except FileNotFoundError:
+        return
+    except (OSError, ValueError, RuntimeError) as e:
+        logger.warning(
+            "_fix_v5_tokenizer_components: unexpected error loading tokenizer.json "
+            "for %s, v5 component fix will not be applied: %s",
+            model_name_or_path,
+            e,
+        )
+        return
+
+    raw_pre = type(raw.pre_tokenizer).__name__ if raw.pre_tokenizer else None
+    loaded_pre = type(backend.pre_tokenizer).__name__ if backend.pre_tokenizer else None
+
+    if raw_pre and loaded_pre and raw_pre != loaded_pre:
+        logger.info(
+            "Fixing v5 tokenizer component mismatch for %s: "
+            "pre_tokenizer %s -> %s, decoder %s -> %s",
+            model_name_or_path,
+            loaded_pre,
+            raw_pre,
+            type(backend.decoder).__name__ if backend.decoder else None,
+            type(raw.decoder).__name__ if raw.decoder else None,
+        )
+        backend.pre_tokenizer = raw.pre_tokenizer
+        backend.decoder = raw.decoder
+
+
+def _fix_v5_add_bos_eos_token(tokenizer, model_name_or_path, revision=None):
+    """Restore add_bos_token/add_eos_token stripped by transformers v5.
+
+    In transformers v5, _from_pretrained() strips add_bos_token and
+    add_eos_token from init kwargs when a tokenizer.json file is present,
+    assuming the tokenizer.json post-processor handles BOS/EOS addition.
+    However, many models (e.g. DeepSeek-V3) have a tokenizer.json whose
+    post-processor does NOT add BOS/EOS, and rely on the add_bos_token flag
+    from tokenizer_config.json instead. This causes silent accuracy regressions.
+
+    This function reads the tokenizer_config.json and restores the values,
+    but only for tokenizer classes that actually supported these flags in v4.
+    Classes like Qwen2Tokenizer did not support add_bos_token/add_eos_token
+    in v4, so restoring them would change behavior.
+    """
+    # In transformers v4, only certain tokenizer classes supported
+    # add_bos_token / add_eos_token as init parameters.  Restoring these
+    # flags for classes that never supported them (e.g. Qwen2Tokenizer)
+    # would incorrectly change tokenization behavior.
+    _V4_CLASSES_WITH_BOS_EOS_FLAGS = frozenset(
+        {
+            "LlamaTokenizer",
+            "LlamaTokenizerFast",
+            "CodeLlamaTokenizer",
+            "CodeLlamaTokenizerFast",
+            "GemmaTokenizer",
+            "GemmaTokenizerFast",
+            "CohereTokenizerFast",
+        }
+    )
+
+    try:
+        config_file = _resolve_local_or_cached_file(
+            model_name_or_path, "tokenizer_config.json", revision
+        )
+        with open(config_file) as f:
+            config = json.load(f)
+    except FileNotFoundError:
+        return
+    except (OSError, json.JSONDecodeError, ValueError) as e:
+        logger.warning(
+            "_fix_v5_add_bos_eos_token: failed to read tokenizer_config.json "
+            "for %s, BOS/EOS token restoration will not be applied: %s",
+            model_name_or_path,
+            e,
+        )
+        return
+
+    tokenizer_class = config.get("tokenizer_class", "")
+    if tokenizer_class not in _V4_CLASSES_WITH_BOS_EOS_FLAGS:
+        logger.debug(
+            "_fix_v5_add_bos_eos_token: skipping %s (tokenizer_class=%s "
+            "did not support add_bos/eos_token in v4)",
+            model_name_or_path,
+            tokenizer_class,
+        )
+        return
+
+    # In v4, Llama/Gemma tokenizers defaulted add_bos_token=True.
+    # When the config omits the key or has null, use the v4 default so that
+    # update_post_processor() doesn't drop BOS/EOS that was there before.
+    _V4_DEFAULTS = {"add_bos_token": True, "add_eos_token": False}
+
+    changed = False
+    for attr in ("add_bos_token", "add_eos_token"):
+        config_val = config.get(attr)
+        if config_val is None:
+            # Key missing or null -> use v4 default for this tokenizer class
+            config_val = _V4_DEFAULTS.get(attr, False)
+        # Fast tokenizers in v4 used tokenizer.json post-processor for EOS —
+        # the add_eos_token Python attribute was set but the post-processor
+        # came from tokenizer.json, not from the attribute.  In v5, the flag is
+        # stripped and both sglang and HF reference end up with add_eos_token=False.
+        # Restoring add_eos_token for fast tokenizers makes sglang diverge from
+        # the HF reference, breaking embedding models like e5-mistral-7b-instruct.
+        if attr == "add_eos_token" and isinstance(tokenizer, PreTrainedTokenizerFast):
+            config_val = _V4_DEFAULTS["add_eos_token"]  # False
+        current_val = getattr(tokenizer, attr, None)
+        if current_val != config_val:
+            logger.info(
+                "Restoring %s=%s for %s (was %s after v5 loading)",
+                attr,
+                config_val,
+                model_name_or_path,
+                current_val,
+            )
+            # Set the private backing attribute (not the property) because
+            # transformers tokenizers expose add_bos/eos_token as properties
+            # that read from the underscore-prefixed attribute.
+            setattr(tokenizer, f"_{attr}", config_val)
+            changed = True
+
+    # Rebuild the post-processor so it respects the restored flags
+    if changed and hasattr(tokenizer, "update_post_processor"):
+        tokenizer.update_post_processor()
+
+
+def _fix_special_tokens_pattern(tokenizer):
+    """Fix https://github.com/huggingface/transformers/pull/42563 which defaults
+    special_tokens_pattern to "cls_sep", inserting None into token IDs when
+    cls_token/sep_token are undefined (e.g. Kimi-VL's TikTokenTokenizer).
+    """
+    pattern = getattr(tokenizer, "special_tokens_pattern", None)
+    if pattern == "cls_sep" and (
+        tokenizer.cls_token_id is None or tokenizer.sep_token_id is None
+    ):
+        tokenizer.special_tokens_pattern = "none"
+
+
+def _apply_post_load_fixes(tokenizer, tokenizer_name, revision):
+    """Apply all post-load patches and return the final tokenizer."""
+    _fix_v5_tokenizer_components(tokenizer, tokenizer_name, revision)
+    _fix_v5_add_bos_eos_token(tokenizer, tokenizer_name, revision)
+
+    if not isinstance(tokenizer, PreTrainedTokenizerFast):
+        warnings.warn(
+            "Using a slow tokenizer. This might cause a significant "
+            "slowdown. Consider using a fast tokenizer instead."
+        )
+
+    patch_mistral_common_tokenizer(tokenizer)
+    _fix_special_tokens_pattern(tokenizer)
+    attach_additional_stop_token_ids(tokenizer)
+    return patch_tokenizer(tokenizer)
+
+
+# ---------------------------------------------------------------------------
+# Public entry point
+# ---------------------------------------------------------------------------
+
+
+_fastokens_patched = False
+
+
+def _ensure_fastokens_patched():
+    """Monkey-patch transformers to use the fastokens backend (once)."""
+    global _fastokens_patched
+    if _fastokens_patched:
+        return
+    try:
+        import fastokens
+    except ImportError:
+        raise ImportError(
+            "The fastokens package is required when --tokenizer-backend=fastokens. "
+            "Install it with: pip install 'sglang[fastokens]'"
+        ) from None
+
+    fastokens.patch_transformers()
+    _fastokens_patched = True
+    logger.info("fastokens backend enabled - transformers patched successfully")
+
+
+def get_tokenizer(
+    tokenizer_name: str,
+    *args,
+    tokenizer_mode: str = "auto",
+    trust_remote_code: bool = False,
+    tokenizer_revision: Optional[str] = None,
+    tokenizer_backend: str = "huggingface",
+    **kwargs,
+) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:
+    """Gets a tokenizer for the given model name via Huggingface."""
+    # Tiktoken format has its own backend — no fastokens patching needed.
+    if tokenizer_name.endswith(".json"):
+        from sglang.srt.tokenizer.tiktoken_tokenizer import TiktokenTokenizer
+
+        return TiktokenTokenizer(tokenizer_name)
+
+    if tokenizer_backend == "fastokens":
+        _ensure_fastokens_patched()
+
+    if tokenizer_mode == "slow":
+        if kwargs.get("use_fast", False):
+            raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.")
+        kwargs["use_fast"] = False
+    elif tokenizer_mode == "auto":
+        # Transformers v5 AutoTokenizer ignores use_fast (always fast), but
+        # some code paths pass kwargs to non-AutoTokenizer loaders where
+        # use_fast still matters. Set explicitly for those fallback paths.
+        if "use_fast" not in kwargs:
+            kwargs["use_fast"] = True
+
+    tokenizer_name = _resolve_tokenizer_name(tokenizer_name, kwargs)
+
+    common_kwargs = dict(
+        trust_remote_code=trust_remote_code,
+        tokenizer_revision=tokenizer_revision,
+        clean_up_tokenization_spaces=False,
+        **kwargs,
+    )
+
+    try:
+        if is_bare_tekken_checkpoint(tokenizer_name, tokenizer_revision):
+            from transformers.tokenization_mistral_common import (
+                MistralCommonTokenizer,
+            )
+
+            logger.info(
+                "Detected bare-tekken checkpoint %s (tekken.json, no "
+                "tokenizer.json); loading via mistral-common MistralCommonTokenizer, "
+                "ignoring tokenizer_backend=%r.",
+                tokenizer_name,
+                tokenizer_backend,
+            )
+
+            tokenizer = MistralCommonTokenizer.from_pretrained(
+                tokenizer_name, revision=tokenizer_revision
+            )
+        else:
+            tokenizer = _auto_tokenizer_from_pretrained(
+                tokenizer_name, *args, **common_kwargs
+            )
+
+            # With fastokens, the patched TokenizersBackend.from_pretrained already
+            # returned a tokenizer whose backend is a fastokens shim. Re-resolving via
+            # the declared class (e.g. Qwen2Tokenizer) would discard that work.
+            if (
+                type(tokenizer).__name__ == _TOKENIZERS_BACKEND
+                and tokenizer_backend != "fastokens"
+            ):
+                tokenizer = _resolve_tokenizers_backend(
+                    tokenizer_name, *args, **common_kwargs
+                )
+
+        return _apply_post_load_fixes(tokenizer, tokenizer_name, tokenizer_revision)
+    except Exception as e:
+        if tokenizer_backend == "fastokens":
+            raise RuntimeError(
+                f"fastokens failed to load tokenizer for {tokenizer_name!r}. "
+                f"This model's tokenizer may not be supported by fastokens — "
+                f"see https://github.com/crusoecloud/fastokens. "
+                f"Re-run without --tokenizer-backend=fastokens to use the default backend."
+            ) from e
+        raise
+
+
+# ---------------------------------------------------------------------------
+# Exported helpers (used by processor.py, etc.)
+# ---------------------------------------------------------------------------
+
+
+def _fix_added_tokens_encoding(tokenizer):
+    """Ensure special tokens encode as single tokens in transformers v5.
+
+    Some model tokenizers (e.g. MiniCPM-V-4) define special tokens like <image>,
+    <slice> as attributes on the tokenizer class with corresponding IDs in the
+    vocabulary (via tokenizer.json's added_tokens). In transformers v5, these
+    tokens may not appear in get_added_vocab() and encode() splits them into
+    subwords, breaking multimodal pipelines that rely on finding them in input_ids.
+
+    This function discovers such tokens by scanning tokenizer attributes, checks
+    if they encode correctly, and re-registers any that don't.
+    """
+
+    # Discover special token strings from tokenizer attributes.
+    # Model tokenizers (e.g. MiniCPMVTokenizerFast) store them as attributes
+    # like im_start="<image>", slice_start="<slice>", etc.
+    def _is_special_token_attr(val):
+        return (
+            isinstance(val, str)
+            and val.startswith("<")
+            and val.endswith(">")
+            and len(val) <= 20
+        )
+
+    candidates = {}
+    for attr in dir(tokenizer):
+        if attr.startswith("_"):
+            continue
+        try:
+            val = getattr(tokenizer, attr)
+        except (AttributeError, TypeError, ValueError):
+            continue
+        if not _is_special_token_attr(val):
+            continue
+        token_id = tokenizer.convert_tokens_to_ids(val)
+        if token_id is not None and token_id != tokenizer.unk_token_id:
+            candidates[val] = token_id
+
+    if not candidates:
+        return
+
+    def _encodes_correctly(token_str, expected_id):
+        try:
+            ids = tokenizer.encode(token_str, add_special_tokens=False)
+            return len(ids) == 1 and ids[0] == expected_id
+        except (ValueError, OverflowError, RuntimeError) as e:
+            logger.debug("Token %s encode check failed: %s", token_str, e)
+            return False
+
+    broken = [
+        tok for tok, eid in candidates.items() if not _encodes_correctly(tok, eid)
+    ]
+
+    if not broken:
+        return
+
+    from transformers import AddedToken
+
+    tokens_to_add = [AddedToken(tok, special=True, normalized=False) for tok in broken]
+    tokenizer.add_tokens(tokens_to_add, special_tokens=True)
+    logger.info(
+        "Re-registered %d special tokens for correct v5 encoding: %s",
+        len(broken),
+        broken[:10],
+    )

python/sglang/srt/utils/hf_transformers_patches.py

@@ -0,0 +1,453 @@
+# 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.
+# ==============================================================================
+"""Monkey-patches on transformers internals.
+
+Mix of backward-compat shims (re-add symbols removed in v5), workarounds
+for transformers v5 bugs, fixes for remote-model-code (trust_remote_code)
+that hasn't been updated for v5 yet, and CI-only patches (e.g. neutralize
+HF API calls to avoid rate limits).
+
+Import this module early (before any ``from_pretrained`` call) to activate
+all patches.  It is safe to import multiple times -- patches are idempotent.
+"""
+
+import inspect
+
+from sglang.srt.utils import logger
+
+_applied = False
+
+
+# ---------------------------------------------------------------------------
+# Public API: apply_all() -- import-time patches (idempotent)
+# ---------------------------------------------------------------------------
+
+
+def apply_all():
+    """Apply all transformers compatibility patches (idempotent).
+
+    Call this once at import time.  It is safe to call multiple times.
+
+    No-op when the ``transformers`` package is not installed -- frontend-only
+    sglang users should not be forced to install transformers just to import
+    the top-level ``sglang`` package.
+    """
+    global _applied
+    if _applied:
+        return
+    try:
+        import transformers  # noqa: F401
+    except ImportError:
+        _applied = True
+        return
+    _applied = True
+
+    # v5.4 patches
+    _patch_flash_attn_availability()
+    _patch_rope_parameters_validation()
+    _patch_removed_symbols()
+    _patch_image_processor_kwargs()
+    _patch_image_process_cuda_tensor()
+    _patch_nemotron_h_pattern()
+
+    # v5 general patches
+    _ensure_clean_up_tokenization_compat()
+    _ensure_is_torch_fx_available_compat()
+
+    # CI-only: neutralize HF API calls inside tokenizer from_pretrained
+    patch_is_base_mistral_in_ci()
+
+    logger.debug("transformers compatibility patches applied")
+
+
+# ---------------------------------------------------------------------------
+# Public API: on-demand helpers (called explicitly by other modules)
+# ---------------------------------------------------------------------------
+
+
+def normalize_rope_scaling_compat(config) -> None:
+    """Ensure rope_scaling dicts have ``"type"`` alongside ``"rope_type"``.
+
+    Transformers v5 standardises rope_scaling to use ``"rope_type"`` and may
+    omit the legacy ``"type"`` key.  Remote-code models (e.g. Kimi-VL) still
+    read ``rope_scaling["type"]``, causing a ``KeyError``.  This helper adds
+    ``"type"`` from ``"rope_type"`` whenever it is missing, recursively across
+    the config and all its sub-configs.
+    """
+
+    def _patch(cfg):
+        rs = getattr(cfg, "rope_scaling", None)
+        if isinstance(rs, dict) and "rope_type" in rs and "type" not in rs:
+            rs["type"] = rs["rope_type"]
+        # Recurse into sub-configs
+        for attr in (
+            "text_config",
+            "llm_config",
+            "language_config",
+            "vision_config",
+            "thinker_config",
+        ):
+            sub = getattr(cfg, attr, None)
+            if sub is not None:
+                _patch(sub)
+
+    _patch(config)
+
+
+def _ensure_gguf_version():
+    """Workaround for transformers v5 bug where is_gguf_available() fails
+    when the gguf package lacks __version__ and metadata lookup also fails,
+    resulting in packaging.version.InvalidVersion: Invalid version: 'N/A'."""
+    try:
+        import gguf
+
+        if not hasattr(gguf, "__version__"):
+            import importlib.metadata
+
+            try:
+                gguf.__version__ = importlib.metadata.version("gguf")
+            except importlib.metadata.PackageNotFoundError:
+                gguf.__version__ = "0.0.0"
+            except (ValueError, OSError, TypeError) as e:
+                logger.warning(
+                    "Failed to determine gguf package version: %s. "
+                    "Falling back to '0.0.0'.",
+                    e,
+                )
+                gguf.__version__ = "0.0.0"
+    except ImportError:
+        pass
+
+
+# ---------------------------------------------------------------------------
+# v5.4 patches (merged from transformers_v54_compat.py)
+# ---------------------------------------------------------------------------
+
+
+def _patch_rope_parameters_validation():
+    """Guard ``standardize_rope_params()`` against missing
+    ``max_position_embeddings``.
+
+    For ``PretrainedConfig``, ``standardize_rope_params()`` accesses
+    ``self.max_position_embeddings`` during ``__post_init__`` before extra
+    kwargs are set as attributes, causing ``AttributeError``.
+
+    Fix: guard ``standardize_rope_params`` against missing
+    ``max_position_embeddings``.
+    """
+    from transformers import PretrainedConfig
+
+    # standardize_rope_params accesses self.max_position_embeddings before
+    # __post_init__ sets extra kwargs — skip when the attribute is absent.
+    if hasattr(PretrainedConfig, "standardize_rope_params"):
+        _orig_standardize = PretrainedConfig.standardize_rope_params
+
+        def _safe_standardize(self):
+            if not hasattr(self, "max_position_embeddings"):
+                return
+            return _orig_standardize(self)
+
+        PretrainedConfig.standardize_rope_params = _safe_standardize
+
+
+def _patch_flash_attn_availability():
+    """Prevent flash-attn-4 from masquerading as flash-attn-2.
+
+    flash-attn-4 registers a bare ``flash_attn`` namespace that makes
+    ``is_flash_attn_2_available()`` return True, but lacks the v2 API.
+    Remote model code (e.g. Kimi-VL) guarded by that check will crash.
+
+    TODO(upstream): model authors should check for specific API symbols.
+    """
+    try:
+        import flash_attn as _fa
+
+        if not hasattr(_fa, "flash_attn_func"):
+            import transformers.utils as _u
+            import transformers.utils.import_utils as _ui
+
+            _ui.is_flash_attn_2_available = lambda: False
+            _u.is_flash_attn_2_available = lambda: False
+    except ImportError:
+        pass
+
+
+def _patch_removed_symbols():
+    """Re-export symbols removed in transformers v5.4.0.
+
+    Remote model code (e.g. DeepSeek-OCR) still imports these.
+    ``check_imports`` in ``dynamic_module_utils.py`` validates imports at
+    config-load time, so these must exist before any ``from_pretrained``.
+
+    Removed symbols:
+    - ``LlamaFlashAttention2`` -- replaced by unified ``LlamaAttention``
+    - ``is_flash_attn_greater_or_equal_2_10`` -- replaced by
+      ``is_flash_attn_greater_or_equal("2.10.0")``
+
+    TODO(upstream): DeepSeek-OCR / deepseek_vl_v2 remote code needs update.
+    """
+    # LlamaFlashAttention2
+    try:
+        import logging
+
+        # Importing modeling_llama triggers a deep import chain:
+        #   modeling_llama -> modeling_utils -> quantizers -> torchao
+        # torchao emits a noisy warning about incompatible torch versions
+        # that is irrelevant here — suppress it during this import.
+        _torchao_logger = logging.getLogger("torchao")
+        _prev_level = _torchao_logger.level
+        _torchao_logger.setLevel(logging.ERROR)
+        try:
+            from transformers.models.llama import modeling_llama
+        finally:
+            _torchao_logger.setLevel(_prev_level)
+
+        if not hasattr(modeling_llama, "LlamaFlashAttention2"):
+            if hasattr(modeling_llama, "LlamaAttention"):
+                modeling_llama.LlamaFlashAttention2 = modeling_llama.LlamaAttention
+    except ImportError:
+        logger.warning(
+            "Could not import transformers.models.llama.modeling_llama; "
+            "LlamaFlashAttention2 compat patch not applied."
+        )
+
+    # is_flash_attn_greater_or_equal_2_10
+    try:
+        import transformers.utils as _u
+
+        if not hasattr(_u, "is_flash_attn_greater_or_equal_2_10"):
+            if hasattr(_u, "is_flash_attn_greater_or_equal"):
+                _u.is_flash_attn_greater_or_equal_2_10 = (
+                    lambda: _u.is_flash_attn_greater_or_equal("2.10.0")
+                )
+            else:
+                _u.is_flash_attn_greater_or_equal_2_10 = lambda: False
+    except ImportError:
+        logger.warning(
+            "Could not import transformers.utils; "
+            "is_flash_attn_greater_or_equal_2_10 compat patch not applied."
+        )
+
+
+def _patch_image_processor_kwargs():
+    """Allow remote image processors that lack ``**kwargs`` in preprocess().
+
+    Transformers v5.4 passes new kwargs (e.g. ``device``) through
+    ``BaseImageProcessor.__call__`` -> ``preprocess()``.  Remote model code
+    (e.g. KimiVL) that defines ``preprocess()`` without ``**kwargs`` will
+    crash with ``TypeError``.
+
+    Fix: wrap ``__call__`` to catch ``TypeError`` and retry with only the
+    kwargs that ``preprocess()`` actually accepts.
+
+    TODO(upstream): KimiVL image_processing_kimi_vl.py needs ``**kwargs``.
+    """
+    try:
+        from transformers.image_processing_utils import BaseImageProcessor
+
+        original = BaseImageProcessor.__call__
+
+        def safe_call(self, images, *args, **kwargs):
+            try:
+                return original(self, images, *args, **kwargs)
+            except TypeError as e:
+                if "unexpected keyword argument" not in str(e):
+                    raise
+                sig = inspect.signature(self.preprocess)
+                params = sig.parameters
+                if any(
+                    p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values()
+                ):
+                    raise
+                dropped = {k for k in kwargs if k not in params}
+                if dropped:
+                    logger.warning(
+                        "Image processor %s.preprocess() does not accept %s; "
+                        "retrying without them. Update the model's image processor "
+                        "to accept **kwargs.",
+                        type(self).__name__,
+                        dropped,
+                    )
+                valid = {k: v for k, v in kwargs.items() if k in params}
+                return original(self, images, *args, **valid)
+
+        BaseImageProcessor.__call__ = safe_call
+    except ImportError:
+        logger.debug(
+            "_patch_image_processor_kwargs: BaseImageProcessor not importable, patch skipped"
+        )
+
+
+def _patch_image_process_cuda_tensor():
+    """Fix ``process_image()`` crashing on CUDA tensors.
+
+    Transformers v5.4's PIL image processing backend calls
+    ``image.numpy()`` on torch tensors, which fails for CUDA tensors.
+    Patch to call ``.cpu().numpy()`` instead.
+
+    TODO(upstream): report to HF transformers.
+    """
+    try:
+        import torch
+        import transformers.image_processing_backends as ipb
+
+        for cls_name in ("PilBackend", "PilImageProcessingMixin"):
+            cls = getattr(ipb, cls_name, None)
+            if cls is None or not hasattr(cls, "process_image"):
+                continue
+            original = cls.process_image
+
+            def patched_process_image(
+                self, image, *args, _orig=original, _Tensor=torch.Tensor, **kwargs
+            ):
+                if isinstance(image, _Tensor) and image.is_cuda:
+                    image = image.cpu()
+                return _orig(self, image, *args, **kwargs)
+
+            cls.process_image = patched_process_image
+    except ImportError:
+        logger.debug(
+            "_patch_image_process_cuda_tensor: required modules not importable, patch skipped"
+        )
+
+
+def _patch_nemotron_h_pattern():
+    """Fix ``_pattern_to_list()`` crashing on ``-`` in hybrid_override_pattern.
+
+    Nemotron-H models (e.g. NVIDIA-Nemotron-Nano-9B-v2) use patterns like
+    ``M-M-M-MM-M-*-...`` where ``-`` denotes an MLP layer.  The upstream
+    ``_pattern_to_list`` tries to map every character and crashes with
+    ``KeyError: '-'``.  We skip ``-`` (and any other unmapped chars)
+    since ``layers_block_type`` only tracks mamba/moe/attention layers.
+    SGLang reads MLP positions from ``hybrid_override_pattern`` directly.
+
+    TODO(upstream): report to HF transformers.
+    """
+    try:
+        from transformers.models.nemotron_h.configuration_nemotron_h import (
+            NemotronHConfig,
+        )
+
+        @staticmethod
+        def _pattern_to_list(pattern: str) -> list:
+            pattern_mapping = {
+                "M": "mamba",
+                "E": "moe",
+                "*": "attention",
+            }
+            return [
+                pattern_mapping[char] for char in pattern if char in pattern_mapping
+            ]
+
+        NemotronHConfig._pattern_to_list = _pattern_to_list
+    except ImportError:
+        logger.debug(
+            "_patch_nemotron_h_pattern: NemotronHConfig not importable, patch skipped"
+        )
+
+
+# ---------------------------------------------------------------------------
+# v5 general patches
+# ---------------------------------------------------------------------------
+
+
+def _ensure_clean_up_tokenization_compat() -> None:
+    """Re-add ``clean_up_tokenization`` removed in transformers v5.
+
+    Remote-code tokenizers (e.g. InternLM2Tokenizer) call
+    ``self.clean_up_tokenization()`` which was a static method on
+    ``PreTrainedTokenizerBase`` in v4 but removed in v5. Patch it back
+    so existing HuggingFace Hub tokenizer code keeps working.
+    """
+    from transformers import PreTrainedTokenizerBase
+
+    if hasattr(PreTrainedTokenizerBase, "clean_up_tokenization"):
+        return
+
+    @staticmethod
+    def clean_up_tokenization(out_string: str) -> str:
+        out_string = (
+            out_string.replace(" .", ".")
+            .replace(" ?", "?")
+            .replace(" !", "!")
+            .replace(" ,", ",")
+            .replace(" ' ", "'")
+            .replace(" n't", "n't")
+            .replace(" 'm", "'m")
+            .replace(" 's", "'s")
+            .replace(" 've", "'ve")
+            .replace(" 're", "'re")
+        )
+        return out_string
+
+    PreTrainedTokenizerBase.clean_up_tokenization = clean_up_tokenization
+
+
+def _ensure_is_torch_fx_available_compat() -> None:
+    """Re-add ``is_torch_fx_available`` removed in transformers v5.
+
+    Remote-code models (e.g. MiniCPM-V) import ``is_torch_fx_available``
+    from ``transformers.utils.import_utils``.  The function was removed
+    in v5.  Patch it back so existing HuggingFace Hub model code keeps
+    working.  torch.fx is always available in PyTorch >= 2.0.
+    """
+    import transformers.utils.import_utils as _import_utils
+
+    if hasattr(_import_utils, "is_torch_fx_available"):
+        return
+
+    _import_utils.is_torch_fx_available = lambda: True
+
+
+# ---------------------------------------------------------------------------
+# CI-only patches
+# ---------------------------------------------------------------------------
+
+_is_base_mistral_patched = False
+
+
+def patch_is_base_mistral_in_ci():
+    """Patch transformers' _patch_mistral_regex to avoid HF API calls in CI.
+
+    transformers defines is_base_mistral as a local function inside
+    _patch_mistral_regex, so it cannot be patched via module attribute.
+    Instead we replace the entire _patch_mistral_regex classmethod with a
+    version that simply returns the tokenizer unchanged.
+
+    In CI this prevents exhausting the 3000 req/5min HF API rate limit.
+
+    TODO(upstream): remove once transformers stops calling model_info()
+    inside _patch_mistral_regex (or removes the method entirely).
+    """
+    global _is_base_mistral_patched
+    if _is_base_mistral_patched:
+        return
+
+    from sglang.srt.environ import envs
+
+    if not envs.SGLANG_IS_IN_CI.get():
+        return
+
+    from transformers import PreTrainedTokenizerFast
+
+    if hasattr(PreTrainedTokenizerFast, "_patch_mistral_regex"):
+
+        @classmethod
+        def _noop_patch_mistral_regex(cls, tokenizer, *args, **kwargs):
+            return tokenizer
+
+        PreTrainedTokenizerFast._patch_mistral_regex = _noop_patch_mistral_regex
+        logger.info("CI: patched _patch_mistral_regex to skip HF API calls")
+
+    _is_base_mistral_patched = True

python/sglang/srt/utils/hf_transformers_utils.py

@@ -0,0 +1,17 @@
+# 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.
+# ==============================================================================
+"""Backward-compatible shim — all code has moved to sglang.srt.utils.hf_transformers."""
+
+from sglang.srt.utils.hf_transformers import *  # noqa: F401, F403
+from sglang.srt.utils.hf_transformers import __all__  # noqa: F401

python/sglang/srt/utils/host_shared_memory.py

@@ -0,0 +1,82 @@
+import logging
+from dataclasses import dataclass
+from multiprocessing import shared_memory
+from pathlib import Path
+from typing import List, Optional
+
+import numpy as np
+import torch
+
+from sglang.srt.distributed.naive_distributed import get_naive_distributed
+from sglang.srt.utils import check_cuda_result
+
+logger = logging.getLogger(__name__)
+
+
+class HostSharedMemoryManager:
+    def __init__(self, base_name: str):
+        self._base_name = Path(base_name)
+        self._operation_index = 0
+        self._records: List[_Record] = []
+
+    def malloc(self, *, shape, dtype):
+        meta_tensor = torch.empty(size=shape, dtype=dtype, device="meta")
+        raw = self._malloc_raw(num_bytes=meta_tensor.nbytes)
+        return raw.view(dtype).view(*shape)
+
+    def _malloc_raw(self, *, num_bytes: int) -> torch.Tensor:
+        import cuda.bindings.runtime as cuda_rt
+
+        self._operation_index += 1
+        shm_name = f"{self._base_name}_op{self._operation_index}"
+
+        # TODO handle dispose
+        if get_naive_distributed().get_rank() == 0:
+            shm = shared_memory.SharedMemory(name=shm_name, create=True, size=num_bytes)
+
+        get_naive_distributed().barrier()
+
+        if get_naive_distributed().get_rank() != 0:
+            shm = shared_memory.SharedMemory(name=shm_name)
+
+        np_array = np.ndarray((num_bytes,), dtype=np.uint8, buffer=shm.buf)
+        tensor = torch.from_numpy(np_array)
+
+        check_cuda_result(
+            cuda_rt.cudaHostRegister(
+                tensor.data_ptr(), num_bytes, cuda_rt.cudaHostRegisterPortable
+            )
+        )
+
+        get_naive_distributed().barrier()
+
+        self._records.append(
+            _Record(
+                shm=shm,
+                np_array=np_array,
+                tensor=tensor,
+            )
+        )
+        return tensor
+
+
+@dataclass
+class _Record:
+    shm: shared_memory.SharedMemory
+    np_array: np.ndarray
+    tensor: torch.Tensor
+
+
+# Can have multi instances if needed
+_instance: Optional[HostSharedMemoryManager] = None
+
+
+def get_host_shared_memory_manager():
+    assert _instance is not None
+    return _instance
+
+
+def set_host_shared_memory_manager(instance: HostSharedMemoryManager):
+    global _instance
+    assert _instance is None
+    _instance = instance

python/sglang/srt/utils/http_middleware_patch.py

@@ -0,0 +1,67 @@
+"""
+Fix @app.middleware("http") whose BaseHTTPMiddleware call_next replaces
+ASGI ``receive``, breaking request.is_disconnected() and preventing
+non-streaming request abort on client disconnect.
+
+patch_app_http_middleware(app) replaces @app.middleware("http") with a
+version whose call_next passes ``receive`` through untouched.
+"""
+
+from __future__ import annotations
+
+from starlette.requests import Request
+
+
+class _SentResponse:
+    """Response proxy returned after the real response was already sent."""
+
+    def __init__(self, status_code: int):
+        self.status_code = status_code
+
+
+class _PureASGIDispatch:
+    """Pure ASGI middleware providing a fixed call_next that passes
+    ``receive`` through untouched (unlike BaseHTTPMiddleware)."""
+
+    def __init__(self, app, dispatch):
+        self.app = app
+        self.dispatch = dispatch
+
+    async def __call__(self, scope, receive, send):
+        if scope["type"] != "http":
+            await self.app(scope, receive, send)
+            return
+
+        request = Request(scope, receive)
+        status_code = 500
+
+        async def call_next(_request):
+            nonlocal status_code
+
+            async def send_and_capture(message):
+                nonlocal status_code
+                if message["type"] == "http.response.start":
+                    status_code = message["status"]
+                await send(message)
+
+            await self.app(scope, receive, send_and_capture)
+            return _SentResponse(status_code)
+
+        await self.dispatch(request, call_next)
+
+
+def patch_app_http_middleware(app):
+    """Replace @app.middleware("http") with a fixed-call_next version."""
+    _orig = app.middleware
+
+    def _fixed(middleware_type):
+        if middleware_type == "http":
+
+            def decorator(fn):
+                app.add_middleware(_PureASGIDispatch, dispatch=fn)
+                return fn
+
+            return decorator
+        return _orig(middleware_type)
+
+    app.middleware = _fixed

python/sglang/srt/utils/json_response.py

@@ -0,0 +1,30 @@
+"""Utilities for JSON serialization in HTTP responses."""
+
+from typing import Any
+
+import orjson
+from fastapi.responses import Response
+
+# Keep response serialization behavior consistent across endpoints:
+# - Support non-string dictionary keys used in some metadata payloads.
+# - Support numpy scalars/arrays without pre-conversion.
+ORJSON_RESPONSE_OPTIONS = orjson.OPT_NON_STR_KEYS | orjson.OPT_SERIALIZE_NUMPY
+
+
+def dumps_json(content: Any) -> bytes:
+    """Serialize content to JSON bytes using SGLang's ORJSON options."""
+    return orjson.dumps(content, option=ORJSON_RESPONSE_OPTIONS)
+
+
+class SGLangORJSONResponse(Response):
+    """ORJSON response with SGLang-specific serialization options."""
+
+    media_type = "application/json"
+
+    def render(self, content: Any) -> bytes:
+        return dumps_json(content)
+
+
+def orjson_response(content: Any, status_code: int = 200) -> Response:
+    """Create a JSON response with stable ORJSON serialization options."""
+    return SGLangORJSONResponse(content=content, status_code=status_code)

python/sglang/srt/utils/log_utils.py

@@ -0,0 +1,74 @@
+from __future__ import annotations
+
+import json
+import logging
+import os
+import socket
+import sys
+from datetime import datetime
+from logging.handlers import TimedRotatingFileHandler
+from typing import List, Optional, Union
+
+import torch.distributed as dist
+
+
+def create_log_targets(
+    *, targets: Optional[List[str]], name_prefix: str
+) -> List[logging.Logger]:
+    if not targets:
+        return [_create_log_target_stdout(name_prefix)]
+    return [_create_log_target(t, name_prefix) for t in targets]
+
+
+def _create_log_target(target: str, name_prefix: str) -> logging.Logger:
+    if target.lower() == "stdout":
+        return _create_log_target_stdout(name_prefix)
+    return _create_log_target_file(target, name_prefix)
+
+
+def _create_log_target_stdout(name_prefix: str) -> logging.Logger:
+    return _create_logger_with_handler(
+        f"{name_prefix}.stdout", logging.StreamHandler(sys.stdout)
+    )
+
+
+def _create_log_target_file(directory: str, name_prefix: str) -> logging.Logger:
+    os.makedirs(directory, exist_ok=True)
+    hostname = socket.gethostname()
+    rank = dist.get_rank() if dist.is_initialized() else 0
+    filename = os.path.join(directory, f"{hostname}_{rank}.log")
+    handler = TimedRotatingFileHandler(
+        filename, when="H", backupCount=0, encoding="utf-8"
+    )
+    return _create_logger_with_handler(
+        f"{name_prefix}.file.{directory}.{hostname}_{rank}", handler
+    )
+
+
+def _create_logger_with_handler(name: str, handler: logging.Handler) -> logging.Logger:
+    logger = logging.getLogger(name)
+    logger.setLevel(logging.INFO)
+    logger.propagate = False
+    if not logger.handlers:
+        handler.setFormatter(
+            logging.Formatter("[%(asctime)s] %(message)s", datefmt="%Y-%m-%d %H:%M:%S")
+        )
+        logger.addHandler(handler)
+    return logger
+
+
+def log_json(
+    loggers: Union[logging.Logger, List[logging.Logger]], event: str, data: dict
+) -> None:
+    log_data = {
+        "timestamp": datetime.now().isoformat(),
+        "event": event,
+        **data,
+    }
+    msg = json.dumps(log_data, ensure_ascii=False)
+
+    if not isinstance(loggers, list):
+        loggers = [loggers]
+
+    for logger in loggers:
+        logger.info(msg)

python/sglang/srt/utils/model_file_verifier.py

@@ -0,0 +1,284 @@
+"""
+Model File Verifier - Verify model file integrity using SHA256 checksums.
+
+Example commands:
+    # Verify using HuggingFace model online metadata
+    python -m sglang.srt.utils.model_file_verifier verify --model-path /path/to/model --model-checksum Qwen/Qwen3-0.6B
+
+    # Verify using locally generated checksum
+    python -m sglang.srt.utils.model_file_verifier generate --model-path <hf-id-or-model-path> --model-checksum checksums.json
+    python -m sglang.srt.utils.model_file_verifier verify --model-path /path/to/model --model-checksum checksums.json
+"""
+
+import argparse
+import fnmatch
+import hashlib
+import json
+import warnings
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+# ======== Data Format ========
+
+
+@dataclass
+class FileInfo:
+    sha256: str
+    size: int
+
+
+@dataclass
+class Manifest:
+    files: Dict[str, FileInfo]
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "Manifest":
+        if "checksums" in data:
+            warnings.warn(
+                "The 'checksums' format is deprecated. "
+                "Please regenerate with the latest version to use the new 'files' format.",
+                DeprecationWarning,
+                stacklevel=3,
+            )
+            return cls(
+                files={
+                    k: FileInfo(sha256=v, size=-1) for k, v in data["checksums"].items()
+                }
+            )
+        return cls(files={k: FileInfo(**v) for k, v in data["files"].items()})
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+
+# ======== Constants ========
+
+
+IGNORE_PATTERNS = [
+    ".DS_Store",
+    "*.lock",
+    ".gitattributes",
+    "LICENSE",
+    "LICENSE.*",
+    "README.md",
+    "README.*",
+    "NOTICE",
+]
+
+
+# ======== Verify ========
+
+
+def verify(*, model_path: str, checksums_source: str, max_workers: int = 4) -> None:
+    model_path = Path(model_path).resolve()
+    expected = _load_checksums(checksums_source)
+    actual = _compute_manifest_from_folder(
+        model_path=model_path,
+        filenames=list(expected.files.keys()),
+        max_workers=max_workers,
+    )
+    _compare_manifests(expected=expected, actual=actual)
+    print(f"[ModelFileVerifier] All {len(expected.files)} files verified successfully.")
+
+
+def _compare_manifests(*, expected: Manifest, actual: Manifest) -> None:
+    errors = []
+    for filename, exp in expected.files.items():
+        if filename not in actual.files:
+            errors.append(f"{filename}: missing (expected size={exp.size})")
+        elif actual.files[filename].sha256 != exp.sha256:
+            act = actual.files[filename]
+            errors.append(
+                f"{filename}: mismatch (expected={exp.sha256[:16]}... size={exp.size}, actual={act.sha256[:16]}... size={act.size})"
+            )
+
+    if errors:
+        raise IntegrityError("Integrity check failed: " + "; ".join(errors))
+
+
+# ======== Generate ========
+
+
+def generate_checksums(
+    *, source: str, output_path: str, max_workers: int = 4
+) -> Manifest:
+    if Path(source).is_dir():
+        model_path = Path(source).resolve()
+        files = _discover_files(model_path)
+        if not files:
+            raise IntegrityError(f"No model files found in {model_path}")
+        manifest = _compute_manifest_from_folder(
+            model_path=model_path, filenames=files, max_workers=max_workers
+        )
+    else:
+        manifest = Manifest(files=_load_file_infos_from_hf(repo_id=source))
+
+    Path(output_path).write_text(
+        json.dumps(manifest.to_dict(), indent=2, sort_keys=True)
+    )
+
+    print(
+        f"[ModelFileVerifier] Generated checksums for {len(manifest.files)} files -> {output_path}"
+    )
+    return manifest
+
+
+def _discover_files(model_path: Path) -> List[str]:
+    return sorted(
+        e.name
+        for e in model_path.iterdir()
+        if e.is_file()
+        and not e.name.startswith(".")
+        and not any(fnmatch.fnmatch(e.name, p) for p in IGNORE_PATTERNS)
+    )
+
+
+# ======== Load Checksums ========
+
+
+def _load_checksums(source: str) -> Manifest:
+    if Path(source).is_file():
+        data = json.loads(Path(source).read_text())
+        return Manifest.from_dict(data)
+    return Manifest(files=_load_file_infos_from_hf(repo_id=source))
+
+
+def _load_file_infos_from_hf(*, repo_id: str) -> Dict[str, FileInfo]:
+    from huggingface_hub import HfFileSystem
+
+    fs = HfFileSystem()
+    files = fs.ls(repo_id, detail=True)
+
+    file_infos = dict(
+        r for r in map(lambda f: _get_filename_and_info_from_hf_file(fs, f), files) if r
+    )
+    if not file_infos:
+        raise IntegrityError(f"No files found in HF repo {repo_id}.")
+
+    return file_infos
+
+
+def _get_filename_and_info_from_hf_file(
+    fs, file_info
+) -> Optional[Tuple[str, FileInfo]]:
+    if file_info.get("type") != "file":
+        return None
+
+    filename = Path(file_info.get("name", "")).name
+    if any(fnmatch.fnmatch(filename, pat) for pat in IGNORE_PATTERNS):
+        return None
+
+    size = file_info.get("size", -1)
+    lfs_info = file_info.get("lfs")
+    if lfs_info and "sha256" in lfs_info:
+        return filename, FileInfo(sha256=lfs_info["sha256"], size=size)
+
+    if "sha256" in file_info:
+        return filename, FileInfo(sha256=file_info["sha256"], size=size)
+
+    content = fs.read_bytes(file_info.get("name", ""))
+    return filename, FileInfo(
+        sha256=hashlib.sha256(content).hexdigest(), size=len(content)
+    )
+
+
+# ======== Compute Checksums ========
+
+
+def _compute_manifest_from_folder(
+    *, model_path: Path, filenames: List[str], max_workers: int
+) -> Manifest:
+    from tqdm import tqdm
+
+    def compute_one(filename: str) -> Tuple[str, Optional[FileInfo]]:
+        full_path = model_path / filename
+        if not full_path.exists():
+            return filename, None
+        sha256 = compute_sha256(file_path=full_path)
+        size = full_path.stat().st_size
+        return filename, FileInfo(sha256=sha256, size=size)
+
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        results = list(
+            tqdm(
+                executor.map(compute_one, filenames),
+                total=len(filenames),
+                desc="Computing checksums",
+            )
+        )
+
+    return Manifest(files={k: v for k, v in results if v is not None})
+
+
+def compute_sha256(*, file_path) -> str:
+    sha256 = hashlib.sha256()
+    with open(file_path, "rb") as f:
+        while chunk := f.read(64 * 1024):
+            sha256.update(chunk)
+    return sha256.hexdigest()
+
+
+# ======== Exceptions ========
+
+
+class IntegrityError(Exception):
+    pass
+
+
+# ======== CLI ========
+
+
+def _add_common_args(parser):
+    parser.add_argument(
+        "--model-path",
+        required=True,
+        help="Local model directory or HuggingFace repo ID",
+    )
+    parser.add_argument(
+        "--model-checksum",
+        required=True,
+        help="Checksums JSON file path",
+    )
+    parser.add_argument(
+        "--workers", type=int, default=4, help="Number of parallel workers"
+    )
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Model File Verifier - Verify model file integrity using checksums"
+    )
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    gen_parser = subparsers.add_parser(
+        "generate", help="Generate checksums.json for a model"
+    )
+    _add_common_args(gen_parser)
+    gen_parser.set_defaults(
+        func=lambda args: generate_checksums(
+            source=args.model_path,
+            output_path=args.model_checksum,
+            max_workers=args.workers,
+        )
+    )
+
+    verify_parser = subparsers.add_parser(
+        "verify", help="Verify model files against checksums"
+    )
+    _add_common_args(verify_parser)
+    verify_parser.set_defaults(
+        func=lambda args: verify(
+            model_path=args.model_path,
+            checksums_source=args.model_checksum,
+            max_workers=args.workers,
+        )
+    )
+
+    args = parser.parse_args()
+    args.func(args)
+
+
+if __name__ == "__main__":
+    main()

python/sglang/srt/utils/msgspec_utils.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+import base64
+import binascii
+from typing import Any
+
+import msgspec
+from pydantic_core import core_schema
+
+
+class Base64Bytes:
+    """Pydantic marker for HTTP JSON base64-encoded bytes fields."""
+
+    def __get_pydantic_core_schema__(self, source_type: Any, handler):
+        return core_schema.no_info_before_validator_function(
+            self._decode_value,
+            handler(source_type),
+        )
+
+    @classmethod
+    def _decode_value(cls, value: Any) -> Any:
+        if isinstance(value, str):
+            try:
+                return base64.b64decode(value, validate=True)
+            except binascii.Error as exc:
+                raise ValueError("Expected base64-encoded bytes") from exc
+
+        if isinstance(value, list):
+            return [cls._decode_value(item) for item in value]
+
+        if isinstance(value, tuple):
+            return tuple(cls._decode_value(item) for item in value)
+
+        return value
+
+
+def msgspec_to_builtins(obj: Any) -> Any:
+    """Recursively convert msgspec structs to dict/list Python builtins."""
+    if isinstance(obj, msgspec.Struct):
+        return {
+            field.name: msgspec_to_builtins(getattr(obj, field.name))
+            for field in msgspec.structs.fields(type(obj))
+        }
+
+    if isinstance(obj, dict):
+        return {key: msgspec_to_builtins(value) for key, value in obj.items()}
+
+    if isinstance(obj, list):
+        return [msgspec_to_builtins(item) for item in obj]
+
+    if isinstance(obj, tuple):
+        return tuple(msgspec_to_builtins(item) for item in obj)
+
+    if isinstance(obj, set):
+        return [msgspec_to_builtins(item) for item in obj]
+
+    return obj
+
+
+def msgspec_struct_pydantic_core_schema(cls: type[msgspec.Struct], handler):
+    fields = {}
+    for struct_field in msgspec.structs.fields(cls):
+        field_schema = handler.generate_schema(struct_field.type)
+        required = (
+            struct_field.default is msgspec.NODEFAULT
+            and struct_field.default_factory is msgspec.NODEFAULT
+        )
+
+        if struct_field.default is not msgspec.NODEFAULT:
+            field_schema = core_schema.with_default_schema(
+                field_schema,
+                default=struct_field.default,
+            )
+        elif struct_field.default_factory is not msgspec.NODEFAULT:
+            field_schema = core_schema.with_default_schema(
+                field_schema,
+                default_factory=struct_field.default_factory,
+            )
+
+        fields[struct_field.name] = core_schema.typed_dict_field(
+            field_schema,
+            required=required,
+        )
+
+    typed_dict_schema = core_schema.typed_dict_schema(
+        fields,
+        cls_name=cls.__name__,
+        extra_behavior="ignore",
+        ref=cls.__name__,
+    )
+
+    def build_struct(value):
+        return value if isinstance(value, cls) else cls(**value)
+
+    dict_to_struct_schema = core_schema.no_info_after_validator_function(
+        build_struct,
+        typed_dict_schema,
+    )
+    return core_schema.json_or_python_schema(
+        json_schema=dict_to_struct_schema,
+        python_schema=core_schema.union_schema(
+            [
+                core_schema.is_instance_schema(cls),
+                dict_to_struct_schema,
+            ],
+            mode="left_to_right",
+        ),
+    )

python/sglang/srt/utils/multi_stream_utils.py

@@ -0,0 +1,60 @@
+# Adapted from trtllm.
+
+from typing import Any, Callable, Optional
+
+import torch
+
+from sglang.srt.runtime_context import get_forward
+
+
+def set_do_multi_stream(enable: bool):
+    get_forward().set("multi_stream", enable)
+
+
+def do_multi_stream() -> bool:
+    return get_forward().multi_stream
+
+
+def with_multi_stream(enable: bool):
+    return get_forward().scoped(multi_stream=enable)
+
+
+def maybe_execute_in_parallel(
+    fn0: Callable,
+    fn1: Callable,
+    events: list[torch.cuda.Event],
+    aux_stream: Optional[torch.cuda.Stream] = None,
+) -> tuple[Any, Any]:
+    """Utility function to run two functions in two cuda streams in parallel. Multi-stream is
+    only enabled when cuda graph is turned on because switch stream has extra host overhead.
+
+    This design is mainly for low latency use case. It needs to be improved for max throughput
+    use case.
+    For simplicity, fn0 and fn1 do not support inputs.
+
+    Args:
+        fn0 (Callable): callable for the default stream
+        fn1 (Callable): callable for the second stream, aux_stream
+        events (list[torch.cuda.Event]): cuda events for callables
+        aux_stream (Optional[torch.cuda.Stream]): the second cuda stream for fn1.
+            Multi-stream is disabled when aux_stream is None.
+
+    Returns:
+        tuple[Any, Any]: the return values of fn0() and fn1()
+    """
+
+    multi_stream = do_multi_stream() and aux_stream is not None
+
+    if multi_stream:
+        events[0].record()
+        result0 = fn0()
+
+        with torch.cuda.stream(aux_stream):
+            events[0].wait()
+            result1 = fn1()
+            events[1].record()
+        events[1].wait()
+    else:
+        result0 = fn0()
+        result1 = fn1()
+    return (result0, result1)

python/sglang/srt/utils/network.py

@@ -0,0 +1,563 @@
+from __future__ import annotations
+
+import ipaddress
+import logging
+import os
+import socket
+import time
+from dataclasses import dataclass
+from typing import Optional, Tuple, Union
+
+import psutil
+import zmq
+
+logger = logging.getLogger(__name__)
+
+
+def get_open_port() -> int:
+    from sglang.srt.environ import envs
+
+    port = envs.SGLANG_PORT.get()
+    if port is not None:
+        while True:
+            if is_port_available(port):
+                return port
+            logger.info("Port %d is already in use, trying port %d", port, port + 1)
+            port += 1
+    sock = try_bind_socket()
+    port = sock.getsockname()[1]
+    sock.close()
+    return port
+
+
+def is_valid_ipv6_address(address: str) -> bool:
+    try:
+        ipaddress.IPv6Address(address)
+        return True
+    except ValueError:
+        return False
+
+
+def find_process_using_port(port: int) -> Optional[psutil.Process]:
+    for conn in psutil.net_connections(kind="inet"):
+        if conn.laddr.port == port:
+            try:
+                return psutil.Process(conn.pid)
+            except psutil.NoSuchProcess:
+                # It could happen by race condition (the proc dies when psutil.Process is called).
+                pass
+
+    return None
+
+
+MAX_VALID_PORT = 65535
+
+
+def wait_port_available(
+    port: int, port_name: str, timeout_s: int = 30, raise_exception: bool = True
+) -> bool:
+    if port < 0 or port > MAX_VALID_PORT:
+        raise ValueError(
+            f"{port_name} has invalid port number {port}. "
+            f"Valid TCP port range is 0-{MAX_VALID_PORT}."
+        )
+
+    error_message = f"{port_name} at {port} is not available"
+    for i in range(timeout_s):
+        if is_port_available(port):
+            return True
+
+        if i > 10 and i % 5 == 0:
+            process = find_process_using_port(port)
+            if process is None:
+                logger.warning(
+                    f"The port {port} is in use, but we could not find the process that uses it."
+                )
+            else:
+                pid = process.pid
+                error_message = f"{port_name} is used by a process already. {process.name()=}' {process.cmdline()=} {process.status()=} {pid=}"
+                logger.info(
+                    f"port {port} is in use. Waiting for {i} seconds for {port_name} to be available. {error_message}"
+                )
+        time.sleep(0.1)
+
+    if raise_exception:
+        raise ValueError(
+            f"{port_name} at {port} is not available in {timeout_s} seconds. {error_message}"
+        )
+    return False
+
+
+def _get_addrinfos_for_bind(host=None, port=0):
+    """Return deduplicated addrinfo tuples for binding (one per address family).
+
+    Args:
+        host: Bind address. None (with AI_PASSIVE) resolves to wildcard
+              addresses (0.0.0.0 / ::) suitable for accepting on all interfaces.
+        port: Port number. 0 lets the OS assign an available ephemeral port.
+
+    Flags:
+        AI_ADDRCONFIG — only return families actually configured on this host.
+        AI_PASSIVE    — return wildcard addresses suitable for bind().
+
+    Falls back to AF_INET if getaddrinfo fails (e.g. DNS misconfiguration).
+    """
+    try:
+        infos = socket.getaddrinfo(
+            host,
+            port,
+            socket.AF_UNSPEC,
+            socket.SOCK_STREAM,
+            0,
+            socket.AI_ADDRCONFIG | socket.AI_PASSIVE,
+        )
+        deduped = []
+        seen_families = set()
+        for info in infos:
+            if info[0] not in seen_families:
+                seen_families.add(info[0])
+                deduped.append(info)
+        # Prefer IPv4 so that callers without an explicit host get consistent
+        # behaviour across platforms (some OSes list IPv6 first).
+        deduped.sort(key=lambda x: (x[0] != socket.AF_INET,))
+        return deduped
+    except socket.gaierror:
+        fallback_host = "0.0.0.0" if host is None else host
+        return [(socket.AF_INET, socket.SOCK_STREAM, 0, "", (fallback_host, port))]
+
+
+def try_bind_socket(host=None, port=0, *, reuse_addr=True, listen=False):
+    """Bind a TCP socket on the first available address family (IPv4/IPv6).
+
+    Iterates over address families returned by _get_addrinfos_for_bind and
+    returns the first socket that successfully binds.
+
+    Args:
+        host: Bind address. None binds to all interfaces (0.0.0.0 / ::).
+        port: Port number. 0 lets the OS assign an available ephemeral port;
+              use sock.getsockname()[1] to retrieve the assigned port.
+        reuse_addr: Set SO_REUSEADDR to allow quick port reuse after close.
+        listen: Call listen(1) after bind, making the socket ready to accept.
+
+    Returns:
+        The bound socket. Caller is responsible for closing it.
+
+    Raises:
+        OSError: If bind fails on all configured address families.
+    """
+    for family, socktype, proto, _, sockaddr in _get_addrinfos_for_bind(host, port):
+        sock = socket.socket(family, socktype, proto)
+        try:
+            if family == socket.AF_INET6:
+                sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
+            if reuse_addr:
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            sock.bind(sockaddr)
+            if listen:
+                sock.listen(1)
+            return sock
+        except (OSError, OverflowError):
+            sock.close()
+    raise OSError(f"Could not bind port {port} on any configured address family")
+
+
+def is_port_available(port):
+    """Return whether a port is available on all configured address families."""
+    try:
+        for family, socktype, proto, _, sockaddr in _get_addrinfos_for_bind(port=port):
+            sock = socket.socket(family, socktype, proto)
+            try:
+                if family == socket.AF_INET6:
+                    sock.setsockopt(socket.IPPROTO_IPV6, socket.IPV6_V6ONLY, 1)
+                sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+                sock.bind(sockaddr)
+            finally:
+                sock.close()
+        return True
+    except (OSError, OverflowError):
+        return False
+
+
+def get_free_port():
+    sock = try_bind_socket()
+    port = sock.getsockname()[1]
+    sock.close()
+    return port
+
+
+def bind_port(port):
+    """Bind to a specific port, assuming it's available."""
+    return try_bind_socket(port=port, listen=True)
+
+
+def get_zmq_socket_on_host(
+    context: zmq.Context,
+    socket_type: zmq.SocketType,
+    host: Optional[str] = None,
+) -> Tuple[int, zmq.Socket]:
+    """Create and configure a ZeroMQ socket.
+
+    Args:
+        context: ZeroMQ context to create the socket from.
+        socket_type: Type of ZeroMQ socket to create.
+        host: Host to bind to, without "tcp://" prefix. Defaults to
+            "127.0.0.1" (localhost-only) to avoid exposing unauthenticated
+            sockets to the network (CVE-2026-3060). Callers that need
+            cross-machine reachability must pass an explicit host.
+
+    Returns:
+        Tuple of (port, socket) where port is the randomly assigned TCP port.
+    """
+    socket = context.socket(socket_type)
+    config_socket(socket, socket_type)
+    if host is None:
+        host = "127.0.0.1"
+    if is_valid_ipv6_address(host):
+        socket.setsockopt(zmq.IPV6, 1)
+        bind_host = f"tcp://[{host}]"
+    else:
+        bind_host = f"tcp://{host}"
+    port = socket.bind_to_random_port(bind_host)
+    return port, socket
+
+
+def config_socket(socket, socket_type: zmq.SocketType):
+    mem = psutil.virtual_memory()
+    total_mem = mem.total / 1024**3
+    available_mem = mem.available / 1024**3
+    if total_mem > 32 and available_mem > 16:
+        buf_size = int(0.5 * 1024**3)
+    else:
+        buf_size = -1
+
+    def set_send_opt():
+        socket.setsockopt(zmq.SNDHWM, 0)
+        socket.setsockopt(zmq.SNDBUF, buf_size)
+
+    def set_recv_opt():
+        socket.setsockopt(zmq.RCVHWM, 0)
+        socket.setsockopt(zmq.RCVBUF, buf_size)
+
+    if socket_type == zmq.PUSH:
+        set_send_opt()
+    elif socket_type == zmq.PULL:
+        set_recv_opt()
+    elif socket_type in [zmq.DEALER, zmq.REQ, zmq.REP, zmq.PAIR]:
+        set_send_opt()
+        set_recv_opt()
+    else:
+        raise ValueError(f"Unsupported socket type: {socket_type}")
+
+
+def get_local_ip_by_nic(interface: str = None) -> Optional[str]:
+    if not (interface := interface or os.environ.get("SGLANG_LOCAL_IP_NIC", None)):
+        return None
+    try:
+        import netifaces
+    except ImportError as e:
+        raise ImportError(
+            "Environment variable SGLANG_LOCAL_IP_NIC requires package netifaces, please install it through 'pip install netifaces'"
+        ) from e
+
+    try:
+        addresses = netifaces.ifaddresses(interface)
+        if netifaces.AF_INET in addresses:
+            for addr_info in addresses[netifaces.AF_INET]:
+                ip = addr_info.get("addr")
+                if ip and ip != "127.0.0.1" and ip != "0.0.0.0":
+                    return ip
+        if netifaces.AF_INET6 in addresses:
+            for addr_info in addresses[netifaces.AF_INET6]:
+                ip = addr_info.get("addr")
+                if ip and not ip.startswith("fe80::") and ip != "::1":
+                    return ip.split("%")[0]
+    except (ValueError, OSError) as e:
+        logger.warning(
+            f"{e} Can not get local ip from NIC. Please verify whether SGLANG_LOCAL_IP_NIC is set correctly."
+        )
+    return None
+
+
+def get_local_ip_by_remote() -> Optional[str]:
+    # Google's public DNS servers, used to discover the local IP.
+    # UDP connect doesn't send packets; it just selects the right source address.
+    # https://developers.google.com/speed/public-dns/docs/using#addresses
+    # Try IPv4 first, then IPv6. getaddrinfo on a literal IP returns exactly
+    # one result, so we unpack directly instead of looping.
+    for dns_host, dns_port in [("8.8.8.8", 80), ("2001:4860:4860::8888", 80)]:
+        try:
+            family, socktype, proto, _, sockaddr = socket.getaddrinfo(
+                dns_host,
+                dns_port,
+                socket.AF_UNSPEC,
+                socket.SOCK_DGRAM,
+                0,
+                socket.AI_ADDRCONFIG,
+            )[0]
+            with socket.socket(family, socktype, proto) as s:
+                s.connect(sockaddr)
+                return s.getsockname()[0]
+        except (socket.gaierror, OSError):
+            continue
+
+    # Fallback: resolve the local hostname to an IP address via /etc/hosts or DNS.
+    # Unreliable — many machines resolve hostname to 127.0.0.1, so we skip loopback.
+    try:
+        hostname = socket.gethostname()
+        ip = socket.getaddrinfo(
+            hostname, None, socket.AF_UNSPEC, 0, 0, socket.AI_ADDRCONFIG
+        )[0][4][0]
+        if ip and ip not in ("127.0.0.1", "0.0.0.0", "::1"):
+            return ip
+    except Exception:
+        pass
+
+    logger.warning("Can not get local ip by remote")
+    return None
+
+
+def get_local_ip_auto(fallback: str = None) -> str:
+    """
+    Automatically detect the local IP address using multiple fallback strategies.
+
+    This function attempts to obtain the local IP address through several methods.
+    If all methods fail, it returns the specified fallback value or raises an exception.
+
+    Args:
+        fallback (str, optional): Fallback IP address to return if all detection
+            methods fail. For server applications, explicitly set this to
+            "0.0.0.0" (IPv4) or "::" (IPv6) to bind to all available interfaces.
+            Defaults to None.
+
+    Returns:
+        str: The detected local IP address, or the fallback value if detection fails.
+
+    Raises:
+        ValueError: If IP detection fails and no fallback value is provided.
+
+    Note:
+        The function tries detection methods in the following order:
+        1. Direct IP detection via get_ip()
+        2. Network interface enumeration via get_local_ip_by_nic()
+        3. Remote connection method via get_local_ip_by_remote()
+    """
+    # Try environment variable
+    host_ip = os.getenv("SGLANG_HOST_IP", "") or os.getenv("HOST_IP", "")
+    if host_ip:
+        return host_ip
+    logger.debug("get_ip failed")
+    # Fallback
+    if ip := get_local_ip_by_nic():
+        return ip
+    logger.debug("get_local_ip_by_nic failed")
+    # Fallback
+    if ip := get_local_ip_by_remote():
+        return ip
+    logger.debug("get_local_ip_by_remote failed")
+    if fallback:
+        return fallback
+    raise ValueError("Can not get local ip")
+
+
+def get_zmq_socket(
+    context: zmq.Context,
+    socket_type: zmq.SocketType,
+    endpoint: Optional[str] = None,
+    bind: bool = True,
+) -> Union[zmq.Socket, Tuple[int, zmq.Socket]]:
+    """Create and configure a ZeroMQ socket.
+
+    Args:
+        context: ZeroMQ context to create the socket from.
+        socket_type: Type of ZeroMQ socket to create.
+        endpoint: Optional endpoint to bind/connect to. If None, binds to a random TCP port.
+        bind: Whether to bind (True) or connect (False) to the endpoint. Ignored if endpoint is None.
+
+    Returns:
+        If endpoint is None: Tuple of (port, socket) where port is the randomly assigned TCP port.
+        If endpoint is provided: The configured ZeroMQ socket.
+    """
+    socket = context.socket(socket_type)
+
+    if endpoint is None:
+        # Bind to random TCP port
+        config_socket(socket, socket_type)
+        port = socket.bind_to_random_port("tcp://*")
+        return port, socket
+    else:
+        if is_zmq_endpoint_ipv6(endpoint):
+            socket.setsockopt(zmq.IPV6, 1)
+
+        config_socket(socket, socket_type)
+
+        if bind:
+            socket.bind(endpoint)
+        else:
+            socket.connect(endpoint)
+
+        return socket
+
+
+def is_zmq_endpoint_ipv6(endpoint: str) -> bool:
+    """Return whether a ZMQ TCP endpoint contains a bracketed IPv6 host."""
+    prefix = "tcp://["
+    if not endpoint.startswith(prefix):
+        return False
+    end = endpoint.find("]", len(prefix))
+    if end == -1:
+        return False
+    return is_valid_ipv6_address(endpoint[len(prefix) : end])
+
+
+def _is_ipv6(host: str) -> bool:
+    """Check whether *host* is a valid IPv6 address (without brackets)."""
+    try:
+        ipaddress.IPv6Address(host)
+        return True
+    except ValueError:
+        return False
+
+
+def _wrap(host: str) -> str:
+    """Wrap an IPv6 address in brackets; pass IPv4/hostname through."""
+    return f"[{host}]" if _is_ipv6(host) else host
+
+
+def _parse_port(s: str) -> int:
+    try:
+        port = int(s)
+    except ValueError:
+        raise ValueError(f"Invalid port number: {s!r}")
+    if not (0 <= port <= 65535):
+        raise ValueError(f"Port out of range (0-65535): {port}")
+    return port
+
+
+@dataclass(frozen=True)
+class NetworkAddress:
+    host: str
+    port: int
+
+    def __post_init__(self):
+        # Auto-strip IPv6 brackets so callers can pass "[::1]" or "::1"
+        if self.host.startswith("[") and self.host.endswith("]"):
+            object.__setattr__(self, "host", self.host[1:-1])
+
+    @property
+    def is_ipv6(self) -> bool:
+        return _is_ipv6(self.host)
+
+    @property
+    def family(self) -> socket.AddressFamily:
+        return socket.AF_INET6 if self.is_ipv6 else socket.AF_INET
+
+    def to_url(self, scheme: str = "http") -> str:
+        """``http://127.0.0.1:30000`` or ``http://[::1]:30000``."""
+        return f"{scheme}://{_wrap(self.host)}:{self.port}"
+
+    def to_tcp(self) -> str:
+        """``tcp://`` endpoint for ZMQ / torch distributed."""
+        return self.to_url("tcp")
+
+    def to_host_port_str(self) -> str:
+        """``host:port`` string for gRPC listen address, session IDs, logs."""
+        return f"{_wrap(self.host)}:{self.port}"
+
+    @staticmethod
+    def resolve_host(host: str) -> str:
+        """Return *host* as-is if it's an IP, otherwise DNS-resolve to one."""
+        try:
+            ipaddress.ip_address(host)
+            return host
+        except ValueError:
+            pass
+        try:
+            return socket.getaddrinfo(
+                host, None, socket.AF_UNSPEC, 0, 0, socket.AI_ADDRCONFIG
+            )[0][4][0]
+        except socket.gaierror as e:
+            raise ValueError(f"Cannot resolve host {host!r}: {e}") from e
+
+    def resolved(self) -> NetworkAddress:
+        """DNS-resolve hostname to IP; return self if already an IP."""
+        ip = self.resolve_host(self.host)
+        return self if ip == self.host else NetworkAddress(ip, self.port)
+
+    def to_bind_tuple(self) -> Tuple[str, int]:
+        """Raw ``(host, port)`` tuple for ``socket.bind()`` / ``socket.connect()``.
+
+        Returns the *unwrapped* host — sockets need the raw address, not
+        the bracketed form.
+        """
+        return (self.host, self.port)
+
+    @staticmethod
+    def parse(addr: str) -> NetworkAddress:
+        """Parse a ``host:port`` string into a ``NetworkAddress``.
+
+        Accepted formats::
+
+            [::1]:8000          → NetworkAddress("::1", 8000)
+            127.0.0.1:8000      → NetworkAddress("127.0.0.1", 8000)
+            my-hostname:8000    → NetworkAddress("my-hostname", 8000)
+
+        IPv6 addresses **must** be bracketed.  Bare ``::1:8000`` is
+        ambiguous and will raise ``ValueError``.
+
+        Raises:
+            ValueError: If the string cannot be unambiguously parsed.
+        """
+        if not addr:
+            raise ValueError("Empty address string")
+
+        # --- Bracketed IPv6: [addr]:port ---
+        if addr.startswith("["):
+            close = addr.find("]")
+            if close == -1:
+                raise ValueError(f"Missing closing bracket in IPv6 address: {addr!r}")
+            host = addr[1:close]
+            if not _is_ipv6(host):
+                raise ValueError(f"Invalid IPv6 address inside brackets: {host!r}")
+            rest = addr[close + 1 :]
+            if not rest.startswith(":") or len(rest) < 2:
+                raise ValueError(
+                    f"Expected ':port' after closing bracket, got: {rest!r}"
+                )
+            return NetworkAddress(host, _parse_port(rest[1:]))
+
+        # --- Plain host:port (IPv4 / hostname) ---
+        if ":" not in addr:
+            raise ValueError(f"Missing port in address (expected host:port): {addr!r}")
+        host, port_str = addr.rsplit(":", 1)
+        if not host:
+            raise ValueError(f"Empty host in address: {addr!r}")
+        # Guard against bare IPv6 slipping through
+        if ":" in host and _is_ipv6(host):
+            raise ValueError(
+                f"Bare IPv6 address without brackets is ambiguous: {addr!r}. "
+                f"Use [{host}]:{port_str} instead."
+            )
+        return NetworkAddress(host, _parse_port(port_str))
+
+    def __str__(self) -> str:
+        return self.to_host_port_str()
+
+    def __repr__(self) -> str:
+        return f"NetworkAddress({self.host!r}, {self.port})"
+
+
+def resolve_base_url(base_url: str, host: str, port: int) -> str:
+    """Base URL a client sends to: ``base_url`` if set, else ``http://host:port``
+    (IPv6-correct via :class:`NetworkAddress`)."""
+    if base_url:
+        return base_url
+    return NetworkAddress(host, port).to_url()
+
+
+def resolve_host_port(base_url: str, host: str, port: int) -> str:
+    """Like :func:`resolve_base_url` but returns the scheme-less ``host:port``
+    form (for gRPC-style endpoints): ``base_url`` if set, else ``host:port``
+    (IPv6-correct via :class:`NetworkAddress`)."""
+    if base_url:
+        return base_url
+    return NetworkAddress(host, port).to_host_port_str()

python/sglang/srt/utils/numa_utils.py

@@ -0,0 +1,429 @@
+import ctypes
+import glob
+import logging
+import math
+import multiprocessing
+import os
+import random
+import shutil
+import subprocess
+import time
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Optional
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import is_cuda
+
+_is_cuda = is_cuda()
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def configure_subprocess(server_args: ServerArgs, gpu_id: int):
+    if envs.SGLANG_NUMA_BIND_V2.get():
+        numa_node = get_numa_node_if_available(server_args, gpu_id)
+        if numa_node is not None:
+            # _numactl_cpu_mem_args returns None (warn/raise) on empty CPU intersection (#26983).
+            numactl_args = _numactl_cpu_mem_args(numa_node, gpu_id)
+            if numactl_args is not None:
+                # Verify numactl can actually apply the binding before we exec it
+                # in front of the interpreter; relax the memory policy if not.
+                numactl_args, probe_err = _probe_numactl_args(numactl_args)
+                if numactl_args is None:
+                    # numactl could not apply even a CPU-only binding (e.g.
+                    # set_mempolicy(2)/sched_setaffinity(2) blocked by seccomp,
+                    # which the read-only get_mempolicy(2) probe in
+                    # _can_set_mempolicy cannot detect). Reuse #26983's failure
+                    # semantics (warn-and-continue, or raise when
+                    # SGLANG_CRASH_ON_NUMA_BIND_FAILURE) with an explicit reason
+                    # carrying the captured stderr: the CPU intersection already
+                    # succeeded here, so the default "no CPU cores allowed"
+                    # message would mislead operators toward the wrong cause.
+                    probe_suffix = f": {probe_err}" if probe_err else ""
+                    _handle_numa_bind_failure(
+                        numa_node,
+                        reason=(
+                            f"numactl could not apply NUMA binding for node "
+                            f"{numa_node} (e.g. set_mempolicy/sched_setaffinity "
+                            f"blocked by seccomp, or cpuset rejects the policy)"
+                            f"{probe_suffix}; skipping NUMA binding for GPU {gpu_id}."
+                        ),
+                    )
+                    yield
+                    return
+                executable, debug_str = _create_numactl_executable(
+                    numactl_args=numactl_args
+                )
+                debug_str += (
+                    f", logical_gpu_id={gpu_id}, "
+                    f"physical_gpu_id={_get_nvml_device_index(gpu_id)}, "
+                    f"CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')}"
+                )
+                with _mp_set_executable(executable=executable, debug_str=debug_str):
+                    yield
+                    return
+    yield
+
+
+def _create_numactl_executable(numactl_args: str):
+    old_executable = os.fsdecode(multiprocessing.spawn.get_executable())
+    script = f'''#!/bin/sh
+exec numactl {numactl_args} {old_executable} "$@"'''
+    path = Path(
+        f"/tmp/sglang_temp_file_{time.time()}_{random.randrange(0, 10000000)}.sh"
+    )
+    path.write_text(script)
+    path.chmod(0o777)
+    return str(path), f"{script=}"
+
+
+@contextmanager
+def _mp_set_executable(executable: str, debug_str: str):
+    start_method = multiprocessing.get_start_method()
+    assert start_method == "spawn", f"{start_method=}"
+
+    old_executable = os.fsdecode(multiprocessing.spawn.get_executable())
+    multiprocessing.spawn.set_executable(executable)
+    logger.debug(f"mp.set_executable {old_executable} -> {executable} ({debug_str})")
+    try:
+        yield
+    finally:
+        assert (
+            os.fsdecode(multiprocessing.spawn.get_executable()) == executable
+        ), f"{multiprocessing.spawn.get_executable()=}"
+        multiprocessing.spawn.set_executable(old_executable)
+        logger.debug(f"mp.set_executable revert to {old_executable}")
+
+
+def _get_nvml_device_index(device_id: int) -> int:
+    # _get_nvml_device_index is an internal PyTorch helper, so fall back to
+    # device_id directly if the helper is unavailable.
+    get_nvml_device_index = getattr(torch.cuda, "_get_nvml_device_index", None)
+    if get_nvml_device_index is None:
+        logger.warning(
+            "torch.cuda._get_nvml_device_index is unavailable; falling back to "
+            f"device_id={device_id} as the NVML device index. This may select "
+            "the wrong physical GPU when CUDA_VISIBLE_DEVICES reorders devices "
+            f"(CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')})."
+        )
+        return device_id
+    return get_nvml_device_index(device_id)
+
+
+def get_numa_node_if_available(server_args: ServerArgs, gpu_id: int) -> Optional[int]:
+    """
+    Returns the NUMA node for the given GPU id. If it is not set in the server_args, it will try to query the NUMA node for the GPU.
+    If the NUMA node is not available, has already been configured externally, or the user lacks permission to set NUMA affinity, it will return None.
+
+    Args:
+        server_args: The server arguments.
+        gpu_id: The GPU id.
+
+    Returns:
+        The NUMA node for the given GPU id or None if it is not available.
+    """
+    if server_args.numa_node is not None:
+        return server_args.numa_node[gpu_id]
+    if _is_numa_available():
+        queried_numa_node = _query_numa_node_for_gpu(gpu_id)
+        if len(queried_numa_node) == 0:
+            return None
+        if len(queried_numa_node) > 1:
+            # get_numa_node_for_gpu could return multiple nodes, we use the first one for now.
+            # I don't think there any hardware configs that would have more than one.
+            logger.warning(
+                f"Multiple NUMA nodes found for GPU {gpu_id}: {queried_numa_node}. Using the first one."
+            )
+        return queried_numa_node[0]
+    return None
+
+
+def get_libnuma():
+    libnuma = None
+
+    for libnuma_so in ["libnuma.so", "libnuma.so.1"]:
+        try:
+            libnuma = ctypes.CDLL(libnuma_so)
+        except OSError as e:
+            logger.debug(f"{e}")
+            libnuma = None
+        if libnuma is not None:
+            break
+    return libnuma
+
+
+def numa_bind_to_node(node: int):
+    libnuma = get_libnuma()
+
+    if libnuma is None or libnuma.numa_available() < 0:
+        logger.warning("numa not available on this system, skip bind action")
+        return
+
+    node_cpus = _node_cpus(node)
+    if node_cpus:
+        allowed_cpus = os.sched_getaffinity(0)
+        target_cpus = node_cpus & allowed_cpus
+        if not target_cpus:
+            _handle_numa_bind_failure(node, allowed_cpus)
+            return
+        os.sched_setaffinity(0, target_cpus)
+    else:
+        libnuma.numa_run_on_node(ctypes.c_int(node))
+    libnuma.numa_set_preferred(ctypes.c_int(node))
+
+
+class _Bitmask(ctypes.Structure):
+    _fields_ = [("size", ctypes.c_ulong), ("maskp", ctypes.POINTER(ctypes.c_ulong))]
+
+
+def _node_cpus(node: int) -> set:
+    libnuma = get_libnuma()
+    if libnuma is None or libnuma.numa_available() < 0:
+        return set()
+    libnuma.numa_allocate_cpumask.restype = ctypes.POINTER(_Bitmask)
+    libnuma.numa_node_to_cpus.argtypes = [ctypes.c_int, ctypes.POINTER(_Bitmask)]
+    libnuma.numa_node_to_cpus.restype = ctypes.c_int
+    libnuma.numa_bitmask_isbitset.argtypes = [ctypes.POINTER(_Bitmask), ctypes.c_uint]
+    libnuma.numa_bitmask_isbitset.restype = ctypes.c_int
+    libnuma.numa_bitmask_free.argtypes = [ctypes.POINTER(_Bitmask)]
+    mask = libnuma.numa_allocate_cpumask()
+    try:
+        if libnuma.numa_node_to_cpus(node, mask) != 0:
+            return set()
+        return {
+            i
+            for i in range(mask.contents.size)
+            if libnuma.numa_bitmask_isbitset(mask, i)
+        }
+    finally:
+        libnuma.numa_bitmask_free(mask)
+
+
+def _numactl_cpu_mem_args(node: int, gpu_id: int) -> Optional[str]:
+    node_cpus = _node_cpus(node)
+    if not node_cpus:
+        return f"--cpunodebind={node} --membind={node}"
+    allowed_cpus = os.sched_getaffinity(0)
+    target_cpus = node_cpus & allowed_cpus
+    if not target_cpus:
+        _handle_numa_bind_failure(node, allowed_cpus, gpu_id)
+        return None
+    if target_cpus == node_cpus:
+        return f"--cpunodebind={node} --membind={node}"
+    cpu_list = ",".join(str(c) for c in sorted(target_cpus))
+    return f"--physcpubind={cpu_list} --membind={node}"
+
+
+def _strip_memory_args(numactl_args: str) -> str:
+    """Return ``numactl_args`` with the ``--membind`` segment removed, keeping
+    only the CPU binding (``--cpunodebind`` / ``--physcpubind``)."""
+    return " ".join(
+        token for token in numactl_args.split() if not token.startswith("--membind")
+    )
+
+
+def _probe_numactl_args(numactl_args: str) -> tuple[Optional[str], str]:
+    """Dry-run ``numactl <args> true`` and fall back to a weaker binding when the
+    kernel rejects the strongest one.
+
+    ``configure_subprocess`` applies NUMA binding by exec-ing ``numactl`` in front
+    of the Python interpreter (see ``_create_numactl_executable``), so a binding
+    that ``numactl`` refuses kills the worker before Python starts, with no
+    traceback. ``_can_set_mempolicy`` only probes ``get_mempolicy(2)`` (read),
+    which does not catch ``set_mempolicy(2)`` being denied (e.g. by a seccomp
+    profile) or a ``--membind`` that the cpuset rejects with ``EINVAL``.
+
+    To avoid that silent crash we probe the requested args and progressively relax
+    the *memory* policy while keeping the CPU binding intact::
+
+        --membind=N  ->  --preferred=N  ->  drop the memory segment
+
+    Returns ``(args, last_stderr)``: ``args`` is the strongest binding that
+    actually runs, or ``None`` if even CPU-only fails (or ``numactl`` is missing /
+    errors out); ``last_stderr`` is the rejection reason numactl printed for the
+    strongest binding that was rejected (empty on success), so the caller can
+    surface it on the total-failure path.
+    """
+
+    def _probe(args: str):
+        """Run ``numactl <args> true``; return ``(succeeded, stderr_text)``."""
+        try:
+            proc = subprocess.run(
+                ["numactl", *args.split(), "true"],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.PIPE,
+                timeout=10,
+            )
+            stderr = proc.stderr.decode("utf-8", errors="replace").strip()
+            if proc.returncode != 0:
+                logger.debug(f"numactl probe for {args!r} rejected: {stderr!r}")
+            return proc.returncode == 0, stderr
+        except Exception as e:
+            # Missing numactl, timeout, etc. Treat as "this binding does not work".
+            logger.debug(f"numactl probe for {args!r} failed: {e}")
+            return False, str(e)
+
+    def _suffix(err: str) -> str:
+        return f": {err}" if err else ""
+
+    # 1. Strongest binding: exactly what was requested.
+    ok, last_err = _probe(numactl_args)
+    if ok:
+        return numactl_args, ""
+
+    # 2. Relax a hard --membind=N to a soft --preferred=N. The memory segment here
+    #    is always a single node, which maps cleanly onto --preferred (single-node
+    #    only). MPOL_PREFERRED is a hint and can succeed where MPOL_BIND is denied.
+    if "--membind=" in numactl_args:
+        preferred_args = numactl_args.replace("--membind=", "--preferred=")
+        ok, _ = _probe(preferred_args)
+        if ok:
+            logger.warning(
+                f"numactl rejected hard memory binding ({numactl_args!r})"
+                f"{_suffix(last_err)}; falling back to soft preferred policy "
+                f"({preferred_args!r})."
+            )
+            return preferred_args, ""
+
+    # 3. Drop the memory segment entirely, keep only the CPU binding.
+    cpu_only_args = _strip_memory_args(numactl_args)
+    if cpu_only_args and cpu_only_args != numactl_args:
+        ok, cpu_err = _probe(cpu_only_args)
+        if ok:
+            logger.warning(
+                f"numactl rejected memory binding ({numactl_args!r})"
+                f"{_suffix(last_err)}; falling back to CPU-only binding "
+                f"({cpu_only_args!r})."
+            )
+            return cpu_only_args, ""
+        last_err = cpu_err
+
+    # 4. Nothing worked.
+    return None, last_err
+
+
+def _handle_numa_bind_failure(
+    node: int,
+    allowed_cpus=None,
+    gpu_id: Optional[int] = None,
+    *,
+    reason: Optional[str] = None,
+) -> None:
+    """Emit the NUMA-bind failure warning, or raise it when
+    ``SGLANG_CRASH_ON_NUMA_BIND_FAILURE`` is set.
+
+    Two call modes:
+      * ``reason is None`` (default): the failure is an empty CPU intersection,
+        so the message reports ``allowed_cpus`` (which must be provided).
+      * ``reason`` provided: the failure is something else (e.g. numactl rejected
+        the binding at runtime); the caller supplies the exact message and
+        ``allowed_cpus`` / ``gpu_id`` are not needed.
+    """
+    if reason is None:
+        gpu_str = f" for GPU {gpu_id}" if gpu_id is not None else ""
+        reason = (
+            f"NUMA node {node} has no CPU cores allowed by the current affinity "
+            f"{sorted(allowed_cpus)}, skipping NUMA binding{gpu_str}."
+        )
+    logger.warning(reason)
+    if envs.SGLANG_CRASH_ON_NUMA_BIND_FAILURE.get():
+        raise RuntimeError(reason)
+
+
+def _can_set_mempolicy() -> bool:
+    """Check if the process has permission to use NUMA memory policy syscalls."""
+    try:
+        libnuma = get_libnuma()
+        if libnuma is None or libnuma.numa_available() < 0:
+            return False
+        mode = ctypes.c_int()
+        ret = libnuma.get_mempolicy(
+            ctypes.byref(mode), None, ctypes.c_ulong(0), None, ctypes.c_ulong(0)
+        )
+        return ret == 0
+    except Exception:
+        return False
+
+
+def _is_numa_available() -> bool:
+    """
+    Check if NUMA is available and not already configured externally.
+    """
+    if not _is_cuda:
+        return False
+
+    # Check if this is a numa system.
+    if not os.path.isdir("/sys/devices/system/node/node1"):
+        return False
+
+    if not shutil.which("numactl") and envs.SGLANG_NUMA_BIND_V2.get():
+        logger.debug(
+            "numactl command not found, skipping NUMA node configuration for GPU. Install numactl (e.g., apt-get install numactl) to enable automatic NUMA binding."
+        )
+        return False
+
+    if not _can_set_mempolicy():
+        logger.warning(
+            "User lacks permission to set NUMA affinity, skipping NUMA node configuration for GPU. If using docker, try adding --cap-add SYS_NICE to your docker run command."
+        )
+        return False
+
+    return True
+
+
+def _query_numa_node_for_gpu(device_id: int):
+    """
+    Get the NUMA node affinity list for a GPU device.
+
+    Args:
+        device_id: CUDA logical device index (post-CUDA_VISIBLE_DEVICES).
+    Returns:
+        List of NUMA node IDs that have affinity with the device.
+    """
+    try:
+        import pynvml
+    except ModuleNotFoundError:
+        logger.warning("pynvml not installed, skipping NUMA node configuration for GPU")
+        return []
+
+    try:
+        pynvml.nvmlInit()
+
+        # device_id is a CUDA logical index. Convert it to the corresponding
+        # NVML index so reordered CUDA_VISIBLE_DEVICES maps to the right GPU.
+        # _get_nvml_device_index takes CUDA_VISIBLE_DEVICES into account.
+        nvml_device_id = _get_nvml_device_index(device_id)
+        handle = pynvml.nvmlDeviceGetHandleByIndex(nvml_device_id)
+        numa_node_count = len(glob.glob("/sys/devices/system/node/node[0-9]*"))
+
+        c_ulong_bits = ctypes.sizeof(ctypes.c_ulong) * 8
+        node_set_size = max(1, math.ceil(numa_node_count / c_ulong_bits))
+        node_set = pynvml.nvmlDeviceGetMemoryAffinity(
+            handle,
+            node_set_size,
+            pynvml.NVML_AFFINITY_SCOPE_NODE,
+        )
+
+        # Decode the bitmask into a list of NUMA node IDs
+        numa_nodes = []
+        for node_id in range(numa_node_count):
+            mask_array_index = node_id // c_ulong_bits
+            mask_bit_index = node_id % c_ulong_bits
+            if node_set[mask_array_index] & (1 << mask_bit_index):
+                numa_nodes.append(node_id)
+        return numa_nodes
+    except pynvml.NVMLError as e:
+        logger.warning(
+            f"NVML error querying memory affinity for GPU {device_id}: {e}, skipping NUMA node configuration for GPU"
+        )
+        return []
+    finally:
+        try:
+            pynvml.nvmlShutdown()
+        except Exception:
+            pass  # Ignore shutdown errors

python/sglang/srt/utils/nvtx_pytorch_hooks.py

@@ -0,0 +1,292 @@
+# 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.
+# ==============================================================================
+"""PyTorch hooks for layerwise NVTX profiling."""
+
+import torch
+import torch.cuda.nvtx as nvtx
+
+
+class PytHooks(object):
+    """This module contains all the code needed to enable forward hooks in a pytorch network.
+
+    To register the hooks for a given network, the user needs to instantiate a PytHook object.
+    Then call the register_hooks method.
+
+    Example:
+
+        my_hook = PytHook()
+        my_hook.register_hooks(my_network_model)
+    """
+
+    def __init__(self):
+        """Initialize module variables
+
+        Returns:
+            None:
+
+        Raises:
+            None:
+        """
+        super().__init__()
+        self.module_to_name_map = {}
+
+    @staticmethod
+    def print_tensor(tensor_obj, prefix, tensor_list=None):
+        """Descends iterators that contains Tensors and prints the Tensor
+
+        Recursive function that descends iterator type arguments until
+        it finds a Tensor object.
+
+        Args:
+            tensor_obj: Could be a Tensor or an iterator type that contains Tensors
+            prefix: String name to assign to the Tensor
+            tensor_list: List to accumulate tensor dimensions
+
+        Returns:
+            List of tensor dimensions
+
+        Raises:
+            None:
+        """
+        if tensor_list is None:
+            tensor_list = []
+
+        if isinstance(tensor_obj, list) or isinstance(tensor_obj, tuple):
+            for ten in tensor_obj:
+                tensor_list = PytHooks.print_tensor(ten, prefix, tensor_list)
+        elif isinstance(tensor_obj, torch.Tensor):
+            tensor_dims = list(tensor_obj.size())
+            tensor_list.append(tensor_dims)
+        return tensor_list
+
+    def process_layer_params(self, module_obj):
+        """Extract the static parameters from LLM and VLM relevant layer types
+
+        Args:
+            module_obj(class): Module state data structure.
+
+        Returns:
+            param_info(dict): Parameter meta_data for the given op.
+
+        Raises:
+            None
+
+        """
+        param_info = {}
+        # Extract parameters for layers commonly used in LLMs and VLMs
+        if (
+            isinstance(module_obj, torch.nn.Conv1d)
+            or isinstance(module_obj, torch.nn.Conv2d)
+            or isinstance(module_obj, torch.nn.Conv3d)
+        ):
+            conv_params = {}
+            conv_params["in_chan"] = module_obj.in_channels
+            conv_params["out_chan"] = module_obj.out_channels
+            conv_params["filter_dim"] = module_obj.kernel_size
+            conv_params["stride"] = module_obj.stride
+            conv_params["padding"] = module_obj.padding
+            conv_params["dilation"] = module_obj.dilation
+            conv_params["transposed"] = module_obj.transposed
+            conv_params["output_padding"] = module_obj.output_padding
+            conv_params["groups"] = module_obj.groups
+            conv_params["padding_mode"] = module_obj.padding_mode
+            param_info = conv_params
+        elif (
+            isinstance(module_obj, torch.nn.ConvTranspose1d)
+            or isinstance(module_obj, torch.nn.ConvTranspose2d)
+            or isinstance(module_obj, torch.nn.ConvTranspose3d)
+        ):
+            convtranspose_params = {}
+            convtranspose_params["in_chan"] = module_obj.in_channels
+            convtranspose_params["out_chan"] = module_obj.out_channels
+            convtranspose_params["filter_dim"] = module_obj.kernel_size
+            convtranspose_params["stride"] = module_obj.stride
+            convtranspose_params["padding"] = module_obj.padding
+            convtranspose_params["dilation"] = module_obj.dilation
+            convtranspose_params["transposed"] = module_obj.transposed
+            convtranspose_params["output_padding"] = module_obj.output_padding
+            convtranspose_params["groups"] = module_obj.groups
+            convtranspose_params["padding_mode"] = module_obj.padding_mode
+            param_info = convtranspose_params
+        elif (
+            isinstance(module_obj, torch.nn.MaxPool1d)
+            or isinstance(module_obj, torch.nn.MaxPool2d)
+            or isinstance(module_obj, torch.nn.MaxPool3d)
+        ):
+
+            def _handle_int_or_tuple(parameter):
+                if isinstance(parameter, tuple):
+                    return list(parameter)
+                elif isinstance(parameter, int):
+                    return [parameter, parameter]
+
+            pooling_params = {}
+            pooling_params["filter_dim"] = _handle_int_or_tuple(module_obj.kernel_size)
+            pooling_params["stride"] = _handle_int_or_tuple(module_obj.stride)
+            pooling_params["padding"] = _handle_int_or_tuple(module_obj.padding)
+            pooling_params["dilation"] = _handle_int_or_tuple(module_obj.dilation)
+            param_info = pooling_params
+        elif (
+            isinstance(module_obj, torch.nn.AvgPool1d)
+            or isinstance(module_obj, torch.nn.AvgPool2d)
+            or isinstance(module_obj, torch.nn.AvgPool3d)
+        ):
+            pooling_params = {}
+            pooling_params["filter_dim"] = [
+                module_obj.kernel_size,
+                module_obj.kernel_size,
+            ]
+            pooling_params["stride"] = [module_obj.stride, module_obj.stride]
+            pooling_params["padding"] = [module_obj.padding, module_obj.padding]
+            pooling_params["ceil_mode"] = module_obj.ceil_mode
+            pooling_params["count_include_pad"] = module_obj.count_include_pad
+            param_info = pooling_params
+        elif (
+            isinstance(module_obj, torch.nn.AdaptiveAvgPool1d)
+            or isinstance(module_obj, torch.nn.AdaptiveAvgPool2d)
+            or isinstance(module_obj, torch.nn.AdaptiveAvgPool3d)
+        ):
+            pooling_params = {}
+            pooling_params["output_size"] = [
+                module_obj.output_size,
+                module_obj.output_size,
+            ]
+            param_info = pooling_params
+        elif isinstance(module_obj, torch.nn.Linear):
+            param_info["in_features"] = module_obj.in_features
+            param_info["out_features"] = module_obj.out_features
+        elif (
+            isinstance(module_obj, torch.nn.BatchNorm1d)
+            or isinstance(module_obj, torch.nn.BatchNorm2d)
+            or isinstance(module_obj, torch.nn.BatchNorm3d)
+        ):
+            param_info["num_features"] = module_obj.num_features
+            param_info["epsilon"] = module_obj.eps
+            param_info["momentum"] = module_obj.momentum
+        elif isinstance(module_obj, torch.nn.ReLU):
+            param_info["in_place"] = module_obj.inplace
+        elif isinstance(module_obj, torch.nn.Dropout):
+            param_info["p"] = module_obj.p
+            param_info["in_place"] = module_obj.inplace
+        elif isinstance(module_obj, torch.nn.Embedding):
+            param_info["num_embeddings"] = module_obj.num_embeddings
+            param_info["embedding_dim"] = module_obj.embedding_dim
+        elif isinstance(
+            module_obj,
+            (
+                torch.nn.Upsample,
+                torch.nn.UpsamplingNearest2d,
+                torch.nn.UpsamplingBilinear2d,
+            ),
+        ):
+            param_info["scale_factor"] = module_obj.scale_factor
+
+        return param_info
+
+    def module_fwd_hook(self, module_obj, in_tensor, out_tensor):
+        """Callback function that ends the NVTX marker
+
+        Records the module name and tensor information
+        Called after the module executes the forward method.
+
+        Args:
+            module_obj: Pointer to the module object
+            in_tensor: Input tensor or list of tensors
+            out_tensor: Output tensor of the resulting forward operator
+
+        Returns:
+            None:
+
+        Raises:
+            None:
+        """
+        nvtx.range_pop()
+        return
+
+    def module_fwd_pre_hook(self, module_obj, in_tensor):
+        """Creates an NVTX marker with the module name in it.
+
+        This function is called before the module executes
+
+        Args:
+            module_obj: Module object data structure - used to get unique module name
+            in_tensor: Input tensor data structure
+
+        Returns:
+            None
+
+        Raises:
+            None
+        """
+        marker_dict = {}
+        module_name = self.module_to_name_map.get(module_obj, "unknown")
+        marker_dict["Module"] = module_name
+
+        ## Get trainable parameters like weights and bias
+        module_params = module_obj.named_parameters(recurse=False)
+        for idx, (param_name, param_obj) in enumerate(module_params):
+            if idx == 0:
+                marker_dict["TrainableParams"] = {}
+            marker_dict["TrainableParams"][param_name] = list(param_obj.size())
+
+        in_tensor_list = PytHooks.print_tensor(in_tensor, "Input")
+        if in_tensor_list:
+            marker_dict["Inputs"] = in_tensor_list
+
+        param_info = self.process_layer_params(module_obj)
+        if param_info:
+            marker_dict["StaticParams"] = param_info
+
+        nvtx.range_push("{}".format(marker_dict))
+
+        return
+
+    def register_hooks(self, network_model, module_prefix="top"):
+        """User level function that activates all the hooks
+
+        The user needs to call this method from the network source code
+        The code descends all the modules in the network and registers their
+        respective hooks.
+
+        Args:
+            network_model: Model object for the network
+            module_prefix: (default: top)
+
+        Returns:
+            None
+
+        Raises:
+            Exception if a module instance is reused
+        """
+        # Module types to skip (simple operations that don't need detailed profiling)
+        skip_types = (
+            torch.nn.Identity,
+            torch.nn.Dropout,
+            torch.nn.Dropout1d,
+            torch.nn.Dropout2d,
+            torch.nn.Dropout3d,
+        )
+
+        for name, module in network_model.named_modules(prefix=module_prefix):
+            # Skip certain module types to reduce profiling overhead
+            if isinstance(module, skip_types):
+                continue
+
+            module.register_forward_pre_hook(self.module_fwd_pre_hook)
+            module.register_forward_hook(self.module_fwd_hook)
+            if module not in self.module_to_name_map:
+                self.module_to_name_map[module] = name
+            else:
+                raise ValueError("Module instance {} is not unique ".format(module))
+        return

python/sglang/srt/utils/nvtx_utils.py

@@ -0,0 +1,119 @@
+# 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.
+# ==============================================================================
+"""Profiler span helpers for hot SGLang code paths.
+
+A span has two independent emitters:
+
+* ``record_function`` -- emitted whenever a torch profiler is active, so spans
+  show up in torch/Perfetto traces for free (no env, no extra package).
+* ``nvtx`` range -- emitted only when the caller opts in via ``nvtx_enabled``
+  (wired to a per-subsystem ``SGLANG_ENABLE_NVTX_*`` gate) and the ``nvtx``
+  package is importable, for Nsight Systems timelines.
+
+Decoupling the two lets every annotation site -- scheduler stages, batch-overlap
+ops, and the speculative-decoding / forward spans -- share one primitive.
+"""
+
+import logging
+from contextlib import ExitStack, contextmanager, nullcontext
+from functools import partial, wraps
+from typing import Optional
+
+import torch
+
+from sglang.srt.environ import envs
+
+logger = logging.getLogger(__name__)
+
+_SCHEDULER_NVTX = envs.SGLANG_ENABLE_NVTX_SCHEDULER.get()
+_OPERATIONS_NVTX = envs.SGLANG_ENABLE_NVTX_OPERATIONS.get()
+
+_nvtx_module = None
+if _SCHEDULER_NVTX or _OPERATIONS_NVTX:
+    try:
+        import nvtx as _nvtx_module  # type: ignore
+    except ImportError:
+        logger.warning(
+            "An SGLANG_ENABLE_NVTX_* flag is set, but the `nvtx` package is "
+            "missing. NVTX markers are disabled; torch profiler spans still emit."
+        )
+
+NVTX_AVAILABLE = _nvtx_module is not None
+# Per-subsystem nvtx gates: emit nvtx ranges only when the flag is set AND the
+# package is importable. The record_function path is independent of both.
+NVTX_SCHEDULER_ENABLED = _SCHEDULER_NVTX and NVTX_AVAILABLE
+NVTX_OPERATIONS_ENABLED = _OPERATIONS_NVTX and NVTX_AVAILABLE
+
+# Default nvtx colors for statically-named spans (only used on the nvtx path).
+_NVTX_COLOR_MAP = {
+    "scheduler.recv_requests": "blue",
+    "scheduler.process_input_requests": "purple",
+    "scheduler.get_next_batch_to_run": "green",
+    "scheduler.run_batch": "red",
+    "scheduler.process_batch_result": "cyan",
+}
+
+_NULL_CONTEXT = nullcontext()
+
+
+@contextmanager
+def _profile_range_impl(
+    debug_name: str, color: Optional[str], record: bool, nvtx_enabled: bool
+):
+    with ExitStack() as stack:
+        if record:
+            stack.enter_context(torch.profiler.record_function(debug_name))
+        if nvtx_enabled:
+            if color is None:
+                color = _NVTX_COLOR_MAP.get(debug_name)
+            stack.enter_context(_nvtx_module.annotate(debug_name, color=color))
+        yield
+
+
+def profile_range(
+    debug_name: str, *, color: Optional[str] = None, nvtx_enabled: bool = False
+):
+    """Context manager emitting a profiler span for ``debug_name``.
+
+    A torch ``record_function`` is emitted whenever a torch profiler is active;
+    an nvtx range is emitted additionally when ``nvtx_enabled`` is true. Returns a
+    shared no-op when neither applies, so off-profile hot paths pay only one
+    ``_profiler_enabled()`` check.
+    """
+    record = torch.autograd._profiler_enabled()
+    if not record and not nvtx_enabled:
+        return _NULL_CONTEXT
+    return _profile_range_impl(debug_name, color, record, nvtx_enabled)
+
+
+def profile_method(
+    debug_name: str, *, color: Optional[str] = None, nvtx_enabled: bool = False
+):
+    """Decorator form of ``profile_range``."""
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with profile_range(debug_name, color=color, nvtx_enabled=nvtx_enabled):
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+# Pre-bound per-subsystem helpers: torch spans always (under a profiler), nvtx
+# ranges only when that subsystem's gate is on.
+scheduler_nvtx_method = partial(profile_method, nvtx_enabled=NVTX_SCHEDULER_ENABLED)
+operations_nvtx_range = partial(profile_range, nvtx_enabled=NVTX_OPERATIONS_ENABLED)

python/sglang/srt/utils/offloader.py

@@ -0,0 +1,583 @@
+import logging
+import os
+from abc import ABC
+from typing import Callable, Generator, List, Optional
+
+import torch
+from torch.func import functional_call
+
+from sglang.srt.distributed.naive_distributed import (
+    NaiveDistributed,
+    get_naive_distributed,
+    set_naive_distributed,
+)
+from sglang.srt.layers.parameter import ModelWeightParameter
+from sglang.srt.runtime_context import get_parallel, get_stream
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import MultiprocessingSerializer, is_pin_memory_available
+from sglang.srt.utils.host_shared_memory import (
+    HostSharedMemoryManager,
+    get_host_shared_memory_manager,
+    set_host_shared_memory_manager,
+)
+
+logger = logging.getLogger(__name__)
+
+_SubmoduleAccessor = Callable[[torch.nn.Module], torch.nn.Module]
+_WhitelistParamNamesCreator = Callable[[torch.nn.Module], List[str]]
+
+
+class BaseOffloader(ABC):
+    def wrap_modules(
+        self,
+        all_modules_generator: Generator[torch.nn.Module, None, None],
+        submodule_accessor: Optional[_SubmoduleAccessor] = None,
+        whitelist_param_names_creator: Optional[_WhitelistParamNamesCreator] = None,
+    ):
+        return list(all_modules_generator)
+
+    def post_init(self):
+        pass
+
+    @property
+    def forbid_copy_engine_usage(self):
+        return False
+
+
+class NoopOffloader(BaseOffloader):
+    pass
+
+
+# For simplicity use singleton, but can surely support multi instance
+_instance: Optional[BaseOffloader] = NoopOffloader()
+
+
+def get_offloader():
+    assert _instance is not None
+    return _instance
+
+
+def set_offloader(instance: BaseOffloader):
+    global _instance
+    _instance = instance
+
+
+def create_offloader_from_server_args(server_args: ServerArgs, dp_rank: int):
+    if server_args.cpu_offload_gb > 0:
+        return OffloaderV1(
+            cpu_offload_max_bytes=int(server_args.cpu_offload_gb * 1024**3)
+        )
+    if server_args.offload_group_size > 0:
+        assert (
+            server_args.cpu_offload_gb == 0
+        ), "V2 offload does not support cpu_offload_gb yet"
+        return OffloaderV2(
+            group_size=server_args.offload_group_size,
+            num_in_group=server_args.offload_num_in_group,
+            prefetch_step=server_args.offload_prefetch_step,
+            mode=server_args.offload_mode,
+            dp_rank=dp_rank,
+            dp_size=server_args.dp_size,
+        )
+    return NoopOffloader()
+
+
+class OffloaderV1(BaseOffloader):
+    def __init__(self, cpu_offload_max_bytes: int):
+        self._cpu_offload_bytes = 0
+        self._cpu_offload_max_bytes = cpu_offload_max_bytes
+
+    def wrap_modules(
+        self,
+        all_modules_generator: Generator[torch.nn.Module, None, None],
+        submodule_accessor: Optional[_SubmoduleAccessor] = None,
+        whitelist_param_names_creator: Optional[_WhitelistParamNamesCreator] = None,
+    ):
+        return [self.maybe_offload_to_cpu(module) for module in all_modules_generator]
+
+    def maybe_offload_to_cpu(self, module: torch.nn.Module) -> torch.nn.Module:
+        if (params := next(module.parameters(), None)) is None:
+            return module
+
+        device = params.device
+
+        if device == torch.device("cpu"):
+            return module
+
+        if self._cpu_offload_bytes >= self._cpu_offload_max_bytes:
+            return module
+
+        pin_memory = is_pin_memory_available()
+        # offload parameters to CPU
+        # use pin_memory if possible, which helps cudagraph capture speed
+        offloaded_parameters = False
+        for p in module.parameters():
+            if self._cpu_offload_bytes >= self._cpu_offload_max_bytes:
+                # we use per-parameter offloading
+                # one module might have some parameters offloaded and some not
+                break
+
+            # `torch.empty_like` does not support `pin_memory` argument
+            cpu_data = torch.empty_strided(
+                size=p.data.size(),
+                stride=p.data.stride(),
+                dtype=p.data.dtype,
+                layout=p.data.layout,
+                device="cpu",
+                pin_memory=pin_memory,
+            )
+            cpu_data.copy_(p.data)
+            p.data = cpu_data
+            self._cpu_offload_bytes += p.data.numel() * p.data.element_size()
+            offloaded_parameters = True
+
+        if offloaded_parameters:
+            original_forward = module.forward
+
+            def forward(*args, **kwargs):
+                module.forward = original_forward
+                device_state = {
+                    # here we blindly call `to(device)`
+                    # if the parameter is already on the device, it will be a no-op
+                    k: v.to(device, non_blocking=True)
+                    for k, v in module.state_dict().items()
+                }
+                output = functional_call(module, device_state, args=args, kwargs=kwargs)
+                module.forward = forward
+                return output
+
+            module.forward = forward
+
+        return module
+
+
+class OffloaderV2(BaseOffloader):
+    def __init__(
+        self,
+        group_size: int,
+        num_in_group: int,
+        prefetch_step: int,
+        mode: str,
+        dp_rank: int,
+        dp_size: int,
+    ):
+        self.group_size = group_size
+        self.num_in_group = num_in_group
+        self.prefetch_step = prefetch_step
+        self.mode = mode
+
+        run_id = os.environ["SGLANG_RUN_ID"]
+
+        # Temporarily init inside Offloader, can move if other modules also need this
+        if self.mode in {"sharded_gpu", "shm_cpu"}:
+
+            assert get_parallel().tp_size == 1, "not yet support tp_size!=1"
+            set_naive_distributed(
+                NaiveDistributed(
+                    rank=dp_rank,
+                    world_size=dp_size,
+                    rendezvous=f"/tmp/{run_id}",
+                )
+            )
+        if self.mode in {"shm_cpu"}:
+            set_host_shared_memory_manager(
+                HostSharedMemoryManager(
+                    base_name=run_id,
+                )
+            )
+
+        self.offloaders = []
+
+    def wrap_modules(
+        self,
+        all_modules_generator: Generator[torch.nn.Module, None, None],
+        submodule_accessor: Optional[_SubmoduleAccessor] = None,
+        whitelist_param_names_creator: Optional[_WhitelistParamNamesCreator] = None,
+    ):
+        assert len(self.offloaders) == 0, "should only call wrap_modules once"
+
+        # The offloader's async prefetch/offload copies run on their own
+        # stream — sharing the models' "alt" overlap stream would serialize
+        # unrelated copy and compute work.
+        alt_stream = get_stream("offload")
+
+        all_modules = []
+        offload_submodules = []
+        for module_index, module in enumerate(all_modules_generator):
+            all_modules.append(module)
+            if module_index % self.group_size >= self.group_size - self.num_in_group:
+                submodule = submodule_accessor(module)
+                whitelist_param_names = whitelist_param_names_creator(submodule)
+                logger.info(
+                    f"[offloader] offload {module_index=} submodule={type(submodule)} params={whitelist_param_names} memory_allocated={torch.cuda.memory_allocated()}"
+                )
+                offload_submodules.append(submodule)
+                self.offloaders.append(
+                    _ModuleOffloader(
+                        mode=self.mode,
+                        module=submodule,
+                        alt_stream=alt_stream,
+                        whitelist_param_names=whitelist_param_names,
+                    )
+                )
+
+        for index, module in enumerate(offload_submodules):
+            _hook_module_forward_for_offloader(
+                index=index,
+                module=module,
+                offloaders=self.offloaders,
+                prefetch_step=self.prefetch_step,
+            )
+
+        return all_modules
+
+    def post_init(self):
+        for offloader in self.offloaders:
+            offloader.post_init()
+
+        for i in range(self.prefetch_step):
+            self.offloaders[i].start_onload()
+
+    @property
+    def forbid_copy_engine_usage(self):
+        return self.mode == "cpu"
+
+
+def _hook_module_forward_for_offloader(index, module, offloaders, prefetch_step):
+    def _on_forward_end():
+        offloaders[(index + prefetch_step) % len(offloaders)].start_onload()
+        offloaders[index].offload()
+
+    _hook_module_forward_raw(
+        module,
+        on_forward_end=_on_forward_end,
+        get_parameter_and_buffer_dicts=lambda: offloaders[
+            index
+        ].wait_and_get_device_tensors(),
+    )
+
+
+def _hook_module_forward_raw(module, on_forward_end, get_parameter_and_buffer_dicts):
+    original_forward = module.forward
+
+    def forward(*args, **kwargs):
+        module.forward = original_forward
+        output = functional_call(
+            module, get_parameter_and_buffer_dicts(), args=args, kwargs=kwargs
+        )
+        on_forward_end()
+        module.forward = forward
+        return output
+
+    module.forward = forward
+
+
+class _ModuleOffloader(ABC):
+    def __init__(
+        self,
+        mode: str,
+        module: torch.nn.Module,
+        alt_stream: torch.cuda.Stream,
+        whitelist_param_names: List[str],
+    ):
+        self.mode = mode
+        self.module = module
+        self.device = next(module.parameters()).device
+        self.alt_stream = alt_stream
+
+        assert self.device != torch.device(
+            "cpu"
+        ), "not handled device=cpu case yet (should skip this tensor)"
+
+        self._device_tensors = None
+        self._load_event = None
+
+        param_dict = dict(self.module.named_parameters())
+        assert all(
+            name in param_dict for name in whitelist_param_names
+        ), f"{whitelist_param_names=} {list(param_dict.keys())=}"
+
+        self._param_offloaders = {
+            name: _BaseParamOffloader.create(mode, module=module, param_name=name)
+            for name in whitelist_param_names
+        }
+
+    def post_init(self):
+        for name, param_offloader in self._param_offloaders.items():
+            param_offloader.post_init()
+
+    def start_onload(self):
+        if torch.cuda.is_current_stream_capturing():
+            self._device_tensors = self._create_device_tensors()
+            self._load_event = None
+            return
+        self.alt_stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(self.alt_stream):
+            self._device_tensors = self._create_device_tensors()
+            self._load_event = torch.cuda.Event()
+            self._load_event.record()
+
+    def offload(self):
+        self._device_tensors = None
+        self._load_event = None
+
+    def wait_and_get_device_tensors(self):
+        assert self._device_tensors is not None
+        if torch.cuda.is_current_stream_capturing():
+            if self._load_event is not None:
+                self._device_tensors = self._create_device_tensors()
+                self._load_event = None
+            return self._device_tensors
+        if self._load_event is not None:
+            self._load_event.wait()
+        return self._device_tensors
+
+    def _create_device_tensors(self):
+        return {k: v.create_device_tensor() for k, v in self._param_offloaders.items()}
+
+
+class _BaseParamOffloader(ABC):
+    @staticmethod
+    def create(mode: str, **kwargs) -> "_BaseParamOffloader":
+        return {
+            "meta": _MetaParamOffloader,
+            "cpu": _CpuParamOffloader,
+            "shm_cpu": _ShmCpuParamOffloader,
+            "sharded_gpu": _ShardedGpuParamOffloader,
+        }[mode](**kwargs)
+
+    def __init__(self, module, param_name):
+        self._module = module
+        self._param_name = param_name
+
+    @property
+    def _param(self):
+        return getattr(self._module, self._param_name)
+
+    def post_init(self):
+        pass
+
+    def create_device_tensor(self):
+        raise NotImplementedError
+
+
+class _MetaParamOffloader(_BaseParamOffloader):
+    """Usually used for debugging."""
+
+    def __init__(self, module, param_name):
+        super().__init__(module, param_name)
+        _move_param_to_meta(module, param_name)
+
+    def create_device_tensor(self):
+        return torch.empty_like(self._param.data, device="cuda")
+
+
+class _CpuParamOffloader(_BaseParamOffloader):
+    def __init__(self, module, param_name):
+        super().__init__(module, param_name)
+        _move_param_to_cpu(self._param, pin_memory=True)
+
+    def create_device_tensor(self):
+        return self._param.to("cuda", non_blocking=True)
+
+
+class _ShmCpuParamOffloader(_BaseParamOffloader):
+    def __init__(self, module, param_name):
+        super().__init__(module, param_name)
+        self._rank = get_naive_distributed().get_rank()
+        self._world_size = get_naive_distributed().get_world_size()
+
+        assert get_parallel().tp_size == 1, "not yet support tp_size!=1"
+        assert (
+            self._param.data.is_contiguous()
+        ), f"not yet support non-contiguous tensor {self._param.shape=} {self._param.stride()=}"
+
+        self.shm_cpu_data = get_host_shared_memory_manager().malloc(
+            shape=self._param.shape, dtype=self._param.dtype
+        )
+
+        if self._rank == 0:
+            self.shm_cpu_data.copy_(self._param.data.to("cpu"))
+            self._param.data = self.shm_cpu_data
+        else:
+            _move_param_to_meta(self._module, self._param_name)
+        get_naive_distributed().barrier()
+
+    def post_init(self):
+        if self._rank == 0:
+            assert (
+                self.shm_cpu_data.data_ptr() == self._param.data.data_ptr()
+            ), f"{self.shm_cpu_data.data_ptr()=} {self._param.data.data_ptr()=} {self.shm_cpu_data=} {self._param.data=}"
+
+        _move_param_to_meta(self._module, self._param_name)
+
+    def create_device_tensor(self):
+        return self.shm_cpu_data.to("cuda", non_blocking=True)
+
+
+def update_param(param, new_tensor):
+    """Update parameter while keeping properties needed by Offloader (e.g. pinned host memory)."""
+
+    if param.device == new_tensor.device:
+        param.data = new_tensor
+    else:
+        assert param.device == torch.device(
+            "cpu"
+        ), f"{param.device=} {new_tensor.device=}"
+        param.data = _create_cpu_data(new_tensor, pin_memory=True)
+
+
+def _move_param_to_cpu(param, pin_memory: bool):
+    param.data = _create_cpu_data(param.data, pin_memory=pin_memory)
+
+
+def _create_cpu_data(data, pin_memory: bool):
+    cpu_data = _empty_strided_like(
+        data,
+        device="cpu",
+        pin_memory=pin_memory,
+    )
+    cpu_data.copy_(data)
+    return cpu_data
+
+
+def _move_param_to_meta(module, param_name):
+    old_param = getattr(module, param_name)
+    old_param_type = type(old_param)
+
+    new_data = old_param.data.to("meta")
+
+    if old_param_type == ModelWeightParameter:
+        # manually checked how `w13_weight` and `w2_weight` are constructed
+        new_param = ModelWeightParameter(
+            data=new_data,
+            **{
+                k: getattr(old_param, k)
+                for k in ["input_dim", "output_dim", "weight_loader"]
+            },
+        )
+    elif old_param_type == torch.nn.Parameter:
+        new_param = torch.nn.Parameter(
+            data=new_data,
+            requires_grad=False,
+        )
+        if hasattr(old_param, "weight_loader"):
+            new_param.weight_loader = old_param.weight_loader
+        else:
+            new_param.weight_loader = lambda *args, **kwargs: None
+    else:
+        raise ValueError(f"Unknown {old_param_type=} {old_param=}")
+
+    setattr(module, param_name, new_param)
+
+
+def _empty_strided_like(x: torch.Tensor, device, pin_memory=False):
+    return torch.empty_strided(
+        size=x.size(),
+        stride=x.stride(),
+        dtype=x.dtype,
+        layout=x.layout,
+        device=device,
+        pin_memory=pin_memory,
+    )
+
+
+# ----------------------------------------- ShardedGpu ------------------------------------------------------
+
+
+# TODO unify with ShmCpu mode
+class _ShardedGpuParamOffloader(_BaseParamOffloader):
+    def __init__(self, module, param_name):
+        super().__init__(module, param_name)
+        self._rank = get_naive_distributed().get_rank()
+        self._world_size = get_naive_distributed().get_world_size()
+
+        assert get_parallel().tp_size == 1, "not yet support tp_size!=1"
+        assert (
+            self._param.data.is_contiguous()
+        ), f"not yet support non-contiguous tensor {self._param.shape=} {self._param.stride()=}"
+
+        if self._rank == 0:
+            _move_param_to_cpu(self._param, pin_memory=True)
+        else:
+            _move_param_to_meta(self._module, self._param_name)
+
+        self.sharded_param_handles = None
+
+    def post_init(self):
+        # check again since it may be changed
+        assert (
+            self._param.data.is_contiguous()
+        ), f"not yet support non-contiguous tensor {self._param.shape=} {self._param.stride()=}"
+
+        scatter_src = self._param.data
+
+        logger.info(
+            f"[offloader] post_init {scatter_src.nbytes=} {scatter_src.dtype=} {scatter_src.shape=} {torch.cuda.memory_allocated()=}"
+        )
+
+        if self._rank == 0:
+            scatter_src = scatter_src.to("cuda")
+        scatter_list = _even_chunk(scatter_src, self._world_size)
+
+        sharded_param = torch.empty(
+            scatter_list[0].shape, dtype=scatter_list[0].dtype, device="cuda"
+        )
+        self.sharded_param_handles = _create_shared_buffer_tensors(
+            local_tensor=sharded_param
+        )
+
+        get_naive_distributed().scatter(
+            sharded_param, scatter_list if self._rank == 0 else None
+        )
+
+        _move_param_to_meta(self._module, self._param_name)
+
+    def create_device_tensor(self):
+        output = _empty_strided_like(self._param, device="cuda")
+        output_chunks = output.chunk(self._world_size)
+
+        for index in range(self._world_size):
+            src_rank = (self._rank + index) % self._world_size
+            src_buf = self.sharded_param_handles[src_rank]
+            output_chunks[src_rank].copy_(src_buf)
+
+        return output
+
+
+def _even_chunk(x: torch.Tensor, chunks: int):
+    assert x.shape[0] % chunks == 0, f"{x.shape=} {chunks=}"
+    return list(x.chunk(chunks))
+
+
+def _create_shared_buffer_tensors(local_tensor: torch.Tensor) -> List[torch.Tensor]:
+    self_rank = get_naive_distributed().get_rank()
+    world_size = get_naive_distributed().get_world_size()
+
+    object_list = get_naive_distributed().all_gather_object(
+        dict(
+            dup_serialized_local_tensor=[
+                (
+                    None
+                    if interesting_rank == self_rank
+                    else MultiprocessingSerializer.serialize(local_tensor)
+                )
+                for interesting_rank in range(world_size)
+            ]
+        )
+    )
+
+    output_tensors = []
+    for output_rank in range(world_size):
+        remote_serialized_tensor = object_list[output_rank][
+            "dup_serialized_local_tensor"
+        ][self_rank]
+        if output_rank == self_rank:
+            assert remote_serialized_tensor is None
+            output_tensors.append(local_tensor)
+        else:
+            output_tensors.append(
+                MultiprocessingSerializer.deserialize(remote_serialized_tensor)
+            )
+
+    return output_tensors

python/sglang/srt/utils/patch_tokenizer.py

@@ -0,0 +1,127 @@
+import logging
+
+from sglang.srt.environ import envs
+
+logger = logging.getLogger(__name__)
+
+
+def patch_tokenizer(tokenizer):
+    if not envs.SGLANG_PATCH_TOKENIZER.get():
+        return tokenizer
+
+    if _is_kimi_tiktoken_tokenizer(tokenizer):
+        logger.info(
+            f"Applying special tokens cache patch for Kimi tokenizer: {type(tokenizer)}"
+        )
+        return _SpecialTokensCachePatcher.patch(tokenizer)
+
+    return tokenizer
+
+
+def unpatch_tokenizer(tokenizer):
+    return _SpecialTokensCachePatcher.unpatch(tokenizer)
+
+
+def _is_kimi_tiktoken_tokenizer(tokenizer):
+    cls = type(tokenizer)
+    class_name = cls.__name__
+    module_name = cls.__module__ or ""
+    return class_name == "TikTokenTokenizer" and "tokenization_kimi" in module_name
+
+
+def decode_without_hf_kwargs(tokenizer, token_ids, skip_special_tokens):
+    if skip_special_tokens:
+        special_ids = getattr(tokenizer, "all_special_ids_set", None)
+        if special_ids is None:
+            special_ids = getattr(tokenizer, "all_special_ids", None)
+        if special_ids is not None:
+            special_ids_set = set(special_ids)
+            token_ids = [tid for tid in token_ids if tid not in special_ids_set]
+    return tokenizer.decode(token_ids)
+
+
+class _SpecialTokensCachePatcher:
+    _PATCHED_FLAG = "_sglang_special_tokens_patched"
+    _CACHED_TOKENS_ATTR = "_sglang_cached_special_tokens"
+    _CACHED_IDS_ATTR = "_sglang_cached_special_ids"
+
+    @classmethod
+    def patch(cls, tokenizer):
+        tokenizer_cls = type(tokenizer)
+
+        if getattr(tokenizer_cls, cls._PATCHED_FLAG, False):
+            return tokenizer
+
+        tokenizer_cls._original_all_special_tokens = (
+            tokenizer_cls.all_special_tokens.fget
+        )
+        tokenizer_cls._original_all_special_ids = tokenizer_cls.all_special_ids.fget
+        tokenizer_cls._original_add_special_tokens = tokenizer_cls.add_special_tokens
+        tokenizer_cls._original_add_tokens = tokenizer_cls.add_tokens
+
+        patched_all_special_tokens = _make_cached_property(
+            cls._CACHED_TOKENS_ATTR, tokenizer_cls._original_all_special_tokens
+        )
+        patched_all_special_ids = _make_cached_property(
+            cls._CACHED_IDS_ATTR, tokenizer_cls._original_all_special_ids
+        )
+
+        def patched_add_special_tokens(self, *args, **kwargs):
+            assert (
+                False
+            ), "Cannot modify special tokens after patch. Call unpatch_tokenizer first."
+
+        def patched_add_tokens(self, new_tokens, special_tokens=False):
+            assert (
+                not special_tokens
+            ), "Cannot add special tokens after patch. Call unpatch_tokenizer first."
+            return tokenizer_cls._original_add_tokens(
+                self, new_tokens, special_tokens=False
+            )
+
+        tokenizer_cls.all_special_tokens = patched_all_special_tokens
+        tokenizer_cls.all_special_ids = patched_all_special_ids
+        tokenizer_cls.add_special_tokens = patched_add_special_tokens
+        tokenizer_cls.add_tokens = patched_add_tokens
+        setattr(tokenizer_cls, cls._PATCHED_FLAG, True)
+
+        return tokenizer
+
+    @classmethod
+    def unpatch(cls, tokenizer):
+        tokenizer_cls = type(tokenizer)
+
+        if not getattr(tokenizer_cls, cls._PATCHED_FLAG, False):
+            return tokenizer
+
+        tokenizer_cls.all_special_tokens = property(
+            tokenizer_cls._original_all_special_tokens
+        )
+        tokenizer_cls.all_special_ids = property(
+            tokenizer_cls._original_all_special_ids
+        )
+        tokenizer_cls.add_special_tokens = tokenizer_cls._original_add_special_tokens
+        tokenizer_cls.add_tokens = tokenizer_cls._original_add_tokens
+
+        del tokenizer_cls._original_all_special_tokens
+        del tokenizer_cls._original_all_special_ids
+        del tokenizer_cls._original_add_special_tokens
+        del tokenizer_cls._original_add_tokens
+        delattr(tokenizer_cls, cls._PATCHED_FLAG)
+
+        for attr in [cls._CACHED_TOKENS_ATTR, cls._CACHED_IDS_ATTR]:
+            if hasattr(tokenizer, attr):
+                delattr(tokenizer, attr)
+
+        logger.info(f"Unpatched special tokens cache for {tokenizer_cls.__name__}")
+        return tokenizer
+
+
+def _make_cached_property(cache_attr, original_fn):
+    @property
+    def cached_prop(self):
+        if getattr(self, cache_attr, None) is None:
+            setattr(self, cache_attr, original_fn(self))
+        return getattr(self, cache_attr)
+
+    return cached_prop

python/sglang/srt/utils/patch_torch.py

@@ -0,0 +1,139 @@
+# 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.
+# ==============================================================================
+from typing import Callable, Union
+
+import torch
+from torch.multiprocessing import reductions
+
+from sglang.srt.utils.common import is_musa, is_npu, torch_release
+
+_is_npu = is_npu()
+_is_musa = is_musa()
+
+if _is_npu:
+    from torch_npu.multiprocessing import reductions as npu_reductions
+
+    def _rebuild_npu_tensor_modified(*args):
+        args = _modify_tuple(args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, npu_verl_to_sglang)
+        return npu_reductions._rebuild_npu_tensor_original(*args)
+
+    def npu_verl_to_sglang(device: int):
+        assert (
+            SGLANG_TP_RANK is not None
+        ), "SGLANG_TP_RANK is not registered. Please call register_sgl_tp_rank() first."
+        return SGLANG_TP_RANK
+
+
+SGLANG_TP_RANK = None
+
+
+def monkey_patch_torch_reductions():
+    """Monkey patching before Torch https://github.com/pytorch/pytorch/pull/149248 is fixed"""
+
+    if not _is_npu:
+        if hasattr(reductions, "_reduce_tensor_original"):
+            return
+        reductions._reduce_tensor_original = reductions.reduce_tensor
+        reductions._rebuild_cuda_tensor_original = reductions.rebuild_cuda_tensor
+
+        reductions.reduce_tensor = _reduce_tensor_modified
+        reductions.rebuild_cuda_tensor = _rebuild_cuda_tensor_modified
+        reductions.init_reductions()
+    else:
+        # FIXME: This is a temp patch for npu as HDK does not support device uuid for now
+        if hasattr(npu_reductions, "_rebuild_npu_tensor_original"):
+            return
+
+        npu_reductions._rebuild_npu_tensor_original = npu_reductions.rebuild_npu_tensor
+        npu_reductions.rebuild_npu_tensor = _rebuild_npu_tensor_modified
+
+
+# The signature has not been changed for years, and we will not need this when the next version is released,
+# so it looks safe to use a constant.
+_REDUCE_TENSOR_ARG_DEVICE_INDEX = 6
+
+
+def register_sgl_tp_rank(rank: int):
+    global SGLANG_TP_RANK
+    SGLANG_TP_RANK = rank
+
+
+def _reduce_tensor_modified(*args, **kwargs):
+    output_fn, output_args = reductions._reduce_tensor_original(*args, **kwargs)
+    output_args = _modify_tuple(
+        output_args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, _device_to_uuid
+    )
+    return output_fn, output_args
+
+
+def _rebuild_cuda_tensor_modified(*args):
+    args = _modify_tuple(args, _REDUCE_TENSOR_ARG_DEVICE_INDEX, _device_from_maybe_uuid)
+    return reductions._rebuild_cuda_tensor_original(*args)
+
+
+def _device_to_uuid(device: int) -> str:
+    return str(torch.cuda.get_device_properties(device).uuid)
+
+
+def _device_from_maybe_uuid(device_maybe_uuid: Union[int, str]) -> int:
+    if isinstance(device_maybe_uuid, int):
+        return device_maybe_uuid
+
+    if isinstance(device_maybe_uuid, str):
+        for device in range(torch.cuda.device_count()):
+            if str(torch.cuda.get_device_properties(device).uuid) == device_maybe_uuid:
+                return device
+        raise Exception("Invalid device_uuid=" + device_maybe_uuid)
+
+    raise Exception(f"Unknown type: {device_maybe_uuid=}")
+
+
+def _modify_tuple(t, index: int, modifier: Callable):
+    return *t[:index], modifier(t[index]), *t[index + 1 :]
+
+
+def monkey_patch_torch_compile():
+    if torch_release < (2, 8):
+        # These things are cacheable by torch.compile. torch.compile just doesn't know it.
+        # This was fixed in PyTorch 2.8, but until then, we monkey patch.
+        import torch._higher_order_ops.auto_functionalize as af
+
+        af.auto_functionalized_v2._cacheable = True
+        af.auto_functionalized._cacheable = True
+
+
+def register_fake_if_exists(op_name):
+    """
+    Decorator factory to conditionally register a fake for a custom op if it exists.
+    Parses op_name (e.g., 'sgl_kernel::gptq_gemm'), checks if the op exists via hasattr
+    on the namespace attribute of torch.ops. Registers the fake if present; otherwise,
+    returns the function unchanged.
+    Args:
+        op_name (str): Full operator name (e.g., 'sgl_kernel::gptq_gemm').
+    Returns:
+        callable: Decorator for the fake function.
+    Example:
+        @register_fake_if_exists('sgl_kernel::gptq_gemm')
+        def fake_gptq_gemm(a, b_q_weight, b_gptq_qzeros, b_gptq_scales, b_g_idx, use_shuffle, bit):
+            return a.new_empty((a.shape[0], b_q_weight.shape[-1]), dtype=a.dtype)
+    """
+
+    def decorator(func):
+        namespace, bare_op = op_name.split("::")
+        ops_namespace = getattr(torch.ops, namespace, None)
+        if ops_namespace and hasattr(ops_namespace, bare_op):
+            torch.library.register_fake(op_name, func)
+        return func
+
+    return decorator

python/sglang/srt/utils/phase_checker.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from enum import IntEnum
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.environ import envs
+
+
+def _phase_repr(phase: int | IntEnum) -> str:
+    if isinstance(phase, IntEnum):
+        return f"{phase.name}({int(phase)})"
+    return str(int(phase))
+
+
+def _host_debug(msg: str) -> None:
+    if envs.SGLANG_PHASE_CHECKER_DEBUG.get():
+        print(msg, flush=True)
+
+
+# debug=True so tl.device_assert below actually raises. Without it the assert
+# is stripped at compile time and only tl.device_print fires (the assert is
+# gated on the TRITON_DEBUG env var by default — see tl.device_assert docstring).
+@triton.jit(debug=True)
+def _phase_check_kernel(
+    phase_ptr,
+    enable_assert_ptr,
+    EXPECT_PHASE: tl.constexpr,
+    NEXT_PHASE: tl.constexpr,
+    CALLER_TAG: tl.constexpr,
+):
+    cur = tl.load(phase_ptr)
+    enable_assert = tl.load(enable_assert_ptr)
+    if enable_assert != 0:
+        if cur != EXPECT_PHASE:
+            # constexpr values get baked into the prefix string at compile time;
+            # only `cur` is runtime.
+            tl.device_print(
+                f"[SimplePhaseChecker FAIL] caller_tag={CALLER_TAG} "
+                f"expect={EXPECT_PHASE} next={NEXT_PHASE} actual=",
+                cur,
+            )
+        tl.device_assert(cur == EXPECT_PHASE, "SimplePhaseChecker: phase mismatch")
+    tl.store(phase_ptr, NEXT_PHASE)
+
+
+class SimplePhaseChecker:
+    """GPU-side state machine for any int-keyed phase sequence."""
+
+    def __init__(self, *, initial_phase: int | IntEnum, device: torch.device) -> None:
+        self._initial_phase = int(initial_phase)
+        self._phase = torch.tensor(
+            self._initial_phase, dtype=torch.int32, device=device
+        )
+        self._enable_assert_device = torch.zeros(1, dtype=torch.int32, device=device)
+        self._caller_tag_registry: dict[str, int] = {}
+        _host_debug(
+            f"[SimplePhaseChecker.__init__] device={device} "
+            f"initial_phase={_phase_repr(initial_phase)} "
+            f"enable_assert=OFF (call enable_assert() after init is done)"
+        )
+
+    def enable_assert(self) -> None:
+        """Reset phase to initial_phase, then enable the device-side assert."""
+        self._reset_to_idle()
+        self._enable_assert_device.fill_(1)
+        _host_debug(f"[SimplePhaseChecker.enable_assert] assert ENABLED")
+
+    def update(
+        self,
+        *,
+        expect_phase: int | IntEnum,
+        next_phase: int | IntEnum,
+        caller_name: str = "",
+    ) -> None:
+        caller_tag = self._resolve_caller_tag(caller_name)
+        _host_debug(
+            f"[SimplePhaseChecker.update] caller={caller_name!r} "
+            f"caller_tag={caller_tag} "
+            f"expect={_phase_repr(expect_phase)} "
+            f"next={_phase_repr(next_phase)} "
+            f"capturing={torch.cuda.is_current_stream_capturing()}"
+        )
+        _phase_check_kernel[(1,)](
+            self._phase,
+            self._enable_assert_device,
+            EXPECT_PHASE=int(expect_phase),
+            NEXT_PHASE=int(next_phase),
+            CALLER_TAG=caller_tag,
+        )
+
+    def _reset_to_idle(self) -> None:
+        self._phase.fill_(self._initial_phase)
+        _host_debug(
+            f"[SimplePhaseChecker._reset_to_idle] phase reset to "
+            f"{self._initial_phase}"
+        )
+
+    def _resolve_caller_tag(self, caller_name: str) -> int:
+        registry = self._caller_tag_registry
+        if caller_name not in registry:
+            registry[caller_name] = len(registry) + 1
+            _host_debug(
+                f"[SimplePhaseChecker] registered caller_tag "
+                f"{registry[caller_name]} <- {caller_name!r}"
+            )
+        return registry[caller_name]

python/sglang/srt/utils/poll_based_barrier.py

@@ -0,0 +1,31 @@
+import torch
+
+from sglang.srt.distributed import get_world_group
+
+
+class PollBasedBarrier:
+    def __init__(self, noop: bool = False):
+        self._noop = noop
+        self._local_arrived = False
+
+    def local_arrive(self):
+        assert not self._local_arrived
+        self._local_arrived = True
+
+    def poll_global_arrived(self) -> bool:
+        global_arrived = self._compute_global_arrived()
+        output = self._local_arrived and global_arrived
+        if output:
+            self._local_arrived = False
+        return output
+
+    def _compute_global_arrived(self) -> bool:
+        local_arrived = self._noop or self._local_arrived
+        global_arrived = torch.tensor(local_arrived)
+        # Can optimize if bottleneck
+        torch.distributed.all_reduce(
+            global_arrived,
+            torch.distributed.ReduceOp.MIN,
+            group=get_world_group().cpu_group,
+        )
+        return global_arrived.item()

python/sglang/srt/utils/profile_merger.py

@@ -0,0 +1,199 @@
+"""Merge Chrome trace files from multiple ranks (TP, DP, PP, EP) into a single trace."""
+
+import glob
+import gzip
+import json
+import logging
+import os
+import re
+from typing import Any, Dict, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+
+class ProfileMerger:
+    """Merge profile traces from all parallelism types: TP, DP, PP, EP."""
+
+    def __init__(self, output_dir: str, profile_id: str):
+        self.output_dir = output_dir
+        self.profile_id = profile_id
+        self.merged_trace_path = os.path.join(
+            output_dir, f"merged-{profile_id}.trace.json.gz"
+        )
+
+        # Rank types in priority order (used for sorting and labeling)
+        self.rank_types = ["tp", "dp", "pp", "ep"]
+
+        # Sort index multipliers: DP (highest) > EP > PP > TP (lowest)
+        # These ensure proper visual ordering in trace viewer
+        self.sort_index_multipliers = {
+            "dp_rank": 100_000_000,
+            "ep_rank": 1_000_000,
+            "pp_rank": 10_000,
+            "tp_rank": 100,
+        }
+
+        # PID threshold for sort_index updates (only update for system PIDs < 1000)
+        self.pid_sort_index_threshold = 1000
+
+    def merge_chrome_traces(self) -> str:
+        """Merge Chrome traces from all ranks into a single trace.
+
+        Returns:
+            Path to merged trace file.
+
+        Raises:
+            ValueError: If no trace files found.
+        """
+        trace_files = self._discover_trace_files()
+        if not trace_files:
+            raise ValueError(f"No trace files found for profile_id: {self.profile_id}")
+
+        logger.info(f"Found {len(trace_files)} trace files to merge")
+
+        merged_trace = {"traceEvents": []}
+        all_device_properties = []
+
+        for trace_file in sorted(trace_files, key=self._get_rank_sort_key):
+            rank_info = self._extract_rank_info(trace_file)
+            logger.info(f"Processing {trace_file} with rank info: {rank_info}")
+
+            output = self._handle_file(trace_file, rank_info)
+
+            merged_trace["traceEvents"].extend(output["traceEvents"])
+
+            if "deviceProperties" in output:
+                all_device_properties.extend(output["deviceProperties"])
+                del output["deviceProperties"]
+
+            for key, value in output.items():
+                if key != "traceEvents" and key not in merged_trace:
+                    merged_trace[key] = value
+
+        if all_device_properties:
+            merged_trace["deviceProperties"] = all_device_properties
+
+        with gzip.open(self.merged_trace_path, "wb") as f:
+            f.write(json.dumps(merged_trace).encode("utf-8"))
+
+        logger.info(f"Merged profile saved to: {self.merged_trace_path}")
+        logger.info(f"Total events merged: {len(merged_trace['traceEvents'])}")
+
+        return self.merged_trace_path
+
+    def _discover_trace_files(self) -> List[str]:
+        """Discover trace files matching profile_id (supports TP/DP/PP/EP formats)."""
+        patterns = [f"{self.profile_id}*.trace.json.gz"]
+
+        trace_files = []
+        for pattern in patterns:
+            search_pattern = os.path.join(self.output_dir, pattern)
+            trace_files.extend(glob.glob(search_pattern))
+
+        trace_files = [
+            f
+            for f in trace_files
+            if not f.endswith(f"merged-{self.profile_id}.trace.json.gz")
+            and not f.endswith("-memory.pickle")
+            and "TP-" in f
+        ]
+        trace_files = list(set(trace_files))
+        return trace_files
+
+    def _extract_rank_info(self, filename: str) -> Dict[str, int]:
+        """Extract rank info (TP/DP/PP/EP) from filename."""
+        basename = os.path.basename(filename)
+        rank_info = {}
+
+        for rank_type in self.rank_types:
+            match = re.search(rf"{rank_type.upper()}-(\d+)", basename)
+            if match:
+                rank_info[f"{rank_type}_rank"] = int(match.group(1))
+
+        return rank_info
+
+    def _create_rank_label(self, rank_info: Dict[str, int]) -> str:
+        parts = []
+        for rank_type in self.rank_types:
+            rank_key = f"{rank_type}_rank"
+            if rank_key in rank_info:
+                parts.append(f"{rank_type.upper()}{rank_info[rank_key]:02d}")
+
+        return f"[{'-'.join(parts)}]" if parts else "[Unknown]"
+
+    def _handle_file(self, path: str, rank_info: Dict[str, int]) -> Dict[str, Any]:
+        logger.info(f"Processing file: {path}")
+
+        try:
+            with gzip.open(path, "rt", encoding="utf-8") as f:
+                trace = json.load(f)
+
+            output = {
+                key: value for key, value in trace.items() if key != "traceEvents"
+            }
+            output["traceEvents"] = self._process_events(
+                trace.get("traceEvents", []), rank_info
+            )
+            return output
+
+        except Exception as e:
+            logger.error(f"Failed to process trace file {path}: {e}")
+            return {"traceEvents": []}
+
+    def _process_events(
+        self, events: List[Dict], rank_info: Dict[str, int]
+    ) -> List[Dict]:
+        """Process events: update sort_index and add rank labels to PIDs."""
+        rank_label = self._create_rank_label(rank_info)
+
+        for event in events:
+            if event.get("name") == "process_sort_index":
+                pid = self._maybe_cast_int(event.get("pid"))
+                if pid is not None and pid < self.pid_sort_index_threshold:
+                    event["args"]["sort_index"] = self._calculate_sort_index(
+                        rank_info, pid
+                    )
+
+            event["pid"] = f"{rank_label} {event['pid']}"
+
+        return events
+
+    def _calculate_sort_index(self, rank_info: Dict[str, int], pid: int) -> int:
+        sort_index = pid
+        for rank_type, multiplier in self.sort_index_multipliers.items():
+            sort_index += rank_info.get(rank_type, 0) * multiplier
+        return sort_index
+
+    def _get_rank_sort_key(self, path: str) -> Tuple[int, int, int, int]:
+        rank_info = self._extract_rank_info(path)
+        return tuple(
+            rank_info.get(f"{rank_type}_rank", 0)
+            for rank_type in ["dp", "ep", "pp", "tp"]
+        )
+
+    def _maybe_cast_int(self, x) -> Optional[int]:
+        try:
+            return int(x)
+        except (ValueError, TypeError):
+            return None
+
+    def get_merge_summary(self) -> Dict[str, Any]:
+        if not os.path.exists(self.merged_trace_path):
+            return {"error": "Merged trace file not found"}
+
+        try:
+            with gzip.open(self.merged_trace_path, "rt") as f:
+                merged_data = json.load(f)
+
+            trace_files = self._discover_trace_files()
+
+            return {
+                "merged_file": self.merged_trace_path,
+                "total_events": len(merged_data.get("traceEvents", [])),
+                "total_files": len(trace_files),
+                "source_files": [os.path.basename(f) for f in trace_files],
+                "profile_id": self.profile_id,
+                "device_properties_count": len(merged_data.get("deviceProperties", [])),
+            }
+        except Exception as e:
+            return {"error": f"Failed to read merged trace: {str(e)}"}

python/sglang/srt/utils/profile_utils.py

@@ -0,0 +1,414 @@
+import logging
+import os
+import time
+from abc import ABC
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Dict, List, Optional
+
+import torch
+
+from sglang.srt.distributed.parallel_state_wrapper import ParallelState
+from sglang.srt.environ import envs
+from sglang.srt.managers.io_struct import ProfileReqOutput
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.runtime_context import get_server_args
+from sglang.srt.utils import is_npu
+from sglang.srt.utils.torch_npu_patch_utils import apply_torch_npu_patches
+
+_is_npu = is_npu()
+if _is_npu:
+    import torch_npu
+
+    patches = [
+        ["profiler.profile", torch_npu.profiler.profile],
+        ["profiler.ProfilerActivity.CUDA", torch_npu.profiler.ProfilerActivity.NPU],
+        ["profiler.ProfilerActivity.CPU", torch_npu.profiler.ProfilerActivity.CPU],
+    ]
+    apply_torch_npu_patches(torch_npu, patches)
+
+logger = logging.getLogger(__name__)
+
+
+def export_cuda_graph_capture_trace(prof_context, *, runner_name: str, tp_rank: int):
+    """Persist a CUDA-graph capture profiler trace (chrome trace) to disk.
+
+    Opt-in via ``SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE`` (no-op otherwise). The
+    capture profiler must have run with ``record_shapes=True`` so the trace can
+    be inspected offline as a per-kernel shape/identity record. The file lands in
+    ``<SGLANG_TORCH_PROFILER_DIR>/graph_capture_profile/`` and is namespaced by
+    runner class and TP rank so concurrent capture passes (e.g. EAGLE3
+    target/draft/draft-extend) and ranks don't overwrite each other.
+    """
+    if not envs.SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE.get():
+        return
+    output_dir = os.path.join(
+        envs.SGLANG_TORCH_PROFILER_DIR.get(), "graph_capture_profile"
+    )
+    os.makedirs(output_dir, exist_ok=True)
+    path = os.path.join(
+        output_dir, f"cuda_graph_capture-{runner_name}-TP-{tp_rank}.json.gz"
+    )
+    prof_context.export_chrome_trace(path)
+    logger.info(f"CUDA graph capture trace saved to: {path}")
+
+
+class ProfileManager:
+    def __init__(self, ps: ParallelState, cpu_group):
+        self.stage_based_trigger = _StageBasedTrigger(
+            on_start=self._do_start,
+            on_stop=self._do_stop,
+        )
+        self.ps = ps
+        self.cpu_group = cpu_group
+        self.first_rank_in_node = ps.gpu_id == get_server_args().base_gpu_id
+        self.profiler_kwargs = None
+        self.profiler = None
+
+    def step(self, forward_mode: ForwardMode):
+        stage = _get_stage_from_forward_mode(forward_mode)
+        if stage is None:
+            return
+
+        self.stage_based_trigger.step(stage=stage)
+
+    def configure(
+        self,
+        *,
+        output_dir: Optional[str],
+        start_step: Optional[int],
+        num_steps: Optional[int],
+        activities: Optional[List[str]],
+        with_stack: Optional[bool],
+        record_shapes: Optional[bool],
+        profile_by_stage: bool,
+        profile_id: str,
+        merge_profiles: bool,
+        profile_prefix: str,
+        profile_stages: Optional[List[str]] = None,
+    ):
+        # not supported yet
+        assert start_step is None
+        assert (
+            profile_by_stage
+        ), "only support profile_by_stage=true now"  # `false` can be easily supported
+        assert not merge_profiles
+
+        if output_dir is None:
+            output_dir = os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp")
+        if activities is None:
+            activities = ["CPU", "GPU"]
+
+        self.profiler_kwargs = dict(
+            activities=activities,
+            with_stack=with_stack,
+            record_shapes=record_shapes,
+            output_dir=output_dir,
+            output_prefix=profile_prefix,
+            profile_id=profile_id,
+        )
+
+        self.stage_based_trigger.configure(
+            num_steps=num_steps,
+            interesting_stages=profile_stages or ["prefill", "decode"],
+        )
+
+        return ProfileReqOutput(success=True, message="Succeeded")
+
+    def manual_start(self):
+        raise NotImplementedError("manually start is only supported yet")
+
+    def manual_stop(self):
+        raise NotImplementedError("manually stop is only supported yet")
+
+    def _do_start(self, stage: Optional[str] = None):
+        logger.info(
+            f"Profiling starts{f' for {stage}' if stage else ''}. "
+            f"Traces will be saved to: {self.profiler_kwargs['output_dir']} "
+            f"(with profile id: {self.profiler_kwargs['profile_id']})",
+        )
+
+        assert self.profiler is None
+        self.profiler = _ProfilerBase.create(
+            **self.profiler_kwargs,
+            ps=self.ps,
+            cpu_group=self.cpu_group,
+            first_rank_in_node=self.first_rank_in_node,
+            output_suffix=f"-{stage}" if stage else "",
+        )
+        self.profiler.start()
+
+    def _do_stop(self):
+        logger.info("Stop profiling...")
+        self.profiler.stop()
+        logger.info(
+            f"Profiling done. Traces are saved to: {self.profiler_kwargs['output_dir']}"
+        )
+        self.profiler = None
+
+
+def _get_stage_from_forward_mode(forward_mode: ForwardMode):
+    if forward_mode.is_prefill():
+        return "prefill"
+    elif forward_mode.is_decode():
+        return "decode"
+    elif forward_mode.is_idle():
+        return None
+    else:
+        raise RuntimeError(f"unsupported profile stage: {forward_mode=}")
+
+
+# ======================================== Stage related ==========================================
+
+
+class _StageBasedTrigger:
+    @dataclass
+    class _StageConfig:
+        target_count: int
+
+    @dataclass
+    class _RunningState:
+        curr_stage: str
+        curr_count: int
+
+    def __init__(self, on_start: Callable, on_stop: Callable):
+        self.on_start = on_start
+        self.on_stop = on_stop
+
+        self.running_state: Optional[_StageBasedTrigger._RunningState] = None
+        # When a stage is in the dict, it means it is being or should be executed
+        self.stage_configs: Dict[str, _StageBasedTrigger._StageConfig] = {}
+
+    def configure(self, num_steps: int, interesting_stages: List[str]):
+        assert self.running_state is None
+        self.stage_configs = {
+            stage: self._StageConfig(target_count=num_steps)
+            for stage in interesting_stages
+        }
+
+    def step(self, stage: str):
+        # Incr counter
+        if (s := self.running_state) is not None:
+            s.curr_count += 1
+
+        # Maybe stop
+        if ((s := self.running_state) is not None) and (
+            (s.curr_count > self.stage_configs[s.curr_stage].target_count)
+            or (stage != s.curr_stage)
+        ):
+            del self.stage_configs[s.curr_stage]
+            self.running_state = None
+            self.on_stop()
+
+        # Maybe start
+        if (self.running_state is None) and (stage in self.stage_configs):
+            self.running_state = self._RunningState(
+                curr_stage=stage,
+                curr_count=0,
+            )
+            self.on_start(stage=stage)
+
+        # Sanity check
+        assert (self.running_state is not None) == (stage in self.stage_configs)
+        if (s := self.running_state) is not None:
+            assert s.curr_stage == stage
+
+
+# ======================================== Concrete profilers ==========================================
+
+
+class _ProfilerBase(ABC):
+    @staticmethod
+    def create(activities, with_stack, record_shapes, **kwargs):
+        inners = []
+        if ("CPU" in activities) or ("GPU" in activities):
+            inners.append(
+                _ProfilerTorch(
+                    **kwargs,
+                    activities=activities,
+                    with_stack=with_stack,
+                    record_shapes=record_shapes,
+                )
+            )
+        if "MEM" in activities:
+            inners.append(_ProfilerMemory(**kwargs))
+        if "CUDA_PROFILER" in activities:
+            inners.append(_ProfilerCudart(**kwargs))
+        if "RPD" in activities:  # for ROCM
+            inners.append(_ProfilerRPD(**kwargs))
+
+        return _ProfilerList(inners)
+
+    def start(self):
+        raise NotImplementedError
+
+    def stop(self):
+        raise NotImplementedError
+
+
+class _ProfilerList(_ProfilerBase):
+    def __init__(self, inners: List[_ProfilerBase]):
+        self.inners = inners
+
+    def start(self):
+        for inner in self.inners:
+            inner.start()
+
+    def stop(self):
+        for inner in self.inners:
+            inner.stop()
+
+
+class _ProfilerConcreteBase(_ProfilerBase):
+    def __init__(
+        self,
+        output_dir: str,
+        output_prefix: str,
+        output_suffix: str,
+        profile_id: str,
+        ps: ParallelState,
+        cpu_group,
+        first_rank_in_node: bool,
+    ):
+        self.output_dir = output_dir
+        self.output_prefix = output_prefix
+        self.output_suffix = output_suffix
+        self.profile_id = profile_id
+        self.ps = ps
+        self.cpu_group = cpu_group
+        self.first_rank_in_node = first_rank_in_node
+
+
+class _ProfilerTorch(_ProfilerConcreteBase):
+    def __init__(self, with_stack: bool, record_shapes: bool, activities, **kwargs):
+        super().__init__(**kwargs)
+        self.with_stack = with_stack
+        self.record_shapes = record_shapes
+        self.activities = activities
+
+    def start(self):
+        activity_map = {
+            "CPU": torch.profiler.ProfilerActivity.CPU,
+            "GPU": torch.profiler.ProfilerActivity.CUDA,
+        }
+        torchprof_activities = [
+            activity_map[a] for a in self.activities if a in activity_map
+        ]
+
+        self.torch_profiler = torch.profiler.profile(
+            activities=torchprof_activities,
+            with_stack=self.with_stack if self.with_stack is not None else True,
+            record_shapes=(
+                self.record_shapes if self.record_shapes is not None else False
+            ),
+            on_trace_ready=(
+                None
+                if not _is_npu
+                else torch_npu.profiler.tensorboard_trace_handler(self.output_dir)
+            ),
+        )
+        self.torch_profiler.start()
+
+    def stop(self):
+        Path(self.output_dir).mkdir(parents=True, exist_ok=True)
+
+        self.torch_profiler.stop()
+        if not _is_npu:
+            # Build filename with only non-zero ranks to maintain backward compatibility
+            filename_parts = [self.profile_id, f"TP-{self.ps.tp_rank}"]
+
+            # Only add other ranks if parallelism is enabled (size > 1)
+            if self.ps.dp_size > 1:
+                filename_parts.append(f"DP-{self.ps.dp_rank}")
+            if self.ps.pp_size > 1:
+                filename_parts.append(f"PP-{self.ps.pp_rank}")
+            if self.ps.moe_ep_size > 1:
+                filename_parts.append(f"EP-{self.ps.moe_ep_rank}")
+
+            filename = (
+                (self.output_prefix + "-" if self.output_prefix else "")
+                + "-".join(filename_parts)
+                + self.output_suffix
+                + ".trace.json.gz"
+            )
+
+            self.torch_profiler.export_chrome_trace(
+                os.path.join(self.output_dir, filename)
+            )
+        torch.distributed.barrier(self.cpu_group)
+
+        # TODO: migrate `_merge_profile_traces`
+
+
+class _ProfilerMemory(_ProfilerConcreteBase):
+    def start(self):
+        torch.cuda.memory._record_memory_history(max_entries=100000)
+
+    def stop(self):
+        Path(self.output_dir).mkdir(parents=True, exist_ok=True)
+
+        memory_profile_path = os.path.join(
+            self.output_dir,
+            str(time.time())
+            + f"-TP-{self.ps.tp_rank}-memory"
+            + self.output_suffix
+            + ".pickle",
+        )
+        torch.cuda.memory._dump_snapshot(memory_profile_path)
+        torch.cuda.memory._record_memory_history(enabled=None)
+
+
+class _ProfilerCudart(_ProfilerConcreteBase):
+    def start(self):
+        if self.first_rank_in_node:
+            logger.info(f"Call cudaProfilerStart")
+            torch.cuda.cudart().cudaProfilerStart()
+
+    def stop(self):
+        if self.first_rank_in_node:
+            logger.info(f"Call cudaProfilerStop")
+            torch.cuda.cudart().cudaProfilerStop()
+
+
+class _ProfilerRPD(_ProfilerConcreteBase):
+    def start(self):
+        Path(self.output_dir).mkdir(parents=True, exist_ok=True)
+
+        from rpdTracerControl import rpdTracerControl
+
+        rpdTracerControl.skipCreate()
+
+        self.rpd_profile_path = os.path.join(
+            self.output_dir,
+            "rpd-" + str(time.time()) + f"-TP-{self.ps.tp_rank}" + ".trace.json.gz",
+        )
+
+        if self.ps.tp_rank == 0:
+            import sqlite3
+
+            from rocpd.schema import RocpdSchema
+
+            if os.path.exists("trace.rpd"):
+                os.unlink("trace.rpd")
+            schema = RocpdSchema()
+            connection = sqlite3.connect("trace.rpd")
+            schema.writeSchema(connection)
+            connection.commit()
+            del connection
+        torch.distributed.barrier(self.cpu_group)
+
+        self.rpd_profiler = rpdTracerControl()
+        self.rpd_profiler.setPythonTrace(True)
+        self.rpd_profiler.start()
+        self.rpd_profiler.rangePush("", "rpd profile range", "")
+
+    def stop(self):
+        self.rpd_profiler.rangePop()
+        self.rpd_profiler.stop()
+        self.rpd_profiler.flush()
+
+        torch.distributed.barrier(self.cpu_group)
+        if self.ps.tp_rank == 0:
+            from sglang.srt.utils.rpd_utils import rpd_to_chrome_trace
+
+            rpd_to_chrome_trace("trace.rpd", self.rpd_profile_path)

python/sglang/srt/utils/request_logger.py

@@ -0,0 +1,316 @@
+# 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.
+# ==============================================================================
+from __future__ import annotations
+
+import dataclasses
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
+
+from sglang.srt.environ import envs
+from sglang.srt.utils.log_utils import create_log_targets, log_json
+
+if TYPE_CHECKING:
+    import fastapi
+
+    from sglang.srt.managers.io_struct import EmbeddingReqInput, GenerateReqInput
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_WHITELISTED_HEADERS = ["x-smg-routing-key"]
+WHITELISTED_HEADERS = _DEFAULT_WHITELISTED_HEADERS + [
+    h.lower() for h in envs.SGLANG_LOG_REQUEST_HEADERS.get()
+]
+
+
+def _extract_whitelisted_headers(
+    request: Optional[fastapi.Request],
+) -> Optional[Dict[str, str]]:
+    if request is None:
+        return None
+    return {h: v for h in WHITELISTED_HEADERS if (v := request.headers.get(h))}
+
+
+class RequestLogger:
+    def __init__(
+        self,
+        log_requests: bool,
+        log_requests_level: int,
+        log_requests_format: str,
+        log_requests_target: Optional[List[str]],
+    ):
+        self.log_requests = log_requests
+        self.log_requests_level = log_requests_level
+        self.log_requests_format = log_requests_format
+        self.log_requests_target = log_requests_target
+
+        self.metadata: Tuple[Optional[int], Optional[Set[str]], Optional[Set[str]]] = (
+            self._compute_metadata()
+        )
+        self.targets = self._setup_targets()
+
+        self.log_exceeded_ms = envs.SGLANG_LOG_REQUEST_EXCEEDED_MS.get()
+
+    def _setup_targets(self) -> List[logging.Logger]:
+        return create_log_targets(
+            targets=self.log_requests_target, name_prefix=__name__
+        )
+
+    def configure(
+        self,
+        log_requests: Optional[bool] = None,
+        log_requests_level: Optional[int] = None,
+        log_requests_format: Optional[str] = None,
+        log_requests_target: Optional[List[str]] = None,
+    ) -> None:
+        if log_requests is not None:
+            self.log_requests = log_requests
+        if log_requests_level is not None:
+            self.log_requests_level = log_requests_level
+        if log_requests_format is not None:
+            self.log_requests_format = log_requests_format
+        if log_requests_target is not None:
+            self.log_requests_target = log_requests_target
+
+        self.metadata = self._compute_metadata()
+        self.targets = self._setup_targets()
+
+    def log_received_request(
+        self,
+        obj: Union[GenerateReqInput, EmbeddingReqInput],
+        tokenizer: Any = None,
+        request: Optional[fastapi.Request] = None,
+    ) -> None:
+        if not self.log_requests:
+            return
+
+        max_length, skip_names, _ = self.metadata
+        headers = _extract_whitelisted_headers(request)
+        if self.log_requests_format == "json":
+            log_data = {
+                "rid": obj.rid,
+                "obj": _transform_data_for_logging(obj, max_length, skip_names),
+            }
+            if headers:
+                log_data["headers"] = headers
+            log_json(self.targets, "request.received", log_data)
+        else:
+            headers_str = f", headers={headers}" if headers else ""
+            self._log(
+                f"Receive: obj={_dataclass_to_string_truncated(obj, max_length, skip_names=skip_names)}{headers_str}"
+            )
+
+        # FIXME: This is a temporary fix to get the text from the input ids.
+        # We should remove this once we have a proper way.
+        if (
+            self.log_requests_level >= 2
+            and obj.text is None
+            and obj.input_ids is not None
+            and tokenizer is not None
+        ):
+            if obj.input_ids and isinstance(obj.input_ids[0], list):
+                # Prefill node warmup while PD disaggregated.
+                decoded = [
+                    tokenizer.decode(_input_ids, skip_special_tokens=False)
+                    for _input_ids in obj.input_ids
+                ]
+            else:
+                decoded = tokenizer.decode(obj.input_ids, skip_special_tokens=False)
+            obj.text = decoded
+
+    def log_openai_received_request(
+        self,
+        obj: Any,
+        request: Optional[fastapi.Request] = None,
+    ) -> None:
+        """Log the raw OpenAI request payload before request adaptation/tokenization."""
+        max_length, _, _ = self.metadata
+        max_length = max_length if max_length is not None else 2048
+        headers = _extract_whitelisted_headers(request)
+
+        if hasattr(obj, "model_dump"):
+            obj_to_log = obj.model_dump(exclude_none=True)
+        else:
+            obj_to_log = obj
+
+        if self.log_requests_format == "json":
+            log_data = {
+                "obj": _transform_data_for_logging(obj_to_log, max_length=max_length),
+            }
+            if headers:
+                log_data["headers"] = headers
+            log_json(self.targets, "request.received.openai", log_data)
+        else:
+            headers_str = f", headers={headers}" if headers else ""
+            self._log(
+                f"Receive OpenAI: obj={_dataclass_to_string_truncated(obj_to_log, max_length)}{headers_str}"
+            )
+
+    def log_finished_request(
+        self,
+        obj: Union[GenerateReqInput, EmbeddingReqInput],
+        out: Any,
+        request: Optional[fastapi.Request] = None,
+    ) -> None:
+        if not self.log_requests:
+            return
+
+        e2e_latency_ms = out["meta_info"].get("e2e_latency", 0) * 1000
+        if self.log_exceeded_ms > 0 and e2e_latency_ms < self.log_exceeded_ms:
+            return
+
+        max_length, skip_names, out_skip_names = self.metadata
+        headers = _extract_whitelisted_headers(request)
+        if self.log_requests_format == "json":
+            log_data = {
+                "rid": obj.rid,
+                "obj": _transform_data_for_logging(obj, max_length, skip_names),
+            }
+            if headers:
+                log_data["headers"] = headers
+            log_data["out"] = _transform_data_for_logging(
+                out, max_length, out_skip_names
+            )
+            log_json(self.targets, "request.finished", log_data)
+        else:
+            obj_str = _dataclass_to_string_truncated(
+                obj, max_length, skip_names=skip_names
+            )
+            out_str = f", out={_dataclass_to_string_truncated(out, max_length, skip_names=out_skip_names)}"
+            headers_str = f", headers={headers}" if headers else ""
+            self._log(f"Finish: obj={obj_str}{headers_str}{out_str}")
+
+    def _compute_metadata(
+        self,
+    ) -> Tuple[Optional[int], Optional[Set[str]], Optional[Set[str]]]:
+        max_length: Optional[int] = None
+        skip_names: Optional[Set[str]] = None
+        out_skip_names: Optional[Set[str]] = None
+        if self.log_requests:
+            if self.log_requests_level == 0:
+                max_length = 1 << 30
+                skip_names = {
+                    "text",
+                    "input_ids",
+                    "input_embeds",
+                    "image_data",
+                    "audio_data",
+                    "video_data",
+                    "mm_data_mooncake",
+                    "lora_path",
+                    "sampling_params",
+                }
+                out_skip_names = {"text", "output_ids", "embedding"}
+            elif self.log_requests_level == 1:
+                max_length = 1 << 30
+                skip_names = {
+                    "text",
+                    "input_ids",
+                    "input_embeds",
+                    "image_data",
+                    "audio_data",
+                    "video_data",
+                    "mm_data_mooncake",
+                    "lora_path",
+                }
+                out_skip_names = {"text", "output_ids", "embedding"}
+            elif self.log_requests_level == 2:
+                max_length = 2048
+            elif self.log_requests_level == 3:
+                max_length = 1 << 30
+            else:
+                raise ValueError(
+                    f"Invalid --log-requests-level: {self.log_requests_level=}"
+                )
+        return max_length, skip_names, out_skip_names
+
+    def _log(self, msg: str) -> None:
+        for target in self.targets:
+            target.info(msg)
+
+
+# TODO unify this w/ `_transform_data_for_logging` if we find performance enough
+def _dataclass_to_string_truncated(
+    data: Any, max_length: int = 2048, skip_names: Optional[Set[str]] = None
+) -> str:
+    if skip_names is None:
+        skip_names = set()
+    if isinstance(data, str):
+        if len(data) > max_length:
+            half_length = max_length // 2
+            return f"{repr(data[:half_length])} ... {repr(data[-half_length:])}"
+        else:
+            return f"{repr(data)}"
+    elif isinstance(data, (list, tuple)):
+        if len(data) > max_length:
+            half_length = max_length // 2
+            return str(data[:half_length]) + " ... " + str(data[-half_length:])
+        else:
+            return str(data)
+    elif isinstance(data, dict):
+        return (
+            "{"
+            + ", ".join(
+                f"'{k}': {_dataclass_to_string_truncated(v, max_length)}"
+                for k, v in data.items()
+                if k not in skip_names
+            )
+            + "}"
+        )
+    elif dataclasses.is_dataclass(data):
+        fields = dataclasses.fields(data)
+        return (
+            f"{data.__class__.__name__}("
+            + ", ".join(
+                f"{f.name}={_dataclass_to_string_truncated(getattr(data, f.name), max_length)}"
+                for f in fields
+                if f.name not in skip_names
+            )
+            + ")"
+        )
+    else:
+        return str(data)
+
+
+def _transform_data_for_logging(
+    data: Any, max_length: int = 2048, skip_names: Optional[Set[str]] = None
+) -> Any:
+    if skip_names is None:
+        skip_names = set()
+    if isinstance(data, str):
+        if len(data) > max_length:
+            half_length = max_length // 2
+            return data[:half_length] + "..." + data[-half_length:]
+        return data
+    elif isinstance(data, (list, tuple)):
+        if len(data) > max_length:
+            half_length = max_length // 2
+            return list(data[:half_length]) + ["..."] + list(data[-half_length:])
+        return [_transform_data_for_logging(v, max_length) for v in data]
+    elif isinstance(data, dict):
+        return {
+            k: _transform_data_for_logging(v, max_length)
+            for k, v in data.items()
+            if k not in skip_names
+        }
+    elif dataclasses.is_dataclass(data):
+        fields = dataclasses.fields(data)
+        return {
+            f.name: _transform_data_for_logging(getattr(data, f.name), max_length)
+            for f in fields
+            if f.name not in skip_names
+        }
+    elif isinstance(data, (int, float, bool, type(None))):
+        return data
+    else:
+        return str(data)

python/sglang/srt/utils/rpd_utils.py

@@ -0,0 +1,450 @@
+# https://raw.githubusercontent.com/ROCm/rocmProfileData/refs/heads/master/tools/rpd2tracing.py
+# commit 92d13a08328625463e9ba944cece82fc5eea36e6
+def rpd_to_chrome_trace(
+    input_rpd, output_json=None, start="0%", end="100%", format="object"
+):
+    import gzip
+    import sqlite3
+
+    if output_json is None:
+        import pathlib
+
+        output_json = pathlib.PurePath(input_rpd).with_suffix(".trace.json.gz")
+
+    connection = sqlite3.connect(input_rpd)
+
+    outfile = gzip.open(output_json, "wt", encoding="utf-8")
+
+    if format == "object":
+        outfile.write('{"traceEvents": ')
+
+    outfile.write("[ {}\n")
+
+    for row in connection.execute("select distinct gpuId from rocpd_op"):
+        try:
+            outfile.write(
+                ',{"name": "process_name", "ph": "M", "pid":"%s","args":{"name":"%s"}}\n'
+                % (row[0], "GPU" + str(row[0]))
+            )
+            outfile.write(
+                ',{"name": "process_sort_index", "ph": "M", "pid":"%s","args":{"sort_index":"%s"}}\n'
+                % (row[0], row[0] + 1000000)
+            )
+        except ValueError:
+            outfile.write("")
+
+    for row in connection.execute("select distinct pid, tid from rocpd_api"):
+        try:
+            outfile.write(
+                ',{"name":"thread_name","ph":"M","pid":"%s","tid":"%s","args":{"name":"%s"}}\n'
+                % (row[0], row[1], "Hip " + str(row[1]))
+            )
+            outfile.write(
+                ',{"name":"thread_sort_index","ph":"M","pid":"%s","tid":"%s","args":{"sort_index":"%s"}}\n'
+                % (row[0], row[1], row[1] * 2)
+            )
+        except ValueError:
+            outfile.write("")
+
+    try:
+        # FIXME - these aren't rendering correctly in chrome://tracing
+        for row in connection.execute("select distinct pid, tid from rocpd_hsaApi"):
+            try:
+                outfile.write(
+                    ',{"name":"thread_name","ph":"M","pid":"%s","tid":"%s","args":{"name":"%s"}}\n'
+                    % (row[0], row[1], "HSA " + str(row[1]))
+                )
+                outfile.write(
+                    ',{"name":"thread_sort_index","ph":"M","pid":"%s","tid":"%s","args":{"sort_index":"%s"}}\n'
+                    % (row[0], row[1], row[1] * 2 - 1)
+                )
+            except ValueError:
+                outfile.write("")
+    except:
+        pass
+
+    rangeStringApi = ""
+    rangeStringOp = ""
+    rangeStringMonitor = ""
+    min_time = connection.execute("select MIN(start) from rocpd_api;").fetchall()[0][0]
+    max_time = connection.execute("select MAX(end) from rocpd_api;").fetchall()[0][0]
+    if min_time is None:
+        raise Exception("Trace file is empty.")
+
+    print("Timestamps:")
+    print(f"\t    first: \t{min_time/1000} us")
+    print(f"\t     last: \t{max_time/1000} us")
+    print(f"\t duration: \t{(max_time-min_time) / 1000000000} seconds")
+
+    start_time = min_time / 1000
+    end_time = max_time / 1000
+
+    if start:
+        if "%" in start:
+            start_time = (
+                (max_time - min_time) * (int(start.replace("%", "")) / 100) + min_time
+            ) / 1000
+        else:
+            start_time = int(start)
+        rangeStringApi = "where rocpd_api.start/1000 >= %s" % (start_time)
+        rangeStringOp = "where rocpd_op.start/1000 >= %s" % (start_time)
+        rangeStringMonitor = "where start/1000 >= %s" % (start_time)
+    if end:
+        if "%" in end:
+            end_time = (
+                (max_time - min_time) * (int(end.replace("%", "")) / 100) + min_time
+            ) / 1000
+        else:
+            end_time = int(end)
+
+        rangeStringApi = (
+            rangeStringApi + " and rocpd_api.start/1000 <= %s" % (end_time)
+            if start != None
+            else "where rocpd_api.start/1000 <= %s" % (end_time)
+        )
+        rangeStringOp = (
+            rangeStringOp + " and rocpd_op.start/1000 <= %s" % (end_time)
+            if start != None
+            else "where rocpd_op.start/1000 <= %s" % (end_time)
+        )
+        rangeStringMonitor = (
+            rangeStringMonitor + " and start/1000 <= %s" % (end_time)
+            if start != None
+            else "where start/1000 <= %s" % (end_time)
+        )
+
+    print("\nFilter: %s" % (rangeStringApi))
+    print(f"Output duration: {(end_time-start_time)/1000000} seconds")
+
+    # Output Ops
+
+    for row in connection.execute(
+        "select A.string as optype, B.string as description, gpuId, queueId, rocpd_op.start/1000.0, (rocpd_op.end-rocpd_op.start) / 1000.0 from rocpd_op INNER JOIN rocpd_string A on A.id = rocpd_op.opType_id INNER Join rocpd_string B on B.id = rocpd_op.description_id %s"
+        % (rangeStringOp)
+    ):
+        try:
+            name = row[0] if len(row[1]) == 0 else row[1]
+            outfile.write(
+                ',{"pid":"%s","tid":"%s","name":"%s","ts":"%s","dur":"%s","ph":"X","args":{"desc":"%s"}}\n'
+                % (row[2], row[3], name, row[4], row[5], row[0])
+            )
+        except ValueError:
+            outfile.write("")
+
+    # Output Graph executions on GPU
+    try:
+        for row in connection.execute(
+            "select graphExec, gpuId, queueId, min(start)/1000.0, (max(end)-min(start))/1000.0, count(*) from rocpd_graphLaunchapi A join rocpd_api_ops B on B.api_id = A.api_ptr_id join rocpd_op C on C.id = B.op_id %s group by api_ptr_id"
+            % (rangeStringMonitor)
+        ):
+            try:
+                outfile.write(
+                    ',{"pid":"%s","tid":"%s","name":"%s","ts":"%s","dur":"%s","ph":"X","args":{"kernels":"%s"}}\n'
+                    % (row[1], row[2], f"Graph {row[0]}", row[3], row[4], row[5])
+                )
+            except ValueError:
+                outfile.write("")
+    except:
+        pass
+
+    # Output apis
+    for row in connection.execute(
+        "select A.string as apiName, B.string as args, pid, tid, rocpd_api.start/1000.0, (rocpd_api.end-rocpd_api.start) / 1000.0, (rocpd_api.end != rocpd_api.start) as has_duration from rocpd_api INNER JOIN rocpd_string A on A.id = rocpd_api.apiName_id INNER Join rocpd_string B on B.id = rocpd_api.args_id %s order by rocpd_api.id"
+        % (rangeStringApi)
+    ):
+        try:
+            if row[0] == "UserMarker":
+                if row[6] == 0:  # instantanuous "mark" messages
+                    outfile.write(
+                        ',{"pid":"%s","tid":"%s","name":"%s","ts":"%s","ph":"i","s":"p","args":{"desc":"%s"}}\n'
+                        % (
+                            row[2],
+                            row[3],
+                            row[1].replace('"', ""),
+                            row[4],
+                            row[1].replace('"', ""),
+                        )
+                    )
+                else:
+                    outfile.write(
+                        ',{"pid":"%s","tid":"%s","name":"%s","ts":"%s","dur":"%s","ph":"X","args":{"desc":"%s"}}\n'
+                        % (
+                            row[2],
+                            row[3],
+                            row[1].replace('"', ""),
+                            row[4],
+                            row[5],
+                            row[1].replace('"', ""),
+                        )
+                    )
+            else:
+                outfile.write(
+                    ',{"pid":"%s","tid":"%s","name":"%s","ts":"%s","dur":"%s","ph":"X","args":{"desc":"%s"}}\n'
+                    % (
+                        row[2],
+                        row[3],
+                        row[0],
+                        row[4],
+                        row[5],
+                        row[1].replace('"', "").replace("\t", ""),
+                    )
+                )
+        except ValueError:
+            outfile.write("")
+
+    # Output api->op linkage
+    for row in connection.execute(
+        "select rocpd_api_ops.id, pid, tid, gpuId, queueId, rocpd_api.end/1000.0 - 2, rocpd_op.start/1000.0 from rocpd_api_ops INNER JOIN rocpd_api on rocpd_api_ops.api_id = rocpd_api.id INNER JOIN rocpd_op on rocpd_api_ops.op_id = rocpd_op.id %s"
+        % (rangeStringApi)
+    ):
+        try:
+            fromtime = row[5] if row[5] < row[6] else row[6]
+            outfile.write(
+                ',{"pid":"%s","tid":"%s","cat":"api_op","name":"api_op","ts":"%s","id":"%s","ph":"s"}\n'
+                % (row[1], row[2], fromtime, row[0])
+            )
+            outfile.write(
+                ',{"pid":"%s","tid":"%s","cat":"api_op","name":"api_op","ts":"%s","id":"%s","ph":"f", "bp":"e"}\n'
+                % (row[3], row[4], row[6], row[0])
+            )
+        except ValueError:
+            outfile.write("")
+
+    try:
+        for row in connection.execute(
+            "select A.string as apiName, B.string as args, pid, tid, rocpd_hsaApi.start/1000.0, (rocpd_hsaApi.end-rocpd_hsaApi.start) / 1000.0 from rocpd_hsaApi INNER JOIN rocpd_string A on A.id = rocpd_hsaApi.apiName_id INNER Join rocpd_string B on B.id = rocpd_hsaApi.args_id %s order by rocpd_hsaApi.id"
+            % (rangeStringApi)
+        ):
+            try:
+                outfile.write(
+                    ',{"pid":"%s","tid":"%s","name":"%s","ts":"%s","dur":"%s","ph":"X","args":{"desc":"%s"}}\n'
+                    % (
+                        row[2],
+                        row[3] + 1,
+                        row[0],
+                        row[4],
+                        row[5],
+                        row[1].replace('"', ""),
+                    )
+                )
+            except ValueError:
+                outfile.write("")
+    except:
+        pass
+
+    #
+    # Counters
+    #
+
+    # Counters should extend to the last event in the trace.  This means they need to have a value at Tend.
+    # Figure out when that is
+
+    T_end = 0
+    for row in connection.execute(
+        "SELECT max(end)/1000 from (SELECT end from rocpd_api UNION ALL SELECT end from rocpd_op)"
+    ):
+        T_end = int(row[0])
+    if end:
+        T_end = end_time
+
+    # Loop over GPU for per-gpu counters
+    gpuIdsPresent = []
+    for row in connection.execute("SELECT DISTINCT gpuId FROM rocpd_op"):
+        gpuIdsPresent.append(row[0])
+
+    for gpuId in gpuIdsPresent:
+        # print(f"Creating counters for: {gpuId}")
+
+        # Create the queue depth counter
+        depth = 0
+        idle = 1
+        for row in connection.execute(
+            'select * from (select rocpd_api.start/1000.0 as ts, "1" from rocpd_api_ops INNER JOIN rocpd_api on rocpd_api_ops.api_id = rocpd_api.id INNER JOIN rocpd_op on rocpd_api_ops.op_id = rocpd_op.id AND rocpd_op.gpuId = %s %s UNION ALL select rocpd_op.end/1000.0, "-1" from rocpd_api_ops INNER JOIN rocpd_api on rocpd_api_ops.api_id = rocpd_api.id INNER JOIN rocpd_op on rocpd_api_ops.op_id = rocpd_op.id AND rocpd_op.gpuId = %s %s) order by ts'
+            % (gpuId, rangeStringOp, gpuId, rangeStringOp)
+        ):
+            try:
+                if idle and int(row[1]) > 0:
+                    idle = 0
+                    outfile.write(
+                        ',{"pid":"%s","name":"Idle","ph":"C","ts":%s,"args":{"idle":%s}}\n'
+                        % (gpuId, row[0], idle)
+                    )
+                if depth == 1 and int(row[1]) < 0:
+                    idle = 1
+                    outfile.write(
+                        ',{"pid":"%s","name":"Idle","ph":"C","ts":%s,"args":{"idle":%s}}\n'
+                        % (gpuId, row[0], idle)
+                    )
+                depth = depth + int(row[1])
+                outfile.write(
+                    ',{"pid":"%s","name":"QueueDepth","ph":"C","ts":%s,"args":{"depth":%s}}\n'
+                    % (gpuId, row[0], depth)
+                )
+            except ValueError:
+                outfile.write("")
+        if T_end > 0:
+            outfile.write(
+                ',{"pid":"%s","name":"Idle","ph":"C","ts":%s,"args":{"idle":%s}}\n'
+                % (gpuId, T_end, idle)
+            )
+            outfile.write(
+                ',{"pid":"%s","name":"QueueDepth","ph":"C","ts":%s,"args":{"depth":%s}}\n'
+                % (gpuId, T_end, depth)
+            )
+
+    # Create SMI counters
+    try:
+        for row in connection.execute(
+            "select deviceId, monitorType, start/1000.0, value from rocpd_monitor %s"
+            % (rangeStringMonitor)
+        ):
+            outfile.write(
+                ',{"pid":"%s","name":"%s","ph":"C","ts":%s,"args":{"%s":%s}}\n'
+                % (row[0], row[1], row[2], row[1], row[3])
+            )
+        # Output the endpoints of the last range
+        for row in connection.execute(
+            "select distinct deviceId, monitorType, max(end)/1000.0, value from rocpd_monitor %s group by deviceId, monitorType"
+            % (rangeStringMonitor)
+        ):
+            outfile.write(
+                ',{"pid":"%s","name":"%s","ph":"C","ts":%s,"args":{"%s":%s}}\n'
+                % (row[0], row[1], row[2], row[1], row[3])
+            )
+    except:
+        print("Did not find SMI data")
+
+    # Create the (global) memory counter
+    """
+    sizes = {}    # address -> size
+    totalSize = 0
+    exp = re.compile("^ptr\((.*)\)\s+size\((.*)\)$")
+    exp2 = re.compile("^ptr\((.*)\)$")
+    for row in connection.execute("SELECT rocpd_api.end/1000.0 as ts, B.string, '1'  FROM rocpd_api INNER JOIN rocpd_string A ON A.id=rocpd_api.apiName_id INNER JOIN rocpd_string B ON B.id=rocpd_api.args_id WHERE A.string='hipFree' UNION ALL SELECT rocpd_api.start/1000.0, B.string, '0' FROM rocpd_api INNER JOIN rocpd_string A ON A.id=rocpd_api.apiName_id INNER JOIN rocpd_string B ON B.id=rocpd_api.args_id WHERE A.string='hipMalloc' ORDER BY ts asc"):
+        try:
+            if row[2] == '0':  #malloc
+                m = exp.match(row[1])
+                if m:
+                    size = int(m.group(2), 16)
+                    totalSize = totalSize + size
+                    sizes[m.group(1)] = size
+                    outfile.write(',{"pid":"0","name":"Allocated Memory","ph":"C","ts":%s,"args":{"depth":%s}}\n'%(row[0],totalSize))
+            else:              #free
+                m = exp2.match(row[1])
+                if m:
+                    try:    # Sometimes free addresses are not valid or listed
+                        size = sizes[m.group(1)]
+                        sizes[m.group(1)] = 0
+                        totalSize = totalSize - size;
+                        outfile.write(',{"pid":"0","name":"Allocated Memory","ph":"C","ts":%s,"args":{"depth":%s}}\n'%(row[0],totalSize))
+                    except KeyError:
+                        pass
+        except ValueError:
+            outfile.write("")
+    if T_end > 0:
+        outfile.write(',{"pid":"0","name":"Allocated Memory","ph":"C","ts":%s,"args":{"depth":%s}}\n'%(T_end,totalSize))
+    """
+
+    # Create "faux calling stack frame" on gpu ops traceS
+    stacks = {}  # Call stacks built from UserMarker entres.     Key is 'pid,tid'
+    currentFrame = {}  # "Current GPU frame" (id, name, start, end).    Key is 'pid,tid'
+
+    class GpuFrame:
+        def __init__(self):
+            self.id = 0
+            self.name = ""
+            self.start = 0
+            self.end = 0
+            self.gpus = []
+            self.totalOps = 0
+
+    # FIXME: include 'start' (in ns) so we can ORDER BY it and break ties?
+    for row in connection.execute(
+        "SELECT '0', start/1000.0, pid, tid, B.string as label, '','','', '' from rocpd_api INNER JOIN rocpd_string A on A.id = rocpd_api.apiName_id AND A.string = 'UserMarker' INNER JOIN rocpd_string B on B.id = rocpd_api.args_id AND rocpd_api.start/1000.0 != rocpd_api.end/1000.0 %s UNION ALL SELECT '1', end/1000.0, pid, tid, B.string as label, '','','', '' from rocpd_api INNER JOIN rocpd_string A on A.id = rocpd_api.apiName_id AND A.string = 'UserMarker' INNER JOIN rocpd_string B on B.id = rocpd_api.args_id AND rocpd_api.start/1000.0 != rocpd_api.end/1000.0 %s UNION ALL SELECT '2', rocpd_api.start/1000.0, pid, tid, '' as label, gpuId, queueId, rocpd_op.start/1000.0, rocpd_op.end/1000.0 from rocpd_api_ops INNER JOIN rocpd_api ON rocpd_api_ops.api_id = rocpd_api.id INNER JOIN rocpd_op ON rocpd_api_ops.op_id = rocpd_op.id %s ORDER BY start/1000.0 asc"
+        % (rangeStringApi, rangeStringApi, rangeStringApi)
+    ):
+        try:
+            key = (row[2], row[3])  # Key is 'pid,tid'
+            if row[0] == "0":  # Frame start
+                if key not in stacks:
+                    stacks[key] = []
+                stacks[key].append((row[1], row[4]))
+                # print(f"0: new api frame: pid_tid={key} -> stack={stacks}")
+
+            elif row[0] == "1":  # Frame end
+                stacks[key].pop()
+                # print(f"1: end api frame: pid_tid={key} -> stack={stacks}")
+
+            elif row[0] == "2":  # API + Op
+                if key in stacks and len(stacks[key]) > 0:
+                    frame = stacks[key][-1]
+                    # print(f"2: Op on {frame} ({len(stacks[key])})")
+                    gpuFrame = None
+                    if key not in currentFrame:  # First op under the current api frame
+                        gpuFrame = GpuFrame()
+                        gpuFrame.id = frame[0]
+                        gpuFrame.name = frame[1]
+                        gpuFrame.start = row[7]
+                        gpuFrame.end = row[8]
+                        gpuFrame.gpus.append((row[5], row[6]))
+                        gpuFrame.totalOps = 1
+                        # print(f"2a: new frame: {gpuFrame.gpus} {gpuFrame.start} {gpuFrame.end} {gpuFrame.end - gpuFrame.start}")
+                    else:
+                        gpuFrame = currentFrame[key]
+                        # Another op under the same frame -> union them (but only if they are butt together)
+                        if (
+                            gpuFrame.id == frame[0]
+                            and gpuFrame.name == frame[1]
+                            and (
+                                abs(row[7] - gpuFrame.end) < 200
+                                or abs(gpuFrame.start - row[8]) < 200
+                            )
+                        ):
+                            if row[7] < gpuFrame.start:
+                                gpuFrame.start = row[7]
+                            if row[8] > gpuFrame.end:
+                                gpuFrame.end = row[8]
+                            if (row[5], row[6]) not in gpuFrame.gpus:
+                                gpuFrame.gpus.append((row[5], row[6]))
+                            gpuFrame.totalOps = gpuFrame.totalOps + 1
+                            # print(f"2c: union frame: {gpuFrame.gpus} {gpuFrame.start} {gpuFrame.end} {gpuFrame.end - gpuFrame.start}")
+
+                        else:  # This is a new frame - dump the last and make new
+                            gpuFrame = currentFrame[key]
+                            for dest in gpuFrame.gpus:
+                                # print(f"2: OUTPUT: dest={dest} time={gpuFrame.start} -> {gpuFrame.end} Duration={gpuFrame.end - gpuFrame.start} TotalOps={gpuFrame.totalOps}")
+                                outfile.write(
+                                    ',{"pid":"%s","tid":"%s","name":"%s","ts":"%s","dur":"%s","ph":"X","args":{"desc":"%s"}}\n'
+                                    % (
+                                        dest[0],
+                                        dest[1],
+                                        gpuFrame.name.replace('"', ""),
+                                        gpuFrame.start - 1,
+                                        gpuFrame.end - gpuFrame.start + 1,
+                                        f"UserMarker frame: {gpuFrame.totalOps} ops",
+                                    )
+                                )
+                            currentFrame.pop(key)
+
+                            # make the first op under the new frame
+                            gpuFrame = GpuFrame()
+                            gpuFrame.id = frame[0]
+                            gpuFrame.name = frame[1]
+                            gpuFrame.start = row[7]
+                            gpuFrame.end = row[8]
+                            gpuFrame.gpus.append((row[5], row[6]))
+                            gpuFrame.totalOps = 1
+                            # print(f"2b: new frame: {gpuFrame.gpus} {gpuFrame.start} {gpuFrame.end} {gpuFrame.end - gpuFrame.start}")
+
+                    currentFrame[key] = gpuFrame
+
+        except ValueError:
+            outfile.write("")
+
+    outfile.write("]\n")
+
+    if format == "object":
+        outfile.write("} \n")
+
+    outfile.close()
+    connection.close()

python/sglang/srt/utils/runai_utils.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/model_executor/model_loader/runai_utils.py
+
+import hashlib
+import logging
+import os
+from pathlib import Path
+
+from sglang.srt.environ import envs
+
+logger = logging.getLogger(__name__)
+
+SUPPORTED_SCHEMES = ["s3://", "gs://", "az://"]
+
+# Design Pattern: Single Metadata Download Before Process Launch
+
+#   1. Engine entrypoint (engine.py) or server arguments post init  (server_args.py):
+#     - Downloads config/tokenizer metadata ONCE before launching subprocesses
+#     - This happens in the main process, avoiding multi-process coordination
+#
+#   2. ModelConfig/HF Utils (model_config.py, hf_transformers_utils.py):
+#     - Use ObjectStorageModel.get_path() to retrieve the cached local path
+#     - NO re-download - just path resolution
+#
+#   3. RunaiModelStreamerLoader (loader.py):
+#     - Calls list_safetensors() which operates directly on the object storage URI
+#     - Streams weights lazily during model loading
+
+#   This avoids file locks, race conditions, and duplicate downloads
+
+
+def list_safetensors(path: str = "") -> list[str]:
+    """
+    List full file names from object path and filter by allow pattern.
+
+    Args:
+        path: The object storage path to list from.
+
+    Returns:
+        list[str]: List of full object storage paths allowed by the pattern
+    """
+    from runai_model_streamer import list_safetensors as runai_list_safetensors
+
+    return runai_list_safetensors(path)
+
+
+def is_runai_obj_uri(model_or_path: str | Path) -> bool:
+    # Cast to str to handle pathlib.Path inputs which lack string methods (like .lower)
+    return str(model_or_path).lower().startswith(tuple(SUPPORTED_SCHEMES))
+
+
+class ObjectStorageModel:
+    """
+    Model loader that uses Runai Model Streamer to load a model.
+
+      Supports object storage (S3, GCS) with lazy weight streaming.
+
+      Configuration (via load_config.model_loader_extra_config):
+          - distributed (bool): Enable distributed streaming
+          - concurrency (int): Number of concurrent downloads
+          - memory_limit (int): Memory limit for streaming buffer
+
+      Note: Metadata files must be pre-downloaded via
+      ObjectStorageModel.download_and_get_path() before instantiation.
+
+    Attributes:
+        dir: The temporary created directory.
+    """
+
+    def __init__(self, url: str) -> None:
+        self.dir = ObjectStorageModel.get_path(url)
+
+        from runai_model_streamer import ObjectStorageModel as RunaiObjectStorageModel
+
+        self._runai_obj = RunaiObjectStorageModel(model_path=url, dst=self.dir)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return self._runai_obj.__exit__(exc_type, exc_val, exc_tb)
+
+    def pull_files(
+        self,
+        allow_pattern: list[str] | None = None,
+        ignore_pattern: list[str] | None = None,
+    ) -> None:
+        """Pull files from object storage into the local cache directory.
+
+        Args:
+            allow_pattern: File patterns to include (e.g. ["*.json"]).
+            ignore_pattern: File patterns to exclude.
+        """
+        self._runai_obj.pull_files(allow_pattern, ignore_pattern)
+
+    @classmethod
+    def download_and_get_path(cls, model_path: str) -> str:
+        """
+        Downloads the model metadata (excluding heavy weights) and returns
+        the local directory path. Safe for concurrent usage by multiple processes
+        """
+        with cls(url=model_path) as downloader:
+            downloader.pull_files(
+                ignore_pattern=[
+                    "*.pt",
+                    "*.safetensors",
+                    "*.bin",
+                    "*.tensors",
+                    "*.pth",
+                ],
+            )
+            cache_dir = downloader.dir
+            logger.info(f"Runai Model : {cache_dir}, metadata ready.")
+        return cache_dir
+
+    @classmethod
+    def get_path(cls, model_path: str) -> str:
+        """
+        Returns the local directory path.
+        """
+        model_hash = hashlib.sha256(str(model_path).encode()).hexdigest()[:16]
+        base_dir = envs.SGLANG_CACHE_DIR.get()
+
+        # Ensure base cache dir exists
+        os.makedirs(os.path.join(base_dir, "model_streamer"), exist_ok=True)
+
+        return os.path.join(
+            base_dir,
+            "model_streamer",
+            model_hash,
+        )

python/sglang/srt/utils/scheduler_status_logger.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+import time
+from typing import TYPE_CHECKING, List, Optional
+
+import torch.distributed as dist
+
+from sglang.srt.environ import envs
+from sglang.srt.utils.log_utils import create_log_targets, log_json
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+
+
+class SchedulerStatusLogger:
+    def __init__(self, targets: List[str], dump_interval: float):
+        self.loggers = create_log_targets(targets=targets, name_prefix=__name__)
+        self.dump_interval = dump_interval
+        self.last_dump_time = 0.0
+        self.rank = dist.get_rank() if dist.is_initialized() else 0
+
+    @staticmethod
+    def maybe_create(enable_metrics: bool) -> Optional[SchedulerStatusLogger]:
+        target = envs.SGLANG_LOG_SCHEDULER_STATUS_TARGET.get()
+        if not target:
+            return None
+
+        if not enable_metrics:
+            raise ValueError(
+                "SGLANG_LOG_SCHEDULER_STATUS_TARGET is set but --enable-metrics "
+                "is not active. Status dumps require --enable-metrics to work."
+            )
+
+        return SchedulerStatusLogger(
+            targets=[t.strip() for t in target.split(",") if t.strip()],
+            dump_interval=envs.SGLANG_LOG_SCHEDULER_STATUS_INTERVAL.get(),
+        )
+
+    def maybe_dump(
+        self, running_batch: ScheduleBatch, waiting_queue: List[Req]
+    ) -> None:
+        now = time.time()
+        if now - self.last_dump_time < self.dump_interval:
+            return
+
+        self.last_dump_time = now
+        log_json(
+            self.loggers,
+            "scheduler.status",
+            {
+                "rank": self.rank,
+                "running_rids": [r.rid for r in running_batch.reqs],
+                "queued_rids": [r.rid for r in waiting_queue],
+            },
+        )

python/sglang/srt/utils/slow_rank_detector.py

@@ -0,0 +1,71 @@
+import logging
+from typing import Any, Dict, List
+
+import torch
+import torch.distributed as dist
+import triton
+
+logger = logging.getLogger(__name__)
+
+
+def execute():
+    if dist.get_rank() == 0:
+        logger.info(f"[slow_rank_detector] Start benchmarking...")
+
+    local_metrics = {
+        bench_name: _compute_local_metric(bench_name) for bench_name in _BENCH_NAMES
+    }
+
+    all_metrics = [None for _ in range(dist.get_world_size())]
+    dist.gather_object(local_metrics, all_metrics if dist.get_rank() == 0 else None)
+
+    if dist.get_rank() == 0:
+        _analyze_metrics(all_metrics)
+
+
+class _GemmExecutor:
+    def __init__(self):
+        self.lhs = torch.randn((8192, 8192), dtype=torch.bfloat16, device="cuda")
+        self.rhs = torch.randn((8192, 8192), dtype=torch.bfloat16, device="cuda")
+
+    def __call__(self):
+        self.lhs @ self.rhs
+
+
+class _ElementwiseExecutor:
+    def __init__(self):
+        self.value = torch.randint(
+            0, 10000, (128 * 1024**2,), dtype=torch.int32, device="cuda"
+        )
+
+    def __call__(self):
+        self.value += 1
+
+
+_EXECUTOR_CLS_OF_BENCH = {
+    "gemm": _GemmExecutor,
+    "elementwise": _ElementwiseExecutor,
+}
+
+_BENCH_NAMES = list(_EXECUTOR_CLS_OF_BENCH.keys())
+
+
+def _compute_local_metric(bench_name):
+    executor = _EXECUTOR_CLS_OF_BENCH[bench_name]()
+    ms = triton.testing.do_bench_cudagraph(executor, return_mode="mean", rep=20)
+    return ms
+
+
+def _analyze_metrics(all_metrics: List[Dict[str, Any]]):
+    for bench_name in _BENCH_NAMES:
+        time_of_rank = torch.tensor([m[bench_name] for m in all_metrics])
+        speed_of_rank = 1 / time_of_rank
+        rel_speed_of_rank = speed_of_rank / speed_of_rank.max()
+        slowest_rel_speed = rel_speed_of_rank.min().item()
+        logger.info(
+            f"[slow_rank_detector] {bench_name=} {slowest_rel_speed=} {rel_speed_of_rank=} {time_of_rank=}"
+        )
+        if slowest_rel_speed < 0.9:
+            logger.warning(
+                "[slow_rank_detector] Some ranks are too slow compared with others"
+            )

python/sglang/srt/utils/stale_shm_cleanup.py

@@ -0,0 +1,123 @@
+"""Self-heal for leaked POSIX shared-memory segments in CI.
+
+SGLang processes are torn down with SIGKILL (kill_process_tree, PDEATHSIG),
+which skips every Python-level unlink path, so /dev/shm segments accumulate
+until the tmpfs is full and the next scheduler init dies with SIGBUS.
+
+Segments created through make_shm_name() embed the creator pid, which lets a
+later server startup safely unlink segments whose creator is gone. The sweep
+only runs in CI (single-tenant runner containers); on shared dev machines a
+pid check against another user's process is not authoritative, so we skip.
+"""
+
+import logging
+import os
+import uuid
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+_SHM_DIR = Path("/dev/shm")
+_SGL_SHM_PREFIX = "sgl_shm"
+
+
+def make_shm_name(kind: str) -> str:
+    """Name a shared-memory segment so cleanup_stale_shm can identify and
+    reclaim it after its creator process dies: sgl_shm_<kind>_<pid>_<rand>."""
+    return f"{_SGL_SHM_PREFIX}_{kind}_{os.getpid()}_{uuid.uuid4().hex[:8]}"
+
+
+def _creator_pid(filename: str) -> int | None:
+    pid = None
+    if filename.startswith(f"{_SGL_SHM_PREFIX}_"):
+        # sgl_shm_<kind>_<pid>_<rand>
+        parts = filename.split("_")
+        if len(parts) >= 4:
+            try:
+                pid = int(parts[-2])
+            except ValueError:
+                return None
+    elif filename.startswith("multi_tokenizer_args_"):
+        try:
+            pid = int(filename.rsplit("_", 1)[-1])
+        except ValueError:
+            return None
+    # os.kill(0, ...) / os.kill(-1, ...) probe process groups, not a process.
+    if pid is not None and pid <= 0:
+        return None
+    return pid
+
+
+def _pid_alive(pid: int) -> bool:
+    try:
+        os.kill(pid, 0)
+        return True
+    except ProcessLookupError:
+        return False
+    except PermissionError:
+        # Process exists but is owned by someone else.
+        return True
+
+
+def cleanup_stale_shm() -> None:
+    """Unlink shared-memory segments whose creator process is dead.
+
+    CI-only: gated on SGLANG_IS_IN_CI because the pid-liveness check is only
+    trustworthy when the container runs one job at a time. Best-effort: never
+    raises, since a failed sweep must not block server startup.
+    """
+    try:
+        _cleanup_stale_shm_impl()
+    except Exception:
+        logger.warning(
+            "cleanup_stale_shm: sweep failed, continuing startup", exc_info=True
+        )
+
+
+def _is_in_ci() -> bool:
+    # Read the env var directly (same semantics as sglang.utils.is_in_ci) so
+    # this module stays import-free and runnable by path from CI scripts
+    # before sglang is installed.
+    return os.environ.get("SGLANG_IS_IN_CI", "false").lower() in ("true", "1")
+
+
+def _cleanup_stale_shm_impl() -> None:
+    if not _is_in_ci():
+        return
+    if not _SHM_DIR.is_dir():
+        return
+
+    removed = 0
+    freed_bytes = 0
+    try:
+        entries = list(_SHM_DIR.iterdir())
+    except OSError as e:
+        logger.warning("cleanup_stale_shm: cannot list %s, skipping: %s", _SHM_DIR, e)
+        return
+    for entry in entries:
+        pid = _creator_pid(entry.name)
+        if pid is None or pid == os.getpid() or _pid_alive(pid):
+            # A recycled pid reads as alive, so pid-reuse degrades to
+            # under-collection (segment leaks), never to deleting a live
+            # segment. Keep that bias when changing this check.
+            continue
+        try:
+            size = entry.stat().st_size
+            entry.unlink()
+            removed += 1
+            freed_bytes += size
+        except FileNotFoundError:
+            pass  # raced with another cleaner
+        except OSError as e:
+            logger.warning("cleanup_stale_shm: failed to remove %s: %s", entry.name, e)
+    if removed:
+        logger.info(
+            "cleanup_stale_shm: removed %d stale segment(s), freed %.1f MiB",
+            removed,
+            freed_bytes / (1 << 20),
+        )
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    cleanup_stale_shm()

python/sglang/srt/utils/tensor_bridge.py

@@ -0,0 +1,228 @@
+# Copied and adapted from: https://github.com/vllm-project/vllm-metal
+# SPDX-License-Identifier: Apache-2.0
+"""Tensor bridge between MLX and PyTorch.
+
+Provides zero-copy conversion when possible using Apple Silicon's unified memory.
+"""
+
+from __future__ import annotations
+
+import logging
+from functools import lru_cache
+from typing import TYPE_CHECKING, Literal
+
+import torch
+
+from sglang.srt.environ import envs
+
+if TYPE_CHECKING:
+    import mlx.core as mx
+
+logger = logging.getLogger(__name__)
+
+_MLX_AVAILABLE: bool = False
+try:
+    import mlx.core as mx  # noqa: F811
+
+    _MLX_AVAILABLE = True
+except ImportError:
+    pass
+
+
+def is_mlx_available() -> bool:
+    """Return True when the ``mlx`` package can be imported."""
+    return _MLX_AVAILABLE
+
+
+@lru_cache(maxsize=1)
+def use_mlx() -> bool:
+    """Return True when the user opted-in via ``SGLANG_USE_MLX=1`` **and** MLX is importable."""
+    return bool(envs.SGLANG_USE_MLX.get()) and _MLX_AVAILABLE
+
+
+# MPS has a 4GB (2^32 bytes) limit for MPSTemporaryNDArray allocations.
+# Metal may allocate multiple temporary buffers internally, so we use a
+# conservative threshold of 1GB to avoid hitting the limit.
+# See: https://github.com/anthropics/vllm-metal/issues/43
+_MPS_SAFE_SIZE_BYTES = 1 << 30  # 1GB
+
+# MLX to PyTorch dtype mapping
+# TODO(perf): float64 is CPU-only in MLX (see ml-explore/mlx#1843).
+# When the target device is GPU/MPS we should auto-downcast float64 → float32
+# to avoid a runtime error; when the target is CPU we can keep float64.
+# For now float64 is omitted from the mapping so it hits the ValueError
+# fallback in mlx_to_torch().
+MLX_TO_TORCH_DTYPE = (
+    {
+        mx.float32: torch.float32,
+        mx.float16: torch.float16,
+        mx.bfloat16: torch.bfloat16,
+        mx.int32: torch.int32,
+        mx.int64: torch.int64,
+        mx.int16: torch.int16,
+        mx.int8: torch.int8,
+        mx.uint8: torch.uint8,
+        mx.bool_: torch.bool,
+    }
+    if _MLX_AVAILABLE
+    else {}
+)
+
+# PyTorch to MLX dtype mapping
+TORCH_TO_MLX_DTYPE = {v: k for k, v in MLX_TO_TORCH_DTYPE.items()}
+
+
+def get_torch_device() -> torch.device:
+    """Get the PyTorch device for Metal/MPS.
+
+    Returns:
+        torch.device for MPS if available, else CPU
+    """
+    if torch.backends.mps.is_available():
+        return torch.device("mps")
+    return torch.device("cpu")
+
+
+def _get_tensor_size_bytes(array: mx.array) -> int:
+    """Calculate the size of an MLX array in bytes.
+
+    Args:
+        array: MLX array
+
+    Returns:
+        Size in bytes
+    """
+    return array.size * array.dtype.size
+
+
+def _is_safe_for_mps(array: mx.array) -> bool:
+    """Check if an array is safe to transfer to MPS without hitting size limits.
+
+    MPS has a 4GB limit for MPSTemporaryNDArray, but Metal may allocate
+    multiple temporary buffers internally. We use a conservative threshold.
+
+    Args:
+        array: MLX array to check
+
+    Returns:
+        True if safe to transfer to MPS, False if should stay on CPU
+    """
+    return _get_tensor_size_bytes(array) < _MPS_SAFE_SIZE_BYTES
+
+
+def torch_to_mlx(tensor: torch.Tensor) -> mx.array:
+    """Convert PyTorch tensor to MLX array.
+
+    Uses numpy as an intermediate to enable zero-copy on unified memory.
+
+    Args:
+        tensor: PyTorch tensor (can be on any device)
+
+    Returns:
+        MLX array with the same data
+    """
+    # Move to CPU if on MPS for numpy conversion
+    if tensor.device.type != "cpu":
+        tensor = tensor.cpu()
+
+    tensor = tensor.detach()
+
+    # Note: numpy does not support bfloat16.
+    if tensor.dtype == torch.bfloat16:
+        return mx.array(tensor)
+
+    return mx.array(tensor.numpy())
+
+
+# TODO(perf): accept a list/batch of arrays and convert them in one pass
+# to reduce the Python ↔ MLX round-trip overhead.
+def mlx_to_torch(
+    array: mx.array,
+    device: torch.device | Literal["mps", "cpu"] | None = None,
+    already_contiguous: bool = False,
+) -> torch.Tensor:
+    """Convert MLX array to PyTorch tensor.
+
+    Uses numpy as an intermediate to enable zero-copy on unified memory.
+
+    Args:
+        array: MLX array
+        device: Target PyTorch device (default: MPS if available)
+        already_contiguous: Skip contiguity check if array is known contiguous
+
+    Returns:
+        PyTorch tensor with the same data
+    """
+    if device is None:
+        device = get_torch_device()
+    elif isinstance(device, str):
+        device = torch.device(device)
+
+    # Use memoryview for zero-copy conversion (bypasses numpy for bfloat16)
+    # reference: https://github.com/ml-explore/mlx/issues/403
+    torch_dtype = MLX_TO_TORCH_DTYPE.get(array.dtype)
+    if torch_dtype is not None:
+        if already_contiguous:
+            # Fast path: skip contiguity check, single eval
+            mx.eval(array)
+            buffer = memoryview(array)
+        else:
+            # MLX views / non-contiguous arrays expose a non-contiguous buffer (or
+            # sometimes no usable buffer), which `torch.frombuffer` can't consume.
+            # Make contiguous first, then eval once
+            array = mx.contiguous(array)
+            mx.eval(array)
+            buffer = memoryview(array)
+
+        tensor = torch.frombuffer(buffer, dtype=torch_dtype).reshape(array.shape)
+    else:
+        # Fallback to numpy path for unsupported dtypes
+        raise ValueError(f"Unsupported MLX dtype: {array.dtype}")
+
+    # Move to target device, but check for MPS size limits first
+    if device.type == "mps":
+        if _is_safe_for_mps(array):
+            tensor = tensor.to(device)
+        else:
+            # Large tensor - keep on CPU to avoid MPS 4GB limit crash
+            # See: https://github.com/anthropics/vllm-metal/issues/43
+            logger.debug(
+                "Tensor too large for MPS (%d bytes > %d limit), keeping on CPU",
+                _get_tensor_size_bytes(array),
+                _MPS_SAFE_SIZE_BYTES,
+            )
+    elif device.type != "cpu":
+        tensor = tensor.to(device)
+
+    return tensor
+
+
+def sync_mlx() -> None:
+    """Synchronize MLX operations.
+
+    Call this before converting MLX arrays to ensure all operations complete.
+    """
+    # Prefer an explicit MLX barrier when available; otherwise force evaluation.
+    # `mx.eval([])` is a no-op, so we evaluate a tiny scalar as a safe fallback.
+    try:
+        mx.synchronize()
+    except (AttributeError, TypeError):
+        mx.eval(mx.array(0, dtype=mx.int32))
+
+
+def sync_torch() -> None:
+    """Synchronize PyTorch MPS operations.
+
+    Call this before converting PyTorch tensors to ensure all operations complete.
+    """
+    if torch.backends.mps.is_available():
+        torch.mps.synchronize()
+
+
+__all__ = [
+    "is_mlx_available",
+    "use_mlx",
+    "mlx_to_torch",
+    "torch_to_mlx",
+    "get_torch_device",
+]

python/sglang/srt/utils/torch_memory_saver_adapter.py

@@ -0,0 +1,112 @@
+import logging
+from abc import ABC
+from contextlib import contextmanager
+
+try:
+    import torch_memory_saver
+
+    _memory_saver = torch_memory_saver.torch_memory_saver
+    import_error = None
+except ImportError as e:
+    import_error = e
+    pass
+
+logger = logging.getLogger(__name__)
+
+
+class TorchMemorySaverAdapter(ABC):
+    @staticmethod
+    def create(enable: bool):
+        if enable and import_error is not None:
+            logger.warning(
+                "enable_memory_saver is enabled, but "
+                "torch-memory-saver is not installed. Please install it "
+                "via `pip3 install torch-memory-saver`. "
+            )
+            raise import_error
+        return (
+            _TorchMemorySaverAdapterReal() if enable else _TorchMemorySaverAdapterNoop()
+        )
+
+    def check_validity(self, caller_name):
+        if not self.enabled:
+            logger.warning(
+                f"`{caller_name}` will not save memory because torch_memory_saver is not enabled. "
+                f"Potential causes: `enable_memory_saver` is false, or torch_memory_saver has installation issues."
+            )
+
+    def configure_subprocess(self):
+        raise NotImplementedError
+
+    def region(self, tag: str, enable_cpu_backup: bool = False):
+        raise NotImplementedError
+
+    def cuda_graph(self, **kwargs):
+        raise NotImplementedError
+
+    def disable(self):
+        raise NotImplementedError
+
+    def pause(self, tag: str):
+        raise NotImplementedError
+
+    def resume(self, tag: str):
+        raise NotImplementedError
+
+    @property
+    def enabled(self):
+        raise NotImplementedError
+
+
+class _TorchMemorySaverAdapterReal(TorchMemorySaverAdapter):
+    """Adapter for TorchMemorySaver with tag-based control"""
+
+    def configure_subprocess(self):
+        return torch_memory_saver.configure_subprocess()
+
+    def region(self, tag: str, enable_cpu_backup: bool = False):
+        return _memory_saver.region(tag=tag, enable_cpu_backup=enable_cpu_backup)
+
+    def cuda_graph(self, **kwargs):
+        return _memory_saver.cuda_graph(**kwargs)
+
+    def disable(self):
+        return _memory_saver.disable()
+
+    def pause(self, tag: str):
+        return _memory_saver.pause(tag=tag)
+
+    def resume(self, tag: str):
+        return _memory_saver.resume(tag=tag)
+
+    @property
+    def enabled(self):
+        return _memory_saver is not None and _memory_saver.enabled
+
+
+class _TorchMemorySaverAdapterNoop(TorchMemorySaverAdapter):
+    @contextmanager
+    def configure_subprocess(self):
+        yield
+
+    @contextmanager
+    def region(self, tag: str, enable_cpu_backup: bool = False):
+        yield
+
+    @contextmanager
+    def cuda_graph(self, **kwargs):
+        yield
+
+    @contextmanager
+    def disable(self):
+        yield
+
+    def pause(self, tag: str):
+        pass
+
+    def resume(self, tag: str):
+        pass
+
+    @property
+    def enabled(self):
+        return False

python/sglang/srt/utils/torch_npu_patch_utils.py

@@ -0,0 +1,17 @@
+from collections.abc import Sequence
+from typing import Any
+
+
+def apply_torch_npu_patches(torch_npu: Any, patches: Sequence[Sequence[Any]]) -> None:
+    """Apply torch_npu patches across old and new torch_npu patch APIs."""
+    if hasattr(torch_npu, "_apply_patches"):
+        torch_npu._apply_patches(patches)
+        return
+
+    if hasattr(torch_npu, "_apply_all_patches"):
+        torch_npu._apply_all_patches()
+        return
+
+    raise AttributeError(
+        "torch_npu must provide either _apply_patches or _apply_all_patches"
+    )

python/sglang/srt/utils/video_decoder.py

@@ -0,0 +1,187 @@
+"""Unified video decoder: torchcodec preferred, decord as fallback."""
+
+import logging
+import os
+
+import numpy as np
+
+logger = logging.getLogger(__name__)
+
+try:
+    from torchcodec.decoders import VideoDecoder
+
+    _BACKEND = "torchcodec"
+except (ImportError, RuntimeError):
+    _BACKEND = "decord"
+
+
+_cuda_backend_enabled: bool | None = None
+
+
+def _try_cuda_backend() -> bool:
+    """Try to enable torchcodec CUDA backend. Caches result after first call."""
+    global _cuda_backend_enabled
+    if _cuda_backend_enabled is not None:
+        return _cuda_backend_enabled
+    try:
+        from torchcodec.decoders import set_cuda_backend
+
+        set_cuda_backend("beta")
+        _cuda_backend_enabled = True
+    except Exception:
+        _cuda_backend_enabled = False
+    return _cuda_backend_enabled
+
+
+class VideoDecoderWrapper:
+    """Unified video decoder that uses torchcodec when available, decord as fallback.
+
+    All frames are returned in NHWC uint8 numpy format for consistency.
+    """
+
+    def __init__(self, source, device: str = "cpu", num_decode_threads: int = 0):
+        """source: file path (str) or video bytes.
+        device: "cpu" or "cuda". GPU decoding only supported with torchcodec.
+        num_decode_threads: number of parallel decoder instances for frame
+            extraction (torchcodec only). 0 = auto (capped at 16),
+            1 = single decoder. Set > 1 to split frame indices across
+            multiple decoders in parallel threads.
+        """
+        self._source = source
+        self._num_decode_threads = num_decode_threads
+        self._source_bytes = source if isinstance(source, bytes) else None
+        self._source_path = source if isinstance(source, str) else None
+        self._tmp_path = None
+        if _BACKEND == "torchcodec":
+            kwargs = {"dimension_order": "NHWC"}
+            if device == "cuda" and _try_cuda_backend():
+                kwargs["device"] = "cuda"
+            self._tc_kwargs = kwargs
+            try:
+                self._decoder = VideoDecoder(source, **kwargs)
+            except RuntimeError:
+                if "device" in kwargs:
+                    logger.warning("CUDA video decoding failed, falling back to CPU.")
+                    kwargs.pop("device")
+                    self._tc_kwargs = kwargs
+                    self._decoder = VideoDecoder(source, **kwargs)
+                else:
+                    raise
+        else:
+            from decord import VideoReader, cpu
+
+            if isinstance(source, bytes):
+                import tempfile
+
+                fd, tmp_path = tempfile.mkstemp(suffix=".mp4")
+                try:
+                    os.write(fd, source)
+                finally:
+                    os.close(fd)
+                self._tmp_path = tmp_path
+                self._decoder = VideoReader(tmp_path, ctx=cpu(0))
+            else:
+                self._decoder = VideoReader(source, ctx=cpu(0))
+
+    def __len__(self):
+        return len(self._decoder)
+
+    def __getitem__(self, idx):
+        """Return single frame as numpy NHWC uint8."""
+        if _BACKEND == "torchcodec":
+            return self._decoder[idx].numpy()
+        else:
+            frame = self._decoder[idx]
+            return frame.asnumpy() if hasattr(frame, "asnumpy") else np.array(frame)
+
+    @property
+    def avg_fps(self) -> float:
+        if _BACKEND == "torchcodec":
+            return self._decoder.metadata.average_fps
+        else:
+            return self._decoder.get_avg_fps()
+
+    def get_frames_at(self, indices: list) -> np.ndarray:
+        """Return frames at given indices as numpy array with shape (N, H, W, C)."""
+        if _BACKEND == "torchcodec":
+            batch = self._decoder.get_frames_at(indices)
+            return batch.data.numpy()
+        else:
+            return self._decoder.get_batch(indices).asnumpy()
+
+    def get_frames_as_tensor(self, indices: list):
+        """Return frames at given indices as a torch tensor (NHWC, uint8, pinned memory)."""
+        import torch
+
+        if (
+            _BACKEND == "torchcodec"
+            and self._num_decode_threads != 1
+            and len(indices) > 1
+        ):
+            num_threads = self._num_decode_threads
+            if num_threads <= 0:
+                num_threads = min(os.cpu_count() or 8, 16)
+            num_threads = min(num_threads, len(indices))
+            if num_threads > 1:
+                return self._parallel_decode(indices, num_threads)
+
+        if _BACKEND == "torchcodec":
+            batch = self._decoder.get_frames_at(indices)
+            return batch.data.pin_memory()
+        else:
+            arr = self._decoder.get_batch(indices).asnumpy()
+            return torch.from_numpy(arr).pin_memory()
+
+    def _parallel_decode(self, indices, num_threads):
+        """Decode frames using multiple VideoDecoder instances in parallel threads."""
+        from concurrent.futures import ThreadPoolExecutor, as_completed
+
+        import torch
+
+        chunks = [list(c) for c in np.array_split(indices, num_threads) if len(c) > 0]
+        source = self._source
+        kwargs = self._tc_kwargs
+
+        def _decode_chunk(chunk):
+            d = VideoDecoder(source, **kwargs)
+            return d.get_frames_at(chunk).data
+
+        with ThreadPoolExecutor(max_workers=len(chunks)) as executor:
+            future_to_idx = {
+                executor.submit(_decode_chunk, chunk): idx
+                for idx, chunk in enumerate(chunks)
+            }
+            results = [None] * len(chunks)
+            for future in as_completed(future_to_idx):
+                idx = future_to_idx[future]
+                results[idx] = future.result()
+
+        return torch.cat(results, dim=0).pin_memory()
+
+    @property
+    def source_bytes(self) -> bytes | None:
+        """Return raw video bytes if available (needed for audio extraction)."""
+        if self._source_bytes is not None:
+            return self._source_bytes
+        path = self._tmp_path or self._source_path
+        if path is not None:
+            if os.path.isfile(path):
+                with open(path, "rb") as f:
+                    return f.read()
+        return None
+
+    def close(self):
+        """Explicitly clean up temporary files."""
+        if self._tmp_path is not None:
+            if os.path.exists(self._tmp_path):
+                os.unlink(self._tmp_path)
+            self._tmp_path = None
+
+    def __del__(self):
+        self.close()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args):
+        self.close()

python/sglang/srt/utils/watchdog.py

@@ -0,0 +1,223 @@
+from __future__ import annotations
+
+import logging
+import os
+import signal
+import sys
+import threading
+import time
+from contextlib import contextmanager
+from multiprocessing import Process
+from typing import Callable, List, Optional
+
+import psutil
+
+from sglang.srt.utils.cudacore_pyspy_dump_utils import pyspy_dump_schedulers
+
+logger = logging.getLogger(__name__)
+
+
+class Watchdog:
+    @staticmethod
+    def create(
+        debug_name: str,
+        watchdog_timeout: Optional[float],
+        soft: bool = False,
+        test_stuck_time: float = 0,
+    ) -> Watchdog:
+        if watchdog_timeout is None:
+            assert (
+                test_stuck_time == 0
+            ), f"stuck tester can be enabled only if soft watchdog is enabled."
+            return _WatchdogNoop()
+        return _WatchdogReal(
+            debug_name=debug_name,
+            watchdog_timeout=watchdog_timeout,
+            soft=soft,
+            test_stuck_time=test_stuck_time,
+        )
+
+    def feed(self):
+        pass
+
+    @contextmanager
+    def disable(self):
+        yield
+
+
+class _WatchdogReal(Watchdog):
+    def __init__(
+        self,
+        debug_name: str,
+        watchdog_timeout: float,
+        soft: bool = False,
+        test_stuck_time: float = 0,
+    ):
+        self._counter = 0
+        self._active = True
+        self._test_stuck_time = test_stuck_time
+        self._test_stuck_triggered = False
+        self._raw = WatchdogRaw(
+            debug_name=debug_name,
+            get_counter=lambda: self._counter,
+            is_active=lambda: self._active,
+            watchdog_timeout=watchdog_timeout,
+            soft=soft,
+        )
+        logger.info(f"Watchdog {self._raw.debug_name} initialized.")
+        if self._test_stuck_time > 0:
+            logger.info(
+                f"Watchdog {self._raw.debug_name} is configured to use {test_stuck_time=}."
+            )
+
+    def feed(self):
+        # Only trigger the test stuck behavior once to avoid blocking server
+        # startup health checks while still testing watchdog timeout detection
+        if self._test_stuck_time > 0 and not self._test_stuck_triggered:
+            self._test_stuck_triggered = True
+            logger.info(
+                f"Watchdog {self._raw.debug_name} start deliberately stuck for {self._test_stuck_time}s"
+            )
+            time.sleep(self._test_stuck_time)
+            logger.info(
+                f"Watchdog {self._raw.debug_name} end deliberately stuck for {self._test_stuck_time}s"
+            )
+
+        self._counter += 1
+
+    @contextmanager
+    def disable(self):
+        assert self._active
+        self._active = False
+        try:
+            yield
+        finally:
+            assert not self._active
+            self._active = True
+
+
+class _WatchdogNoop(Watchdog):
+    pass
+
+
+class WatchdogRaw:
+    def __init__(
+        self,
+        debug_name: str,
+        get_counter: Callable[[], int],
+        is_active: Callable[[], bool],
+        watchdog_timeout: float,
+        soft: bool = False,
+        dump_info: Optional[Callable[[], str]] = None,
+    ):
+        self.debug_name = debug_name
+        self.get_counter = get_counter
+        self.is_active = is_active
+        self.watchdog_timeout = watchdog_timeout
+        self.soft = soft
+        self.dump_info = dump_info
+
+        self.parent_process = psutil.Process().parent()
+        t = threading.Thread(target=self._watchdog_thread, daemon=True)
+        t.start()
+
+    def _watchdog_thread(self):
+        try:
+            while True:
+                self._watchdog_once()
+        except Exception as e:
+            logger.error(
+                f"{self.debug_name} watchdog thread crashed: {e}", exc_info=True
+            )
+
+    def _watchdog_once(self):
+        watchdog_last_counter = 0
+        watchdog_last_time = time.perf_counter()
+
+        while True:
+            current = time.perf_counter()
+            if self.is_active():
+                current_counter = self.get_counter()
+                if watchdog_last_counter == current_counter:
+                    if current > watchdog_last_time + self.watchdog_timeout:
+                        break
+                else:
+                    watchdog_last_counter = current_counter
+                    watchdog_last_time = current
+            time.sleep(self.watchdog_timeout / 2)
+
+        if self.dump_info is not None and (info_msg := self.dump_info()):
+            logger.error(f"{self.debug_name} debug info:\n{info_msg}")
+
+        pyspy_dump_schedulers()
+        logger.error(
+            f"{self.debug_name} watchdog timeout "
+            f"({self.watchdog_timeout=}, {self.soft=})"
+        )
+        print(file=sys.stderr, flush=True)
+        print(file=sys.stdout, flush=True)
+
+        if not self.soft:
+            # Wait for some time so that the parent process can print the error.
+            time.sleep(5)
+            self.parent_process.send_signal(signal.SIGQUIT)
+
+
+class SubprocessWatchdog:
+    """Monitors subprocess liveness and triggers SIGQUIT when a crash is detected.
+
+    When a subprocess crashes (e.g., NCCL timeout causing C++ std::terminate()),
+    Python exception handlers never run, leaving the main process as a zombie
+    service. This watchdog polls subprocess liveness in a daemon thread and
+    sends SIGQUIT to trigger proper cleanup.
+
+    See: https://github.com/sgl-project/sglang/issues/18421
+    """
+
+    def __init__(
+        self,
+        processes: List[Process],
+        process_names: Optional[List[str]] = None,
+        interval: float = 1.0,
+    ):
+        self._processes = processes
+        self._names = process_names or [f"process_{i}" for i in range(len(processes))]
+        self._interval = interval
+        self._stop_event = threading.Event()
+        self._thread: Optional[threading.Thread] = None
+
+    def start(self) -> None:
+        if self._thread is not None or not self._processes:
+            return
+        self._thread = threading.Thread(
+            target=self._monitor_loop, daemon=True, name="subprocess-watchdog"
+        )
+        self._thread.start()
+
+    def stop(self) -> None:
+        self._stop_event.set()
+        if self._thread is not None:
+            self._thread.join(timeout=self._interval * 2)
+            self._thread = None
+
+    def _monitor_loop(self) -> None:
+        try:
+            while not self._stop_event.wait(self._interval):
+                if self._check_processes():
+                    return
+        except Exception as e:
+            logger.error(f"SubprocessWatchdog thread crashed: {e}", exc_info=True)
+
+    def _check_processes(self) -> bool:
+        for proc, name in zip(self._processes, self._names):
+            if proc.is_alive() or proc.exitcode == 0:
+                continue
+
+            logger.error(
+                f"Subprocess {name} (pid={proc.pid}) crashed "
+                f"with exit code {proc.exitcode}. "
+                f"Triggering SIGQUIT for cleanup..."
+            )
+            os.kill(os.getpid(), signal.SIGQUIT)
+            return True
+        return False

python/sglang/srt/utils/weight_checker.py

@@ -0,0 +1,296 @@
+import hashlib
+import logging
+import time
+from typing import Dict, Iterable, NamedTuple, Optional, Set
+
+import torch
+import torch.distributed as dist
+from pydantic import BaseModel, ConfigDict
+
+from sglang.srt.managers.mm_utils import tensor_hash
+from sglang.srt.utils.weight_checker_comparator import (
+    CHUNK_NUMEL,
+    ComparableWeight,
+    RawComparable,
+    compare_weights,
+    select_comparable_weight,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class _StrictBaseModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class ParallelismInfo(_StrictBaseModel):
+    tp_rank: int
+    tp_size: int
+    dp_rank: int
+    dp_size: int
+    pp_rank: int
+    pp_size: int
+    rank: int
+    size: int
+
+
+class ChecksumInfo(_StrictBaseModel):
+    checksums: Dict[str, str]
+    per_gpu_checksum: str
+    parallelism_info: ParallelismInfo
+
+
+class CheckEntry(NamedTuple):
+    name: str
+    should_compare: bool
+    comparable: ComparableWeight
+
+
+class QuantizedWeight(NamedTuple):
+    comparable_cls: type[ComparableWeight]
+    scale_name: str
+
+
+_NON_PERSISTENT_BUFFER_PATTERNS = (
+    "cos_sin_cache",
+    "inv_freq",
+    "freqs_cis",
+    "_weight_fp32",
+)
+
+
+def _is_non_persistent_buffer_name(name: str) -> bool:
+    return any(pat in name for pat in _NON_PERSISTENT_BUFFER_PATTERNS)
+
+
+class WeightChecker:
+    def __init__(self, model_runner):
+        self._model_runner = model_runner
+        self._snapshot_tensors = None
+
+    def handle(self, action: str, allow_quant_error: bool = False) -> Optional[Dict]:
+        logger.info(
+            f"[WeightChecker] handle action={action} allow_quant_error={allow_quant_error}"
+        )
+        if action == "snapshot":
+            return self._snapshot()
+        elif action == "reset_tensors":
+            return self._reset_tensors()
+        elif action == "compare":
+            return self._compare(allow_quant_error=allow_quant_error)
+        elif action == "checksum":
+            return self._compute_checksum()
+        else:
+            raise Exception(f"Unsupported {action=}")
+
+    def _snapshot(self):
+        named_tensors = [
+            (name, param.data.detach().cpu()) for name, param in self._model_state()
+        ]
+        self._snapshot_tensors = dict(named_tensors)
+        assert len(self._snapshot_tensors) == len(
+            named_tensors
+        ), f"should not have duplicated tensor name"
+
+    def _reset_tensors(self):
+        for name, param in self._model_state():
+            if _is_non_persistent_buffer_name(name):
+                continue
+            param.copy_(_random_like(param))
+
+    def _compare(self, allow_quant_error: bool = False):
+        assert self._snapshot_tensors is not None
+
+        quantized_set = _build_quantized_set(self._model_runner.model)
+        skip_compare_names = {
+            name
+            for name, param in self._model_state()
+            if getattr(param, "_skip_weight_check", False)
+        }
+        _check_tensors(
+            expect_tensors=_build_check_entries(
+                self._snapshot_tensors, skip_compare_names, quantized_set
+            ),
+            actual_tensors=_build_check_entries(
+                dict(self._model_state()), skip_compare_names, quantized_set
+            ),
+            allow_quant_error=allow_quant_error,
+        )
+
+    def _compute_checksum(self) -> Dict:
+        torch.cuda.synchronize()
+        start = time.perf_counter()
+
+        quantized_set = _build_quantized_set(self._model_runner.model)
+        skip_compare_names = {
+            name
+            for name, param in self._model_state()
+            if getattr(param, "_skip_weight_check", False)
+        }
+
+        # Hash the dequantized weight so two (qweight, scale) pairs with the same
+        # bf16 hash equal.
+        checksums = {}
+        for name, should_compare, comparable in _build_check_entries(
+            dict(self._model_state()), skip_compare_names, quantized_set
+        ):
+            if should_compare:
+                checksums[name] = _hash_tensor(comparable.dequantize().data)
+
+        h = hashlib.sha256()
+        for name in sorted(checksums):
+            h.update(name.encode())
+            h.update(checksums[name].encode())
+        overall = h.hexdigest()
+
+        torch.cuda.synchronize()
+        elapsed = time.perf_counter() - start
+        logger.info(
+            f"[WeightChecker] checksum computed for {len(checksums)} tensors in {elapsed:.3f}s"
+        )
+
+        info = ChecksumInfo(
+            checksums=checksums,
+            per_gpu_checksum=overall,
+            parallelism_info=self._parallelism_info(),
+        )
+        return info.model_dump()
+
+    def _parallelism_info(self) -> ParallelismInfo:
+        mr = self._model_runner
+        return ParallelismInfo(
+            tp_rank=mr.tp_rank,
+            tp_size=mr.tp_size,
+            dp_rank=mr.dp_rank if mr.dp_rank is not None else 0,
+            dp_size=mr.dp_size,
+            pp_rank=mr.pp_rank,
+            pp_size=mr.pp_size,
+            rank=dist.get_rank() if dist.is_initialized() else 0,
+            size=dist.get_world_size() if dist.is_initialized() else 1,
+        )
+
+    def _model_state(self):
+        yield from self._model_runner.model.named_parameters()
+        yield from self._model_runner.model.named_buffers()
+
+
+def _hash_tensor(t: torch.Tensor) -> str:
+    return f"{tensor_hash(t):016x}"
+
+
+def _check_tensors(
+    expect_tensors: Iterable[CheckEntry],
+    actual_tensors: Iterable[CheckEntry],
+    allow_quant_error: bool = False,
+):
+    good_names = []
+    error_messages = []
+    info_messages = []
+
+    for (expect_name, should_compare, expect_comparable), (
+        actual_name,
+        actual_should_compare,
+        actual_comparable,
+    ) in zip(expect_tensors, actual_tensors, strict=True):
+        assert expect_name == actual_name, f"{expect_name=} {actual_name=}"
+        assert (
+            should_compare == actual_should_compare
+        ), f"{should_compare=} {actual_should_compare=}"
+        name = expect_name
+
+        try:
+            equal, max_abs_err, mean_abs_err, num_exceed = compare_weights(
+                expect_comparable, actual_comparable
+            )
+        except Exception as e:
+            e.add_note(
+                f"when handling {name=} expect={expect_comparable!r} actual={actual_comparable!r}"
+            )
+            raise
+        if equal:
+            good_names.append(name)
+            continue
+        msg = (
+            f"name={name} "
+            f"max_abs_err={max_abs_err} "
+            f"mean_abs_err={mean_abs_err} "
+            f"num_exceed={num_exceed} "
+            f"expect={expect_comparable!r} actual={actual_comparable!r} "
+        )
+        if not should_compare:
+            info_messages.append(msg)
+        elif allow_quant_error and num_exceed == 0:
+            info_messages.append(msg + "(within quantization ULP tolerance)")
+        else:
+            error_messages.append(msg)
+
+    logger.info(f"[check_tensors] equal tensors: {good_names}")
+    if len(info_messages) > 0:
+        logger.info(f"[check_tensors] info: {info_messages}")
+    if len(error_messages) > 0:
+        raise Exception(f"check tensor equality failed:\n" + "\n".join(error_messages))
+
+
+def _random_like(t: torch.Tensor):
+    device = t.device
+    shape = t.shape
+    dtype = t.dtype
+
+    if dtype.is_floating_point:
+        out = torch.empty(shape, device=device, dtype=dtype)
+        for chunk in out.view(-1).split(CHUNK_NUMEL):
+            chunk.copy_(
+                torch.rand(chunk.shape, device=device, dtype=torch.float32).to(dtype)
+            )
+        return out
+
+    if dtype == torch.bool:
+        return torch.rand(shape, device=device) > 0.5
+
+    info = torch.iinfo(dtype)
+    return torch.randint(
+        low=int(info.min), high=int(info.max), size=shape, device=device, dtype=dtype
+    )
+
+
+def _build_quantized_set(model) -> Dict[str, QuantizedWeight]:
+    """Run the router over the model: {weight_name: QuantizedWeight} for each
+    quantized weight; weights absent from the set compare raw."""
+    quantized_set = {}
+    for module_name, module in model.named_modules():
+        comparable_cls = select_comparable_weight(getattr(module, "quant_method", None))
+        if comparable_cls is None:
+            continue
+        prefix = f"{module_name}." if module_name else ""
+        own = {name for name, _ in module.named_parameters(recurse=False)}
+        for name in own:
+            scale = name.replace("weight", "weight_scale_inv")
+            if name.endswith("weight") and scale in own:
+                quantized_set[prefix + name] = QuantizedWeight(
+                    comparable_cls, prefix + scale
+                )
+    return quantized_set
+
+
+def _build_check_entries(
+    raw: Dict[str, torch.Tensor],
+    skip_compare_names: Set[str],
+    quantized_set: Optional[Dict[str, QuantizedWeight]] = None,
+) -> Iterable[CheckEntry]:
+    """Yields a CheckEntry per weight; quantized weights consume their scale, everything
+    else is raw."""
+    skip_compare_names = set(skip_compare_names)
+    quantized_set = quantized_set or {}
+    scale_names = {qw.scale_name for qw in quantized_set.values()}
+
+    for name, tensor in raw.items():
+        if name in scale_names:
+            continue  # compared via its weight's comparable
+        if name in quantized_set:
+            qw = quantized_set[name]
+            yield CheckEntry(name, True, qw.comparable_cls(tensor, raw[qw.scale_name]))
+        else:
+            should_compare = name not in skip_compare_names and (
+                not _is_non_persistent_buffer_name(name)
+            )
+            yield CheckEntry(name, should_compare, RawComparable(tensor))

python/sglang/srt/utils/weight_checker_comparator.py

@@ -0,0 +1,168 @@
+from typing import Iterable, NamedTuple, Optional, Tuple
+
+import torch
+
+from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod, Fp8MoEMethod
+from sglang.srt.layers.quantization.fp8_utils import (
+    block_quant_dequant,
+    inverse_transform_scale_ue8m0,
+)
+from sglang.srt.layers.quantization.modelopt_quant import (
+    ModelOptFp4LinearMethod,
+    ModelOptNvFp4FusedMoEMethod,
+)
+
+# chunk to avoid too high GPU memory peak
+CHUNK_NUMEL = 64 * 1024 * 1024
+
+
+class CompareResult(NamedTuple):
+    equal: bool
+    max_abs_err: float
+    mean_abs_err: float
+    num_exceed: int  # elements past the combined per-side tolerance
+
+
+class ComparableWeight:
+    """Base comparable-weight class; one subclass per precision or raw tensor."""
+
+    @staticmethod
+    def _quant_ulp(w_q: torch.Tensor) -> torch.Tensor:
+        """Per-element ULP of w_q in its own dtype."""
+        finfo = torch.finfo(w_q.dtype)
+        x = w_q.to(torch.float32).abs()
+        # frexp: x = m * 2^e, m in [0.5, 1), so 2^(e-1) is x's binade base.
+        _, exponent = torch.frexp(x)
+        binade = torch.exp2((exponent - 1).to(torch.float32))
+        # Zeros and subnormals share the spacing of the smallest normal binade.
+        binade = binade.masked_fill(x < finfo.smallest_normal, finfo.smallest_normal)
+        return binade * finfo.eps
+
+    def iter_chunks(self) -> Iterable[Tuple[torch.Tensor, Optional[torch.Tensor]]]:
+        raise NotImplementedError
+
+    def dequantize(self, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
+        raise NotImplementedError
+
+
+class Fp8BlockComparable(ComparableWeight):
+    """Deepseek-style FP8 quantization."""
+
+    def __init__(self, w_q: torch.Tensor, w_s: torch.Tensor):
+        self.w_q = w_q
+        self.w_s = w_s
+
+    def __repr__(self) -> str:
+        return f"fp8_block(shape={tuple(self.w_q.shape)} dtype={self.w_q.dtype})"
+
+    @staticmethod
+    def _normalize_scale(w_q: torch.Tensor, w_s: torch.Tensor) -> torch.Tensor:
+        if w_s.dtype == torch.int32:
+            w_s = inverse_transform_scale_ue8m0(w_s, mn=w_q.shape[-2])
+            # ue8m0 packing aligns k to a multiple of 4; drop the padding blocks.
+            w_s = w_s[..., : -(-w_q.shape[-1] // 128)]
+        return w_s.to(torch.float32)
+
+    @staticmethod
+    def _infer_block_size(w_q: torch.Tensor, w_s: torch.Tensor) -> list:
+        k, s_k = w_q.shape[-1], w_s.shape[-1]
+        assert k % s_k == 0, f"cannot infer block size from {w_q.shape=} {w_s.shape=}"
+        block = k // s_k
+        return [block, block]
+
+    @staticmethod
+    def _iter_quant_chunks(w_q: torch.Tensor, w_s: torch.Tensor, block_n: int):
+        """Yields block-row-aligned (q_slice, s_slice) pairs of bounded size."""
+        q3 = w_q.reshape(-1, *w_q.shape[-2:])
+        s3 = w_s.reshape(-1, *w_s.shape[-2:])
+        n, k = q3.shape[-2:]
+        rows = max(block_n, CHUNK_NUMEL // k // block_n * block_n)
+        for b in range(q3.shape[0]):
+            for r0 in range(0, n, rows):
+                r1 = min(r0 + rows, n)
+                yield q3[b, r0:r1], s3[b, r0 // block_n : -(-r1 // block_n)]
+
+    def _scale_and_block_size(self):
+        s = self._normalize_scale(self.w_q, self.w_s)
+        return s, self._infer_block_size(self.w_q, s)
+
+    def iter_chunks(self):
+        s, block_size = self._scale_and_block_size()
+        for q, s_chunk in self._iter_quant_chunks(self.w_q, s, block_size[0]):
+            q, s_chunk = q.cuda(), s_chunk.cuda()
+            yield (
+                block_quant_dequant(q, s_chunk, block_size, dtype=torch.bfloat16),
+                block_quant_dequant(
+                    self._quant_ulp(q), s_chunk, block_size, dtype=torch.float32
+                ),
+            )
+
+    def dequantize(self, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
+        s, block_size = self._scale_and_block_size()
+        return block_quant_dequant(self.w_q, s, block_size, dtype=dtype)
+
+
+class RawComparable(ComparableWeight):
+    """Bitwise equal compare on raw tensor."""
+
+    def __init__(self, tensor: torch.Tensor):
+        self.tensor = tensor
+
+    def __repr__(self) -> str:
+        return f"raw(shape={tuple(self.tensor.shape)} dtype={self.tensor.dtype})"
+
+    def iter_chunks(self):
+        flat = self.tensor.reshape(-1)
+        for start in range(0, flat.numel(), CHUNK_NUMEL):
+            yield flat[start : start + CHUNK_NUMEL].cuda(), None
+
+    def dequantize(self, dtype: torch.dtype = torch.bfloat16) -> torch.Tensor:
+        return self.tensor
+
+
+def compare_weights(
+    expect: ComparableWeight, actual: ComparableWeight
+) -> CompareResult:
+    """Chunked element-wise compare in ComparableWeight space."""
+    equal = True
+    max_abs_err = torch.zeros((), dtype=torch.float32)
+    sum_abs_err = 0.0
+    num_exceed = 0
+    numel = 0
+    for (expect_dq, expect_tol), (actual_dq, actual_tol) in zip(
+        expect.iter_chunks(), actual.iter_chunks(), strict=True
+    ):
+        assert (
+            expect_dq.shape == actual_dq.shape
+        ), f"{expect_dq.shape=} {actual_dq.shape=}"
+        numel += expect_dq.numel()
+        abs_diff = (actual_dq.float() - expect_dq.float()).abs()
+        if torch.all(abs_diff == 0):
+            continue
+        equal = False
+        # |actual_dq - expect_dq| ≤ |actual_dq - w| + |expect_dq - w| ≤ actual_tol + expect_tol
+        tol = (
+            0.0 if expect_tol is None or actual_tol is None else expect_tol + actual_tol
+        )
+        max_abs_err = torch.maximum(max_abs_err, abs_diff.max().cpu())
+        sum_abs_err += abs_diff.sum().item()
+        # `~(diff <= tol)` instead of `diff > tol` so NaN counts as exceeding.
+        num_exceed += int((~(abs_diff <= tol)).sum())
+    return CompareResult(
+        equal, max_abs_err.item(), sum_abs_err / max(numel, 1), num_exceed
+    )
+
+
+def select_comparable_weight(quant_method) -> Optional[type]:
+    """Map a module's quant_method to its ComparableWeight. None means raw (bitwise equal) compare."""
+    if (
+        isinstance(quant_method, (Fp8LinearMethod, Fp8MoEMethod))
+        and quant_method.block_quant
+        and not quant_method.use_mxfp8
+    ):
+        return Fp8BlockComparable
+    if isinstance(quant_method, (ModelOptFp4LinearMethod, ModelOptNvFp4FusedMoEMethod)):
+        raise NotImplementedError(
+            f"weight checker has no ComparableWeight for {type(quant_method).__name__}"
+        )
+    return None