[Misc] Clean up python/sglang package structure (#35062)
This commit is contained in:
+24
-16
@@ -1,18 +1,26 @@
|
||||
# Code Structure
|
||||
|
||||
- `eval`: The evaluation utilities.
|
||||
- `lang`: The frontend language.
|
||||
- `multimodal_gen`: Inference framework for accelerated image/video generation.
|
||||
- `srt`: The backend engine for running local models. (SRT = SGLang Runtime).
|
||||
- `test`: The test utilities.
|
||||
- `api.py`: The public APIs.
|
||||
- `bench_offline_throughput.py`: Benchmark the performance in the offline mode.
|
||||
- `bench_one_batch.py`: Benchmark the latency of running a single static batch without a server.
|
||||
- `bench_one_batch_server.py`: Benchmark the latency of running a single batch with a server.
|
||||
- `bench_serving.py`: Benchmark online serving with dynamic requests.
|
||||
- `check_env.py`: Check the environment variables and dependencies.
|
||||
- `global_config.py`: The global configs and constants.
|
||||
- `launch_server.py`: The entry point for launching a local server.
|
||||
- `profiler.py`: The profiling entry point to send profile requests.
|
||||
- `utils.py`: Common utilities.
|
||||
- `version.py`: Version info.
|
||||
## Folders
|
||||
|
||||
- `benchmark/`: Benchmark implementations and dataset utilities.
|
||||
- `cli/`: Command-line interface commands and entrypoints.
|
||||
- `kernels/`: Kernel interfaces, implementations, selection, and debugging utilities shared by the runtimes.
|
||||
- `lang/`: Deprecated language frontend that is no longer actively maintained.
|
||||
- `multimodal_gen/`: Core runtime for image, video, and audio generation models, most of which are diffusion models.
|
||||
- `srt/`: Core runtime for autoregressive language models. (SRT = SGLang Runtime.)
|
||||
- `test/`: Shared test and evaluation utilities.
|
||||
|
||||
## Files
|
||||
|
||||
- `README.md`: This package structure overview.
|
||||
- `__init__.py`: Package initialization and public Python APIs.
|
||||
- `bench_offline_throughput.py`: Deprecated wrapper for the offline throughput benchmark.
|
||||
- `bench_one_batch.py`: Deprecated wrapper for the one-batch benchmark.
|
||||
- `bench_one_batch_server.py`: Deprecated wrapper for the server-based one-batch benchmark.
|
||||
- `bench_serving.py`: Deprecated wrapper for the online serving benchmark.
|
||||
- `check_env.py`: Environment and dependency diagnostics.
|
||||
- `compile_deep_gemm.py`: DeepGEMM kernel precompilation entrypoint.
|
||||
- `launch_server.py`: Compatibility entrypoint for launching an inference server.
|
||||
- `profiler.py`: Client entrypoint for collecting server profiling traces.
|
||||
- `utils.py`: Common package utilities.
|
||||
- `version.py`: Public package version resolution.
|
||||
|
||||
+25
-44
@@ -1,48 +1,28 @@
|
||||
# SGLang public APIs
|
||||
"""SGLang public API."""
|
||||
|
||||
import platform as _platform
|
||||
import sys as _sys
|
||||
|
||||
# sglang.srt.environ must run before the rest of this file's imports
|
||||
# (hf_transformers_patches, lang.api, ...), which pull in torch and
|
||||
# FlashInfer: those claim these cache dirs early, and the first value set is
|
||||
# the one that sticks. Safe here -- environ has no heavy dependency (no torch).
|
||||
from sglang.srt.environ import redirect_third_party_caches
|
||||
from sglang.srt.environ import (
|
||||
redirect_third_party_caches as _redirect_third_party_caches,
|
||||
)
|
||||
|
||||
redirect_third_party_caches()
|
||||
|
||||
# Install stubs early for platforms where certain dependencies are unavailable
|
||||
# (e.g. macOS/MPS has no triton, and torch.mps lacks Stream / set_device /
|
||||
# get_device_properties). This must run before any downstream imports.
|
||||
import platform as _platform
|
||||
import sys as _sys
|
||||
_redirect_third_party_caches()
|
||||
|
||||
if _sys.platform == "darwin" and _platform.machine() == "arm64":
|
||||
try:
|
||||
import torch as _torch
|
||||
from sglang._platform_stubs import install_platform_stubs as _install_platform_stubs
|
||||
|
||||
if _torch.backends.mps.is_available():
|
||||
from sglang._triton_stub import install as _install_triton_stub
|
||||
|
||||
_install_triton_stub()
|
||||
del _install_triton_stub
|
||||
|
||||
from sglang._mps_stub import install as _install_mps_stub
|
||||
|
||||
_install_mps_stub()
|
||||
del _install_mps_stub
|
||||
del _torch
|
||||
except ImportError:
|
||||
pass
|
||||
del _platform
|
||||
del _sys
|
||||
_install_platform_stubs()
|
||||
|
||||
from sglang.srt.utils.hf_transformers_patches import apply_all as _apply_hf_patches
|
||||
|
||||
_apply_hf_patches()
|
||||
del _apply_hf_patches
|
||||
|
||||
# Frontend Language APIs
|
||||
from sglang.global_config import global_config
|
||||
from sglang.lang.api import (
|
||||
Engine,
|
||||
Runtime,
|
||||
assistant,
|
||||
assistant_begin,
|
||||
@@ -71,8 +51,9 @@ from sglang.lang.choices import (
|
||||
token_length_normalized,
|
||||
unconditional_likelihood_normalized,
|
||||
)
|
||||
from sglang.lang.global_config import global_config
|
||||
|
||||
# Lazy import some libraries
|
||||
# Lazy backend clients
|
||||
from sglang.utils import LazyImport
|
||||
from sglang.version import __version__
|
||||
|
||||
@@ -82,13 +63,21 @@ LiteLLM = LazyImport("sglang.lang.backend.litellm", "LiteLLM")
|
||||
OpenAI = LazyImport("sglang.lang.backend.openai", "OpenAI")
|
||||
VertexAI = LazyImport("sglang.lang.backend.vertexai", "VertexAI")
|
||||
|
||||
# Runtime Engine APIs
|
||||
# Runtime API
|
||||
ServerArgs = LazyImport("sglang.srt.server_args", "ServerArgs")
|
||||
Engine = LazyImport("sglang.srt.entrypoints.engine", "Engine")
|
||||
|
||||
__all__ = [
|
||||
"Anthropic",
|
||||
"Crusoe",
|
||||
"Engine",
|
||||
"LiteLLM",
|
||||
"OpenAI",
|
||||
"Runtime",
|
||||
"RuntimeEndpoint",
|
||||
"ServerArgs",
|
||||
"VertexAI",
|
||||
"__version__",
|
||||
"assistant",
|
||||
"assistant_begin",
|
||||
"assistant_end",
|
||||
@@ -98,6 +87,8 @@ __all__ = [
|
||||
"gen_int",
|
||||
"gen_string",
|
||||
"get_server_info",
|
||||
"global_config",
|
||||
"greedy_token_selection",
|
||||
"image",
|
||||
"select",
|
||||
"separate_reasoning",
|
||||
@@ -105,20 +96,10 @@ __all__ = [
|
||||
"system",
|
||||
"system_begin",
|
||||
"system_end",
|
||||
"token_length_normalized",
|
||||
"unconditional_likelihood_normalized",
|
||||
"user",
|
||||
"user_begin",
|
||||
"user_end",
|
||||
"video",
|
||||
"RuntimeEndpoint",
|
||||
"greedy_token_selection",
|
||||
"token_length_normalized",
|
||||
"unconditional_likelihood_normalized",
|
||||
"ServerArgs",
|
||||
"Anthropic",
|
||||
"Crusoe",
|
||||
"LiteLLM",
|
||||
"OpenAI",
|
||||
"VertexAI",
|
||||
"global_config",
|
||||
"__version__",
|
||||
]
|
||||
|
||||
@@ -1,19 +1,155 @@
|
||||
"""Stub implementations for APIs missing from ``torch.mps``.
|
||||
"""Install compatibility stubs needed by SGLang on Apple Silicon MPS.
|
||||
|
||||
``torch.mps`` lacks several APIs that ``torch.cuda`` provides (``Stream``,
|
||||
``set_device``, ``get_device_properties``, …). Rather than scattering
|
||||
``hasattr`` / ``getattr`` guards throughout the codebase, we monkey-patch
|
||||
``torch.mps`` once at startup so that generic device-agnostic code paths
|
||||
just work.
|
||||
just work. Triton is also stubbed because it is unavailable on macOS.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import functools
|
||||
import importlib
|
||||
import platform
|
||||
import sys
|
||||
import types
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
|
||||
class _StubBase:
|
||||
"""A base class that any mock attribute can safely be subclassed from.
|
||||
|
||||
Used when external code does ``class Foo(triton.runtime.KernelInterface):``.
|
||||
"""
|
||||
|
||||
def __init_subclass__(cls, **kwargs):
|
||||
super().__init_subclass__(**kwargs)
|
||||
|
||||
|
||||
class _MockModule(types.ModuleType):
|
||||
"""A module whose every attribute is itself a ``_MockModule``.
|
||||
|
||||
When called (e.g. ``@triton.jit``), it acts as a pass-through decorator so
|
||||
that kernel *definitions* are syntactically valid even though they will never
|
||||
be compiled.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str):
|
||||
super().__init__(name)
|
||||
self.__path__: list[str] = [] # make it look like a package
|
||||
self.__package__ = name
|
||||
self.__file__ = __file__
|
||||
self._children: dict[str, object] = {}
|
||||
# Set __spec__ so that importlib.util.find_spec() works on cached modules
|
||||
self.__spec__ = importlib.machinery.ModuleSpec(name, None, is_package=True)
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
"""Handle attribute access by creating and returning a child _MockModule."""
|
||||
if name.startswith("__") and name.endswith("__"):
|
||||
raise AttributeError(name)
|
||||
full = f"{self.__name__}.{name}"
|
||||
if full in sys.modules:
|
||||
return sys.modules[full]
|
||||
# If the name looks like a class (CamelCase / uppercase), return a
|
||||
# stub class that can be used as a base class for inheritance.
|
||||
if name[0:1].isupper():
|
||||
stub_cls = type(name, (_StubBase,), {"__module__": self.__name__})
|
||||
self._children[name] = stub_cls
|
||||
return stub_cls
|
||||
child = _MockModule(full)
|
||||
sys.modules[full] = child
|
||||
self._children[name] = child
|
||||
return child
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
# Direct decorator usage: @triton.jit (receives the function)
|
||||
if len(args) == 1 and callable(args[0]) and not kwargs:
|
||||
return args[0]
|
||||
|
||||
# Parameterised decorator: @triton.jit(...) → returns a decorator
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
|
||||
return _decorator
|
||||
|
||||
def __instancecheck__(self, instance):
|
||||
"""Return False for all instance checks against the mock."""
|
||||
return False
|
||||
|
||||
def __contains__(self, item):
|
||||
"""Return False for all membership checks."""
|
||||
return False
|
||||
|
||||
def __iter__(self):
|
||||
return iter([])
|
||||
|
||||
def __len__(self):
|
||||
return 0
|
||||
|
||||
def __bool__(self):
|
||||
return False
|
||||
|
||||
def __repr__(self):
|
||||
return f"<triton-stub {self.__name__!r}>"
|
||||
|
||||
|
||||
def _cdiv(a: int, b: int) -> int:
|
||||
"""Ceiling division – mirrors ``triton.cdiv``."""
|
||||
return -(a // -b)
|
||||
|
||||
|
||||
def _next_power_of_2(n: int) -> int:
|
||||
"""Mirrors ``triton.next_power_of_2``."""
|
||||
return 1 << (n - 1).bit_length() if n > 0 else 1
|
||||
|
||||
|
||||
class _Config:
|
||||
"""Minimal stand-in for ``triton.Config`` used in ``@triton.autotune``."""
|
||||
|
||||
def __init__(self, kwargs=None, num_warps=4, num_stages=2, **extra):
|
||||
self.kwargs = kwargs or {}
|
||||
self.num_warps = num_warps
|
||||
self.num_stages = num_stages
|
||||
|
||||
|
||||
class _TritonFinder:
|
||||
"""A meta-path finder that intercepts all ``import triton.*`` statements.
|
||||
|
||||
When Python encounters ``import triton.backends.compiler``, it walks the
|
||||
dotted path and tries to import each component. Our mock module's
|
||||
``__getattr__`` handles *attribute* access, but the import machinery uses
|
||||
``importlib`` finders, not attribute access, for sub-module resolution.
|
||||
This finder bridges that gap by creating ``_MockModule`` instances for any
|
||||
``triton.*`` sub-module that isn't already in ``sys.modules``.
|
||||
"""
|
||||
|
||||
def find_spec(self, fullname, path=None, target=None):
|
||||
"""PEP 451 meta-path finder for ``triton.*`` sub-modules."""
|
||||
if fullname == "triton" or fullname.startswith("triton."):
|
||||
if fullname in sys.modules:
|
||||
return getattr(sys.modules[fullname], "__spec__", None)
|
||||
# Create and register the mock so the import machinery finds it
|
||||
mod = _MockModule(fullname)
|
||||
sys.modules[fullname] = mod
|
||||
parts = fullname.rsplit(".", 1)
|
||||
if len(parts) == 2:
|
||||
parent_name, child_name = parts
|
||||
parent = sys.modules.get(parent_name)
|
||||
if parent is not None:
|
||||
setattr(parent, child_name, mod)
|
||||
return mod.__spec__
|
||||
return None
|
||||
|
||||
|
||||
def _make_mock(name: str) -> _MockModule:
|
||||
"""Create a ``_MockModule`` and register it in ``sys.modules``."""
|
||||
mod = _MockModule(name)
|
||||
sys.modules[name] = mod
|
||||
return mod
|
||||
|
||||
|
||||
class Stream:
|
||||
"""Minimal stand-in for ``torch.cuda.Stream``.
|
||||
|
||||
@@ -233,16 +369,90 @@ def _patch_non_blocking() -> None:
|
||||
torch.Tensor.copy_ = _patched_copy_
|
||||
|
||||
|
||||
_installed = False
|
||||
_platform_stubs_installed = False
|
||||
|
||||
|
||||
def install() -> None:
|
||||
"""Patch ``torch.mps`` with the stubs above. Safe to call multiple times."""
|
||||
global _installed
|
||||
if _installed:
|
||||
def install_platform_stubs() -> None:
|
||||
"""Install the Triton and MPS compatibility stubs when they are needed."""
|
||||
global _platform_stubs_installed
|
||||
if _platform_stubs_installed:
|
||||
return
|
||||
|
||||
import torch
|
||||
if sys.platform != "darwin" or platform.machine() != "arm64":
|
||||
return
|
||||
|
||||
try:
|
||||
import torch
|
||||
except ImportError:
|
||||
return
|
||||
|
||||
if not torch.backends.mps.is_available():
|
||||
return
|
||||
|
||||
if "triton" not in sys.modules and importlib.util.find_spec("triton") is None:
|
||||
# Register the meta-path finder first so later ``import triton.X``
|
||||
# statements are handled by the stub.
|
||||
sys.meta_path.insert(0, _TritonFinder())
|
||||
|
||||
triton = _make_mock("triton")
|
||||
triton.__version__ = "3.0.0"
|
||||
triton.cdiv = _cdiv
|
||||
triton.next_power_of_2 = _next_power_of_2
|
||||
triton.Config = _Config
|
||||
|
||||
# triton.language (commonly imported as ``tl``)
|
||||
tl = _make_mock("triton.language")
|
||||
|
||||
class _constexpr:
|
||||
"""Stand-in for ``tl.constexpr`` as an annotation and value wrapper."""
|
||||
|
||||
def __init__(self, value=None):
|
||||
self.value = value
|
||||
|
||||
def __repr__(self):
|
||||
return f"constexpr({self.value!r})"
|
||||
|
||||
tl.constexpr = _constexpr
|
||||
triton.language = tl
|
||||
|
||||
# triton.language.extra.libdevice
|
||||
extra = _make_mock("triton.language.extra")
|
||||
tl.extra = extra
|
||||
libdevice = _make_mock("triton.language.extra.libdevice")
|
||||
extra.libdevice = libdevice
|
||||
|
||||
# triton.runtime.jit (JITFunction is used in isinstance checks)
|
||||
runtime = _make_mock("triton.runtime")
|
||||
triton.runtime = runtime
|
||||
jit_mod = _make_mock("triton.runtime.jit")
|
||||
|
||||
class _JITFunction:
|
||||
"""Dummy type for ``triton.runtime.jit.JITFunction`` checks."""
|
||||
|
||||
pass
|
||||
|
||||
jit_mod.JITFunction = _JITFunction
|
||||
runtime.jit = jit_mod
|
||||
|
||||
# triton.runtime.driver
|
||||
driver = _make_mock("triton.runtime.driver")
|
||||
runtime.driver = driver
|
||||
|
||||
# triton.testing
|
||||
testing = _make_mock("triton.testing")
|
||||
triton.testing = testing
|
||||
|
||||
# triton.tools / triton.tools.tensor_descriptor
|
||||
tools = _make_mock("triton.tools")
|
||||
triton.tools = tools
|
||||
td = _make_mock("triton.tools.tensor_descriptor")
|
||||
tools.tensor_descriptor = td
|
||||
|
||||
# triton.backends / triton.backends.compiler
|
||||
backends = _make_mock("triton.backends")
|
||||
triton.backends = backends
|
||||
compiler = _make_mock("triton.backends.compiler")
|
||||
backends.compiler = compiler
|
||||
|
||||
mps = torch.mps
|
||||
# Only patch attributes that are actually missing
|
||||
@@ -267,4 +477,4 @@ def install() -> None:
|
||||
|
||||
_patch_non_blocking()
|
||||
|
||||
_installed = True
|
||||
_platform_stubs_installed = True
|
||||
@@ -1,228 +0,0 @@
|
||||
"""
|
||||
Mock triton module for platforms where triton is not available (e.g., macOS/MPS).
|
||||
|
||||
This module provides stub implementations of triton APIs so that modules which
|
||||
import triton at the top level can be loaded without error. The actual triton
|
||||
kernels are never executed on these platforms – alternative backends (e.g. SDPA
|
||||
for MPS) are used instead.
|
||||
|
||||
Usage – call ``install()`` **before** any ``import triton`` in the process:
|
||||
|
||||
from sglang._triton_stub import install
|
||||
install()
|
||||
"""
|
||||
|
||||
import sys
|
||||
import types
|
||||
|
||||
|
||||
class _StubBase:
|
||||
"""A base class that any mock attribute can safely be subclassed from.
|
||||
|
||||
Used when external code does ``class Foo(triton.runtime.KernelInterface):``.
|
||||
"""
|
||||
|
||||
def __init_subclass__(cls, **kwargs):
|
||||
super().__init_subclass__(**kwargs)
|
||||
|
||||
|
||||
class _MockModule(types.ModuleType):
|
||||
"""A module whose every attribute is itself a ``_MockModule``.
|
||||
|
||||
When called (e.g. ``@triton.jit``), it acts as a pass-through decorator so
|
||||
that kernel *definitions* are syntactically valid even though they will never
|
||||
be compiled.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str):
|
||||
super().__init__(name)
|
||||
self.__path__: list[str] = [] # make it look like a package
|
||||
self.__package__ = name
|
||||
self.__file__ = __file__
|
||||
self._children: dict[str, object] = {}
|
||||
# Set __spec__ so that importlib.util.find_spec() works on cached modules
|
||||
import importlib
|
||||
|
||||
self.__spec__ = importlib.machinery.ModuleSpec(name, None, is_package=True)
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
"""Handle attribute access by creating and returning a child _MockModule."""
|
||||
if name.startswith("__") and name.endswith("__"):
|
||||
raise AttributeError(name)
|
||||
full = f"{self.__name__}.{name}"
|
||||
if full in sys.modules:
|
||||
return sys.modules[full]
|
||||
# If the name looks like a class (CamelCase / uppercase), return a
|
||||
# stub class that can be used as a base class for inheritance.
|
||||
if name[0:1].isupper():
|
||||
stub_cls = type(name, (_StubBase,), {"__module__": self.__name__})
|
||||
self._children[name] = stub_cls
|
||||
return stub_cls
|
||||
child = _MockModule(full)
|
||||
sys.modules[full] = child
|
||||
self._children[name] = child
|
||||
return child
|
||||
|
||||
def __call__(self, *args, **kwargs):
|
||||
# Direct decorator usage: @triton.jit (receives the function)
|
||||
if len(args) == 1 and callable(args[0]) and not kwargs:
|
||||
return args[0]
|
||||
|
||||
# Parameterised decorator: @triton.jit(...) → returns a decorator
|
||||
def _decorator(fn):
|
||||
return fn
|
||||
|
||||
return _decorator
|
||||
|
||||
def __instancecheck__(self, instance):
|
||||
"""Return False for all instance checks against the mock."""
|
||||
return False
|
||||
|
||||
def __contains__(self, item):
|
||||
"""Return False for all membership checks."""
|
||||
return False
|
||||
|
||||
def __iter__(self):
|
||||
return iter([])
|
||||
|
||||
def __len__(self):
|
||||
return 0
|
||||
|
||||
def __bool__(self):
|
||||
return False
|
||||
|
||||
def __repr__(self):
|
||||
return f"<triton-stub {self.__name__!r}>"
|
||||
|
||||
|
||||
def _cdiv(a: int, b: int) -> int:
|
||||
"""Ceiling division – mirrors ``triton.cdiv``."""
|
||||
return -(a // -b)
|
||||
|
||||
|
||||
def _next_power_of_2(n: int) -> int:
|
||||
"""Mirrors ``triton.next_power_of_2``."""
|
||||
return 1 << (n - 1).bit_length() if n > 0 else 1
|
||||
|
||||
|
||||
class _Config:
|
||||
"""Minimal stand-in for ``triton.Config`` used in ``@triton.autotune``."""
|
||||
|
||||
def __init__(self, kwargs=None, num_warps=4, num_stages=2, **extra):
|
||||
self.kwargs = kwargs or {}
|
||||
self.num_warps = num_warps
|
||||
self.num_stages = num_stages
|
||||
|
||||
|
||||
class _TritonFinder:
|
||||
"""A meta-path finder that intercepts all ``import triton.*`` statements.
|
||||
|
||||
When Python encounters ``import triton.backends.compiler``, it walks the
|
||||
dotted path and tries to import each component. Our mock module's
|
||||
``__getattr__`` handles *attribute* access, but the import machinery uses
|
||||
``importlib`` finders, not attribute access, for sub-module resolution.
|
||||
This finder bridges that gap by creating ``_MockModule`` instances for any
|
||||
``triton.*`` sub-module that isn't already in ``sys.modules``.
|
||||
"""
|
||||
|
||||
def find_spec(self, fullname, path=None, target=None):
|
||||
"""PEP 451 meta-path finder for ``triton.*`` sub-modules."""
|
||||
if fullname == "triton" or fullname.startswith("triton."):
|
||||
if fullname in sys.modules:
|
||||
return getattr(sys.modules[fullname], "__spec__", None)
|
||||
# Create and register the mock so the import machinery finds it
|
||||
mod = _MockModule(fullname)
|
||||
sys.modules[fullname] = mod
|
||||
parts = fullname.rsplit(".", 1)
|
||||
if len(parts) == 2:
|
||||
parent_name, child_name = parts
|
||||
parent = sys.modules.get(parent_name)
|
||||
if parent is not None:
|
||||
setattr(parent, child_name, mod)
|
||||
return mod.__spec__
|
||||
return None
|
||||
|
||||
|
||||
def _make_mock(name: str) -> _MockModule:
|
||||
"""Create a ``_MockModule`` and register it in ``sys.modules``."""
|
||||
mod = _MockModule(name)
|
||||
sys.modules[name] = mod
|
||||
return mod
|
||||
|
||||
|
||||
def install() -> None:
|
||||
"""Register a mock ``triton`` package in *sys.modules*.
|
||||
|
||||
This is a no-op if a real ``triton`` is already importable.
|
||||
"""
|
||||
if "triton" in sys.modules:
|
||||
return
|
||||
# Check whether a real triton exists before installing the stub.
|
||||
import importlib.util
|
||||
|
||||
if importlib.util.find_spec("triton") is not None:
|
||||
return
|
||||
|
||||
# Register the meta-path finder FIRST so that any ``import triton.X``
|
||||
# during the rest of install() (or later) is handled.
|
||||
sys.meta_path.insert(0, _TritonFinder())
|
||||
|
||||
triton = _make_mock("triton")
|
||||
triton.__version__ = "3.0.0"
|
||||
triton.cdiv = _cdiv
|
||||
triton.next_power_of_2 = _next_power_of_2
|
||||
triton.Config = _Config
|
||||
|
||||
# triton.language (commonly imported as ``tl``)
|
||||
tl = _make_mock("triton.language")
|
||||
|
||||
class _constexpr:
|
||||
"""Stand-in for ``tl.constexpr`` – works as both annotation and value wrapper."""
|
||||
|
||||
def __init__(self, value=None):
|
||||
self.value = value
|
||||
|
||||
def __repr__(self):
|
||||
return f"constexpr({self.value!r})"
|
||||
|
||||
tl.constexpr = _constexpr
|
||||
triton.language = tl
|
||||
|
||||
# triton.language.extra.libdevice
|
||||
extra = _make_mock("triton.language.extra")
|
||||
tl.extra = extra
|
||||
libdevice = _make_mock("triton.language.extra.libdevice")
|
||||
extra.libdevice = libdevice
|
||||
|
||||
# triton.runtime.jit (JITFunction used in isinstance checks)
|
||||
runtime = _make_mock("triton.runtime")
|
||||
triton.runtime = runtime
|
||||
jit_mod = _make_mock("triton.runtime.jit")
|
||||
|
||||
class _JITFunction:
|
||||
"""Dummy so ``isinstance(fn, triton.runtime.jit.JITFunction)`` works."""
|
||||
|
||||
pass
|
||||
|
||||
jit_mod.JITFunction = _JITFunction
|
||||
runtime.jit = jit_mod
|
||||
|
||||
# triton.runtime.driver (used by fla/utils.py)
|
||||
driver = _make_mock("triton.runtime.driver")
|
||||
runtime.driver = driver
|
||||
|
||||
# triton.testing
|
||||
testing = _make_mock("triton.testing")
|
||||
triton.testing = testing
|
||||
|
||||
# triton.tools / triton.tools.tensor_descriptor
|
||||
tools = _make_mock("triton.tools")
|
||||
triton.tools = tools
|
||||
td = _make_mock("triton.tools.tensor_descriptor")
|
||||
tools.tensor_descriptor = td
|
||||
|
||||
# triton.backends / triton.backends.compiler (used by torch._inductor)
|
||||
backends = _make_mock("triton.backends")
|
||||
triton.backends = backends
|
||||
compiler = _make_mock("triton.backends.compiler")
|
||||
backends.compiler = compiler
|
||||
@@ -1,315 +0,0 @@
|
||||
# Adapt from https://github.com/fw-ai/llm_eval_meta
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import os
|
||||
import pickle
|
||||
import re
|
||||
import shutil
|
||||
from collections import defaultdict
|
||||
from dataclasses import dataclass
|
||||
|
||||
import httpx
|
||||
import numpy as np
|
||||
import openai
|
||||
from datasets import load_dataset
|
||||
from openai import AsyncOpenAI
|
||||
from tqdm import tqdm
|
||||
|
||||
# Mapping providers to their clients and models
|
||||
provider_to_models = {
|
||||
"b10": {
|
||||
"8b": "meta-llama/Llama-3.1-8B-Instruct",
|
||||
"70b": "meta-llama/Llama-3.1-70B-Instruct",
|
||||
"405b": "meta-llama/Llama-3.1-405B-Instruct",
|
||||
},
|
||||
"oai": {
|
||||
"8b": "meta-llama/Llama-3.1-8B-Instruct",
|
||||
"70b": "meta-llama/Llama-3.1-70B-Instruct",
|
||||
"405b": "meta-llama/Llama-3.1-405B-Instruct",
|
||||
},
|
||||
"sgl": {
|
||||
"8b": "meta-llama/Llama-3.1-8B-Instruct",
|
||||
"70b": "meta-llama/Llama-3.1-70B-Instruct",
|
||||
"405b": "meta-llama/Llama-3.1-405B-Instruct",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
async def fetch_responses(
|
||||
client, prompt, semaphore, index, provider, model_size, output_dir, max_tokens
|
||||
):
|
||||
output_file = os.path.join(output_dir, f"response_{index}.pkl")
|
||||
if os.path.exists(output_file):
|
||||
print(f"File {output_file} already exists, skipping.")
|
||||
return
|
||||
|
||||
async with semaphore:
|
||||
response = await client.completions.create(
|
||||
model=provider_to_models[provider][model_size],
|
||||
prompt=prompt,
|
||||
temperature=0.0,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
if isinstance(response, openai.BadRequestError):
|
||||
with open(output_file, "wb") as f:
|
||||
pickle.dump("bad_response", f)
|
||||
assert isinstance(response, openai.types.completion.Completion)
|
||||
# Save response to a file
|
||||
with open(output_file, "wb") as f:
|
||||
pickle.dump(response, f)
|
||||
|
||||
|
||||
TASK_TO_MAX_TOKENS = {
|
||||
"evals__mmlu__details": 1,
|
||||
"evals__mmlu__0_shot__cot__details": 1024,
|
||||
# Official meta uses 1024, but a small % (.05) of questions are answered correctly after relaxing
|
||||
"evals__mmlu_pro__details": 2048,
|
||||
"evals__gsm8k__details": 1024,
|
||||
}
|
||||
|
||||
TASK_TO_EVAL_SET = {
|
||||
"mmlu": "evals__mmlu__details",
|
||||
"mmlu_cot": "evals__mmlu__0_shot__cot__details",
|
||||
"mmlu_pro": "evals__mmlu_pro__details",
|
||||
"gsm8k": "evals__gsm8k__details",
|
||||
}
|
||||
|
||||
|
||||
class CustomAsyncHTTPXClient(httpx.AsyncClient):
|
||||
async def send(self, request: httpx.Request, *args, **kwargs) -> httpx.Response:
|
||||
request.url = httpx.URL(
|
||||
f"https://model-{os.getenv('MODEL_ID')}.api.baseten.co/development/predict"
|
||||
)
|
||||
return await super().send(request, *args, **kwargs)
|
||||
|
||||
|
||||
def get_client(provider):
|
||||
if provider not in "b10":
|
||||
if os.getenv("OPENAI_API_KEY") is None:
|
||||
os.environ["OPENAI_API_KEY"] = "EMPTY"
|
||||
return {
|
||||
"oai": AsyncOpenAI(base_url="http://127.0.0.1:8000/v1/"),
|
||||
"b10": AsyncOpenAI(
|
||||
api_key=f"Api-Key {os.getenv('OPENAI_API_KEY')}",
|
||||
base_url=f"https://model-{os.getenv('MODEL_ID')}.api.baseten.co/development/predict",
|
||||
http_client=CustomAsyncHTTPXClient(),
|
||||
),
|
||||
"sgl": AsyncOpenAI(base_url="http://127.0.0.1:30000/v1/"),
|
||||
}[provider]
|
||||
|
||||
|
||||
# Define the benchmark function
|
||||
async def benchmark(args):
|
||||
ds = load_dataset(
|
||||
"meta-llama/Llama-3.1-405B-Instruct-evals",
|
||||
f"Llama-3.1-405B-Instruct-{TASK_TO_EVAL_SET[args.task]}",
|
||||
)
|
||||
semaphore = asyncio.Semaphore(args.concurrency) # Limit to 16 concurrent tasks
|
||||
|
||||
if args.num_examples is None:
|
||||
args.num_examples = len(ds["latest"]["input_final_prompts"])
|
||||
prompts = ds["latest"]["input_final_prompts"][: args.num_examples]
|
||||
|
||||
# Create the output directory if it does not exist
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
tasks = []
|
||||
# Create the tasks with tqdm progress bar
|
||||
max_tokens = TASK_TO_MAX_TOKENS[TASK_TO_EVAL_SET[args.task]]
|
||||
client = get_client(args.provider)
|
||||
for idx, prompt in enumerate(tqdm(prompts, desc="Creating tasks")):
|
||||
tasks.append(
|
||||
asyncio.create_task(
|
||||
fetch_responses(
|
||||
client,
|
||||
f"<|begin_of_text|>{prompt[0]}",
|
||||
semaphore,
|
||||
idx,
|
||||
args.provider,
|
||||
args.model_size,
|
||||
args.output_dir,
|
||||
max_tokens=max_tokens,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
# Run the tasks with tqdm progress bar
|
||||
for future in tqdm(
|
||||
asyncio.as_completed(tasks), total=len(tasks), desc="Processing tasks"
|
||||
):
|
||||
await future
|
||||
|
||||
|
||||
def get_mmlu_answer(response):
|
||||
if response is not None:
|
||||
return response.choices[0].text.lstrip().rstrip().upper().replace(".", "")
|
||||
return None
|
||||
|
||||
|
||||
def get_mmlu_cot_answer(response):
|
||||
pattern = r"The best answer is (.+)\.?"
|
||||
match = re.search(pattern, response.choices[0].text)
|
||||
if match:
|
||||
return match.group(1).replace(".", "").replace("*", "")
|
||||
|
||||
pattern = r"the best answer is (.+)\.?"
|
||||
match = re.search(pattern, response.choices[0].text)
|
||||
if match:
|
||||
return match.group(1).replace(".", "")
|
||||
|
||||
pattern = r"The correct answer is (.+)\.?"
|
||||
match = re.search(pattern, response.choices[0].text)
|
||||
if match:
|
||||
return match.group(1).replace(".", "")
|
||||
|
||||
pattern = r"the correct answer is (.+)\.?"
|
||||
match = re.search(pattern, response.choices[0].text)
|
||||
if match:
|
||||
return match.group(1).replace(".", "")
|
||||
|
||||
|
||||
def get_answer_gsm8k(response):
|
||||
pattern = r"The final answer is (.+)\.?"
|
||||
match = re.search(pattern, response.choices[0].text)
|
||||
if match:
|
||||
s = match.group(1)
|
||||
for ok_symbol in ["%", "$"]:
|
||||
s = s.replace(ok_symbol, "")
|
||||
return s
|
||||
|
||||
|
||||
TASK_TO_ANSWER_EXTRACTOR = {
|
||||
"evals__mmlu__details": get_mmlu_answer,
|
||||
"evals__mmlu__0_shot__cot__details": get_mmlu_cot_answer,
|
||||
"evals__gsm8k__details": get_answer_gsm8k,
|
||||
"evals__mmlu_pro__details": get_mmlu_cot_answer,
|
||||
}
|
||||
|
||||
|
||||
def get_dataset_from_task(task, response_path, model_size):
|
||||
ds_405b = load_dataset(
|
||||
f"meta-llama/Llama-3.1-405B-Instruct-evals",
|
||||
f"Llama-3.1-405B-Instruct-{task}",
|
||||
)
|
||||
ds_405b_hash_order = [x[0] for x in ds_405b["latest"]["input_final_prompts_hash"]]
|
||||
|
||||
if "70b" in model_size or "8b" in model_size:
|
||||
if "70" in model_size:
|
||||
ref_model_ds = load_dataset(
|
||||
f"meta-llama/Llama-3.1-70B-Instruct-evals",
|
||||
f"Llama-3.1-70B-Instruct-{task}",
|
||||
)
|
||||
else:
|
||||
ref_model_ds = load_dataset(
|
||||
f"meta-llama/Llama-3.1-8B-Instruct-evals",
|
||||
f"Llama-3.1-8B-Instruct-{task}",
|
||||
)
|
||||
|
||||
hash_to_row = {}
|
||||
for row in ref_model_ds["latest"]:
|
||||
hash_to_row[row["input_final_prompts_hash"][0]] = row
|
||||
reordered_rows = []
|
||||
for prompt_hash in ds_405b_hash_order:
|
||||
reordered_rows.append(hash_to_row[prompt_hash])
|
||||
ref_model_ds["latest"] = reordered_rows
|
||||
return ref_model_ds
|
||||
|
||||
return ds_405b
|
||||
|
||||
|
||||
def analyze(task, response_path, model_size):
|
||||
ds = get_dataset_from_task(task, response_path, model_size)
|
||||
|
||||
responses = []
|
||||
total = len(ds["latest"])
|
||||
|
||||
for i in range(0, total):
|
||||
response = pickle.load(
|
||||
open(os.path.join(response_path, f"response_{i}.pkl"), "rb")
|
||||
)
|
||||
responses.append(response)
|
||||
|
||||
@dataclass
|
||||
class Stats:
|
||||
correct: int = 0
|
||||
total: int = 0
|
||||
meta_correct: int = 0
|
||||
|
||||
average: float = None
|
||||
|
||||
subtask_name_to_stats = defaultdict(lambda: Stats())
|
||||
|
||||
for response, ds_row in zip(responses, ds["latest"]):
|
||||
model_answer = TASK_TO_ANSWER_EXTRACTOR[task](response)
|
||||
|
||||
subtask = ds_row["subtask_name"]
|
||||
|
||||
is_eval_correct = model_answer in ds_row["input_correct_responses"]
|
||||
if is_eval_correct:
|
||||
subtask_name_to_stats[subtask].correct += 1
|
||||
|
||||
if ds_row["is_correct"]:
|
||||
subtask_name_to_stats[subtask].meta_correct += 1
|
||||
|
||||
subtask_name_to_stats[subtask].total += 1
|
||||
|
||||
micro_stats = Stats()
|
||||
for subtask, stats in subtask_name_to_stats.items():
|
||||
stats.average = stats.correct / stats.total
|
||||
stats.meta_average = stats.meta_correct / stats.total
|
||||
|
||||
micro_stats.correct += stats.correct
|
||||
micro_stats.total += stats.total
|
||||
micro_stats.meta_correct += stats.meta_correct
|
||||
|
||||
micro_stats.average = micro_stats.correct / micro_stats.total
|
||||
micro_stats.meta_average = micro_stats.meta_correct / micro_stats.total
|
||||
|
||||
print("Macro average", np.mean([x.average for x in subtask_name_to_stats.values()]))
|
||||
print(
|
||||
"Meta Macro average",
|
||||
np.mean([x.meta_average for x in subtask_name_to_stats.values()]),
|
||||
)
|
||||
print("Micro average", micro_stats.average)
|
||||
print("Meta Micro average", micro_stats.meta_average)
|
||||
|
||||
|
||||
# Entry point for the script
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Script to run model with specified parameters."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model-size",
|
||||
type=str,
|
||||
default="8b",
|
||||
help="Size of the model (e.g., 8b or 70b)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--provider",
|
||||
type=str,
|
||||
default="sgl",
|
||||
help="Provider name (e.g., sgl, oai, b10)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--task",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Task (e.g., mmlu, mmlu_cot, mmlu_pro, gsm8k)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-examples", type=int, default=None, help="Number of examples to process"
|
||||
)
|
||||
parser.add_argument("--concurrency", type=int, default=16)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default="tmp-output-dir",
|
||||
help="Directory to save responses",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
asyncio.run(benchmark(args))
|
||||
analyze(TASK_TO_EVAL_SET[args.task], args.output_dir, args.model_size)
|
||||
shutil.rmtree("tmp-output-dir", ignore_errors=True)
|
||||
@@ -1,164 +0,0 @@
|
||||
import argparse
|
||||
import asyncio
|
||||
import os
|
||||
import pickle
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
import openai
|
||||
import torch
|
||||
from bert_score import BERTScorer
|
||||
from datasets import load_dataset
|
||||
from tqdm import tqdm
|
||||
|
||||
|
||||
def get_client(api_url: str) -> openai.AsyncOpenAI:
|
||||
if os.getenv("OPENAI_API_KEY") is None:
|
||||
os.environ["OPENAI_API_KEY"] = "EMPTY"
|
||||
return openai.AsyncOpenAI(base_url=api_url)
|
||||
|
||||
|
||||
def get_dataset():
|
||||
return load_dataset("bigai-nlco/LooGLE", "longdep_qa", split="test")
|
||||
|
||||
|
||||
async def fetch_response(
|
||||
client: openai.AsyncOpenAI,
|
||||
context: str,
|
||||
question: str,
|
||||
semaphore: asyncio.Semaphore,
|
||||
index: int,
|
||||
model: str,
|
||||
output_dir: Path,
|
||||
):
|
||||
output_file = output_dir / f"response_{index}.pkl"
|
||||
if output_file.exists():
|
||||
return
|
||||
|
||||
prompt = (
|
||||
"Please answer the question based on the long texts below.\n"
|
||||
f"{context}\n"
|
||||
f"Question: {question}\n"
|
||||
"Answer:"
|
||||
)
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": prompt},
|
||||
]
|
||||
|
||||
async with semaphore:
|
||||
try:
|
||||
response = await client.chat.completions.create(
|
||||
model=model,
|
||||
messages=messages,
|
||||
temperature=0.0,
|
||||
max_tokens=512,
|
||||
)
|
||||
except openai.BadRequestError as e:
|
||||
with open(output_file, "wb") as f:
|
||||
pickle.dump({"error": str(e)}, f)
|
||||
return
|
||||
|
||||
with open(output_file, "wb") as f:
|
||||
pickle.dump(response, f)
|
||||
|
||||
|
||||
async def benchmark(args):
|
||||
dataset = get_dataset()
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
client = get_client(args.api_url)
|
||||
semaphore = asyncio.Semaphore(args.max_concurrency)
|
||||
|
||||
tasks: List[asyncio.Task] = []
|
||||
for idx, ex in enumerate(dataset):
|
||||
if idx >= args.num_prompts:
|
||||
break
|
||||
tasks.append(
|
||||
asyncio.create_task(
|
||||
fetch_response(
|
||||
client,
|
||||
ex["context"],
|
||||
ex["question"],
|
||||
semaphore,
|
||||
idx,
|
||||
args.model,
|
||||
output_dir,
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
for _ in tqdm(
|
||||
asyncio.as_completed(tasks), total=len(tasks), desc="Running benchmark"
|
||||
):
|
||||
await _
|
||||
|
||||
|
||||
def analyse(args):
|
||||
dataset = get_dataset()
|
||||
output_dir = Path(args.output_dir)
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
scorer = BERTScorer(lang="en", device=device)
|
||||
|
||||
hyps: List[str] = []
|
||||
refs: List[str] = []
|
||||
for idx, ex in enumerate(tqdm(dataset, desc="Loading responses")):
|
||||
if idx >= args.num_prompts:
|
||||
break
|
||||
pkl_file = output_dir / f"response_{idx}.pkl"
|
||||
if not pkl_file.exists():
|
||||
raise FileNotFoundError(pkl_file)
|
||||
|
||||
response = pickle.load(open(pkl_file, "rb"))
|
||||
if isinstance(response, dict) and "error" in response:
|
||||
continue
|
||||
|
||||
hyps.append(response.choices[0].message.content.strip())
|
||||
refs.append(ex["answer"])
|
||||
|
||||
if not hyps:
|
||||
print("No valid responses to score!")
|
||||
return
|
||||
|
||||
batch_size = 64
|
||||
all_f1: List[float] = []
|
||||
for i in tqdm(range(0, len(hyps), batch_size), desc="Scoring batches"):
|
||||
h_batch = hyps[i : i + batch_size]
|
||||
r_batch = refs[i : i + batch_size]
|
||||
_, _, f1_scores = scorer.score(h_batch, r_batch, verbose=False)
|
||||
all_f1.extend([float(x) for x in f1_scores])
|
||||
|
||||
avg = sum(all_f1) / len(all_f1)
|
||||
print(f"Average BERTScore (F1): {avg:.2%}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Run benchmark and evaluation in one go."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--api-url",
|
||||
default="http://127.0.0.1:30000/v1",
|
||||
help="OpenAI‑compatible API base URL",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="meta-llama/Llama-4-Maverick-17B-128E-Instruct",
|
||||
help="Model name or ID, only used for model name",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-concurrency", type=int, default=144, help="Maximum concurrent requests"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir", default="tmp-output-dir", help="Directory for cached responses"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-prompts", type=int, default=10000, help="Number of prompts to run"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
asyncio.run(benchmark(args))
|
||||
|
||||
analyse(args)
|
||||
@@ -12,7 +12,7 @@ def _wrap_debug_kernel(func: F, op_name: str | None = None) -> F:
|
||||
return func
|
||||
|
||||
try:
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
except Exception:
|
||||
return func
|
||||
|
||||
|
||||
@@ -84,7 +84,7 @@ import msgspec
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.registry import register_kernel
|
||||
from sglang.kernels.spec import (
|
||||
CapabilityRequirement,
|
||||
|
||||
@@ -4,13 +4,13 @@ from typing import TYPE_CHECKING, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
from .utils import make_name
|
||||
|
||||
@@ -4,8 +4,8 @@ from typing import Optional, Union
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import cache_once
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils import get_device_capability, is_musa
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@ from typing import Callable, Optional, Tuple, Union
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
|
||||
try:
|
||||
if os.environ.get("SGLANG_INKLING_FA4_USE_PIP") == "1":
|
||||
|
||||
@@ -9,7 +9,7 @@ from typing import Callable, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.ops.attention.flash_attention_v4 import (
|
||||
_flash_attn_import_error,
|
||||
_flash_attn_varlen_func,
|
||||
|
||||
@@ -13,13 +13,13 @@ from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
@@ -14,8 +14,8 @@ from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
@@ -10,8 +10,8 @@ from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
@@ -7,7 +7,6 @@ import torch
|
||||
import tvm_ffi
|
||||
from tvm_ffi import Module
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
@@ -15,6 +14,7 @@ from sglang.kernels.jit.utils import (
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
|
||||
|
||||
class AllReduceAlgo(enum.Enum):
|
||||
|
||||
@@ -4,8 +4,8 @@ from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
@@ -2,7 +2,7 @@ import torch
|
||||
import triton # type: ignore
|
||||
import triton.language as tl # type: ignore
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
|
||||
@@ -38,7 +38,7 @@ from cutlass.cute import experimental as cute_ext
|
||||
from cutlass.cute.nvgpu import tcgen05
|
||||
from cutlass.cute.runtime import from_dlpack, make_fake_stream
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.utils import is_sm100_supported
|
||||
from sglang.srt.utils.common import direct_register_custom_op
|
||||
|
||||
|
||||
@@ -32,7 +32,7 @@ from cutlass._mlir.dialects import llvm
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
from cutlass.utils import get_smem_capacity_in_bytes
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.utils import get_device_sm
|
||||
from sglang.srt.utils.common import direct_register_custom_op
|
||||
|
||||
|
||||
@@ -11,13 +11,13 @@ from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.utils.common import direct_register_custom_op
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
@@ -11,13 +11,13 @@ from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
@@ -4,8 +4,8 @@ from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.utils.common import is_sm120_supported
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
|
||||
@@ -3,8 +3,8 @@ from __future__ import annotations
|
||||
import logging
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import torch
|
||||
|
||||
@@ -5,13 +5,13 @@ from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
@@ -15,8 +15,8 @@ from __future__ import annotations
|
||||
import logging
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import torch
|
||||
|
||||
@@ -7,8 +7,8 @@ import torch
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import cache_once, is_arch_support_pdl, load_jit
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.ops.moe import moe_route_radix
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
@@ -4,8 +4,8 @@ from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sgl_kernel.scalar_type import ScalarType
|
||||
|
||||
@@ -4,8 +4,8 @@ from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
@@ -4,8 +4,8 @@ from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit, make_cpp_args
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sgl_kernel.scalar_type import ScalarType
|
||||
|
||||
@@ -4,8 +4,8 @@ from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from tvm_ffi.module import Module
|
||||
|
||||
@@ -4,13 +4,13 @@ from typing import TYPE_CHECKING, Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
@@ -10,13 +10,13 @@ from typing import TYPE_CHECKING, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import (
|
||||
cache_once,
|
||||
is_arch_support_pdl,
|
||||
load_jit,
|
||||
make_cpp_args,
|
||||
)
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.utils.custom_op import register_custom_op
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
@@ -2,8 +2,8 @@ from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.jit.utils import cache_once, load_jit
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import torch
|
||||
|
||||
@@ -3,9 +3,9 @@
|
||||
import re
|
||||
from typing import Callable, List, Optional, Union
|
||||
|
||||
from sglang.global_config import global_config
|
||||
from sglang.lang.backend.base_backend import BaseBackend
|
||||
from sglang.lang.choices import ChoicesSamplingMethod, token_length_normalized
|
||||
from sglang.lang.global_config import global_config
|
||||
from sglang.lang.ir import (
|
||||
SglExpr,
|
||||
SglExprList,
|
||||
|
||||
@@ -8,10 +8,10 @@ from typing import Dict, List, Optional, Union
|
||||
import aiohttp
|
||||
import requests
|
||||
|
||||
from sglang.global_config import global_config
|
||||
from sglang.lang.backend.base_backend import BaseBackend
|
||||
from sglang.lang.chat_template import get_chat_template, get_chat_template_by_model_path
|
||||
from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod
|
||||
from sglang.lang.global_config import global_config
|
||||
from sglang.lang.interpreter import StreamExecutor
|
||||
from sglang.lang.ir import (
|
||||
REGEX_BOOL,
|
||||
|
||||
@@ -1,12 +1,8 @@
|
||||
"""Global configurations"""
|
||||
|
||||
# FIXME: deprecate this file and move all usage to sglang.srt.environ or sglang.__init__.py
|
||||
"""Global configuration for the frontend language API."""
|
||||
|
||||
|
||||
class GlobalConfig:
|
||||
"""
|
||||
Store some global constants.
|
||||
"""
|
||||
"""Store mutable process-wide frontend settings."""
|
||||
|
||||
def __init__(self):
|
||||
# Verbosity level
|
||||
@@ -27,3 +23,5 @@ class GlobalConfig:
|
||||
|
||||
|
||||
global_config = GlobalConfig()
|
||||
|
||||
__all__ = ["GlobalConfig", "global_config"]
|
||||
@@ -14,7 +14,7 @@ from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
import tqdm
|
||||
|
||||
from sglang.global_config import global_config
|
||||
from sglang.lang.global_config import global_config
|
||||
from sglang.lang.ir import (
|
||||
SglCommitLazy,
|
||||
SglConcateAndAppend,
|
||||
|
||||
@@ -5,8 +5,8 @@ import inspect
|
||||
import warnings
|
||||
from typing import List, Optional, Union
|
||||
|
||||
from sglang.global_config import global_config
|
||||
from sglang.lang.choices import ChoicesSamplingMethod
|
||||
from sglang.lang.global_config import global_config
|
||||
|
||||
REGEX_INT = r"[-+]?[0-9]+[ \n]*"
|
||||
REGEX_FLOAT = r"[-+]?[0-9]*\.?[0-9]+[ \n]*"
|
||||
|
||||
@@ -5,6 +5,7 @@ import os
|
||||
import sys
|
||||
import warnings
|
||||
|
||||
from sglang.srt.plugins import load_plugins
|
||||
from sglang.srt.server_args import prepare_server_args
|
||||
from sglang.srt.utils import kill_process_tree
|
||||
from sglang.srt.utils.common import suppress_noisy_warnings
|
||||
@@ -61,8 +62,6 @@ if __name__ == "__main__":
|
||||
stacklevel=1,
|
||||
)
|
||||
|
||||
from sglang.srt.plugins import load_plugins
|
||||
|
||||
load_plugins()
|
||||
|
||||
server_args = prepare_server_args(sys.argv[1:])
|
||||
|
||||
@@ -12,7 +12,7 @@ if TYPE_CHECKING:
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import wrap_method_with_debug_kernel_once
|
||||
from sglang.kernels.kernel_api_logging import wrap_method_with_debug_kernel_once
|
||||
from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
|
||||
|
||||
|
||||
|
||||
@@ -8,7 +8,7 @@ from typing import Any
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
|
||||
@@ -10,7 +10,7 @@ import torch.distributed as dist
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from sglang.kernel_api_logging import wrap_method_with_debug_kernel_once
|
||||
from sglang.kernels.kernel_api_logging import wrap_method_with_debug_kernel_once
|
||||
from sglang.multimodal_gen.runtime.distributed import (
|
||||
divide,
|
||||
get_tp_group,
|
||||
|
||||
@@ -4,7 +4,7 @@ from typing import Optional, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.ops.diffusion.triton.rotary import apply_rotary_embedding
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
|
||||
|
||||
@@ -10,7 +10,7 @@ from typing import Any, Callable, List, Optional
|
||||
import torch
|
||||
from torch.library import Library
|
||||
|
||||
from sglang.kernel_api_logging import debug_torch_op
|
||||
from sglang.kernels.kernel_api_logging import debug_torch_op
|
||||
from sglang.multimodal_gen.runtime.platforms import current_platform
|
||||
|
||||
|
||||
|
||||
@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Iterable, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.utils.common import is_npu
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
@@ -18,7 +18,7 @@ from typing import TYPE_CHECKING, Callable, List, Optional, Union
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.ops.attention.utils import (
|
||||
assert_buffer_fits,
|
||||
create_flashinfer_kv_indices_triton,
|
||||
|
||||
@@ -12,7 +12,7 @@ import torch
|
||||
from torch import nn
|
||||
from torch.nn.parameter import Parameter, UninitializedParameter
|
||||
|
||||
from sglang.kernel_api_logging import wrap_method_with_debug_kernel_once
|
||||
from sglang.kernels.kernel_api_logging import wrap_method_with_debug_kernel_once
|
||||
from sglang.srt.distributed import (
|
||||
divide,
|
||||
get_tp_group,
|
||||
|
||||
@@ -5,7 +5,7 @@ from typing import NamedTuple, Optional
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.dp_attention import (
|
||||
get_dp_global_num_tokens,
|
||||
|
||||
@@ -102,14 +102,14 @@ def awq_dequantize_func():
|
||||
|
||||
return awq_dequantize
|
||||
elif _is_hip:
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.ops.quantization.awq_triton import (
|
||||
awq_dequantize_triton as awq_dequantize,
|
||||
)
|
||||
|
||||
return debug_kernel_api(awq_dequantize, op_name="DeepseekCommon.awq_dequantize")
|
||||
elif _is_npu:
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.ops.quantization.awq_triton import (
|
||||
awq_dequantize_decomposition as awq_dequantize,
|
||||
)
|
||||
|
||||
@@ -26,7 +26,7 @@ import triton.language as tl
|
||||
from torch import nn
|
||||
from transformers import PretrainedConfig
|
||||
|
||||
from sglang.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.kernel_api_logging import debug_kernel_api
|
||||
from sglang.kernels.ops.communication.all_reduce import (
|
||||
fused_parallel_qknorm,
|
||||
get_fused_parallel_qknorm_max_occupancy,
|
||||
|
||||
@@ -6,7 +6,7 @@ from typing import Any, Callable, List, Optional, TypeVar, Union, overload
|
||||
import torch
|
||||
import torch.library
|
||||
|
||||
from sglang.kernel_api_logging import debug_torch_op
|
||||
from sglang.kernels.kernel_api_logging import debug_torch_op
|
||||
|
||||
F = TypeVar("F", bound=Callable)
|
||||
|
||||
|
||||
@@ -33,7 +33,7 @@ import torch.nn.functional as F
|
||||
from PIL import Image
|
||||
|
||||
from sglang.benchmark.serving import run_benchmark
|
||||
from sglang.global_config import global_config
|
||||
from sglang.lang.global_config import global_config
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.utils import (
|
||||
get_bool_env_var,
|
||||
|
||||
Reference in New Issue
Block a user