From 198a7b2fc990bccd5d1fbeba7a5fe2ef992d5f8e Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Mon, 17 Aug 2026 14:24:35 -0700 Subject: [PATCH] [Misc] Clean up python/sglang package structure (#35062) --- python/sglang/README.md | 40 ++- python/sglang/__init__.py | 69 ++-- .../{_mps_stub.py => _platform_stubs.py} | 228 ++++++++++++- python/sglang/_triton_stub.py | 228 ------------- python/sglang/eval/llama3_eval.py | 315 ------------------ python/sglang/eval/loogle_eval.py | 164 --------- .../aot/python/sgl_kernel/debug_utils.py | 2 +- python/sglang/kernels/fused_op.py | 2 +- .../{ => kernels}/kernel_api_logging.py | 0 .../kernels/ops/attention/dsv4/fp8_wo_a.py | 2 +- .../ops/attention/flash_attention_v3.py | 2 +- .../ops/attention/flash_attention_v4.py | 2 +- .../ops/attention/flash_attention_v4_sm120.py | 2 +- .../ops/attention/fused_store_index_cache.py | 2 +- .../ops/attention/qprep_bf16_fp8_sm90.py | 2 +- .../sparse_mla_q8kv8_prefill_sm90.py | 2 +- .../kernels/ops/communication/all_reduce.py | 2 +- .../ops/diffusion/timestep_embedding.py | 2 +- .../ops/diffusion/triton/rmsnorm_onepass.py | 2 +- .../kernels/ops/gemm/cutedsl_bf16_gemm.py | 2 +- .../ops/gemm/cutedsl_dsv3_fused_a_gemm.py | 2 +- .../kernels/ops/gemm/dsv3_fused_a_gemm.py | 2 +- .../kernels/ops/gemm/dsv3_router_gemm.py | 2 +- .../kernels/ops/gemm/fp8_blockwise_gemm.py | 2 +- python/sglang/kernels/ops/kvcache/hicache.py | 2 +- python/sglang/kernels/ops/layernorm/norm.py | 2 +- .../kernels/ops/mamba/transfer_mamba.py | 2 +- .../sglang/kernels/ops/moe/moe_fused_gate.py | 2 +- .../kernels/ops/moe/moe_wna16_marlin.py | 2 +- .../ops/quantization/awq_marlin_repack.py | 2 +- .../kernels/ops/quantization/gptq_marlin.py | 2 +- .../ops/quantization/gptq_marlin_repack.py | 2 +- .../ops/quantization/per_token_group_quant.py | 2 +- .../per_token_group_quant_8bit_v2.py | 2 +- .../ops/speculative/ngram_embedding.py | 2 +- python/sglang/lang/api.py | 2 +- .../sglang/lang/backend/runtime_endpoint.py | 2 +- python/sglang/{ => lang}/global_config.py | 10 +- python/sglang/lang/interpreter.py | 2 +- python/sglang/lang/ir.py | 2 +- python/sglang/launch_server.py | 3 +- .../attention/backends/attention_backend.py | 2 +- .../runtime/layers/custom_op.py | 2 +- .../multimodal_gen/runtime/layers/linear.py | 2 +- .../runtime/layers/rotary_embedding/utils.py | 2 +- .../multimodal_gen/runtime/layers/utils.py | 2 +- .../srt/layers/attention/base_attn_backend.py | 2 +- .../layers/attention/flashinfer_backend.py | 2 +- python/sglang/srt/layers/linear.py | 2 +- .../layers/moe/token_dispatcher/flashinfer.py | 2 +- .../srt/models/deepseek_common/utils.py | 4 +- python/sglang/srt/models/minimax_m2.py | 2 +- python/sglang/srt/utils/custom_op.py | 2 +- python/sglang/test/test_utils.py | 2 +- 54 files changed, 319 insertions(+), 830 deletions(-) rename python/sglang/{_mps_stub.py => _platform_stubs.py} (50%) delete mode 100644 python/sglang/_triton_stub.py delete mode 100644 python/sglang/eval/llama3_eval.py delete mode 100644 python/sglang/eval/loogle_eval.py rename python/sglang/{ => kernels}/kernel_api_logging.py (100%) rename python/sglang/{ => lang}/global_config.py (77%) diff --git a/python/sglang/README.md b/python/sglang/README.md index de0a7189f..e9801d2ac 100644 --- a/python/sglang/README.md +++ b/python/sglang/README.md @@ -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. diff --git a/python/sglang/__init__.py b/python/sglang/__init__.py index 37ae1d38e..3fbec34a6 100644 --- a/python/sglang/__init__.py +++ b/python/sglang/__init__.py @@ -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__", ] diff --git a/python/sglang/_mps_stub.py b/python/sglang/_platform_stubs.py similarity index 50% rename from python/sglang/_mps_stub.py rename to python/sglang/_platform_stubs.py index b463ff6be..ccb3c2c15 100644 --- a/python/sglang/_mps_stub.py +++ b/python/sglang/_platform_stubs.py @@ -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"" + + +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 diff --git a/python/sglang/_triton_stub.py b/python/sglang/_triton_stub.py deleted file mode 100644 index b2e252bf1..000000000 --- a/python/sglang/_triton_stub.py +++ /dev/null @@ -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"" - - -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 diff --git a/python/sglang/eval/llama3_eval.py b/python/sglang/eval/llama3_eval.py deleted file mode 100644 index 4a3c736de..000000000 --- a/python/sglang/eval/llama3_eval.py +++ /dev/null @@ -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) diff --git a/python/sglang/eval/loogle_eval.py b/python/sglang/eval/loogle_eval.py deleted file mode 100644 index 895362cd1..000000000 --- a/python/sglang/eval/loogle_eval.py +++ /dev/null @@ -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) diff --git a/python/sglang/kernels/aot/python/sgl_kernel/debug_utils.py b/python/sglang/kernels/aot/python/sgl_kernel/debug_utils.py index 57e176b15..820aa53e8 100644 --- a/python/sglang/kernels/aot/python/sgl_kernel/debug_utils.py +++ b/python/sglang/kernels/aot/python/sgl_kernel/debug_utils.py @@ -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 diff --git a/python/sglang/kernels/fused_op.py b/python/sglang/kernels/fused_op.py index 441bd3107..332274152 100644 --- a/python/sglang/kernels/fused_op.py +++ b/python/sglang/kernels/fused_op.py @@ -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, diff --git a/python/sglang/kernel_api_logging.py b/python/sglang/kernels/kernel_api_logging.py similarity index 100% rename from python/sglang/kernel_api_logging.py rename to python/sglang/kernels/kernel_api_logging.py diff --git a/python/sglang/kernels/ops/attention/dsv4/fp8_wo_a.py b/python/sglang/kernels/ops/attention/dsv4/fp8_wo_a.py index 39cb66e85..a170c58ab 100644 --- a/python/sglang/kernels/ops/attention/dsv4/fp8_wo_a.py +++ b/python/sglang/kernels/ops/attention/dsv4/fp8_wo_a.py @@ -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 diff --git a/python/sglang/kernels/ops/attention/flash_attention_v3.py b/python/sglang/kernels/ops/attention/flash_attention_v3.py index 0d7cbf5ec..898f60088 100644 --- a/python/sglang/kernels/ops/attention/flash_attention_v3.py +++ b/python/sglang/kernels/ops/attention/flash_attention_v3.py @@ -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 diff --git a/python/sglang/kernels/ops/attention/flash_attention_v4.py b/python/sglang/kernels/ops/attention/flash_attention_v4.py index 613c6daf0..499b956d2 100644 --- a/python/sglang/kernels/ops/attention/flash_attention_v4.py +++ b/python/sglang/kernels/ops/attention/flash_attention_v4.py @@ -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": diff --git a/python/sglang/kernels/ops/attention/flash_attention_v4_sm120.py b/python/sglang/kernels/ops/attention/flash_attention_v4_sm120.py index 460ad3ab3..d534618df 100644 --- a/python/sglang/kernels/ops/attention/flash_attention_v4_sm120.py +++ b/python/sglang/kernels/ops/attention/flash_attention_v4_sm120.py @@ -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, diff --git a/python/sglang/kernels/ops/attention/fused_store_index_cache.py b/python/sglang/kernels/ops/attention/fused_store_index_cache.py index cb93b352c..500a2d97c 100644 --- a/python/sglang/kernels/ops/attention/fused_store_index_cache.py +++ b/python/sglang/kernels/ops/attention/fused_store_index_cache.py @@ -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 diff --git a/python/sglang/kernels/ops/attention/qprep_bf16_fp8_sm90.py b/python/sglang/kernels/ops/attention/qprep_bf16_fp8_sm90.py index a8557c1af..66a7b8589 100644 --- a/python/sglang/kernels/ops/attention/qprep_bf16_fp8_sm90.py +++ b/python/sglang/kernels/ops/attention/qprep_bf16_fp8_sm90.py @@ -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 diff --git a/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py b/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py index 355e0dde0..734a78c63 100644 --- a/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py +++ b/python/sglang/kernels/ops/attention/sparse_mla_q8kv8_prefill_sm90.py @@ -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: diff --git a/python/sglang/kernels/ops/communication/all_reduce.py b/python/sglang/kernels/ops/communication/all_reduce.py index 54cd7757f..65c5d1474 100644 --- a/python/sglang/kernels/ops/communication/all_reduce.py +++ b/python/sglang/kernels/ops/communication/all_reduce.py @@ -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): diff --git a/python/sglang/kernels/ops/diffusion/timestep_embedding.py b/python/sglang/kernels/ops/diffusion/timestep_embedding.py index c78423845..7aaedd877 100644 --- a/python/sglang/kernels/ops/diffusion/timestep_embedding.py +++ b/python/sglang/kernels/ops/diffusion/timestep_embedding.py @@ -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 diff --git a/python/sglang/kernels/ops/diffusion/triton/rmsnorm_onepass.py b/python/sglang/kernels/ops/diffusion/triton/rmsnorm_onepass.py index 065205381..3c14a30c9 100644 --- a/python/sglang/kernels/ops/diffusion/triton/rmsnorm_onepass.py +++ b/python/sglang/kernels/ops/diffusion/triton/rmsnorm_onepass.py @@ -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 diff --git a/python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py b/python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py index 073935744..45553b7a8 100644 --- a/python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py +++ b/python/sglang/kernels/ops/gemm/cutedsl_bf16_gemm.py @@ -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 diff --git a/python/sglang/kernels/ops/gemm/cutedsl_dsv3_fused_a_gemm.py b/python/sglang/kernels/ops/gemm/cutedsl_dsv3_fused_a_gemm.py index 7e89986d2..0df7e9470 100644 --- a/python/sglang/kernels/ops/gemm/cutedsl_dsv3_fused_a_gemm.py +++ b/python/sglang/kernels/ops/gemm/cutedsl_dsv3_fused_a_gemm.py @@ -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 diff --git a/python/sglang/kernels/ops/gemm/dsv3_fused_a_gemm.py b/python/sglang/kernels/ops/gemm/dsv3_fused_a_gemm.py index 7652d1b0d..daf809270 100644 --- a/python/sglang/kernels/ops/gemm/dsv3_fused_a_gemm.py +++ b/python/sglang/kernels/ops/gemm/dsv3_fused_a_gemm.py @@ -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: diff --git a/python/sglang/kernels/ops/gemm/dsv3_router_gemm.py b/python/sglang/kernels/ops/gemm/dsv3_router_gemm.py index 05ed46d72..6a49bbdbb 100644 --- a/python/sglang/kernels/ops/gemm/dsv3_router_gemm.py +++ b/python/sglang/kernels/ops/gemm/dsv3_router_gemm.py @@ -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: diff --git a/python/sglang/kernels/ops/gemm/fp8_blockwise_gemm.py b/python/sglang/kernels/ops/gemm/fp8_blockwise_gemm.py index e7176fa4c..826f34cf0 100644 --- a/python/sglang/kernels/ops/gemm/fp8_blockwise_gemm.py +++ b/python/sglang/kernels/ops/gemm/fp8_blockwise_gemm.py @@ -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 diff --git a/python/sglang/kernels/ops/kvcache/hicache.py b/python/sglang/kernels/ops/kvcache/hicache.py index 1cce02fdb..66c5c0feb 100644 --- a/python/sglang/kernels/ops/kvcache/hicache.py +++ b/python/sglang/kernels/ops/kvcache/hicache.py @@ -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 diff --git a/python/sglang/kernels/ops/layernorm/norm.py b/python/sglang/kernels/ops/layernorm/norm.py index 046afbe48..c3a8e948f 100644 --- a/python/sglang/kernels/ops/layernorm/norm.py +++ b/python/sglang/kernels/ops/layernorm/norm.py @@ -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 diff --git a/python/sglang/kernels/ops/mamba/transfer_mamba.py b/python/sglang/kernels/ops/mamba/transfer_mamba.py index 2f7c42c23..6cdc1a713 100644 --- a/python/sglang/kernels/ops/mamba/transfer_mamba.py +++ b/python/sglang/kernels/ops/mamba/transfer_mamba.py @@ -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 diff --git a/python/sglang/kernels/ops/moe/moe_fused_gate.py b/python/sglang/kernels/ops/moe/moe_fused_gate.py index 112f501d0..5e9c8b23a 100644 --- a/python/sglang/kernels/ops/moe/moe_fused_gate.py +++ b/python/sglang/kernels/ops/moe/moe_fused_gate.py @@ -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: diff --git a/python/sglang/kernels/ops/moe/moe_wna16_marlin.py b/python/sglang/kernels/ops/moe/moe_wna16_marlin.py index 3e3e25d0d..1defc8a92 100644 --- a/python/sglang/kernels/ops/moe/moe_wna16_marlin.py +++ b/python/sglang/kernels/ops/moe/moe_wna16_marlin.py @@ -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 diff --git a/python/sglang/kernels/ops/quantization/awq_marlin_repack.py b/python/sglang/kernels/ops/quantization/awq_marlin_repack.py index 8e8707709..87faeaf81 100644 --- a/python/sglang/kernels/ops/quantization/awq_marlin_repack.py +++ b/python/sglang/kernels/ops/quantization/awq_marlin_repack.py @@ -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 diff --git a/python/sglang/kernels/ops/quantization/gptq_marlin.py b/python/sglang/kernels/ops/quantization/gptq_marlin.py index a980c1960..9b52594f0 100644 --- a/python/sglang/kernels/ops/quantization/gptq_marlin.py +++ b/python/sglang/kernels/ops/quantization/gptq_marlin.py @@ -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 diff --git a/python/sglang/kernels/ops/quantization/gptq_marlin_repack.py b/python/sglang/kernels/ops/quantization/gptq_marlin_repack.py index 251f3b91f..bf1541a82 100644 --- a/python/sglang/kernels/ops/quantization/gptq_marlin_repack.py +++ b/python/sglang/kernels/ops/quantization/gptq_marlin_repack.py @@ -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 diff --git a/python/sglang/kernels/ops/quantization/per_token_group_quant.py b/python/sglang/kernels/ops/quantization/per_token_group_quant.py index 8d499a347..ca608e425 100644 --- a/python/sglang/kernels/ops/quantization/per_token_group_quant.py +++ b/python/sglang/kernels/ops/quantization/per_token_group_quant.py @@ -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: diff --git a/python/sglang/kernels/ops/quantization/per_token_group_quant_8bit_v2.py b/python/sglang/kernels/ops/quantization/per_token_group_quant_8bit_v2.py index 70d92e96d..7ad66bfa1 100644 --- a/python/sglang/kernels/ops/quantization/per_token_group_quant_8bit_v2.py +++ b/python/sglang/kernels/ops/quantization/per_token_group_quant_8bit_v2.py @@ -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: diff --git a/python/sglang/kernels/ops/speculative/ngram_embedding.py b/python/sglang/kernels/ops/speculative/ngram_embedding.py index 3b2de30c0..300cca030 100644 --- a/python/sglang/kernels/ops/speculative/ngram_embedding.py +++ b/python/sglang/kernels/ops/speculative/ngram_embedding.py @@ -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 diff --git a/python/sglang/lang/api.py b/python/sglang/lang/api.py index aa691bb9f..a4f09966d 100644 --- a/python/sglang/lang/api.py +++ b/python/sglang/lang/api.py @@ -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, diff --git a/python/sglang/lang/backend/runtime_endpoint.py b/python/sglang/lang/backend/runtime_endpoint.py index 9b453f72f..c29b82735 100644 --- a/python/sglang/lang/backend/runtime_endpoint.py +++ b/python/sglang/lang/backend/runtime_endpoint.py @@ -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, diff --git a/python/sglang/global_config.py b/python/sglang/lang/global_config.py similarity index 77% rename from python/sglang/global_config.py rename to python/sglang/lang/global_config.py index fcd65b5ed..b63060c80 100644 --- a/python/sglang/global_config.py +++ b/python/sglang/lang/global_config.py @@ -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"] diff --git a/python/sglang/lang/interpreter.py b/python/sglang/lang/interpreter.py index 2e17c9cbe..90dd41857 100644 --- a/python/sglang/lang/interpreter.py +++ b/python/sglang/lang/interpreter.py @@ -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, diff --git a/python/sglang/lang/ir.py b/python/sglang/lang/ir.py index 5a808e191..45cb6c859 100644 --- a/python/sglang/lang/ir.py +++ b/python/sglang/lang/ir.py @@ -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]*" diff --git a/python/sglang/launch_server.py b/python/sglang/launch_server.py index 8091692fd..5e47dd22b 100644 --- a/python/sglang/launch_server.py +++ b/python/sglang/launch_server.py @@ -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:]) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py index 86e0ea6e7..41a080258 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/attention_backend.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/layers/custom_op.py b/python/sglang/multimodal_gen/runtime/layers/custom_op.py index eb745b6fe..98016b155 100644 --- a/python/sglang/multimodal_gen/runtime/layers/custom_op.py +++ b/python/sglang/multimodal_gen/runtime/layers/custom_op.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/layers/linear.py b/python/sglang/multimodal_gen/runtime/layers/linear.py index 4e6657388..07b69cba2 100644 --- a/python/sglang/multimodal_gen/runtime/layers/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/linear.py @@ -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, diff --git a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py index 9fa7d8de1..559703c71 100644 --- a/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py +++ b/python/sglang/multimodal_gen/runtime/layers/rotary_embedding/utils.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/layers/utils.py b/python/sglang/multimodal_gen/runtime/layers/utils.py index 1feeb3f36..2454eae78 100644 --- a/python/sglang/multimodal_gen/runtime/layers/utils.py +++ b/python/sglang/multimodal_gen/runtime/layers/utils.py @@ -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 diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index 614d63216..b9eb964c3 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/python/sglang/srt/layers/attention/base_attn_backend.py @@ -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: diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 28b0d2205..185943e0c 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -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, diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index 57a8e4459..6cf18264b 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -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, diff --git a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py index 35759b919..577a9bf8b 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/flashinfer.py @@ -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, diff --git a/python/sglang/srt/models/deepseek_common/utils.py b/python/sglang/srt/models/deepseek_common/utils.py index f56f47ece..43a26b275 100644 --- a/python/sglang/srt/models/deepseek_common/utils.py +++ b/python/sglang/srt/models/deepseek_common/utils.py @@ -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, ) diff --git a/python/sglang/srt/models/minimax_m2.py b/python/sglang/srt/models/minimax_m2.py index 29adf3ef4..0b655a872 100644 --- a/python/sglang/srt/models/minimax_m2.py +++ b/python/sglang/srt/models/minimax_m2.py @@ -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, diff --git a/python/sglang/srt/utils/custom_op.py b/python/sglang/srt/utils/custom_op.py index 720776501..7895393b2 100644 --- a/python/sglang/srt/utils/custom_op.py +++ b/python/sglang/srt/utils/custom_op.py @@ -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) diff --git a/python/sglang/test/test_utils.py b/python/sglang/test/test_utils.py index cdd8f8c55..3d88230d7 100644 --- a/python/sglang/test/test_utils.py +++ b/python/sglang/test/test_utils.py @@ -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,