commit 466c38a2d43837c82f6c09c55dbd69eb7a8f2336 Author: hailin Date: Tue May 27 11:42:20 2025 +0800 first commit diff --git a/sglang/.clang-format-ignore b/sglang/.clang-format-ignore new file mode 100644 index 0000000..15c76cc --- /dev/null +++ b/sglang/.clang-format-ignore @@ -0,0 +1 @@ +sgl-kernel/3rdparty/tensorrt_llm/* diff --git a/sglang/.devcontainer/Dockerfile b/sglang/.devcontainer/Dockerfile new file mode 100644 index 0000000..0c061cd --- /dev/null +++ b/sglang/.devcontainer/Dockerfile @@ -0,0 +1,35 @@ +From lmsysorg/sglang:dev + +# Create non-root user with specified UID and GID +# NOTE: Replace with your own UID and GID. This is a workaround from https://github.com/microsoft/vscode-remote-release/issues/49#issuecomment-489060908. +ARG HOST_UID=1003 +ARG HOST_GID=1003 +RUN groupadd -g $HOST_GID devuser && \ + useradd -m -u $HOST_UID -g $HOST_GID -s /bin/zsh devuser + +# Give devuser sudo access +RUN apt-get update && apt-get install -y sudo && \ + echo "devuser ALL=(ALL) NOPASSWD:ALL" > /etc/sudoers.d/devuser && \ + rm -rf /var/lib/apt/lists/* && \ + apt-get clean + +# Set up oh-my-zsh for devuser +RUN cp -r /root/.oh-my-zsh /home/devuser/.oh-my-zsh && \ + cp /root/.zshrc /home/devuser/.zshrc && \ + cp /root/.vimrc /home/devuser/.vimrc && \ + cp /root/.tmux.conf /home/devuser/.tmux.conf && \ + sed -i 's|/root/.oh-my-zsh|/home/devuser/.oh-my-zsh|g' /home/devuser/.zshrc && \ + chown -R devuser:devuser /home/devuser/ + +# Set workspace directory and ownership +WORKDIR /sgl-workspace/sglang +RUN chown -R devuser:devuser /sgl-workspace + +# Switch to devuser +USER devuser + +# Install uv +RUN curl -LsSf https://astral.sh/uv/install.sh | sh + +# Install rust +RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y diff --git a/sglang/.devcontainer/devcontainer.json b/sglang/.devcontainer/devcontainer.json new file mode 100644 index 0000000..5767aa2 --- /dev/null +++ b/sglang/.devcontainer/devcontainer.json @@ -0,0 +1,24 @@ +{ + "name": "sglang", + "build": { + "dockerfile": "Dockerfile" + }, + "remoteUser": "devuser", + "customizations": { + "vscode": { + "extensions": [ + // Python development + "ms-python.python", + "charliermarsh.ruff", + // Rust development + "rust-lang.rust-analyzer", + "tamasfe.even-better-toml" + ] + } + }, + "forwardPorts": [], + "runArgs": [ + "--gpus", + "all" + ] +} diff --git a/sglang/.editorconfig b/sglang/.editorconfig new file mode 100644 index 0000000..030a729 --- /dev/null +++ b/sglang/.editorconfig @@ -0,0 +1,25 @@ +# https://editorconfig.org/ + +root = true + +[*] +charset = utf-8 +end_of_line = lf +indent_style = space +indent_size = 4 +trim_trailing_whitespace = true +insert_final_newline = true + +[*.{json,yaml,yml}] +indent_size = 2 + +[*.md] +indent_size = 2 +x-soft-wrap-text = true + +[*.rst] +indent_size = 4 +x-soft-wrap-text = true + +[Makefile] +indent_style = tab diff --git a/sglang/.gitignore b/sglang/.gitignore new file mode 100644 index 0000000..75e29fa --- /dev/null +++ b/sglang/.gitignore @@ -0,0 +1,229 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class + +# C extensions +*.so + +# Distribution / packaging +.Python +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +wheels/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +*.py,cover +.hypothesis/ +.pytest_cache/ +cover/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 +db.sqlite3-journal + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/_build/ + +# PyBuilder +.pybuilder/ +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +# For a library or package, you might want to ignore these files since the code is +# intended to run in multiple environments; otherwise, check them in: +# .python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don't work, or not +# install all needed dependencies. +#Pipfile.lock + +# poetry +# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control. +# This is especially recommended for binary packages to ensure reproducibility, and is more +# commonly ignored for libraries. +# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control +#poetry.lock + +# pdm +# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control. +#pdm.lock +# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it +# in version control. +# https://pdm.fming.dev/#use-with-ide +.pdm.toml + +# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm +__pypackages__/ + +# Celery stuff +celerybeat-schedule +celerybeat.pid + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# pytype static type analyzer +.pytype/ + +# Cython debug symbols +cython_debug/ + +# PyCharm +# JetBrains specific template is maintained in a separate JetBrains.gitignore that can +# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore +# and can be added to the global gitignore or merged into this file. For a more nuclear +# option (not recommended) you can uncomment the following to ignore the entire idea folder. +.idea/ + +# MacOS +.DS_Store + +# Vim +*.swp + +# Documentation +docs/_build + +# SGL +benchmark/mmlu/data +benchmark/mmlu/data.tar +benchmark/llava_bench/images +benchmark/llava_bench/mme_pack +*.jsonl +tmp*.txt + +# Plots +*.png +*.pdf + +# personnal +work_dirs/ +*.csv + +!logo.png + +# Prerequisites +*.d + +# Compiled Object files +*.slo +*.lo +*.o +*.obj + +# Precompiled Headers +*.gch +*.pch + +# Compiled Dynamic libraries +*.so +*.dylib +*.dll + +# Fortran module files +*.mod +*.smod + +# Compiled Static libraries +*.lai +*.la +*.a +*.lib + +# Executables +*.exe +*.out +*.app + +compile_commands.json + +*.iml + +# VSCode +.vscode + +1 diff --git a/sglang/.gitmodules b/sglang/.gitmodules new file mode 100644 index 0000000..265d8d9 --- /dev/null +++ b/sglang/.gitmodules @@ -0,0 +1,12 @@ +[submodule "sgl-kernel/3rdparty/cutlass"] + path = sgl-kernel/3rdparty/cutlass + url = https://github.com/NVIDIA/cutlass.git +[submodule "sgl-kernel/3rdparty/cccl"] + path = sgl-kernel/3rdparty/cccl + url = https://github.com/NVIDIA/cccl.git +[submodule "sgl-kernel/3rdparty/flashinfer"] + path = sgl-kernel/3rdparty/flashinfer + url = https://github.com/flashinfer-ai/flashinfer.git +[submodule "sgl-kernel/3rdparty/deepgemm"] + path = sgl-kernel/3rdparty/deepgemm + url = https://github.com/deepseek-ai/DeepGEMM diff --git a/sglang/.isort.cfg b/sglang/.isort.cfg new file mode 100644 index 0000000..8350951 --- /dev/null +++ b/sglang/.isort.cfg @@ -0,0 +1,3 @@ +[settings] +profile=black +known_first_party=sglang diff --git a/sglang/.pre-commit-config.yaml b/sglang/.pre-commit-config.yaml new file mode 100644 index 0000000..eb30735 --- /dev/null +++ b/sglang/.pre-commit-config.yaml @@ -0,0 +1,48 @@ +default_stages: [pre-commit, pre-push, manual] + +repos: + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v5.0.0 + hooks: + - id: check-symlinks + - id: destroyed-symlinks + - id: trailing-whitespace + - id: end-of-file-fixer + - id: check-yaml + args: [--allow-multiple-documents] + - id: check-toml + - id: check-ast + - id: check-added-large-files + - id: check-merge-conflict + - id: check-shebang-scripts-are-executable + - id: detect-private-key + - id: debug-statements + - id: no-commit-to-branch + - repo: https://github.com/PyCQA/isort + rev: 5.13.2 + hooks: + - id: isort + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.11.2 + hooks: + - id: ruff + args: [--select=F401, --fixable=F401] + files: ^(benchmark/|docs/|examples/) + exclude: \.ipynb$ + - repo: https://github.com/psf/black + rev: 24.10.0 + hooks: + - id: black-jupyter + - repo: https://github.com/pre-commit/mirrors-clang-format + rev: v18.1.8 + hooks: + - id: clang-format + types_or: [c++, cuda] + args: [--style=file, --verbose] + - repo: https://github.com/kynan/nbstripout + rev: 0.8.1 + hooks: + - id: nbstripout + args: + - '--keep-output' + - '--extra-keys=metadata.kernelspec metadata.language_info.version' diff --git a/sglang/3rdparty/amd/profiling/PROFILING.md b/sglang/3rdparty/amd/profiling/PROFILING.md new file mode 100644 index 0000000..7e15ec8 --- /dev/null +++ b/sglang/3rdparty/amd/profiling/PROFILING.md @@ -0,0 +1,425 @@ +## Profiling SGLang Infer System with AMD GPUs +This AppNote describes the SGLang profiling technical, code augment and running steps for systems with AMD Instinct GPUs, nevertheless the same procedure may work with Nvidia GPUs too. +Examples and steps are provided in detail, to facilitate easy reproduce and use to localize performance problem towards optimizations. +Two primary methods are covered: +- [RPD](https://github.com/ROCm/rocmProfileData.git) +- [PyTorch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) + +### Profiling SGLang Infer System with RPD Profiler +RPD profiler is a low-overhead cross-platform profiler. Therefore, the same RPD code augment not only works for profiling on ROCm/AMD GPUs, but also works for profiling on CUDA/Nvidia GPUs as well. To do RPD profiling on SGLang repository, please use scripts and patch files included in this directory and follow the steps below: +1. Install RPD with rpd.patch applied during installation using install_rpd.sh, both files are in this directory. + +install_rpd.sh + +```bash +# download and install RPD +apt update && apt install -y sqlite3 libsqlite3-dev libfmt-dev + +# install rpd module +git clone https://github.com/ROCmSoftwarePlatform/rocmProfileData +cd rocmProfileData +git checkout 976899e9c6dbc6dd2bccf770818e4e44125590ac +git apply rpd.patch +make && make install +cd rocpd_python && python setup.py install && cd .. +cd rpd_tracer && make clean;make install && python setup.py install && cd .. +``` + +rpd.patch + +```bash +diff --git a/rpd_tracer/Makefile b/rpd_tracer/Makefile +index e9d9feb..b2e9e1a 100644 +--- a/rpd_tracer/Makefile ++++ b/rpd_tracer/Makefile +@@ -16,7 +16,7 @@ ifneq (,$(HIP_PATH)) + $(info Building with roctracer) + RPD_LIBS += -L/opt/rocm/lib -lroctracer64 -lroctx64 -lamdhip64 -lrocm_smi64 + RPD_INCLUDES += -I/opt/rocm/include -I/opt/rocm/include/roctracer -I/opt/rocm/include/hsa +- RPD_SRCS += RoctracerDataSource.cpp RocmSmiDataSource.cpp ++ RPD_SRCS += RoctracerDataSource.cpp + RPD_INCLUDES += -D__HIP_PLATFORM_AMD__ + endif +``` +2. Add loadTracer.sh file included in this directory to /sglang/python/sglang. + +loadTracer.sh + +```bash +#!/bin/bash +################################################################################ +# Copyright (c) 2021 - 2023 Advanced Micro Devices, Inc. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +################################################################################ +OUTPUT_FILE="trace.rpd" + +if [ "$1" = "-o" ] ; then + OUTPUT_FILE=$2 + shift + shift +fi + +if [ -e ${OUTPUT_FILE} ] ; then + rm ${OUTPUT_FILE} +fi + +python3 -m rocpd.schema --create ${OUTPUT_FILE} +if [ $? != 0 ] ; then + echo "Error: Could not create rpd file. Please run 'python setup.py install' from the rocpd_python dir" + exit +fi + +export RPDT_FILENAME=${OUTPUT_FILE} +export RPDT_AUTOSTART=0 +LD_PRELOAD=librocm-smi_64:librpd_tracer.so "$@" +``` +3. Apply patch (provided in this directory) with "git apply rpd_profile_server_enable.patch" if the main profiling purpose is to get info on gpu kernels as well as limited cpu activity info. + +#### Common Notes 1 +Please note that although we are doing TP=8 in the example, we purposely only log RPD profiling on 2 ranks in the patch file (i.e.tp_rank=0/1) for profiling/visualization convenience, as even Perfetto streaming mode can only load maximal 8GB json file for visualization. With 2 ranks logged in RPD profiling, we could still check whether there are issues among ranks (e.g. load imbalance issue, nccl issue), and at the same time, we could log relatively longer time duration before the json file generated from RPD file hits 8GB size. + +rpd_profile_server_enable.patch + +```bash +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index 62d1ff9..9021c01 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -71,6 +71,8 @@ from sglang.srt.utils import ( + suppress_other_loggers, + ) + from sglang.utils import get_exception_traceback ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() + + logger = logging.getLogger(__name__) + +@@ -245,6 +247,7 @@ class Scheduler: + ], + with_stack=True, + ) ++ self.rpd = rpdTracerControl() + + @torch.inference_mode() + def event_loop(self): +@@ -1027,15 +1030,24 @@ class Scheduler: + def start_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.start() ++ #self.profiler.start() #block pytorch profiler for rpd profiler enabling ++ if self.tp_rank == 0 or self.tp_rank == 1: ++ self.rpd.start() ++ self.rpd.rangePush("", "rpd profile range", "") ++ logger.info("rpd is enabled") + + def stop_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.stop() +- self.profiler.export_chrome_trace( +- self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" +- ) ++ #self.profiler.stop() ++ #self.profiler.export_chrome_trace( ++ # self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" ++ #) ++ if self.tp_rank ==0 or self.tp_rank ==1: ++ self.rpd.rangePop() ++ self.rpd.stop() ++ self.rpd.flush() ++ logger.info("rpd is done") + logger.info("Profiler is done") +``` + +#### Advanced Debugging with RPD Profiler +Sometimes, we want to use rpd profiler to capture more CPU and python activities in order to debug some challenging issues (e.g. root cause of load imbalance across gpu processes, root cause of bubbles, etc). Only in such cases, we need to apply patch "git apply rpd_profile_server_enable_wCPU_activities.patch", where 3 files are modified. + +rpd_profile_server_enable_wCPU_activities.patch + +```bash +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index 62d1ff9..2edb427 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -71,6 +71,8 @@ from sglang.srt.utils import ( + suppress_other_loggers, + ) + from sglang.utils import get_exception_traceback ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() + + logger = logging.getLogger(__name__) + +@@ -245,6 +247,7 @@ class Scheduler: + ], + with_stack=True, + ) ++ self.rpd = rpdTracerControl() + + @torch.inference_mode() + def event_loop(self): +@@ -1027,15 +1030,26 @@ class Scheduler: + def start_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.start() ++ #self.profiler.start() ++ logger.info("torch profiler is disabled") ++ if self.tp_rank == 0 or self.tp_rank == 1: ++ self.rpd.setPythonTrace(True) ++ self.rpd.start() ++ self.rpd.rangePush("", "scheduler", "") ++ logger.info("rpd is enabled inside scheduler profiling") + + def stop_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.stop() +- self.profiler.export_chrome_trace( +- self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" +- ) ++ #self.profiler.stop() ++ #self.profiler.export_chrome_trace( ++ # self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" ++ #) ++ if self.tp_rank ==0 or self.tp_rank ==1: ++ self.rpd.rangePop() ++ self.rpd.stop() ++ self.rpd.flush() ++ logger.info("rpd is done inside scheduler") + logger.info("Profiler is done") + + +diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py +index 2621ccd..181df85 100644 +--- a/python/sglang/srt/managers/tokenizer_manager.py ++++ b/python/sglang/srt/managers/tokenizer_manager.py +@@ -58,6 +58,10 @@ from sglang.srt.sampling.sampling_params import SamplingParams + from sglang.srt.server_args import PortArgs, ServerArgs + from sglang.srt.utils import is_generation_model, is_multimodal_model + ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() ++ ++ + asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) + + logger = logging.getLogger(__name__) +@@ -514,10 +518,20 @@ class TokenizerManager: + self.send_to_scheduler.send_pyobj(req) + + def start_profile(self): ++ rpd = rpdTracerControl() ++ rpd.setPythonTrace(True) ++ rpd.start() ++ rpd.rangePush("", "tokenizer_manager", "") ++ logger.info("tokenizer_manager rpd profiling started!") + req = ProfileReq.START_PROFILE + self.send_to_scheduler.send_pyobj(req) + + def stop_profile(self): ++ rpd = rpdTracerControl() ++ rpd.rangePop() ++ rpd.stop() ++ rpd.flush() ++ logger.info("rpd profiling is done inside tokenizer_manager!") + req = ProfileReq.STOP_PROFILE + self.send_to_scheduler.send_pyobj(req) + +diff --git a/python/sglang/srt/server.py b/python/sglang/srt/server.py +index 7111c93..2bd722c 100644 +--- a/python/sglang/srt/server.py ++++ b/python/sglang/srt/server.py +@@ -30,6 +30,8 @@ import threading + import time + from http import HTTPStatus + from typing import Dict, List, Optional, Union ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() + + # Fix a bug of Python threading + setattr(threading, "_register_atexit", lambda *args, **kwargs: None) +@@ -152,6 +154,11 @@ async def flush_cache(): + @app.post("/start_profile") + async def start_profile(): + """Start profiling.""" ++ rpd = rpdTracerControl() ++ rpd.setPythonTrace(True) ++ rpd.start() ++ rpd.rangePush("", "server rpd profile range", "") ++ logger.info("rpd profiling started in server.py!") + tokenizer_manager.start_profile() + return Response( + content="Start profiling.\n", +@@ -164,6 +171,11 @@ async def start_profile(): + async def stop_profile(): + """Stop profiling.""" + tokenizer_manager.stop_profile() ++ rpd = rpdTracerControl() ++ rpd.rangePop() ++ rpd.stop() ++ rpd.flush() ++ logger.info("rpd profiling is done in server.py!") + return Response( + content="Stop profiling. This will take some time.\n", + status_code=200, +``` + +4. As an example for grok1 profiling, we create a dummy_grok1 directory with config.json (see content below) inside this directory and copy this directory to the right path for "--model-path" if you want to use the example server.sh file provided. +```bash +cat ../dummy_grok1/config.json +{ + "architectures": [ + "Grok1ModelForCausalLM" + ], + "embedding_multiplier_scale": 78.38367176906169, + "output_multiplier_scale": 0.5773502691896257, + "vocab_size": 131072, + "hidden_size": 6144, + "intermediate_size": 32768, + "max_position_embeddings": 8192, + "num_experts_per_tok": 2, + "num_local_experts": 8, + "num_attention_heads": 48, + "num_hidden_layers": 64, + "num_key_value_heads": 8, + "head_dim": 128, + "rms_norm_eps": 1e-05, + "rope_theta": 10000.0, + "model_type": "mixtral", + "torch_dtype": "bfloat16" +} +``` +5. Launch server with rpd enabled script ./server.sh in one terminal inside the docker container. + +#### Common Notes 2 +- Remember to change model-path to the correct path +- loadTracer.sh is needed to conduct profiling +- SGLANG_TORCH_PROFILER_DIR is used for default torch profiler +- Do not use loadTracer.sh if you are using the torch profiler, simply use python3 -m sglang.launch_server. + + +server.sh + +```bash +#!/bin/bash + +# export SGLANG_TORCH_PROFILER_DIR=/data/sglang/ +export SGLANG_TORCH_PROFILER_DIR=/sgl-workspace/sglang/profile/ + +# Get the current timestamp +TIMESTAMP=$(date +"%Y%m%d_%H%M%S") + +# Define the log file with a timestamp +LOGFILE="sglang_server_log_$TIMESTAMP.json" + +# Run the Python command and save the output to the log file +loadTracer.sh python3 -m sglang.launch_server \ + --model-path /sgl-workspace/sglang/dummy_grok1 \ + --tokenizer-path Xenova/grok-1-tokenizer \ + --load-format dummy \ + --quantization fp8 \ + --tp 8 \ + --port 30000 \ + --disable-radix-cache 2>&1 | tee "$LOGFILE" +``` +6. Open another terminal for the same docker container, and run the rpd enabled ./client.sh after you see "The server is fired up and is ready to roll!" message from server side terminal. + +#### Common Notes 3 +- Use curl http://localhost:30000/start_profile & curl http://localhost:30000/stop_profile to control the start and end of profiling. Check sglang/python/sglang/srt/managers/scheduler.py for more details. +- Please don't use RPD profiler together with PyTorch profiler to avoid interference. +- The rocmProfileData/tools/rpd2tracing.py file is used to generate json file from RPD file. + +client.sh + +```bash +#!/bin/bash + +# Start profiling via API +curl http://localhost:30000/start_profile -H "Content-Type: application/json" + +# Benchmark serving using sglang with random dataset and tokenizer +# Define the log file with a timestamp +TIMESTAMP=$(date +%Y%m%d_%H%M%S) +LOGFILE="sglang_client_log_$TIMESTAMP.json" + +# Run the benchmark with specified parameters and save logs +python3 -m sglang.bench_serving \ + --backend sglang \ + --tokenizer Xenova/grok-1-tokenizer \ + --dataset-name random \ + --random-input 1024\ + --random-output 1024 \ + --num-prompts 120 \ + --request-rate 8 \ + --output-file online.jsonl 2>&1 | tee "$LOGFILE" + +# Stop profiling via API +curl http://localhost:30000/stop_profile -H "Content-Type: application/json" + +# Convert tracing file to csv & json +sqlite3 trace.rpd ".mode csv" ".header on" ".output trace.csv" "select * from top;" ".output stdout" +python3 ./rocmProfileData/tools/rpd2tracing.py trace.rpd trace.json +``` +7. Follow [Perfetto docs](https://perfetto.dev/docs/visualization/large-traces) to visualize large json files. Try to adjust parameters so that the trace.json file size is less than 9GB. + +### Profiling SGLang Infer System with PyTorch Profiler + +Please use the steps as follows: + +1. Apply the patch torch_profiler.patch. Note that you can modify "if self.tp_rank == 0" in the patch to allow more ranks be recorded in profiling. + +torch_profiler.patch +```bash +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index 62d1ff9..6ecd78c 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -240,7 +240,6 @@ class Scheduler: + ) + self.profiler = torch.profiler.profile( + activities=[ +- torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + with_stack=True, +@@ -1033,9 +1032,11 @@ class Scheduler: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") + self.profiler.stop() +- self.profiler.export_chrome_trace( +- self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" +- ) ++ if self.tp_rank == 0: ++ with open(f"stats_repro_{int(time.time())}.txt", "w") as f: ++ print(self.profiler.key_averages(group_by_input_shape=True).table(sort_by="cuda_time_total", row_limit=-1), file=f) ++ print("Profiling stats done.") ++ + logger.info("Profiler is done") +``` + +2. Create the model path directory and copy it to the right path for "--model-path" if you want to use the server.sh file provided. + +3. Modify the included server.sh by removing "loadTracer.sh" before python command and launch script ./server.sh in one terminal inside the docker container. + +4. Similar to step 6 in RPD profiling section, but remove the last 2 lines in client.sh, which converted rpd file into csv and json files. Run modified client.sh for PyTorch profiling. +------- +- [Torch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) diff --git a/sglang/3rdparty/amd/profiling/client.sh b/sglang/3rdparty/amd/profiling/client.sh new file mode 100644 index 0000000..150ea9f --- /dev/null +++ b/sglang/3rdparty/amd/profiling/client.sh @@ -0,0 +1,27 @@ +#!/bin/bash + +# Start profiling via API +curl http://localhost:30000/start_profile -H "Content-Type: application/json" + +# Benchmark serving using sglang with random dataset and tokenizer +# Define the log file with a timestamp +TIMESTAMP=$(date +%Y%m%d_%H%M%S) +LOGFILE="sglang_client_log_$TIMESTAMP.json" + +# Run the benchmark with specified parameters and save logs +python3 -m sglang.bench_serving \ + --backend sglang \ + --tokenizer Xenova/grok-1-tokenizer \ + --dataset-name random \ + --random-input 1024\ + --random-output 1024 \ + --num-prompts 240 \ + --request-rate 8 \ + --output-file online.jsonl 2>&1 | tee "$LOGFILE" + +# Stop profiling via API +curl http://localhost:30000/stop_profile -H "Content-Type: application/json" + +# Convert tracing file to csv & json +sqlite3 trace.rpd ".mode csv" ".header on" ".output trace.csv" "select * from top;" ".output stdout" +python3 /sgl-workspace/rocmProfileData/tools/rpd2tracing.py trace.rpd trace.json diff --git a/sglang/3rdparty/amd/profiling/install_rpd.sh b/sglang/3rdparty/amd/profiling/install_rpd.sh new file mode 100644 index 0000000..d1b04b9 --- /dev/null +++ b/sglang/3rdparty/amd/profiling/install_rpd.sh @@ -0,0 +1,10 @@ +# download and install RPD +apt update && apt install -y sqlite3 libsqlite3-dev libfmt-dev + +# install rpd module +git clone https://github.com/ROCmSoftwarePlatform/rocmProfileData +cd rocmProfileData +git apply rpd.patch +make && make install +cd rocpd_python && python setup.py install && cd .. +cd rpd_tracer && make clean;make install && python setup.py install && cd .. diff --git a/sglang/3rdparty/amd/profiling/loadTracer.sh b/sglang/3rdparty/amd/profiling/loadTracer.sh new file mode 100644 index 0000000..8a95a33 --- /dev/null +++ b/sglang/3rdparty/amd/profiling/loadTracer.sh @@ -0,0 +1,43 @@ +#!/bin/bash +################################################################################ +# Copyright (c) 2021 - 2023 Advanced Micro Devices, Inc. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +################################################################################ +OUTPUT_FILE="trace.rpd" + +if [ "$1" = "-o" ] ; then + OUTPUT_FILE=$2 + shift + shift +fi + +if [ -e ${OUTPUT_FILE} ] ; then + rm ${OUTPUT_FILE} +fi + +python3 -m rocpd.schema --create ${OUTPUT_FILE} +if [ $? != 0 ] ; then + echo "Error: Could not create rpd file. Please run 'python setup.py install' from the rocpd_python dir" + exit +fi + +export RPDT_FILENAME=${OUTPUT_FILE} +export RPDT_AUTOSTART=0 +LD_PRELOAD=librocm-smi_64:librpd_tracer.so "$@" diff --git a/sglang/3rdparty/amd/profiling/rpd.patch b/sglang/3rdparty/amd/profiling/rpd.patch new file mode 100644 index 0000000..8791765 --- /dev/null +++ b/sglang/3rdparty/amd/profiling/rpd.patch @@ -0,0 +1,12 @@ +diff --git a/rpd_tracer/Makefile b/rpd_tracer/Makefile +index e9d9feb..b2e9e1a 100644 +--- a/rpd_tracer/Makefile ++++ b/rpd_tracer/Makefile +@@ -16,7 +16,7 @@ ifneq (,$(HIP_PATH)) + $(info Building with roctracer) + RPD_LIBS += -L/opt/rocm/lib -lroctracer64 -lroctx64 -lamdhip64 -lrocm_smi64 + RPD_INCLUDES += -I/opt/rocm/include -I/opt/rocm/include/roctracer -I/opt/rocm/include/hsa +- RPD_SRCS += RoctracerDataSource.cpp RocmSmiDataSource.cpp ++ RPD_SRCS += RoctracerDataSource.cpp + RPD_INCLUDES += -D__HIP_PLATFORM_AMD__ + endif diff --git a/sglang/3rdparty/amd/profiling/rpd_profile_server_enable.patch b/sglang/3rdparty/amd/profiling/rpd_profile_server_enable.patch new file mode 100644 index 0000000..3cd3915 --- /dev/null +++ b/sglang/3rdparty/amd/profiling/rpd_profile_server_enable.patch @@ -0,0 +1,49 @@ +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index 62d1ff9..9021c01 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -71,6 +71,8 @@ from sglang.srt.utils import ( + suppress_other_loggers, + ) + from sglang.utils import get_exception_traceback ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() + + logger = logging.getLogger(__name__) + +@@ -245,6 +247,7 @@ class Scheduler: + ], + with_stack=True, + ) ++ self.rpd = rpdTracerControl() + + @torch.inference_mode() + def event_loop(self): +@@ -1027,15 +1030,24 @@ class Scheduler: + def start_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.start() ++ #self.profiler.start() #block pytorch profiler for rpd profiler enabling ++ if self.tp_rank == 0 or self.tp_rank == 1: ++ self.rpd.start() ++ self.rpd.rangePush("", "rpd profile range", "") ++ logger.info("rpd is enabled") + + def stop_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.stop() +- self.profiler.export_chrome_trace( +- self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" +- ) ++ #self.profiler.stop() ++ #self.profiler.export_chrome_trace( ++ # self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" ++ #) ++ if self.tp_rank ==0 or self.tp_rank ==1: ++ self.rpd.rangePop() ++ self.rpd.stop() ++ self.rpd.flush() ++ logger.info("rpd is done") + logger.info("Profiler is done") diff --git a/sglang/3rdparty/amd/profiling/rpd_profile_server_enable_wCPU_activities.patch b/sglang/3rdparty/amd/profiling/rpd_profile_server_enable_wCPU_activities.patch new file mode 100644 index 0000000..5416f4d --- /dev/null +++ b/sglang/3rdparty/amd/profiling/rpd_profile_server_enable_wCPU_activities.patch @@ -0,0 +1,126 @@ +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index 62d1ff9..2edb427 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -71,6 +71,8 @@ from sglang.srt.utils import ( + suppress_other_loggers, + ) + from sglang.utils import get_exception_traceback ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() + + logger = logging.getLogger(__name__) + +@@ -245,6 +247,7 @@ class Scheduler: + ], + with_stack=True, + ) ++ self.rpd = rpdTracerControl() + + @torch.inference_mode() + def event_loop(self): +@@ -1027,15 +1030,26 @@ class Scheduler: + def start_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.start() ++ #self.profiler.start() ++ logger.info("torch profiler is disabled") ++ if self.tp_rank == 0 or self.tp_rank == 1: ++ self.rpd.setPythonTrace(True) ++ self.rpd.start() ++ self.rpd.rangePush("", "scheduler", "") ++ logger.info("rpd is enabled inside scheduler profiling") + + def stop_profile(self) -> None: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") +- self.profiler.stop() +- self.profiler.export_chrome_trace( +- self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" +- ) ++ #self.profiler.stop() ++ #self.profiler.export_chrome_trace( ++ # self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" ++ #) ++ if self.tp_rank ==0 or self.tp_rank ==1: ++ self.rpd.rangePop() ++ self.rpd.stop() ++ self.rpd.flush() ++ logger.info("rpd is done inside scheduler") + logger.info("Profiler is done") + + +diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py +index 2621ccd..181df85 100644 +--- a/python/sglang/srt/managers/tokenizer_manager.py ++++ b/python/sglang/srt/managers/tokenizer_manager.py +@@ -58,6 +58,10 @@ from sglang.srt.sampling.sampling_params import SamplingParams + from sglang.srt.server_args import PortArgs, ServerArgs + from sglang.srt.utils import is_generation_model, is_multimodal_model + ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() ++ ++ + asyncio.set_event_loop_policy(uvloop.EventLoopPolicy()) + + logger = logging.getLogger(__name__) +@@ -514,10 +518,20 @@ class TokenizerManager: + self.send_to_scheduler.send_pyobj(req) + + def start_profile(self): ++ rpd = rpdTracerControl() ++ rpd.setPythonTrace(True) ++ rpd.start() ++ rpd.rangePush("", "tokenizer_manager", "") ++ logger.info("tokenizer_manager rpd profiling started!") + req = ProfileReq.START_PROFILE + self.send_to_scheduler.send_pyobj(req) + + def stop_profile(self): ++ rpd = rpdTracerControl() ++ rpd.rangePop() ++ rpd.stop() ++ rpd.flush() ++ logger.info("rpd profiling is done inside tokenizer_manager!") + req = ProfileReq.STOP_PROFILE + self.send_to_scheduler.send_pyobj(req) + +diff --git a/python/sglang/srt/server.py b/python/sglang/srt/server.py +index 7111c93..2bd722c 100644 +--- a/python/sglang/srt/server.py ++++ b/python/sglang/srt/server.py +@@ -30,6 +30,8 @@ import threading + import time + from http import HTTPStatus + from typing import Dict, List, Optional, Union ++from rpdTracerControl import rpdTracerControl ++rpdTracerControl.skipCreate() + + # Fix a bug of Python threading + setattr(threading, "_register_atexit", lambda *args, **kwargs: None) +@@ -152,6 +154,11 @@ async def flush_cache(): + @app.post("/start_profile") + async def start_profile(): + """Start profiling.""" ++ rpd = rpdTracerControl() ++ rpd.setPythonTrace(True) ++ rpd.start() ++ rpd.rangePush("", "server rpd profile range", "") ++ logger.info("rpd profiling started in server.py!") + tokenizer_manager.start_profile() + return Response( + content="Start profiling.\n", +@@ -164,6 +171,11 @@ async def start_profile(): + async def stop_profile(): + """Stop profiling.""" + tokenizer_manager.stop_profile() ++ rpd = rpdTracerControl() ++ rpd.rangePop() ++ rpd.stop() ++ rpd.flush() ++ logger.info("rpd profiling is done in server.py!") + return Response( + content="Stop profiling. This will take some time.\n", + status_code=200, diff --git a/sglang/3rdparty/amd/profiling/server.sh b/sglang/3rdparty/amd/profiling/server.sh new file mode 100644 index 0000000..f877e6c --- /dev/null +++ b/sglang/3rdparty/amd/profiling/server.sh @@ -0,0 +1,20 @@ +#!/bin/bash + +# export SGLANG_TORCH_PROFILER_DIR=/data/sglang/ +export SGLANG_TORCH_PROFILER_DIR=/sgl-workspace/sglang/profile/ + +# Get the current timestamp +TIMESTAMP=$(date +"%Y%m%d_%H%M%S") + +# Define the log file with a timestamp +LOGFILE="sglang_server_log_$TIMESTAMP.json" + +# Run the Python command and save the output to the log file +loadTracer.sh python3 -m sglang.launch_server \ + --model-path /sgl-workspace/sglang/dummy_grok1 \ + --tokenizer-path Xenova/grok-1-tokenizer \ + --load-format dummy \ + --quantization fp8 \ + --tp 8 \ + --port 30000 \ + --disable-radix-cache 2>&1 | tee "$LOGFILE" diff --git a/sglang/3rdparty/amd/profiling/torch_profiler.patch b/sglang/3rdparty/amd/profiling/torch_profiler.patch new file mode 100644 index 0000000..40f5574 --- /dev/null +++ b/sglang/3rdparty/amd/profiling/torch_profiler.patch @@ -0,0 +1,25 @@ +diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py +index 62d1ff9..6ecd78c 100644 +--- a/python/sglang/srt/managers/scheduler.py ++++ b/python/sglang/srt/managers/scheduler.py +@@ -240,7 +240,6 @@ class Scheduler: + ) + self.profiler = torch.profiler.profile( + activities=[ +- torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + with_stack=True, +@@ -1033,9 +1032,11 @@ class Scheduler: + if self.profiler is None: + raise RuntimeError("Profiler is not enabled.") + self.profiler.stop() +- self.profiler.export_chrome_trace( +- self.torch_profiler_trace_dir + "/" + str(time.time()) + ".trace.json.gz" +- ) ++ if self.tp_rank == 0: ++ with open(f"stats_repro_{int(time.time())}.txt", "w") as f: ++ print(self.profiler.key_averages(group_by_input_shape=True).table(sort_by="cuda_time_total", row_limit=-1), file=f) ++ print("Profiling stats done.") ++ + logger.info("Profiler is done") diff --git a/sglang/3rdparty/amd/tuning/TUNING.md b/sglang/3rdparty/amd/tuning/TUNING.md new file mode 100644 index 0000000..0638041 --- /dev/null +++ b/sglang/3rdparty/amd/tuning/TUNING.md @@ -0,0 +1,118 @@ +## Tuning SGLang Infer System with AMD GPUs +This AppNote describes the SGLang performance tuning technical, code harness and running steps for systems with AMD Instinct GPUs. +Harness code, examples and steps are provided in detail, to facilitate easy reproduce & use to tune performance towards workloads. +Three primary runtime areas are covered: + +## 1. Triton Kernels +To maximize Triton kernel efficiency, several strategies can be employed: + +### Key Environment Variables: +- **num_stages**: Adjusts the number of pipeline stages to optimize kernel efficiency based on the specific type of operations (e.g., General Matrix Multiplication - GEMM). +- **waves_per_eu**: Controls the usage of Vector General Purpose Registers (VGPR) to enhance occupancy, thereby improving latency or throughput. +- **BLOCK_M, BLOCK_N, BLOCK_K**: Tunable tile sizes that assist in balancing memory transfer and computational efficiency. +- **matrix_instr_nonkdim**: Optimizes the usage of Matrix-Fused Multiply-Add (MFMA) instructions for specific kernel types, such as Flash Attention. +- **OPTIMIZE_EPILOGUE**: An environment variable that can be set to `1` to enhance performance by eliminating the `convert_layout` operation in the kernel's epilogue. +```python +@triton.autotune(configs=[ + triton.Config({'waves_per_eu': 1}, num_warps=4, num_stages=1), + triton.Config({'waves_per_eu': 1}, num_warps=8, num_stages=1), + triton.Config({'waves_per_eu': 1}, num_warps=16, num_stages=1), + triton.Config({'waves_per_eu': 2}, num_warps=4, num_stages=1), + triton.Config({'waves_per_eu': 2}, num_warps=8, num_stages=1), + triton.Config({'waves_per_eu': 2}, num_warps=16, num_stages=1), + triton.Config({'waves_per_eu': 4}, num_warps=4, num_stages=1), + triton.Config({'waves_per_eu': 4}, num_warps=8, num_stages=1), + triton.Config({'waves_per_eu': 4}, num_warps=16, num_stages=1), + ], key=['BLOCK_N', 'NUM_TOKEN_BLKS'], use_cuda_graph=True) +@triton.jit +def _triton_kernel_funtion(): + ... +``` +## 2. Torch Tunable Operations +**TunableOp** is a feature in PyTorch that allows for the definition and optimization of custom kernels with tunable parameters. This feature is particularly useful for enhancing the performance of kernels by experimenting with different configurations. + +### Key Environment Variables: +1. **PYTORCH_TUNABLEOP_ENABLED**: + - Default: `0` + - Set to `1` to enable TunableOp. + +2. **PYTORCH_TUNABLEOP_TUNING**: + - Default: `1` + - Set to `0` to disable tuning. If a tuned entry is not found, it will run the tuning step and record the entry when PYTORCH_TUNABLEOP_ENABLED is enabled. + +3. **PYTORCH_TUNABLEOP_VERBOSE**: + - Default: `0` + - Set to `1` to enable verbose output for TunableOp. + +### Usage Example: +To enable TunableOp and tuning, and optionally enable verbose mode, you can run the following command in your terminal: + +```bash +#Tuning +PYTORCH_TUNABLEOP_ENABLED=1 PYTORCH_TUNABLEOP_TUNING=1 your_script.sh + +#Inference with tuning op +PYTORCH_TUNABLEOP_ENABLED=1 PYTORCH_TUNABLEOP_TUNING=0 your_script.sh + +#Print out the log +PYTORCH_TUNABLEOP_ENABLED=1 PYTORCH_TUNABLEOP_TUNING=0 PYTORCH_TUNABLEOP_VERBOSE=1 your_script.sh + +``` +## 3. Torch Compilation + + +The following are suggestions for optimizing matrix multiplication (GEMM) and convolution (conv) operations in PyTorch using Inductor, a part of the PyTorch compilation framework. The goal is to leverage Triton to achieve better performance. + +To tune Triton kernels with GEMM and convolution ops (conv), use the `torch.compile` function with the max-autotune mode. This benchmarks a predefined list of Triton configurations and selects the fastest one for each shape. + +### Key Configurations: +1. **Max Autotune**: + - Set `torch._inductor.config.max_autotune = True` or `TORCHINDUCTOR_MAX_AUTOTUNE=1`. + +2. **Fine-Grained Control**: + - Enable GEMM tuning: `torch._inductor.config.max_autotune_gemm = True`. + - Enable tuning for pointwise/reduction ops: `torch._inductor.config.max_autotune.pointwise = True`. + +3. **Backend Selection**: + - Use `torch._inductor.max_autotune_gemm_backends` to limit backends to TRITON for better performance. + +4. **Freezing for Inference**: + - Use `torch._inductor.config.freezing=True` to enable constant folding optimizations. + +5. **Debugging**: + - Set `TORCH_COMPILE_DEBUG=1` to extract Triton kernels generated by Inductor. + +### Example Code Block: +```bash +#Gemm Tuning +TORCHINDUCTOR_MAX_AUTOTUNE=1 TORCHINDUCTOR_COORDINATE_DESCENT_TUNING=1 your_script.sh + +#Specify your backend to TRITON for Gemm Tuning +TORCHINDUCTOR_MAX_AUTOTUNE=1 TORCHINDUCTOR_COORDINATE_DESCENT_TUNING=1 TORCHINDUCTOR_MAX_AUTOTUNE_GEMM_BACKENDS=TRITON your_script.sh + +#Inference with large improvement on AMD GPU +TORCHINDUCTOR_FREEZING=1 your_script.sh +``` +## 4. Fused MOE kernel +To maximize moe kernel efficiency, need to use below scripts to find out the best launch configuration + +### Key parameters: +- **--model**: what moe model type to do tuning, it will automatically decide the size of d_model, model_intermediate_size, num_layers +- **--tp-size**: simulate the whole model run configuration to set the dimension size using tp correctly +- **--batch**: M dimension size of moe kernel, for prefill moe kernel the value is batch*input_len, for decode moe kernel the value is batch +- **--dtype**: computation type + +```bash +#Tuning +#for example, we have one case like this "python3 -m sglang.bench_latency --model dummy_grok1/ --load-format dummy --tokenizer-path Xenova/grok-1-tokenizer --tp 8 --batch-size 32 --input 1024 --output 8 --attention-backend triton --sampling-backend pytorch --quantization fp8" to run, it defined batch-size 32 input lenth 1024 and output length 8, from "--batch" in moe view point, the prefill batch is 32*1024 = 32768, the decode batch is 32*1(only one output token generated in each run). +#so we can tune decode moe use below command +python benchmark_moe_rocm.py --model grok1 --tp-size 8 --dtype float8 --batch "32" +# and use this command to tune prefill moe +python benchmark_moe_rocm.py --model grok1 --tp-size 8 --dtype float8 --batch "32768" +``` + +## Reference + +For more detailed information on tuning SGLang performance with AMD GPUs, please refer to the following link: + +[ROCm Documentation: Triton Kernel Performance Optimization](https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/workload.html#triton-kernel-performance-optimization) diff --git a/sglang/3rdparty/amd/tuning/benchmark_moe_rocm.py b/sglang/3rdparty/amd/tuning/benchmark_moe_rocm.py new file mode 100644 index 0000000..5aff8c0 --- /dev/null +++ b/sglang/3rdparty/amd/tuning/benchmark_moe_rocm.py @@ -0,0 +1,380 @@ +import argparse +import json +import os +import sys + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl +from tqdm import tqdm +from transformers import AutoConfig + +from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( + fused_moe, + get_config_file_name, +) + +padding_size = 128 if bool(int(os.getenv("MOE_PADDING", "0"))) else 0 + + +def main(model, tp_size, dtype: str, batches): + method = fused_moe + + for bs in batches: + run_grid(int(bs), model=model, method=method, tp_size=tp_size, dtype=dtype) + + +def prune_configs(M, N, K, configs): + pruned_configs = [] + elemBytes_a = 1 # [DV Note] Hard-coded for float16 (2 bytes) + elemBytes_b = 1 # [DV Note] Hard-coded for float16 (2 bytes) + + mfma = 16 if M < 32 or N < 32 else 32 + + # TODO (zhanglx): figure out the boundary between large and small gemms + large_gemm = False + if M >= 2048 and N >= 2048: + large_gemm = True + + for config in configs: + BLOCK_SIZE_M = config.get("BLOCK_SIZE_M") + BLOCK_SIZE_N = config.get("BLOCK_SIZE_N") + BLOCK_SIZE_K = config.get("BLOCK_SIZE_K") + num_warps = config.get("num_warps") + matrix_instr_nonkdim = config.get("matrix_instr_nonkdim") + # kpack = config.get("kpack") + if matrix_instr_nonkdim > mfma: + continue + if mfma == 4 and BLOCK_SIZE_K < 64: + continue + # some layouts could not work properly in case + # number elements per thread is less 1 + if BLOCK_SIZE_M * BLOCK_SIZE_N < 64: + continue + SPLIT_K = 1 # config.get("SPLIT_K") + GROUP_M = config.get("GROUP_SIZE_M") + if matrix_instr_nonkdim > BLOCK_SIZE_M or matrix_instr_nonkdim > BLOCK_SIZE_N: + continue + if matrix_instr_nonkdim >= M and matrix_instr_nonkdim != BLOCK_SIZE_M: + continue + if matrix_instr_nonkdim >= N and matrix_instr_nonkdim != BLOCK_SIZE_N: + continue + # Skip BLOCK_SIZE that is too large compare to M/N + # unless BLOCK_SIZE is already small enough + if M * 2 < BLOCK_SIZE_M and BLOCK_SIZE_M != 16: + continue + if N * 2 < BLOCK_SIZE_N and BLOCK_SIZE_N != 16: + continue + # skip large split_k when not necessary + if SPLIT_K != 1 and not need_split_k(M, N, K): + continue + # skip split_k that leads to EVEN_K = false + leap = SPLIT_K * BLOCK_SIZE_K + modv = K % leap + if modv != 0: + continue + # skip large GROUP_M + if GROUP_M * BLOCK_SIZE_M > M and GROUP_M != 1: + continue + # out of shared memory resource + # TODO (zhanglx): This does not consider the LDS usage in the epilogue + LDS = ( + BLOCK_SIZE_K * BLOCK_SIZE_M * elemBytes_a + + BLOCK_SIZE_K * BLOCK_SIZE_N * elemBytes_b + ) + if LDS > 65536: + continue + # Skip small block sizes and num_warps for large gemm + # For fp16 and f8, we want to only use BLOCK_SIZE >= 64 + if large_gemm: + if BLOCK_SIZE_M < 64 or BLOCK_SIZE_N < 64: + continue + if BLOCK_SIZE_K < 64: + continue + if num_warps < 4: + continue + + pruned_configs.append(config) + + return pruned_configs + + +def union_of_list_of_dicts(l1, l2): + result = [] + temp_list = l1.copy() + temp_list.extend(l2) + for myDict in temp_list: + if myDict not in result: + result.append(myDict) + + return result + + +def run_grid(bs, model, method, tp_size, dtype: str): + + config = AutoConfig.from_pretrained(model) + + top_k = config.num_experts_per_tok + d_model = config.hidden_size + model_intermediate_size = config.intermediate_size + num_layers = config.num_hidden_layers + hidden_states_dtype = config.torch_dtype + + if config.num_experts_per_tok: + if config.architectures[0] == "Grok1ModelForCausalLM": + num_total_experts = config.num_experts + else: + num_total_experts = config.num_local_experts + else: + raise ValueError(f"Unsupported Mixtral model {model}") + + # tp_size = 2 + num_warmup_calls = 10 + num_calls = 30 + + num_warmup_trials = 1 + num_trials = 1 + + full_configs = [] + + block_m_range = [16, 32, 64, 128, 256] + block_n_range = [16, 32, 64, 128, 256] + block_k_range = [32, 64, 128, 256] # MUST >= 32 + num_warps_range = [1, 2, 4, 8] + group_m_range = [1, 4, 8, 16, 32] + # For now we see better perf with num_stages=0 for all gemm configs we care + # But keep this explicit so that we do not forget we may need to set it to + # other values in the future + num_stage_range = [2] + waves_per_eu_range = [0, 1, 2, 4, 8] + # Remove 32 because of triton compiling error + matrix_instr_nonkdim_range = [16] + kpack_range = [1, 2] + + for block_size_m in block_m_range: + for block_size_n in block_n_range: + for block_size_k in block_k_range: + for group_size_m in group_m_range: + for num_warps in num_warps_range: + for num_stages in num_stage_range: + for waves_per_eu in waves_per_eu_range: + for matrix_instr_nonkdim in matrix_instr_nonkdim_range: + for kpack in kpack_range: + full_configs.append( + { + "BLOCK_SIZE_M": block_size_m, + "BLOCK_SIZE_N": block_size_n, + "BLOCK_SIZE_K": block_size_k, + "GROUP_SIZE_M": group_size_m, + "num_warps": num_warps, + "num_stages": num_stages, + "waves_per_eu": waves_per_eu, + "matrix_instr_nonkdim": matrix_instr_nonkdim, + "kpack": kpack, + } + ) + + M1 = bs * 2 + N1 = model_intermediate_size * 2 // tp_size + K1 = d_model + prune_configs_1 = prune_configs(M1, N1, K1, full_configs) + + M2 = bs * 2 + N2 = d_model + K2 = model_intermediate_size // tp_size + prune_configs_2 = prune_configs(M2, N2, K2, full_configs) + + configs = union_of_list_of_dicts(prune_configs_1, prune_configs_2) + + print( + f"{bs=} || {len(full_configs)=} | {len(prune_configs_1)=} | \ + {len(prune_configs_2)=} | {len(configs)=}" + ) + + best_config = None + best_time_us = 1e20 + + print(f"{tp_size=} {bs=}") + + for config in tqdm(configs): + # warmup + try: + print(config) + for _ in range(num_warmup_trials): + run_timing( + num_calls=num_warmup_calls, + bs=bs, + d_model=d_model, + num_total_experts=num_total_experts, + top_k=top_k, + tp_size=tp_size, + model_intermediate_size=model_intermediate_size, + method=method, + config=config, + dtype=dtype, + hidden_states_dtype=hidden_states_dtype, + ) + except triton.runtime.autotuner.OutOfResources: + continue + + # trial + for _ in range(num_trials): + kernel_dur_ms = run_timing( + num_calls=num_calls, + bs=bs, + d_model=d_model, + num_total_experts=num_total_experts, + top_k=top_k, + tp_size=tp_size, + model_intermediate_size=model_intermediate_size, + method=method, + config=config, + dtype=dtype, + hidden_states_dtype=hidden_states_dtype, + ) + + kernel_dur_us = 1000 * kernel_dur_ms + model_dur_ms = kernel_dur_ms * num_layers + + if kernel_dur_us < best_time_us: + best_config = config + best_time_us = kernel_dur_us + + tqdm.write( + f"{kernel_dur_us=:.1f} {model_dur_ms=:.1f}" + f" {bs=} {tp_size=} {top_k=} {num_total_experts=} " + f"{d_model=} {model_intermediate_size=} {num_layers=}" + ) + + print("best_time_us", best_time_us) + print("best_config", best_config) + + # holds Dict[str, Dict[str, int]] + filename = get_config_file_name( + num_total_experts, + model_intermediate_size // tp_size, + "float8" if dtype == "float8" else None, + ) + print(f"writing config to file {filename}") + existing_content = {} + if os.path.exists(filename): + with open(filename, "r") as f: + existing_content = json.load(f) + existing_content[str(bs)] = best_config + with open(filename, "w") as f: + json.dump(existing_content, f, indent=4) + f.write("\n") + + +def run_timing( + num_calls: int, + bs: int, + d_model: int, + num_total_experts: int, + top_k: int, + tp_size: int, + model_intermediate_size: int, + method, + config, + dtype: str, + hidden_states_dtype, +) -> float: + shard_intermediate_size = model_intermediate_size // tp_size + + hidden_states = torch.rand( + (bs, d_model), + device="cuda:0", + dtype=hidden_states_dtype, + ) + + w1 = torch.rand( + (num_total_experts, 2 * shard_intermediate_size, d_model + padding_size), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + w2 = torch.rand( + (num_total_experts, d_model, shard_intermediate_size + padding_size), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + w1_scale = None + w2_scale = None + a1_scale = None + a2_scale = None + + if dtype == "float8": + w1 = w1.to(torch.float8_e4m3fnuz) + w2 = w2.to(torch.float8_e4m3fnuz) + w1_scale = torch.ones( + num_total_experts, device=hidden_states.device, dtype=torch.float32 + ) + w2_scale = torch.ones( + num_total_experts, device=hidden_states.device, dtype=torch.float32 + ) + a1_scale = torch.ones(1, device=hidden_states.device, dtype=torch.float32) + a2_scale = torch.ones(1, device=hidden_states.device, dtype=torch.float32) + + gating_output = F.softmax( + torch.rand( + (num_calls, bs, num_total_experts), + device=hidden_states.device, + dtype=torch.float32, + ), + dim=-1, + ) + + ################################## + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + + start_event.record() + for i in range(num_calls): + hidden_states = method( + hidden_states=hidden_states, + w1=w1, + w2=w2, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + gating_output=gating_output[0], + topk=top_k, + renormalize=True, + inplace=True, + override_config=config, + use_fp8=dtype == "float8", + ) + + end_event.record() + end_event.synchronize() + + dur_ms = start_event.elapsed_time(end_event) / num_calls + return dur_ms + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + prog="benchmark_mixtral_moe", + description="Benchmark and tune the fused_moe kernel", + ) + parser.add_argument( + "--dtype", + type=str, + default="auto", + choices=["float8", "float16", "bfloat16"], + help="Data type used for fused_moe kernel computations", + ) + parser.add_argument("--model", type=str, default="hpcai-tech/grok-1") + + parser.add_argument("--tp-size", type=int, default=2, help="Tensor paralleli size") + parser.add_argument("-b", "--batches", type=str) + + args = parser.parse_args() + + batches = args.batches.split(",") + + sys.exit(main(args.model, args.tp_size, args.dtype, batches)) diff --git a/sglang/LICENSE b/sglang/LICENSE new file mode 100644 index 0000000..9c42268 --- /dev/null +++ b/sglang/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright 2023-2024 SGLang Team + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/sglang/Makefile b/sglang/Makefile new file mode 100644 index 0000000..23d1ddf --- /dev/null +++ b/sglang/Makefile @@ -0,0 +1,43 @@ +.PHONY: check-deps install-deps format update help + +# Show help for each target +help: + @echo "Available targets:" + @grep -E '^[a-zA-Z0-9_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[36m%-20s\033[0m %s\n", $$1, $$2}' + +check-deps: ## Check and install required Python formatting dependencies + @command -v isort >/dev/null 2>&1 || (echo "Installing isort..." && pip install isort) + @command -v black >/dev/null 2>&1 || (echo "Installing black..." && pip install black) + +install-deps: ## Install Python formatting tools (isort and black) + pip install isort black + +format: check-deps ## Format modified Python files using isort and black + @echo "Formatting modified Python files..." + git diff --name-only --diff-filter=M | grep '\.py$$' | xargs -I {} sh -c 'isort {} && black {}' + +FILES_TO_UPDATE = docker/Dockerfile.rocm \ + python/pyproject.toml \ + python/sglang/version.py \ + docs/developer/setup_github_runner.md \ + docs/start/install.md + +update: ## Update version numbers across project files. Usage: make update + @if [ -z "$(filter-out $@,$(MAKECMDGOALS))" ]; then \ + echo "Version required. Usage: make update "; \ + exit 1; \ + fi + @OLD_VERSION=$$(grep "version" python/sglang/version.py | cut -d '"' -f2); \ + NEW_VERSION=$(filter-out $@,$(MAKECMDGOALS)); \ + echo "Updating version from $$OLD_VERSION to $$NEW_VERSION"; \ + for file in $(FILES_TO_UPDATE); do \ + if [ "$(shell uname)" = "Darwin" ]; then \ + sed -i '' -e "s/$$OLD_VERSION/$$NEW_VERSION/g" $$file; \ + else \ + sed -i -e "s/$$OLD_VERSION/$$NEW_VERSION/g" $$file; \ + fi \ + done; \ + echo "Version update complete" + +%: + @: diff --git a/sglang/README.md b/sglang/README.md new file mode 100644 index 0000000..bcb6771 --- /dev/null +++ b/sglang/README.md @@ -0,0 +1,75 @@ +
+logo + +[![PyPI](https://img.shields.io/pypi/v/sglang)](https://pypi.org/project/sglang) +![PyPI - Downloads](https://img.shields.io/pypi/dm/sglang) +[![license](https://img.shields.io/github/license/sgl-project/sglang.svg)](https://github.com/sgl-project/sglang/tree/main/LICENSE) +[![issue resolution](https://img.shields.io/github/issues-closed-raw/sgl-project/sglang)](https://github.com/sgl-project/sglang/issues) +[![open issues](https://img.shields.io/github/issues-raw/sgl-project/sglang)](https://github.com/sgl-project/sglang/issues) +[![](https://img.shields.io/badge/Gurubase-(experimental)-006BFF)](https://gurubase.io/g/sglang) + +
+ +-------------------------------------------------------------------------------- + +| [**Blog**](https://lmsys.org/blog/2024-07-25-sglang-llama3/) +| [**Documentation**](https://docs.sglang.ai/) +| [**Join Slack**](https://slack.sglang.ai/) +| [**Join Bi-Weekly Development Meeting**](https://meeting.sglang.ai/) +| [**Roadmap**](https://github.com/sgl-project/sglang/issues/4042) +| [**Slides**](https://github.com/sgl-project/sgl-learning-materials?tab=readme-ov-file#slides) | + +## News +- [2025/03] Supercharge DeepSeek-R1 Inference on AMD Instinct MI300X ([AMD blog](https://rocm.blogs.amd.com/artificial-intelligence/DeepSeekR1-Part2/README.html)) +- [2025/03] SGLang Joins PyTorch Ecosystem: Efficient LLM Serving Engine ([PyTorch blog](https://pytorch.org/blog/sglang-joins-pytorch/)) +- [2025/02] Unlock DeepSeek-R1 Inference Performance on AMD Instinct™ MI300X GPU ([AMD blog](https://rocm.blogs.amd.com/artificial-intelligence/DeepSeekR1_Perf/README.html)) +- [2025/01] 🔥 SGLang provides day one support for DeepSeek V3/R1 models on NVIDIA and AMD GPUs with DeepSeek-specific optimizations. ([instructions](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3), [AMD blog](https://www.amd.com/en/developer/resources/technical-articles/amd-instinct-gpus-power-deepseek-v3-revolutionizing-ai-development-with-sglang.html), [10+ other companies](https://x.com/lmsysorg/status/1887262321636221412)) +- [2024/12] 🔥 v0.4 Release: Zero-Overhead Batch Scheduler, Cache-Aware Load Balancer, Faster Structured Outputs ([blog](https://lmsys.org/blog/2024-12-04-sglang-v0-4/)). +- [2024/09] v0.3 Release: 7x Faster DeepSeek MLA, 1.5x Faster torch.compile, Multi-Image/Video LLaVA-OneVision ([blog](https://lmsys.org/blog/2024-09-04-sglang-v0-3/)). +- [2024/07] v0.2 Release: Faster Llama3 Serving with SGLang Runtime (vs. TensorRT-LLM, vLLM) ([blog](https://lmsys.org/blog/2024-07-25-sglang-llama3/)). + +
+More + +- [2024/10] The First SGLang Online Meetup ([slides](https://github.com/sgl-project/sgl-learning-materials?tab=readme-ov-file#the-first-sglang-online-meetup)). +- [2024/02] SGLang enables **3x faster JSON decoding** with compressed finite state machine ([blog](https://lmsys.org/blog/2024-02-05-compressed-fsm/)). +- [2024/01] SGLang provides up to **5x faster inference** with RadixAttention ([blog](https://lmsys.org/blog/2024-01-17-sglang/)). +- [2024/01] SGLang powers the serving of the official **LLaVA v1.6** release demo ([usage](https://github.com/haotian-liu/LLaVA?tab=readme-ov-file#demo)). + +
+ +## About +SGLang is a fast serving framework for large language models and vision language models. +It makes your interaction with models faster and more controllable by co-designing the backend runtime and frontend language. +The core features include: + +- **Fast Backend Runtime**: Provides efficient serving with RadixAttention for prefix caching, zero-overhead CPU scheduler, continuous batching, token attention (paged attention), speculative decoding, tensor parallelism, chunked prefill, structured outputs, and quantization (FP8/INT4/AWQ/GPTQ). +- **Flexible Frontend Language**: Offers an intuitive interface for programming LLM applications, including chained generation calls, advanced prompting, control flow, multi-modal inputs, parallelism, and external interactions. +- **Extensive Model Support**: Supports a wide range of generative models (Llama, Gemma, Mistral, QWen, DeepSeek, LLaVA, etc.), embedding models (e5-mistral, gte, mcdse) and reward models (Skywork), with easy extensibility for integrating new models. +- **Active Community**: SGLang is open-source and backed by an active community with industry adoption. + +## Getting Started +- [Install SGLang](https://docs.sglang.ai/start/install.html) +- [Quick Start](https://docs.sglang.ai/backend/send_request.html) +- [Backend Tutorial](https://docs.sglang.ai/backend/openai_api_completions.html) +- [Frontend Tutorial](https://docs.sglang.ai/frontend/frontend.html) +- [Contribution Guide](https://docs.sglang.ai/references/contribution_guide.html) + +## Benchmark and Performance +Learn more in the release blogs: [v0.2 blog](https://lmsys.org/blog/2024-07-25-sglang-llama3/), [v0.3 blog](https://lmsys.org/blog/2024-09-04-sglang-v0-3/), [v0.4 blog](https://lmsys.org/blog/2024-12-04-sglang-v0-4/) + +## Roadmap +[Development Roadmap (2025 H1)](https://github.com/sgl-project/sglang/issues/4042) + +## Adoption and Sponsorship +The project has been deployed to large-scale production, generating trillions of tokens every day. +It is supported by the following institutions: AMD, Atlas Cloud, Baseten, Cursor, DataCrunch, Etched, Hyperbolic, Iflytek, Jam & Tea Studios, LinkedIn, LMSYS, Meituan, Nebius, Novita AI, NVIDIA, RunPod, Stanford, UC Berkeley, UCLA, xAI, and 01.AI. + +logo + +## Contact Us + +For enterprises interested in adopting or deploying SGLang at scale, including technical consulting, sponsorship opportunities, or partnership inquiries, please contact us at contact@sglang.ai. + +## Acknowledgment and Citation +We learned the design and reused code from the following projects: [Guidance](https://github.com/guidance-ai/guidance), [vLLM](https://github.com/vllm-project/vllm), [LightLLM](https://github.com/ModelTC/lightllm), [FlashInfer](https://github.com/flashinfer-ai/flashinfer), [Outlines](https://github.com/outlines-dev/outlines), and [LMQL](https://github.com/eth-sri/lmql). Please cite the paper, [SGLang: Efficient Execution of Structured Language Model Programs](https://arxiv.org/abs/2312.07104), if you find the project useful. diff --git a/sglang/assets/logo.png b/sglang/assets/logo.png new file mode 100644 index 0000000..2a8bc25 Binary files /dev/null and b/sglang/assets/logo.png differ diff --git a/sglang/benchmark/bench_in_batch_prefix/bench_in_batch_prefix.py b/sglang/benchmark/bench_in_batch_prefix/bench_in_batch_prefix.py new file mode 100644 index 0000000..86648e5 --- /dev/null +++ b/sglang/benchmark/bench_in_batch_prefix/bench_in_batch_prefix.py @@ -0,0 +1,130 @@ +# Benchmark with lots of common prefixes. Used to benchmark prefix caching performance. +# +# Launch a server: +# python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --log-level-http warning + +import random +import string +import time + +from tqdm import tqdm +from transformers import AutoTokenizer + +import sglang as sgl +from sglang import set_default_backend +from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint + + +def generate_random_string(token_length: int) -> str: + random_string = "".join( + random.choices(string.ascii_letters + string.digits, k=token_length * 100) + ) + tokenized_output = tokenizer.encode(random_string, add_special_tokens=False)[ + :token_length + ] + + if len(tokenized_output) < token_length: + tokenized_output = tokenized_output + [tokenizer.pad_token_id] * ( + token_length - len(tokenized_output) + ) + + decoded_string = tokenizer.decode(tokenized_output, skip_special_tokens=False) + return decoded_string + + +def generate_unique_prefix(base_text, index): + return str(index) + base_text[len(str(index)) :] + + +@sgl.function +def text_qa(s, question, gen_len): + s += "Q: " + question + "\n" + s += "A:" + sgl.gen("answer", stop="\n", temperature=0, max_tokens=gen_len) + + +def prepare_prompts(num_prefix, num_samples_per_prefix, prefix_length, suffix_length): + base_prefix = generate_random_string(prefix_length) + + tot_input_len = 0 + all_prompts = [] + for i in tqdm(range(num_prefix), desc="prepare prompts"): + unique_prefix = generate_unique_prefix(base_prefix, i) + prompt_list = [] + for j in range(num_samples_per_prefix): + suffix = generate_random_string(suffix_length) + prompt = unique_prefix + suffix + prompt_list.append(prompt) + tot_input_len += len(tokenizer.encode(prompt)) + all_prompts.append(prompt_list) + return all_prompts, tot_input_len + + +def test_batch_by_batch(all_prompts, gen_len): + backend.flush_cache() + + tot_time = 0 + for i in range(len(all_prompts)): + tic = time.time() + text_qa.run_batch( + list(zip(all_prompts[i], [gen_len] * len(all_prompts[i]))), + ) + tot_time += time.time() - tic + + return tot_time + + +def test_batch_by_batch_with_hint(all_prompts, gen_len): + backend.flush_cache() + + tot_time = 0 + for i in range(len(all_prompts)): + tic = time.time() + # Send a hint to cache the prefix + text_qa.run_batch(list(zip(all_prompts[i][:1], [gen_len]))) + # Send the batch + text_qa.run_batch(list(zip(all_prompts[i], [gen_len] * len(all_prompts[i])))) + + tot_time += time.time() - tic + + return tot_time + + +def test_send_all(all_prompts, gen_len): + backend.flush_cache() + + all_prompts = [x for prompt_list in all_prompts for x in prompt_list] + + tic = time.time() + text_qa.run_batch( + list(zip(all_prompts, [gen_len] * len(all_prompts))), + ) + tot_time = time.time() - tic + + return tot_time + + +if __name__ == "__main__": + tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer") + backend = RuntimeEndpoint("http://127.0.0.1:30000") + set_default_backend(backend) + + random.seed(0) + num_prefix = 10 + num_samples_per_prefix = 32 + prefix_length = 1024 + suffix_length = 128 + gen_len = 1 + all_prompts, tot_input_len = prepare_prompts( + num_prefix, num_samples_per_prefix, prefix_length, suffix_length + ) + + print(f"Total input token length: {tot_input_len}\n") + + cost = test_batch_by_batch(all_prompts, gen_len) + print(f"Latency of test_batch_by_batch : {cost:.4f} s\n") + + cost = test_batch_by_batch_with_hint(all_prompts, gen_len) + print(f"Latency of test_batch_by_batch_with_hint: {cost:.4f} s\n") + + cost = test_send_all(all_prompts, gen_len) + print(f"Latency of test_send_all : {cost:.4f} s\n") diff --git a/sglang/benchmark/benchmark_vllm_060/README.md b/sglang/benchmark/benchmark_vllm_060/README.md new file mode 100644 index 0000000..b480dab --- /dev/null +++ b/sglang/benchmark/benchmark_vllm_060/README.md @@ -0,0 +1,89 @@ +## How to reproduce the benchmark results for SGLang v0.3.0 compared to vLLM v0.6.0 + +In short, with multi step enabled, in online scenarios that we benchmarked, the Median TTFT of vLLM is **3 times** that of SGLang, and the Median ITL is **10 times** that of SGLang. Lower Median TTFT and ITL are better. vLLM's multi-step optimization did not improve throughput while ensuring lower Median TTFT and ITL. Also, under maximum throughput benchmark, if vLLM does not set gpu util to 0.95 separately and uses the default configuration instead, its maximum throughput is **lower** than that of SGLang. + +## Online benchmark results + +### Llama 3.1 8B Instruct 1 x A100 80G + +| RPS | Num prompts | Engine | Median E2E Latency | Median TTFT | Median TPOT | Median ITL | +|------|-------------|--------|--------------------|-------------|-------------|------------| +| 4 | 1200 | SGLang | 1564.17 | **31.98** | 13.17 | **11.93** | +| 4 | 1200 | vLLM | 1691.97 | **100.48** | 14.14 | **129.32** | +| 8 | 2400 | SGLang | 2175.02 | **35.68** | 17.85 | **14.41** | +| 8 | 2400 | vLLM | 2137.16 | **120.39** | 17.09 | **158.63** | + +### Llama 3.1 70B Insruct 4 x H100 80G + +| RPS | Num Prompts | Engine | Median E2E Latency | Median TTFT | Median TPOT | Median ITL | +|------|-------------|--------|--------------------|-------------|-------------|------------| +| 4 | 1200 | SGLang | 3005.24 | **53.94** | 25.03 | **21.67** | +| 4 | 1200 | vLLM | 2915.60 | **179.15** | 23.58 | **231.23** | +| 8 | 2400 | SGLang | 4064.98 | **58.11** | 33.07 | **24.45** | +| 8 | 2400 | vLLM | 3752.38 | **207.12** | 29.15 | **275.32** | + +## Offline benchmark results + +### Llama 3.1 8B Instruct 1 x A100 80G + +| RPS | Num Prompts | Engine | Request throughput | Output token throughput | +|------|-------------|--------|--------------------|-------------------------| +| inf | 5000 | SGLang | 22.03 | **4281.51** | +| inf | 5000 | vLLM | 21.27 | **4132.37** | + +### Llama 3.1 70B Insruct 4 x H100 80G + +| RPS | Num Prompts | Engine | Request throughput | Output token throughput | +|------|-------------|--------|--------------------|-------------------------| +| inf | 5000 | SGLang | 19.84 | **3856.01** | +| inf | 5000 | vLLM | 19.04 | **3700.64** | + +## Installation + +```bash +# install sglang v0.3.0 +pip install --upgrade pip +pip install "sglang[all]"==0.3.0 +pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.4/ + +# install vllm v0.6.0 +pip install vllm==0.6.0 +``` + +## Notes + +We referred to the reproduction method in https://github.com/vllm-project/vllm/issues/8176, and added the `--num-scheduler-steps 10` parameter when starting the vLLM server. The `gpu_memory_utilization` of vLLM is by default 0.9 at both TP 1 and TP 4, while SGLang's `mem_frac` is 0.88 at TP 1 and 0.85 at TP 4, so we manually set it to 0.88 at TP 4. + +## Online benchmarks + +```bash +# Llama 3.1 8B Instruct on 1 x A100 +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --enable-torch-compile --disable-radix-cache +python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-3.1-8B-Instruct --disable-log-requests --num-scheduler-steps 10 --max_model_len 4096 + +# Llama 3.1 70B Instruct on 4 x H100 +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-70B-Instruct --disable-radix-cache --tp 4 +python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-3.1-70B-Instruct --disable-log-requests --num-scheduler-steps 10 --tensor 4 --max_model_len 4096 + +# bench serving +python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompts 1200 --request-rate 4 +python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompts 2400 --request-rate 8 +python3 -m sglang.bench_serving --backend vllm --dataset-name sharegpt --num-prompts 1200 --request-rate 4 +python3 -m sglang.bench_serving --backend vllm --dataset-name sharegpt --num-prompts 2400 --request-rate 8 +``` + +## Offline benchmarks + +```bash +# Llama 3.1 8B Instruct on 1 x A100 +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --enable-torch-compile --disable-radix-cache +python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-3.1-8B-Instruct --disable-log-requests --num-scheduler-steps 10 --max_model_len 4096 + +# Llama 3.1 70B Instruct on 4 x H100 +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-70B-Instruct --disable-radix-cache --tp 4 --mem-frac 0.88 +python -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-3.1-70B-Instruct --disable-log-requests --num-scheduler-steps 10 --tensor 4 --max_model_len 4096 + +# bench serving +python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompts 5000 +python3 -m sglang.bench_serving --backend vllm --dataset-name sharegpt --num-prompts 5000 +``` diff --git a/sglang/benchmark/blog_v0_2/405b_sglang.sh b/sglang/benchmark/blog_v0_2/405b_sglang.sh new file mode 100644 index 0000000..4918537 --- /dev/null +++ b/sglang/benchmark/blog_v0_2/405b_sglang.sh @@ -0,0 +1,24 @@ +# Create dummy weights: +# 1. Create a folder `~/llama-3.1-405b-fp8-dummy` and create `config.json` and tokenizer under this folder. +# 2. Get `config.json`` from ./config.md +# 3. Download the tokenizer +# wget https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8/resolve/main/tokenizer.json +# wget https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8/resolve/main/tokenizer_config.json + +# Launch sglang +# python -m sglang.launch_server --model-path ~/llama-3.1-405b-fp8-dummy/ --load-format dummy --tp 8 --quantization fp8 --disable-radix --mem-frac 0.87 + +# offline +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 3000 --random-input 1024 --random-output 1024 > sglang_log11 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 4000 --random-input 1024 --random-output 512 > sglang_log12 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 800 --random-input 4096 --random-output 2048 > sglang_log13 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 1500 --random-input 4096 --random-output 1024 > sglang_log14 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 6000 --random-input 256 --random-output 512 > sglang_log15 +python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompt 2000 > sglang_log21 + +# online +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 300 --request-rate 1 --random-input 1024 --random-output 1024 > sglang_log31 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 600 --request-rate 2 --random-input 1024 --random-output 1024 > sglang_log32 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 1200 --request-rate 4 --random-input 1024 --random-output 1024 > sglang_log33 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 2400 --request-rate 8 --random-input 1024 --random-output 1024 > sglang_log34 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompt 3200 --request-rate 16 --random-input 1024 --random-output 1024 > sglang_log35 diff --git a/sglang/benchmark/blog_v0_2/405b_trt.sh b/sglang/benchmark/blog_v0_2/405b_trt.sh new file mode 100644 index 0000000..1950bc9 --- /dev/null +++ b/sglang/benchmark/blog_v0_2/405b_trt.sh @@ -0,0 +1,17 @@ +# Launch trtllm +# https://github.com/sgl-project/tensorrt-demo + +# offline +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 3000 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log11 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 4000 --random-input 1024 --random-output 512 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log12 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 800 --random-input 4096 --random-output 2048 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log13 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 1500 --random-input 4096 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log14 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 6000 --random-input 256 --random-output 512 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log15 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name sharegpt --num-prompt 2000 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log21 + +# online +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 300 --request-rate 1 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log31 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 600 --request-rate 2 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log32 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 1200 --request-rate 4 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log33 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 2400 --request-rate 8 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log34 +python3 ../../python/sglang/bench_serving.py --backend trt --dataset-name random --num-prompt 3200 --request-rate 16 --random-input 1024 --random-output 1024 --model /root/Meta-Llama-3-8B-Instruct > trtllm_log35 diff --git a/sglang/benchmark/blog_v0_2/405b_vllm.sh b/sglang/benchmark/blog_v0_2/405b_vllm.sh new file mode 100644 index 0000000..d1ef048 --- /dev/null +++ b/sglang/benchmark/blog_v0_2/405b_vllm.sh @@ -0,0 +1,24 @@ +# Create dummy weights: +# 1. Create a folder `~/llama-3.1-405b-fp8-dummy` and create `config.json` and tokenizer under this folder. +# 2. Get `config.json`` from ./config.md +# 3. Download the tokenizer +# wget https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8/resolve/main/tokenizer.json +# wget https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-quantized.w8a8/resolve/main/tokenizer_config.json + +# Launch vllm +# python3 -m vllm.entrypoints.openai.api_server --model ~/llama-3.1-405b-fp8-dummy/ --load-format dummy --disable-log-requests --tensor-parallel-size 8 --max-model-len 10000 + +# offline +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 3000 --random-input 1024 --random-output 1024 > vllm_log11 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 4000 --random-input 1024 --random-output 512 > vllm_log12 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 800 --random-input 4096 --random-output 2048 > vllm_log13 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 1500 --random-input 4096 --random-output 1024 > vllm_log14 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 6000 --random-input 256 --random-output 512 > vllm_log15 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name sharegpt --num-prompt 2000 > vllm_log21 + +# online +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 300 --request-rate 1 --random-input 1024 --random-output 1024 > vllm_log31 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 600 --request-rate 2 --random-input 1024 --random-output 1024 > vllm_log32 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 1200 --request-rate 4 --random-input 1024 --random-output 1024 > vllm_log33 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 2400 --request-rate 8 --random-input 1024 --random-output 1024 > vllm_log34 +python3 ../../python/sglang/bench_serving.py --backend vllm --dataset-name random --num-prompt 3200 --request-rate 16 --random-input 1024 --random-output 1024 > vllm_log35 diff --git a/sglang/benchmark/blog_v0_2/README.md b/sglang/benchmark/blog_v0_2/README.md new file mode 100644 index 0000000..7448554 --- /dev/null +++ b/sglang/benchmark/blog_v0_2/README.md @@ -0,0 +1,164 @@ +# How to reproduce the benchmark results of SGLang + +## Prerequisite + +### Install the latest SGLang + +```bash +git clone https://github.com/sgl-project/sglang.git +cd sglang +git checkout v0.2.7 + +pip install --upgrade pip +pip install -e "python[all]" + +pip install flashinfer -i https://flashinfer.ai/whl/cu121/torch2.3/ +``` + +### Set up ulimit and HF_TOKEN + +```bash +ulimit -n 65535 +# Change the token to a real and usable one, with access permissions for the Llama 3 models. +export HF_TOKEN=hf_token +``` + +### Launch the server + +```bash +# Meta-Llama-3.1-8B-Instruct +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --enable-torch-compile --disable-radix-cache + +# Meta-Llama-3.1-70B-Instruct +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-70B-Instruct --disable-radix-cache --tp 8 + +# Meta-Llama-3-70B-Instruct-FP8 +python -m sglang.launch_server --model-path neuralmagic/Meta-Llama-3-70B-Instruct-FP8 --disable-radix-cache --tp 8 +``` + +## Benchmark + +### Hardware Requirements + +- 8B models: Single NVIDIA A100 80GB GPU +- 70B models: 8 x NVIDIA A100 80GB GPUs with Tensor Parallelism (TP) 8 +- 70B FP8 models: 8 x NVIDIA H100 GPUs with Tensor Parallelism (TP) 8 + +Please ensure you have the appropriate hardware before running the benchmarks. + +#### Offline benchmark + +```bash +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 4000 --random-input 1024 --random-output 1024 --output-file offline.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 5000 --random-input 1024 --random-output 512 --output-file offline.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 1000 --random-input 4096 --random-output 2048 --output-file offline.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 2000 --random-input 4096 --random-output 1024 --output-file offline.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 6000 --random-input 256 --random-output 512 --output-file offline.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name sharegpt --num-prompts 3000 --output-file offline.jsonl +cat offline.jsonl | cut -d':' -f12 | cut -d',' -f1 +``` + +#### Online benchmark + +```bash +python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 300 --request-rate 1 --output-file online.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 600 --request-rate 2 --output-file online.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 1200 --request-rate 4 --output-file online.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 2400 --request-rate 8 --output-file online.jsonl +python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 3200 --request-rate 16 --output-file online.jsonl +cat online.jsonl | cut -d':' -f9 | cut -d',' -f1 +``` + +## Other + +We tried using vLLM 0.5.3.post1, but it often crashes under high loads, and it seems to have similar or worse performance compared to vLLM 0.5.2 from our partial benchmarking, so we are using the older version, vLLM 0.5.2. + +Preparation for TensorRT LLM can refer to https://github.com/sgl-project/tensorrt-demo. Specifically, we used a batch size of 512, a max input length of 8192, and a max number of tokens of 8192. The instance count for preprocessing and postprocessing in Triton Server is 16. + +```bash +# vLLM +pip install vllm==0.5.2 +pip install jsonschema==4.21.1 + +# Meta-Llama-3-8B-Instruct +python -m vllm.entrypoints.openai.api_server --model meta-llama/Meta-Llama-3-8B-Instruct --disable-log-requests + +# meta-llama/Meta-Llama-3-70B-Instruct +python -m vllm.entrypoints.openai.api_server --model meta-llama/Meta-Llama-3-70B-Instruct --disable-log-requests --tensor 8 + +# neuralmagic/Meta-Llama-3-70B-Instruct-FP8 +python -m vllm.entrypoints.openai.api_server --model neuralmagic/Meta-Llama-3-70B-Instruct-FP8 --disable-log-requests --tensor 8 +``` + +```bash +wget https://raw.githubusercontent.com/sgl-project/sglang/main/python/sglang/bench_serving.py +``` + +```bash +# vLLM Offline + +python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 4000 --random-input 1024 --random-output 1024 --output-file offline_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 5000 --random-input 1024 --random-output 512 --output-file offline_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 1000 --random-input 4096 --random-output 2048 --output-file offline_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 2000 --random-input 4096 --random-output 1024 --output-file offline_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --num-prompts 6000 --random-input 256 --random-output 512 --output-file offline_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name sharegpt --num-prompts 3000 --output-file offline_vllm.jsonl +cat offline_vllm.jsonl | cut -d':' -f12 | cut -d',' -f1 +``` + +```bash +# vLLM Online + +python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 300 --request-rate 1 --output-file online_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 600 --request-rate 2 --output-file online_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 1200 --request-rate 4 --output-file online_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 2400 --request-rate 8 --output-file online_vllm.jsonl +python3 bench_serving.py --backend vllm --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 3200 --request-rate 16 --output-file online_vllm.jsonl +cat online_vllm.jsonl | cut -d':' -f9 | cut -d',' -f1 +``` + +```bash +# TensorRT LLM Offline 8B + +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --num-prompts 4000 --random-input 1024 --random-output 1024 --output-file offline_trt_8b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --num-prompts 5000 --random-input 1024 --random-output 512 --output-file offline_trt_8b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --num-prompts 1000 --random-input 4096 --random-output 2048 --output-file offline_trt_8b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --num-prompts 2000 --random-input 4096 --random-output 1024 --output-file offline_trt_8b.jsonl +python3 bench_serving.py --backend trt --dataset-name random --num-prompts 6000 --random-input 256 --random-output 512 --output-file offline_trt_8b.jsonl --model meta-llama/Meta-Llama-3-8B-Instruct +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name sharegpt --num-prompts 3000 --output-file offline_trt_8b.jsonl +cat offline_trt_8b.jsonl | cut -d':' -f12 | cut -d',' -f1 +``` + +```bash +# TensorRT LLM Online 8B + +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 300 --request-rate 1 --output-file online_trt_8b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 600 --request-rate 2 --output-file online_trt_8b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 1200 --request-rate 4 --output-file online_trt_8b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 2400 --request-rate 8 --output-file online_trt_8b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 3200 --request-rate 16 --output-file online_trt_8b.jsonl +cat online_trt_8b.jsonl | cut -d':' -f9 | cut -d',' -f1 +``` + +```bash +# TensorRT LLM Offline 70B + +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --num-prompts 4000 --random-input 1024 --random-output 1024 --output-file offline_trt_70b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --num-prompts 5000 --random-input 1024 --random-output 512 --output-file offline_trt_70b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --num-prompts 1000 --random-input 4096 --random-output 2048 --output-file offline_trt_70b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --num-prompts 2000 --random-input 4096 --random-output 1024 --output-file offline_trt_70b.jsonl +python3 bench_serving.py --backend trt --dataset-name random --num-prompts 6000 --random-input 256 --random-output 512 --output-file offline_trt_70b.jsonl --model meta-llama/Meta-Llama-3-70B-Instruct +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name sharegpt --num-prompts 3000 --output-file offline_trt_70b.jsonl +cat offline_trt_70b.jsonl | cut -d':' -f12 | cut -d',' -f1 +``` + +```bash +# TensorRT LLM Online 70B + +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 300 --request-rate 1 --output-file online_trt_70b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 600 --request-rate 2 --output-file online_trt_70b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 1200 --request-rate 4 --output-file online_trt_70b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 2400 --request-rate 8 --output-file online_trt_70b.jsonl +python3 bench_serving.py --backend trt --model meta-llama/Meta-Llama-3-70B-Instruct --dataset-name random --random-input 1024 --random-output 1024 --num-prompts 3200 --request-rate 16 --output-file online_trt_70b.jsonl +cat online_trt_70b.jsonl | cut -d':' -f9 | cut -d',' -f1 +``` diff --git a/sglang/benchmark/blog_v0_2/config.md b/sglang/benchmark/blog_v0_2/config.md new file mode 100644 index 0000000..3faf600 --- /dev/null +++ b/sglang/benchmark/blog_v0_2/config.md @@ -0,0 +1,100 @@ +### used for TensorRT LLM + +``` +{ + "architecture": "LlamaForCausalLM", + "dtype": "float16", + "logits_dtype": "float32", + "vocab_size": 128256, + "max_position_embeddings": 8192, + "hidden_size": 16384, + "num_hidden_layers": 126, + "num_attention_heads": 128, + "num_key_value_heads": 16, + "head_size": 128, + "qk_layernorm": false, + "hidden_act": "silu", + "intermediate_size": 53248, + "norm_epsilon": 1e-05, + "position_embedding_type": "rope_gpt_neox", + "use_parallel_embedding": false, + "embedding_sharding_dim": 0, + "share_embedding_table": false, + "mapping": { + "world_size": 8, + "tp_size": 8, + "pp_size": 1, + "gpus_per_node": 8 + }, + "quantization": { + "quant_algo": "FP8", + "kv_cache_quant_algo": null, + "group_size": 128, + "smoothquant_val": null, + "has_zero_point": false, + "pre_quant_scale": false, + "exclude_modules": [ + "lm_head" + ] + }, + "kv_dtype": "float16", + "rotary_scaling": null, + "residual_mlp": false, + "moe_normalization_mode": null, + "rotary_base": 500000.0, + "moe_num_experts": 0, + "moe_top_k": 0, + "moe_tp_mode": 2, + "attn_bias": false, + "disable_weight_only_quant_plugin": false, + "mlp_bias": false +} +``` + +### used for vLLM and SGLang + +``` +{ + "_name_or_path": "dummy_fp8", + "architectures": [ + "LlamaForCausalLM" + ], + "attention_bias": false, + "attention_dropout": 0.0, + "bos_token_id": 128000, + "eos_token_id": 128009, + "hidden_act": "silu", + "hidden_size": 16384, + "initializer_range": 0.02, + "intermediate_size": 53248, + "mlp_bias": false, + "model_type": "llama", + "num_attention_heads": 128, + "num_hidden_layers": 126, + "num_key_value_heads": 8, + "pretraining_tp": 1, + "quantization_config": { + "activation_scheme": "static", + "ignored_layers": [ + "lm_head" + ], + "quant_method": "fp8" + }, + "rope_scaling": { + "factor": 8.0, + "low_freq_factor": 1.0, + "high_freq_factor": 4.0, + "original_max_position_embeddings": 8192, + "rope_type": "llama3" + }, + "max_position_embeddings": 131072, + "rms_norm_eps": 1e-05, + "rope_scaling": null, + "rope_theta": 500000.0, + "tie_word_embeddings": false, + "torch_dtype": "bfloat16", + "transformers_version": "4.41.1", + "use_cache": true, + "vocab_size": 128256 +} +``` diff --git a/sglang/benchmark/deepseek_v3/README.md b/sglang/benchmark/deepseek_v3/README.md new file mode 100644 index 0000000..8863142 --- /dev/null +++ b/sglang/benchmark/deepseek_v3/README.md @@ -0,0 +1,285 @@ +# DeepSeek V3 Support + +The SGLang and DeepSeek teams collaborated to get DeepSeek V3 FP8 running on NVIDIA and AMD GPUs **from day one**. SGLang also supports [MLA optimization](https://lmsys.org/blog/2024-09-04-sglang-v0-3/#deepseek-multi-head-latent-attention-mla-throughput-optimizations) and [DP attention](https://lmsys.org/blog/2024-12-04-sglang-v0-4/#data-parallelism-attention-for-deepseek-models), making SGLang one of the best open-source LLM engines for running DeepSeek models. SGLang is the inference engine recommended by the official [DeepSeek team](https://github.com/deepseek-ai/DeepSeek-V3/tree/main?tab=readme-ov-file#62-inference-with-sglang-recommended). + +Special thanks to Meituan's Search & Recommend Platform Team and Baseten's Model Performance Team for implementing the model, and DataCrunch for providing GPU resources. + +For optimizations made on the DeepSeek series models regarding SGLang, please refer to [DeepSeek Model Optimizations in SGLang](https://docs.sglang.ai/references/deepseek.html). + +## Installation & Launch + +If you encounter errors when starting the server, ensure the weights have finished downloading. It's recommended to download them beforehand or restart multiple times until all weights are downloaded. + +### Using Docker (Recommended) + +```bash +# Pull latest image +# https://hub.docker.com/r/lmsysorg/sglang/tags +docker pull lmsysorg/sglang:latest + +# Launch +docker run --gpus all --shm-size 32g -p 30000:30000 -v ~/.cache/huggingface:/root/.cache/huggingface --ipc=host --network=host --privileged lmsysorg/sglang:latest \ + python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code --port 30000 +``` + +If you are using RDMA, please note that: + +1. `--network host` and `--privileged` are required by RDMA. If you don't need RDMA, you can remove them. +2. You may need to set `NCCL_IB_GID_INDEX` if you are using RoCE, for example: `export NCCL_IB_GID_INDEX=3`. + +Add [performance optimization options](#performance-optimization-options) as needed. + +### Using pip + +```bash +# Installation +pip install "sglang[all]>=0.4.3" --find-links https://flashinfer.ai/whl/cu124/torch2.5/flashinfer-python + +# Launch +python3 -m sglang.launch_server --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code +``` + +Add [performance optimization options](#performance-optimization-options) as needed. + + + +### Performance Optimization Options + +[MLA optimizations](https://lmsys.org/blog/2024-09-04-sglang-v0-3/#deepseek-multi-head-latent-attention-mla-throughput-optimizations) are enabled by default. Here are some optional optimizations can be enabled as needed. + +- [Data Parallelism Attention](https://lmsys.org/blog/2024-12-04-sglang-v0-4/#data-parallelism-attention-for-deepseek-models): For high QPS scenarios, add the `--enable-dp-attention` argument to boost throughput. +- [Torch.compile Optimization](https://lmsys.org/blog/2024-09-04-sglang-v0-3/#torchcompile-latency-optimizations): Add `--enable-torch-compile` argument to enable it. This will take some time while server starts. The maximum batch size for torch.compile optimization can be controlled with `--torch-compile-max-bs`. It's recommended to set it between `1` and `8`. (e.g., `--torch-compile-max-bs 8`) + +### Example: Sending requests with OpenAI API + +```python3 +import openai +client = openai.Client( + base_url="http://127.0.0.1:30000/v1", api_key="EMPTY") + +# Chat completion +response = client.chat.completions.create( + model="default", + messages=[ + {"role": "system", "content": "You are a helpful AI assistant"}, + {"role": "user", "content": "List 3 countries and their capitals."}, + ], + temperature=0, + max_tokens=64, +) +print(response) +``` + +### Example: Serving with two H20\*8 nodes + +For example, there are two H20 nodes, each with 8 GPUs. The first node's IP is `10.0.0.1`, and the second node's IP is `10.0.0.2`. Please **use the first node's IP** for both commands. + +If the command fails, try setting the `GLOO_SOCKET_IFNAME` parameter. For more information, see [Common Environment Variables](https://pytorch.org/docs/stable/distributed.html#common-environment-variables). + +If the multi nodes support NVIDIA InfiniBand and encounter hanging issues during startup, consider adding the parameter `export NCCL_IB_GID_INDEX=3`. For more information, see [this](https://github.com/sgl-project/sglang/issues/3516#issuecomment-2668493307). + +```bash +# node 1 +python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3 --tp 16 --dist-init-addr 10.0.0.1:5000 --nnodes 2 --node-rank 0 --trust-remote-code + +# node 2 +python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3 --tp 16 --dist-init-addr 10.0.0.1:5000 --nnodes 2 --node-rank 1 --trust-remote-code +``` + +If you have two H100 nodes, the usage is similar to the aforementioned H20. + +> **Note that the launch command here does not enable Data Parallelism Attention or `torch.compile` Optimization**. For optimal performance, please refer to the command options in [Performance Optimization Options](#option_args). + +### Example: Serving with two H200\*8 nodes and docker + +There are two H200 nodes, each with 8 GPUs. The first node's IP is `192.168.114.10`, and the second node's IP is `192.168.114.11`. Configure the endpoint to expose it to another Docker container using `--host 0.0.0.0` and `--port 40000`, and set up communications with `--dist-init-addr 192.168.114.10:20000`. +A single H200 with 8 devices can run DeepSeek V3, the dual H200 setup is just to demonstrate multi-node usage. + +```bash +# node 1 +docker run --gpus all \ + --shm-size 32g \ + --network=host \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --name sglang_multinode1 \ + -it \ + --rm \ + --env "HF_TOKEN=$HF_TOKEN" \ + --ipc=host \ + lmsysorg/sglang:latest \ + python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3 --tp 16 --dist-init-addr 192.168.114.10:20000 --nnodes 2 --node-rank 0 --trust-remote-code --host 0.0.0.0 --port 40000 +``` + +```bash +# node 2 +docker run --gpus all \ + --shm-size 32g \ + --network=host \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --name sglang_multinode2 \ + -it \ + --rm \ + --env "HF_TOKEN=$HF_TOKEN" \ + --ipc=host \ + lmsysorg/sglang:latest \ + python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3 --tp 16 --dist-init-addr 192.168.114.10:20000 --nnodes 2 --node-rank 1 --trust-remote-code --host 0.0.0.0 --port 40000 +``` + +To ensure functionality, we include a test from a client Docker container. + +```bash +docker run --gpus all \ + --shm-size 32g \ + --network=host \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --name sglang_multinode_client \ + -it \ + --rm \ + --env "HF_TOKEN=$HF_TOKEN" \ + --ipc=host \ + lmsysorg/sglang:latest \ + python3 -m sglang.bench_serving --backend sglang --dataset-name random --random-input 1 --random-output 512 --random-range-ratio 1 --num-prompts 1 --host 0.0.0.0 --port 40000 --output-file "deepseekv3_multinode.jsonl" +``` + +> **Note that the launch command here does not enable Data Parallelism Attention or `torch.compile` Optimization**. For optimal performance, please refer to the command options in [Performance Optimization Options](#option_args). + +### Example: Serving with four A100\*8 nodes + +To serve DeepSeek-V3 with A100 GPUs, we need to convert the [FP8 model checkpoints](https://huggingface.co/deepseek-ai/DeepSeek-V3) to BF16 with [script](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/fp8_cast_bf16.py) mentioned [here](https://github.com/deepseek-ai/DeepSeek-V3/blob/main/inference/fp8_cast_bf16.py) first. + +Since the BF16 model is over 1.3 TB, we need to prepare four A100 nodes, each with 8 80GB GPUs. Assume the first node's IP is `10.0.0.1`, and the converted model path is `/path/to/DeepSeek-V3-BF16`, we can have following commands to launch the server. + +```bash +# node 1 +python3 -m sglang.launch_server --model-path /path/to/DeepSeek-V3-BF16 --tp 32 --dist-init-addr 10.0.0.1:5000 --nnodes 4 --node-rank 0 --trust-remote-code --host 0.0.0.0 --port 30000 + +# node 2 +python3 -m sglang.launch_server --model-path /path/to/DeepSeek-V3-BF16 --tp 32 --dist-init-addr 10.0.0.1:5000 --nnodes 4 --node-rank 1 --trust-remote-code + +# node 3 +python3 -m sglang.launch_server --model-path /path/to/DeepSeek-V3-BF16 --tp 32 --dist-init-addr 10.0.0.1:5000 --nnodes 4 --node-rank 2 --trust-remote-code + +# node 4 +python3 -m sglang.launch_server --model-path /path/to/DeepSeek-V3-BF16 --tp 32 --dist-init-addr 10.0.0.1:5000 --nnodes 4 --node-rank 3 --trust-remote-code +``` + +> **Note that the launch command here does not enable Data Parallelism Attention or `torch.compile` Optimization**. For optimal performance, please refer to the command options in [Performance Optimization Options](#option_args). + +Then we can benchmark the accuracy and latency by accessing the first node's exposed port with the following example commands. + +```bash +# bench accuracy +python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319 --host http://10.0.0.1 --port 30000 + +# bench latency +python3 -m sglang.bench_one_batch_server --model None --base-url http://10.0.0.1:30000 --batch-size 1 --input-len 128 --output-len 128 +``` + + +### Example: Serving with 8 A100/A800 with AWQ Quantization + +AWQ does not support BF16, so add the `--dtype half` flag if AWQ is used for quantization. One example is as follows: + +```bash +python3 -m sglang.launch_server --model cognitivecomputations/DeepSeek-R1-AWQ --tp 8 --trust-remote-code --dtype half +``` + + +### Example: Serving with 16 A100/A800 with int8 Quantization + +There are block-wise and per-channel quantization methods, and the quantization parameters have already been uploaded to Huggingface. One example is as follows: + +- [meituan/DeepSeek-R1-Block-INT8](https://huggingface.co/meituan/DeepSeek-R1-Block-INT8) +- [meituan/DeepSeek-R1-Channel-INT8](https://huggingface.co/meituan/DeepSeek-R1-Channel-INT8) + +Assuming that master node IP is `MASTER_IP`, checkpoint path is `/path/to/DeepSeek-R1-INT8` and port=5000, we can have following commands to launch the server: +```bash +#master +python3 -m sglang.launch_server \ + --model meituan/DeepSeek-R1-Block-INT8 --tp 16 --dist-init-addr \ + MASTER_IP:5000 --nnodes 2 --node-rank 0 --trust-remote-code --enable-torch-compile --torch-compile-max-bs 8 +#cluster +python3 -m sglang.launch_server \ + --model meituan/DeepSeek-R1-Block-INT8 --tp 16 --dist-init-addr \ + MASTER_IP:5000 --nnodes 2 --node-rank 1 --trust-remote-code --enable-torch-compile --torch-compile-max-bs 8 +``` + +> **Note that the launch command here enables `torch.compile` Optimization**. For optimal performance, please refer to the command options in [Performance Optimization Options](#option_args). + +Then on the **master node**, supposing the ShareGPT data is located at `/path/to/ShareGPT_V3_unfiltered_cleaned_split.json`, you can run the following commands to benchmark the launched server: + +```bash +# bench accuracy +python3 benchmark/gsm8k/bench_sglang.py --num-questions 1319 + +# bench serving +python3 -m sglang.bench_serving --dataset-path /path/to/ShareGPT_V3_unfiltered_cleaned_split.json --dataset-name random --random-input 128 --random-output 128 --num-prompts 1000 --request-rate 128 --random-range-ratio 1.0 +``` + +> **Note: using `--parallel 200` can accelerate accuracy benchmarking**. + +### Example: Serving with 32 L40S with int8 Quantization + +Running with per-channel quantization model: + +- [meituan/DeepSeek-R1-Channel-INT8](https://huggingface.co/meituan/DeepSeek-R1-Channel-INT8) + +Assuming that master node IP is `MASTER_IP`, checkpoint path is `/path/to/DeepSeek-R1-Channel-INT8` and port=5000, we can have following commands to launch the server: + +```bash +#master +python3 -m sglang.launch_server --model meituan/DeepSeek-R1-Channel-INT8 --tp 32 --quantization w8a8_int8 \ + --dist-init-addr MASTER_IP:5000 --nnodes 4 --node-rank 0 --trust-remote \ + --enable-torch-compile --torch-compile-max-bs 32 +#cluster +python3 -m sglang.launch_server --model meituan/DeepSeek-R1-Channel-INT8 --tp 32 --quantization w8a8_int8 \ + --dist-init-addr MASTER_IP:5000 --nnodes 4 --node-rank 1 --trust-remote \ + --enable-torch-compile --torch-compile-max-bs 32 +python3 -m sglang.launch_server --model meituan/DeepSeek-R1-Channel-INT8 --tp 32 --quantization w8a8_int8 \ + --dist-init-addr MASTER_IP:5000 --nnodes 4 --node-rank 2 --trust-remote \ + --enable-torch-compile --torch-compile-max-bs 32 +python3 -m sglang.launch_server --model meituan/DeepSeek-R1-Channel-INT8 --tp 32 --quantization w8a8_int8 \ + --dist-init-addr MASTER_IP:5000 --nnodes 4 --node-rank 3 --trust-remote \ + --enable-torch-compile --torch-compile-max-bs 32 +``` + +The benchmarking method is the same as describted in the previous [16 x A100](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-16-a100a800-with-int8-quantization) example. + +### Example: Serving on any cloud or Kubernetes with SkyPilot + +SkyPilot helps find cheapest available GPUs across any cloud or existing Kubernetes clusters and launch distributed serving with a single command. See details [here](https://github.com/skypilot-org/skypilot/tree/master/llm/deepseek-r1). + +To serve on multiple nodes: + +```bash +git clone https://github.com/skypilot-org/skypilot.git +# Serve on 2 H100/H200x8 nodes +sky launch -c r1 llm/deepseek-r1/deepseek-r1-671B.yaml --retry-until-up +# Serve on 4 A100x8 nodes +sky launch -c r1 llm/deepseek-r1/deepseek-r1-671B-A100.yaml --retry-until-up +``` + +#### Troubleshooting + +If you encounter the following error with fp16/bf16 checkpoint: + +```bash +ValueError: Weight output_partition_size = 576 is not divisible by weight quantization block_n = 128. +``` + +edit your `config.json` and remove the `quantization_config` block. For example: + +```json +"quantization_config": { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "weight_block_size": [128, 128] +}, +``` + +Removing this block typically resolves the error. For more details, see the discussion in [sgl-project/sglang#3491](https://github.com/sgl-project/sglang/issues/3491#issuecomment-2650779851). + +## DeepSeek V3 Optimization Plan + +https://github.com/sgl-project/sglang/issues/2591 diff --git a/sglang/benchmark/dspy/README.md b/sglang/benchmark/dspy/README.md new file mode 100644 index 0000000..a2b213a --- /dev/null +++ b/sglang/benchmark/dspy/README.md @@ -0,0 +1,51 @@ +## Install + +``` +pip3 install dspy-ai +``` + +Turn off cache at https://github.com/stanfordnlp/dspy/blob/34d8420383ec752037aa271825c1d3bf391e1277/dsp/modules/cache_utils.py#L10. +``` +cache_turn_on = False +``` + +or set the environment variable + +``` +export DSP_CACHEBOOL=false +``` + +## Benchmark SGLang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_dspy_intro.py --backend sglang +``` + + +## Benchmark TGI +``` +docker run --name tgi --rm -ti --gpus all --network host \ + -v /home/ubuntu/model_weights/Llama-2-7b-chat-hf:/Llama-2-7b-chat-hf \ + ghcr.io/huggingface/text-generation-inference:1.3.0 \ + --model-id /Llama-2-7b-chat-hf --num-shard 1 --trust-remote-code \ + --max-input-length 2048 --max-total-tokens 4096 \ + --port 24000 +``` + +``` +python3 bench_dspy_intro.py --backend tgi +``` + + + +## Benchmark vLLM +``` +python3 -m vllm.entrypoints.openai.api_server --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_dspy_intro.py --backend vllm +``` diff --git a/sglang/benchmark/dspy/bench_dspy_intro.py b/sglang/benchmark/dspy/bench_dspy_intro.py new file mode 100644 index 0000000..2b0936e --- /dev/null +++ b/sglang/benchmark/dspy/bench_dspy_intro.py @@ -0,0 +1,192 @@ +""" +Adapted from +https://github.com/stanfordnlp/dspy/blob/34d8420383ec752037aa271825c1d3bf391e1277/intro.ipynb#L9 +""" + +import argparse + +import dspy +from dspy.datasets import HotPotQA + + +class BasicQA(dspy.Signature): + """Answer questions with short factoid answers.""" + + question = dspy.InputField() + answer = dspy.OutputField(desc="often between 1 and 5 words") + + +class GenerateAnswer(dspy.Signature): + """Answer questions with short factoid answers.""" + + context = dspy.InputField(desc="may contain relevant facts") + question = dspy.InputField() + answer = dspy.OutputField(desc="often between 1 and 5 words") + + +class RAG(dspy.Module): + def __init__(self, num_passages=3): + super().__init__() + + self.retrieve = dspy.Retrieve(k=num_passages) + self.generate_answer = dspy.ChainOfThought(GenerateAnswer) + + def forward(self, question): + context = self.retrieve(question).passages + prediction = self.generate_answer(context=context, question=question) + return dspy.Prediction(context=context, answer=prediction.answer) + + +def main(args): + # lm = dspy.OpenAI(model='gpt-3.5-turbo') + if args.backend == "tgi": + lm = dspy.HFClientTGI( + model="meta-llama/Llama-2-7b-chat-hf", + port=args.port, + url="http://localhost", + ) + elif args.backend == "sglang": + lm = dspy.HFClientSGLang( + model="meta-llama/Llama-2-7b-chat-hf", + port=args.port, + url="http://localhost", + ) + elif args.backend == "vllm": + lm = dspy.HFClientVLLM( + model="meta-llama/Llama-2-7b-chat-hf", + port=args.port, + url="http://localhost", + ) + else: + raise ValueError(f"Invalid backend: {args.backend}") + + colbertv2_wiki17_abstracts = dspy.ColBERTv2( + url="http://20.102.90.50:2017/wiki17_abstracts" + ) + dspy.settings.configure(lm=lm, rm=colbertv2_wiki17_abstracts) + + # Load the dataset. + dataset = HotPotQA( + train_seed=1, train_size=20, eval_seed=2023, dev_size=args.dev_size, test_size=0 + ) + + # Tell DSPy that the 'question' field is the input. Any other fields are labels and/or metadata. + trainset = [x.with_inputs("question") for x in dataset.train] + devset = [x.with_inputs("question") for x in dataset.dev] + + print(len(trainset), len(devset)) + + train_example = trainset[0] + print(f"Question: {train_example.question}") + print(f"Answer: {train_example.answer}") + + dev_example = devset[18] + print(f"Question: {dev_example.question}") + print(f"Answer: {dev_example.answer}") + print(f"Relevant Wikipedia Titles: {dev_example.gold_titles}") + + print( + f"For this dataset, training examples have input keys {train_example.inputs().keys()} and label keys {train_example.labels().keys()}" + ) + print( + f"For this dataset, dev examples have input keys {dev_example.inputs().keys()} and label keys {dev_example.labels().keys()}" + ) + + # Define the predictor. + generate_answer = dspy.Predict(BasicQA) + + # Call the predictor on a particular input. + pred = generate_answer(question=dev_example.question) + + # Print the input and the prediction. + print(f"Question: {dev_example.question}") + print(f"Predicted Answer: {pred.answer}") + + lm.inspect_history(n=1) + + # Define the predictor. Notice we're just changing the class. The signature BasicQA is unchanged. + generate_answer_with_chain_of_thought = dspy.ChainOfThought(BasicQA) + + # Call the predictor on the same input. + pred = generate_answer_with_chain_of_thought(question=dev_example.question) + + # Print the input, the chain of thought, and the prediction. + print(f"Question: {dev_example.question}") + print(f"Thought: {pred.rationale.split('.', 1)[1].strip()}") + print(f"Predicted Answer: {pred.answer}") + + retrieve = dspy.Retrieve(k=3) + topK_passages = retrieve(dev_example.question).passages + + print( + f"Top {retrieve.k} passages for question: {dev_example.question} \n", + "-" * 30, + "\n", + ) + + for idx, passage in enumerate(topK_passages): + print(f"{idx+1}]", passage, "\n") + + retrieve("When was the first FIFA World Cup held?").passages[0] + + from dspy.teleprompt import BootstrapFewShot + + # Validation logic: check that the predicted answer is correct. + # Also check that the retrieved context does actually contain that answer. + def validate_context_and_answer(example, pred, trace=None): + answer_EM = dspy.evaluate.answer_exact_match(example, pred) + answer_PM = dspy.evaluate.answer_passage_match(example, pred) + return answer_EM and answer_PM + + # Set up a basic teleprompter, which will compile our RAG program. + teleprompter = BootstrapFewShot(metric=validate_context_and_answer) + + # Compile! + compiled_rag = teleprompter.compile(RAG(), trainset=trainset) + + # Ask any question you like to this simple RAG program. + my_question = "What castle did David Gregory inherit?" + + # Get the prediction. This contains `pred.context` and `pred.answer`. + pred = compiled_rag(my_question) + + # Print the contexts and the answer. + print(f"Question: {my_question}") + print(f"Predicted Answer: {pred.answer}") + print(f"Retrieved Contexts (truncated): {[c[:200] + '...' for c in pred.context]}") + + from dspy.evaluate.evaluate import Evaluate + + # Set up the `evaluate_on_hotpotqa` function. We'll use this many times below. + evaluate_on_hotpotqa = Evaluate( + devset=devset, + num_threads=args.num_threads, + display_progress=True, + display_table=5, + ) + + # Evaluate the `compiled_rag` program with the `answer_exact_match` metric. + metric = dspy.evaluate.answer_exact_match + evaluate_on_hotpotqa(compiled_rag, metric=metric) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--port", type=int) + parser.add_argument("--num-threads", type=int, default=32) + parser.add_argument("--dev-size", type=int, default=150) + parser.add_argument( + "--backend", type=str, choices=["sglang", "tgi", "vllm"], default="sglang" + ) + args = parser.parse_args() + + if args.port is None: + default_port = { + "vllm": 21000, + "lightllm": 22000, + "tgi": 24000, + "sglang": 30000, + } + args.port = default_port.get(args.backend, None) + + main(args) diff --git a/sglang/benchmark/generative_agents/README.md b/sglang/benchmark/generative_agents/README.md new file mode 100644 index 0000000..393a9ce --- /dev/null +++ b/sglang/benchmark/generative_agents/README.md @@ -0,0 +1,38 @@ +## Download the dataset + +``` +wget -O agent_calls.jsonl https://drive.google.com/uc?export=download&id=19qLpD45e9JGTKF2cUjJJegwzSUEZEKht +``` + +## Run benchmark + +Ensure that this benchmark is run in a serial manner (using --parallel 1) to preserve any potential dependencies between requests. + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-events 1000 --parallel 1 +``` + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-events 1000 --backend vllm --parallel 1 +``` + +### Benchmark guidance +``` +python3 bench_other.py --num-events 1000 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +### Benchmark lmql + +``` +python3 bench_other.py --num-events 1000 --backend lmql --parallel 1 +``` diff --git a/sglang/benchmark/generative_agents/agent_functions.py b/sglang/benchmark/generative_agents/agent_functions.py new file mode 100644 index 0000000..9785fc1 --- /dev/null +++ b/sglang/benchmark/generative_agents/agent_functions.py @@ -0,0 +1,300 @@ +import sglang as sgl + +# here are the top five agent functions contributing ~70% LLM calls +# reference: https://github.com/joonspk-research/generative_agents/ + + +@sgl.function +def poignancy_event(s, persona_name, persona_iss, event): + s += "Here is a brief description of " + persona_name + ".\n" + s += persona_iss + "\n" + s += "On the scale of 1 to 10, where 1 is purely mundane (e.g., brushing teeth, making bed) and 10 is extremely poignant (e.g., a break up, college acceptance), rate the likely poignancy of the following event for" + s += persona_name + ".\n\n" + s += "Event: " + event + s += "Rate (return a number between 1 to 10):" + s += sgl.gen(name="Rate", max_tokens=2) + + +def poignancy_event_prompt(persona_name, persona_iss, event): + # return prompt and max_tokens + s = "" + s += "Here is a brief description of " + persona_name + ".\n" + s += persona_iss + "\n" + s += "On the scale of 1 to 10, where 1 is purely mundane (e.g., brushing teeth, making bed) and 10 is extremely poignant (e.g., a break up, college acceptance), rate the likely poignancy of the following event for" + s += persona_name + ".\n\n" + s += "Event: " + event + s += "Rate (return a number between 1 to 10):" + return {"prompt": s, "max_tokens": 2, "stop": None} + + +@sgl.function +def generate_event_triple(s, persona_name, action): + s += """Task: Turn the input into (subject, predicate, object). +Input: Sam Johnson is eating breakfast. +Output: (Dolores Murphy, eat, breakfast) +--- +Input: Joon Park is brewing coffee. +Output: (Joon Park, brew, coffee) +--- +Input: Jane Cook is sleeping. +Output: (Jane Cook, is, sleep) +--- +Input: Michael Bernstein is writing email on a computer. +Output: (Michael Bernstein, write, email) +--- +Input: Percy Liang is teaching students in a classroom. +Output: (Percy Liang, teach, students) +--- +Input: Merrie Morris is running on a treadmill. +Output: (Merrie Morris, run, treadmill) +---""" + s += persona_name + "is" + action + ".\n" + s += "(" + persona_name + "," + s += sgl.gen(name="Triple", max_tokens=20, stop=")") + + +def generate_event_triple_prompt(persona_name, action): + s = "" + s += """Task: Turn the input into (subject, predicate, object). +Input: Sam Johnson is eating breakfast. +Output: (Dolores Murphy, eat, breakfast) +--- +Input: Joon Park is brewing coffee. +Output: (Joon Park, brew, coffee) +--- +Input: Jane Cook is sleeping. +Output: (Jane Cook, is, sleep) +--- +Input: Michael Bernstein is writing email on a computer. +Output: (Michael Bernstein, write, email) +--- +Input: Percy Liang is teaching students in a classroom. +Output: (Percy Liang, teach, students) +--- +Input: Merrie Morris is running on a treadmill. +Output: (Merrie Morris, run, treadmill) +---""" + s += persona_name + "is" + action + ".\n" + s += "(" + persona_name + "," + return {"prompt": s, "max_tokens": 20, "stop": ")"} + + +@sgl.function +def generate_pronunciatio(s, action): + s += "Convert an action description to an emoji (important: use two or less emojis).\n" + s += "Action description: " + action + ".\n" + s += "Emoji:" + sgl.gen(name="Emoji", max_tokens=6) + + +def generate_pronunciatio_prompt(action): + s = "" + s += "Convert an action description to an emoji (important: use two or less emojis).\n" + s += "Action description: " + action + ".\n" + s += "Emoji:" + return {"prompt": s, "max_tokens": 6, "stop": None} + + +@sgl.function +def action_location_sector( + s, + persona_name, + living_sector, + living_sector_areas, + current_sector, + current_sector_areas, + daily_plan, + sector_options, + current_action, + next_action, +): + s += """Task -- choose an appropriate area from the area options for a task at hand. +Sam Kim lives in {Sam Kim's house} that has Sam Kim's room, bathroom, kitchen. +Sam Kim is currently in {Sam Kim's house} that has Sam Kim's room, bathroom, kitchen. +Area options: {Sam Kim's house, The Rose and Crown Pub, Hobbs Cafe, Oak Hill College, Johnson Park, Harvey Oak Supply Store, The Willows Market and Pharmacy}. +* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. +* Must be one of the "Area options," verbatim. +For taking a walk, Sam Kim should go to the following area: {Johnson Park} +--- +Jane Anderson lives in {Oak Hill College Student Dormatory} that has Jane Anderson's room. +Jane Anderson is currently in {Oak Hill College} that has a classroom, library +Area options: {Oak Hill College Student Dormatory, The Rose and Crown Pub, Hobbs Cafe, Oak Hill College, Johnson Park, Harvey Oak Supply Store, The Willows Market and Pharmacy}. +* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. +* Must be one of the "Area options," verbatim. +For eating dinner, Jane Anderson should go to the following area: {Hobbs Cafe} +---""" + s += ( + persona_name + + " lives in " + + living_sector + + " that has " + + living_sector_areas + + ".\n" + ) + s += ( + persona_name + + " is currently in " + + current_sector + + " that has " + + current_sector_areas + + ".\n" + ) + s += daily_plan + ".\n" + s += "Area options: " + sector_options + ".\n" + s += """* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. +* Must be one of the "Area options," verbatim.\n""" + s += ( + persona_name + + " is " + + current_action + + ". For " + + next_action + + ", " + + persona_name + + " should go to the following area: {" + ) + s += sgl.gen(name="Location", max_tokens=10, stop="}") + + +def action_location_sector_prompt( + persona_name, + living_sector, + living_sector_areas, + current_sector, + current_sector_areas, + daily_plan, + sector_options, + current_action, + next_action, +): + s = "" + s += """Task -- choose an appropriate area from the area options for a task at hand. +Sam Kim lives in {Sam Kim's house} that has Sam Kim's room, bathroom, kitchen. +Sam Kim is currently in {Sam Kim's house} that has Sam Kim's room, bathroom, kitchen. +Area options: {Sam Kim's house, The Rose and Crown Pub, Hobbs Cafe, Oak Hill College, Johnson Park, Harvey Oak Supply Store, The Willows Market and Pharmacy}. +* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. +* Must be one of the "Area options," verbatim. +For taking a walk, Sam Kim should go to the following area: {Johnson Park} +--- +Jane Anderson lives in {Oak Hill College Student Dormatory} that has Jane Anderson's room. +Jane Anderson is currently in {Oak Hill College} that has a classroom, library +Area options: {Oak Hill College Student Dormatory, The Rose and Crown Pub, Hobbs Cafe, Oak Hill College, Johnson Park, Harvey Oak Supply Store, The Willows Market and Pharmacy}. +* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. +* Must be one of the "Area options," verbatim. +For eating dinner, Jane Anderson should go to the following area: {Hobbs Cafe} +---""" + s += ( + persona_name + + " lives in " + + living_sector + + " that has " + + living_sector_areas + + ".\n" + ) + s += ( + persona_name + + " is currently in " + + current_sector + + " that has " + + current_sector_areas + + ".\n" + ) + s += daily_plan + ".\n" + s += "Area options: " + sector_options + ".\n" + s += """* Stay in the current area if the activity can be done there. Only go out if the activity needs to take place in another place. +* Must be one of the "Area options," verbatim.\n""" + s += ( + persona_name + + " is " + + current_action + + ". For " + + next_action + + ", " + + persona_name + + " should go to the following area: {" + ) + return {"prompt": s, "max_tokens": 10, "stop": "}"} + + +@sgl.function +def action_location_object( + s, persona_name, target_sector, target_sector_areas, current_action, next_action +): + s += """ +Jane Anderson is in kitchen in Jane Anderson's house. +Jane Anderson is going to Jane Anderson's house that has the following areas: {kitchen, bedroom, bathroom} +Stay in the current area if the activity can be done there. Never go into other people's rooms unless necessary. +For cooking, Jane Anderson should go to the following area in Jane Anderson's house: +Answer: {kitchen} +--- +Tom Watson is in common room in Tom Watson's apartment. +Tom Watson is going to Hobbs Cafe that has the following areas: {cafe} +Stay in the current area if the activity can be done there. Never go into other people's rooms unless necessary. +For getting coffee, Tom Watson should go to the following area in Hobbs Cafe: +Answer: {cafe} +---""" + s += ( + persona_name + + " is going to " + + target_sector + + " that has the following areas: {" + + target_sector_areas + + "}\n" + ) + s += """* Stay in the current area if the activity can be done there. +* NEVER go into other people's rooms unless necessary.""" + s += ( + persona_name + + " is " + + current_action + + ". For " + + next_action + + ", " + + persona_name + + "should go to the following area in " + + target_sector + ) + s += " (MUST pick one of {" + target_sector_areas + "}):\n" + s += "Answer: {" + sgl.gen(name="Area", max_tokens=5, stop="}") + + +def action_location_object_prompt( + persona_name, target_sector, target_sector_areas, current_action, next_action +): + s = "" + s += """ +Jane Anderson is in kitchen in Jane Anderson's house. +Jane Anderson is going to Jane Anderson's house that has the following areas: {kitchen, bedroom, bathroom} +Stay in the current area if the activity can be done there. Never go into other people's rooms unless necessary. +For cooking, Jane Anderson should go to the following area in Jane Anderson's house: +Answer: {kitchen} +--- +Tom Watson is in common room in Tom Watson's apartment. +Tom Watson is going to Hobbs Cafe that has the following areas: {cafe} +Stay in the current area if the activity can be done there. Never go into other people's rooms unless necessary. +For getting coffee, Tom Watson should go to the following area in Hobbs Cafe: +Answer: {cafe} +---""" + s += ( + persona_name + + " is going to " + + target_sector + + " that has the following areas: {" + + target_sector_areas + + "}\n" + ) + s += """* Stay in the current area if the activity can be done there. +* NEVER go into other people's rooms unless necessary.""" + s += ( + persona_name + + " is " + + current_action + + ". For " + + next_action + + ", " + + persona_name + + "should go to the following area in " + + target_sector + ) + s += " (MUST pick one of {" + target_sector_areas + "}):\n" + s += "Answer: {" + return {"prompt": s, "max_tokens": 5, "stop": "}"} diff --git a/sglang/benchmark/generative_agents/bench_other.py b/sglang/benchmark/generative_agents/bench_other.py new file mode 100644 index 0000000..48f6ebc --- /dev/null +++ b/sglang/benchmark/generative_agents/bench_other.py @@ -0,0 +1,80 @@ +import argparse +import json +import time + +from agent_functions import ( + action_location_object_prompt, + action_location_sector_prompt, + generate_event_triple_prompt, + generate_pronunciatio_prompt, + poignancy_event_prompt, +) +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_events] + mapping = { + "poignancy_event": poignancy_event_prompt, + "generate_event_triple": generate_event_triple_prompt, + "generate_pronunciatio": generate_pronunciatio_prompt, + "action_location_sector": action_location_sector_prompt, + "action_location_object": action_location_object_prompt, + } + + arguments = [mapping[k](**v) for l in lines for k, v in l.items()] + states = [] + + # Select backend + call_generate = get_call_generate(args) + + def get_one_answer(arg): + answer = call_generate(**arg, temperature=0) + states.append(answer) + + async def get_one_answer_async(arg): + answer = await call_generate(**arg, temperature=0) + states.append(answer) + + tic = time.time() + # we always sequentially execute agent calls to maintain its dependency + if args.backend != "lmql": + for arg in tqdm(arguments): + get_one_answer(arg) + else: + import asyncio + + loop = asyncio.get_event_loop() + for arg in tqdm(arguments): + loop.run_until_complete(get_one_answer_async(arg)) + latency = time.time() - tic + + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "Generative Agents", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + # to pack weighted functions as a single agent + "num_requests": len(arguments) / len(mapping), + "other": { + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="agent_calls.jsonl") + parser.add_argument("--num-events", type=int, default=10) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/generative_agents/bench_sglang.py b/sglang/benchmark/generative_agents/bench_sglang.py new file mode 100644 index 0000000..b42a32b --- /dev/null +++ b/sglang/benchmark/generative_agents/bench_sglang.py @@ -0,0 +1,74 @@ +import argparse +import json +import time + +from agent_functions import ( + action_location_object, + action_location_sector, + generate_event_triple, + generate_pronunciatio, + poignancy_event, +) + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_events] + mapping = { + "poignancy_event": poignancy_event, + "generate_event_triple": generate_event_triple, + "generate_pronunciatio": generate_pronunciatio, + "action_location_sector": action_location_sector, + "action_location_object": action_location_object, + } + arguments = [{mapping[k]: v for k, v in l.items()} for l in lines] + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + states = [] + # Run requests + tic = time.time() + for a in arguments: + # only a single key in the dict + for func, arg in a.items(): + result = func.run(**arg) + result.sync() + states.append(result) + latency = time.time() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "Generative Agents", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + # to pack weighted functions as a single agent + "num_requests": len(arguments) / len(mapping), + "other": { + "num_events": args.num_events, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="agent_calls.jsonl") + parser.add_argument("--num-events", type=int, default=10) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/gsm8k/README.md b/sglang/benchmark/gsm8k/README.md new file mode 100644 index 0000000..c110f53 --- /dev/null +++ b/sglang/benchmark/gsm8k/README.md @@ -0,0 +1,47 @@ +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 200 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-questions 200 --backend vllm +``` + + +### Benchmark lightllm +``` +# A10G +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 +``` + +``` +python3 bench_other.py --num-questions 200 --backend lightllm +``` + + +### Benchmark guidance +``` +python3 bench_other.py --num-questions 200 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + + +### Benchmark lmql +``` +CUDA_VISIBLE_DEVICES=0,1 lmql serve-model meta-llama/Llama-2-7b-chat-hf --cuda --port 23000 +``` + +``` +python3 bench_other.py --num-questions 100 --backend lmql --parallel 2 +``` diff --git a/sglang/benchmark/gsm8k/bench_other.py b/sglang/benchmark/gsm8k/bench_other.py new file mode 100644 index 0000000..a8bbcfb --- /dev/null +++ b/sglang/benchmark/gsm8k/bench_other.py @@ -0,0 +1,151 @@ +import argparse +import ast +import asyncio +import json +import re +import time +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import download_and_cache_file, dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_one_example(lines, i, include_answer): + ret = "Question: " + lines[i]["question"] + "\nAnswer:" + if include_answer: + ret += " " + lines[i]["answer"] + return ret + + +def get_few_shot_examples(lines, k): + ret = "" + for i in range(k): + ret += get_one_example(lines, i, True) + "\n\n" + return ret + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def main(args): + # Select backend + call_generate = get_call_generate(args) + + # Read data + url = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl" + filename = download_and_cache_file(url) + lines = list(read_jsonl(filename)) + + # Construct prompts + num_questions = args.num_questions + num_shots = args.num_shots + few_shot_examples = get_few_shot_examples(lines, num_shots) + + questions = [] + labels = [] + for i in range(len(lines[:num_questions])): + questions.append(get_one_example(lines, i, False)) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + + states = [None] * len(labels) + + # Run requests + if args.backend != "lmql": + # Use thread pool + def get_one_answer(i): + answer = call_generate( + prompt=few_shot_examples + questions[i], + temperature=0, + max_tokens=256, + stop=["Question", "Assistant:", "<|separator|>"], + ) + states[i] = answer + + tic = time.time() + if args.parallel == 1: + for i in tqdm(range(len(questions))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(questions)))), + total=len(questions), + ) + ) + + else: + # Use asyncio + async def batched_call(batch_size): + for i in range(0, len(questions), batch_size): + tasks = [] + for q in questions[i : i + batch_size]: + tasks.append( + call_generate( + few_shot_examples + q, + temperature=0, + max_tokens=256, + stop="Question", + ) + ) + rets = await asyncio.gather(*tasks) + for j in range(len(rets)): + states[i + j] = rets[j] + + tic = time.time() + asyncio.run(batched_call(batch_size=args.parallel)) + latency = time.time() - tic + + preds = [] + for i in range(len(states)): + preds.append(get_answer_value(states[i])) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + + # Print results + print(f"Accuracy: {acc:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Latency: {latency:.3f} s") + + # Dump results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-shots", type=int, default=5) + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/gsm8k/bench_sglang.py b/sglang/benchmark/gsm8k/bench_sglang.py new file mode 100644 index 0000000..b6bdbef --- /dev/null +++ b/sglang/benchmark/gsm8k/bench_sglang.py @@ -0,0 +1,141 @@ +import argparse +import ast +import json +import os +import re +import time + +import numpy as np + +from sglang.api import set_default_backend +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import download_and_cache_file, dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_one_example(lines, i, include_answer): + ret = "Question: " + lines[i]["question"] + "\nAnswer:" + if include_answer: + ret += " " + lines[i]["answer"] + return ret + + +def get_few_shot_examples(lines, k): + ret = "" + for i in range(k): + ret += get_one_example(lines, i, True) + "\n\n" + return ret + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def main(args): + # Select backend + set_default_backend(select_sglang_backend(args)) + + # Read data + data_path = args.data_path + url = "https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl" + if not os.path.isfile(data_path): + data_path = download_and_cache_file(url) + lines = list(read_jsonl(data_path)) + + # Construct prompts + num_questions = args.num_questions + num_shots = args.num_shots + few_shot_examples = get_few_shot_examples(lines, num_shots) + + questions = [] + labels = [] + for i in range(len(lines[:num_questions])): + questions.append(get_one_example(lines, i, False)) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q} for q in questions] + + ##################################### + ######### SGL Program Begin ######### + ##################################### + + import sglang as sgl + + @sgl.function + def few_shot_gsm8k(s, question): + s += few_shot_examples + question + s += sgl.gen( + "answer", max_tokens=512, stop=["Question", "Assistant:", "<|separator|>"] + ) + + ##################################### + ########## SGL Program End ########## + ##################################### + + # Run requests + tic = time.time() + states = few_shot_gsm8k.run_batch( + arguments, + temperature=0, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.time() - tic + + preds = [] + for i in range(len(states)): + preds.append(get_answer_value(states[i]["answer"])) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + + # Compute speed + num_output_tokens = sum( + s.get_meta_info("answer")["completion_tokens"] for s in states + ) + output_throughput = num_output_tokens / latency + + # Print results + print(f"Accuracy: {acc:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Latency: {latency:.3f} s") + print(f"Output throughput: {output_throughput:.3f} token/s") + + # Dump results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-shots", type=int, default=5) + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/hellaswag/README.md b/sglang/benchmark/hellaswag/README.md new file mode 100644 index 0000000..cb7e653 --- /dev/null +++ b/sglang/benchmark/hellaswag/README.md @@ -0,0 +1,47 @@ +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 200 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-questions 200 --backend vllm +``` + + +### Benchmark lightllm +``` +# A10G +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 +``` + +``` +python3 bench_other.py --num-questions 200 --backend lightllm +``` + + +### Benchmark guidance +``` +CUDA_VISIBLE_DEVICES=0,1 python3 bench_other.py --num-questions 200 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + + +### Benchmark lmql +``` +lmql serve-model meta-llama/Llama-2-7b-chat-hf --cuda --port 23000 +``` + +``` +python3 bench_other.py --num-questions 200 --backend lmql --port 23000 --parallel 1 +``` diff --git a/sglang/benchmark/hellaswag/bench_other.py b/sglang/benchmark/hellaswag/bench_other.py new file mode 100644 index 0000000..04be456 --- /dev/null +++ b/sglang/benchmark/hellaswag/bench_other.py @@ -0,0 +1,118 @@ +import argparse +import asyncio +import json +import time +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_select +from sglang.utils import download_and_cache_file, read_jsonl + + +def get_one_example(lines, i, include_answer): + ret = lines[i]["activity_label"] + ": " + lines[i]["ctx"] + " " + if include_answer: + ret += lines[i]["endings"][lines[i]["label"]] + return ret + + +def get_few_shot_examples(lines, k): + ret = "" + for i in range(k): + ret += get_one_example(lines, i, True) + "\n\n" + return ret + + +def main(args): + # Select backend + call_select = get_call_select(args) + + # Read data + url = "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_val.jsonl" + filename = download_and_cache_file(url) + lines = list(read_jsonl(filename)) + + # Construct prompts + num_questions = args.num_questions + num_shots = args.num_shots + few_shot_examples = get_few_shot_examples(lines, num_shots) + + questions = [] + choices = [] + labels = [] + for i in range(len(lines[:num_questions])): + questions.append(get_one_example(lines, i, False)) + choices.append(lines[i]["endings"]) + labels.append(lines[i]["label"]) + + preds = [None] * len(labels) + + # Run requests + if args.backend != "lmql": + # Use thread pool + def get_one_answer(i): + preds[i] = call_select( + context=few_shot_examples + questions[i], choices=choices[i] + ) + + tic = time.time() + if args.parallel == 1: + for i in tqdm(range(len(questions))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(questions)))), + total=len(questions), + ) + ) + else: + # Use asyncio + async def batched_call(batch_size): + for i in range(0, len(questions), batch_size): + tasks = [] + for q, c in zip( + questions[i : i + batch_size], choices[i : i + batch_size] + ): + tasks.append(call_select(context=few_shot_examples + q, choices=c)) + rets = await asyncio.gather(*tasks) + for j in range(len(rets)): + preds[i + j] = rets[j] + + tic = time.time() + asyncio.run(batched_call(batch_size=args.parallel)) + + latency = time.time() - tic + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + print(f"Latency: {latency:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + with open(args.result_file, "a") as fout: + value = { + "task": "hellaswag", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-shots", type=int, default=20) + parser.add_argument("--data-path", type=str, default="hellaswag_val.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/hellaswag/bench_sglang.py b/sglang/benchmark/hellaswag/bench_sglang.py new file mode 100644 index 0000000..798521f --- /dev/null +++ b/sglang/benchmark/hellaswag/bench_sglang.py @@ -0,0 +1,109 @@ +import argparse +import json +import os +import time + +import numpy as np + +from sglang.api import set_default_backend +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import download_and_cache_file, read_jsonl + + +def get_one_example(lines, i, include_answer): + ret = lines[i]["activity_label"] + ": " + lines[i]["ctx"] + " " + if include_answer: + ret += lines[i]["endings"][lines[i]["label"]] + return ret + + +def get_few_shot_examples(lines, k): + ret = "" + for i in range(k): + ret += get_one_example(lines, i, True) + "\n\n" + return ret + + +def main(args): + # Select backend + set_default_backend(select_sglang_backend(args)) + + # Read data + data_path = args.data_path + url = "https://raw.githubusercontent.com/rowanz/hellaswag/master/data/hellaswag_val.jsonl" + if not os.path.isfile(data_path): + data_path = download_and_cache_file(url) + lines = list(read_jsonl(data_path)) + + # Construct prompts + num_questions = args.num_questions + num_shots = args.num_shots + few_shot_examples = get_few_shot_examples(lines, num_shots) + + questions = [] + choices = [] + labels = [] + for i in range(len(lines[:num_questions])): + questions.append(get_one_example(lines, i, False)) + choices.append(lines[i]["endings"]) + labels.append(lines[i]["label"]) + arguments = [{"question": q, "choices": c} for q, c in zip(questions, choices)] + + ##################################### + ######### SGL Program Begin ######### + ##################################### + + import sglang as sgl + + @sgl.function + def few_shot_hellaswag(s, question, choices): + s += few_shot_examples + question + s += sgl.select("answer", choices=choices) + + ##################################### + ########## SGL Program End ########## + ##################################### + + # Run requests + tic = time.time() + rets = few_shot_hellaswag.run_batch( + arguments, + temperature=0, + num_threads=args.parallel, + progress_bar=True, + ) + preds = [choices[i].index(rets[i]["answer"]) for i in range(len(rets))] + latency = time.time() - tic + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + print(f"Latency: {latency:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + with open(args.result_file, "a") as fout: + value = { + "task": "hellaswag", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-shots", type=int, default=20) + parser.add_argument("--data-path", type=str, default="hellaswag_val.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/hicache/README.md b/sglang/benchmark/hicache/README.md new file mode 100644 index 0000000..f2aa6ff --- /dev/null +++ b/sglang/benchmark/hicache/README.md @@ -0,0 +1,91 @@ +## Run synthetic multi-turn benchmark + +``` +# SGLang server with radix cache disabled +python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 --disable-radix-cache + +# SGLang server with radix cache on and first-come-first-serve policy +python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 --schedule-policy fcfs + +# The default SGLang server with radix cache on and long-prefix-match policy +python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 + +# SGLang server with hierarchical radix cache enabled +python -m sglang.launch_server --model-path Qwen/Qwen2.5-14B-Instruct --port 30000 --enable-hierarchical-cache + +``` + +``` +python bench_multiturn.py --model-path Qwen/Qwen2.5-14B-Instruct +``` + +Note: The performance gain of hierarchical caching depends on the ratio of reusable tokens to GPU memory capacity. The more tokens to be reused, the larger the model, and the more constrained the GPU memory size, the greater the benefit one can expect from hierarchical caching. + + +# Benchmark with more datasets +## Download Dataset +```bash +./download.sh {sharegpt|ultragpt|loogle|nextqa|all} +``` +This script will automatically download the required dataset to the current working directory + +## Multiturn Benchmark +### Supported Datasets +- sharegpt +- ultrachat +- loogle +### Example Usage: +```bash +python3 bench_serving.py --model mistralai/Mistral-7B-Instruct-v0.3 --backend sglang \ +--dataset-path longdep_qa.json --dataset-name loogle --request-rate 10 --num-prompts 10 \ +--port 8001 --enable-multiturn --disable-shuffle +``` +This uses `mistralai/Mistral-7B-Instruct-v0.3` model with `sglang` as backend. The dataset +is `longdep_qa.json`. We send `10 conversations` with `10 req/s` to port 8001. We enable +multiturn chat without shuffling the order of conversations (i.e. following the original +order in the dataset file). + +### Note: +The requests of multiple conversations are sent in a round robin fashion. +For example, if we have 3 conversations A, B, C whose rounds are `[2, 3, 4]` correspondingly, +multiturn chat will send the requests to the backend in the following order: `[A1, B1, C1, A2, B2, C2, B3, C3, C4]` +This has implications on the cache reuse patterns: the cache reuse distance is the largest +under this request pattern (which means a prefix-aware local scheduler in the backend can +yield the most benefit compared to a FIFO scheduler) + +## Shared Prefix Benchmark +### Supported Datasets +- loogle +### Example Usage: +```bash +python3 bench_serving.py --model mistralai/Mistral-7B-Instruct-v0.3 --backend sglang \ +--dataset-path longdep_qa.json --dataset-name loogle --request-rate 10 --num-prompts 10 \ +--port 8001 --enable-shared-prefix --disable-shuffle +``` +### Note: +Shared Prefix benchmark sends the questions for the same prompt together. For example, +if we have 3 shared prefix A, B, C, which have [2, 3, 4] questions correspondingly, +the shared prefix benchmark will send the requests to the +backend in the following order: `[A+Q1, A+Q2, B+Q1, B+Q2, B+Q3, C+Q1, C+Q2, C+Q3]`. + + +## Multi Modality Benchmark (WIP) +### Supported Datasets: +- nextqa +### Example Usage: +```bash +Server: +python3 -m sglang.launch_server --model-path lmms-lab/LLaVA-NeXT-Video-7B --tp 2 --dp 1 --port 8001 \ +--host 0.0.0.0 --mem-fraction-static 0.9 --tokenizer-path llava-hf/llava-1.5-7b-hf \ +--json-model-override-args "{\"architectures\": [\"LlavaVidForCausalLM\"], \"model_type\":\"llava\", \"mm_spatial_pool_stride\":2}" + +Client: +python3 bench_serving.py --model lmms-lab/LLaVA-NeXT-Video-7B --backend sglang --dataset-path \ +NExTVideo --dataset-name nextqa --request-rate 10 --num-prompts 1 --disable-shuffle --port 8001 \ --enable-multiturn --max-frames 16 --tokenizer llava-hf/llava-1.5-7b-hf --fixed-output-len 2048 +``` +Note: for the server args, `tokenizer-path`, overriding architecture are necessary. + +## Supported Backend +- sglang (oai) +- vllm (oai) +- lmdeploy (oai) diff --git a/sglang/benchmark/hicache/bench_multiturn.py b/sglang/benchmark/hicache/bench_multiturn.py new file mode 100644 index 0000000..6bd0bd9 --- /dev/null +++ b/sglang/benchmark/hicache/bench_multiturn.py @@ -0,0 +1,395 @@ +import argparse +import asyncio +import json +import queue +import random +import threading +import time +from datetime import datetime +from typing import Optional + +import aiohttp +import requests +from tqdm.asyncio import tqdm + +from sglang.bench_serving import ( + RequestFuncOutput, + get_tokenizer, + remove_prefix, + sample_random_requests, +) + + +def parse_args(): + parser = argparse.ArgumentParser( + description="Script to benchmark concurrent requests to a server." + ) + parser.add_argument( + "--num-clients", + type=int, + default=256, + help="Number of concurrent clients", + ) + parser.add_argument( + "--max-parallel", + type=int, + default=128, + help="Maximum number of parallel requests", + ) + parser.add_argument( + "--request-length", + type=int, + default=512, + help="Length of each new request", + ) + parser.add_argument( + "--output-length", + type=int, + default=64, + help="Length of each output", + ) + parser.add_argument( + "--num-rounds", + type=int, + default=5, + help="Number of rounds per client", + ) + parser.add_argument( + "--distribution", + type=str, + default="poisson", + choices=["poisson", "uniform"], + help="Distribution type for request intervals (poisson or uniform)", + ) + parser.add_argument( + "--request-rate", + type=float, + default=1.0, + help="Average number of requests per second", + ) + parser.add_argument( + "--host", + type=str, + default="localhost", + help="Server hostname or IP (default: localhost)", + ) + parser.add_argument( + "--port", + type=int, + default=30000, + help="Server port (default: 30000)", + ) + parser.add_argument( + "--model-path", + type=str, + default="meta-llama/Llama-3.1-8B-Instruct", + help="model path compatible with Hugging Face Transformers", + ) + parser.add_argument( + "--dataset-path", + type=str, + default="", + help="local dataset to sample tokens from", + ) + parser.add_argument( + "--log-file", + type=str, + default="performance_metrics.jsonl", + help="File to log performance metrics", + ) + return parser.parse_args() + + +async def async_request_sglang_generate( + payload, + url, + pbar: Optional[tqdm] = None, +): + """ + Sends a streaming request to the server. Gathers text token-by-token. + """ + async with aiohttp.ClientSession() as session: + headers = {} + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + output = RequestFuncOutput() + + try: + async with session.post(url=url, json=payload, headers=headers) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + if data["text"]: + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + generated_text = data["text"] + + output.generated_text = generated_text + output.success = True + output.latency = latency + else: + output.error = response.reason or "" + output.success = False + except Exception as e: + output.success = False + output.error = str(e) + print(f"Request failed: {e}") + + if pbar: + pbar.update(1) + return output + + +def gen_payload(prompt, output_len): + payload = { + "text": prompt, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": output_len, + "ignore_eos": True, + }, + "stream": True, + "lora_path": "", + "return_logprob": False, + "logprob_start_len": -1, + } + return payload + + +def log_to_jsonl_file(data, file_path="performance_metrics.jsonl"): + """Append the data with a timestamp to the specified JSONL file.""" + timestamped_data = {"timestamp": datetime.now().isoformat(), **data} + try: + with open(file_path, "a") as file: + file.write( + json.dumps(timestamped_data) + "\n" + ) # Write as a single line in JSONL format + except IOError as e: + print(f"Error writing to JSONL file: {e}") + + +class ReadyQueue: + """ + Thread-safe queue that can pop requests in different orders based on given policy. + """ + + def __init__(self, init_requests=None, policy="random"): + self.lock = threading.Lock() + self.requests = init_requests or [] + self.policy = policy + + def append(self, item): + with self.lock: + self.requests.append(item) + + def pop(self): + with self.lock: + if not self.requests: + return None + if self.policy == "random": + index = random.randrange(len(self.requests)) + return self.requests.pop(index) + elif self.policy == "fifo": + return self.requests.pop(0) + else: + # todo, varying thinking time of clients + raise ValueError(f"{self.policy} not implemented") + + +class WorkloadGenerator: + def __init__(self, args): + # Construct the base URL for requests + self.url = f"http://{args.host}:{args.port}/generate" + + self.tokenizer = get_tokenizer(args.model_path) + self.distribution = args.distribution + self.request_rate = args.request_rate + self.start_time = None + self.finished_time = None + + self.sent_requests = 0 + self.completed_requests = 0 + + self.candidate_inputs = sample_random_requests( + input_len=args.request_length, + output_len=args.output_length, + num_prompts=args.num_clients * args.num_rounds, + range_ratio=1.0, + tokenizer=self.tokenizer, + dataset_path=args.dataset_path, + ) + self.candidate_inputs = [i[0] for i in self.candidate_inputs] + + init_requests = [ + (i, gen_payload(self.candidate_inputs[i], args.output_length)) + for i in range(args.num_clients) + ] + self.client_records = { + i: {"round": 0, "history": init_requests[i][1]["text"]} + for i in range(args.num_clients) + } + self.ready_queue = ReadyQueue(init_requests=init_requests) + self.candidate_inputs = self.candidate_inputs[args.num_clients :] + + self.response_queue = queue.Queue() + self.pbar = tqdm(total=args.num_clients * args.num_rounds) + self.performance_metrics = {"ttft": [], "latency": []} + + async def handle_request(self, item): + try: + client_id, payload = item + response = await async_request_sglang_generate(payload, self.url, self.pbar) + if self.pbar.n == self.pbar.total: + self.finished_time = time.time() + self.response_queue.put((client_id, response)) + except Exception as e: + print(f"Request failed: {e}") + + def request_sender(self): + async def request_loop(): + while True: + if self.sent_requests - self.completed_requests < args.max_parallel: + new_request = self.ready_queue.pop() + if new_request: + asyncio.create_task(self.handle_request(new_request)) + self.sent_requests += 1 + else: + await asyncio.sleep(0.05) + continue + + if self.pbar.n == self.pbar.total: + break + + # Calculate Poisson-distributed wait time + if self.distribution == "poisson": + sleep_time = random.expovariate(self.request_rate) + elif self.distribution == "uniform": + avg_interval = ( + 1.0 / self.request_rate if self.request_rate > 0 else 1.0 + ) + sleep_time = random.uniform(0, 2 * avg_interval) + else: + raise ValueError("Invalid distribution type") + await asyncio.sleep(sleep_time) # Wait before sending the next request + + # Create and run the event loop for asynchronous requests + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + loop.run_until_complete(request_loop()) + loop.close() + + def response_handler(self): + while True: + try: + client_id, response = self.response_queue.get( + timeout=10 + ) # Block until response is available + if not response.success: + raise ValueError(f"Request failed with error: {response.error}") + self.client_records[client_id]["history"] += response.generated_text + self.client_records[client_id]["round"] += 1 + self.performance_metrics["ttft"].append(response.ttft) + self.performance_metrics["latency"].append(response.latency) + self.completed_requests += 1 + + if self.client_records[client_id]["round"] < args.num_rounds: + self.client_records[client_id][ + "history" + ] += self.candidate_inputs.pop() + self.ready_queue.append( + ( + client_id, + gen_payload( + self.client_records[client_id]["history"], + args.output_length, + ), + ) + ) + except queue.Empty: + if self.pbar.n == self.pbar.total: + break + + def run(self): + request_thread = threading.Thread(target=self.request_sender, daemon=True) + response_thread = threading.Thread(target=self.response_handler, daemon=True) + + self.start_time = time.time() + request_thread.start() + response_thread.start() + + request_thread.join() + response_thread.join() + self.pbar.close() + + performance_data = { + "summary": { + "total_requests": len(self.performance_metrics["ttft"]), + "request_rate": self.request_rate, + "average_ttft": sum(self.performance_metrics["ttft"]) + / len(self.performance_metrics["ttft"]), + "p90_ttft": sorted(self.performance_metrics["ttft"])[ + int(0.9 * len(self.performance_metrics["ttft"])) + ], + "median_ttft": sorted(self.performance_metrics["ttft"])[ + len(self.performance_metrics["ttft"]) // 2 + ], + "average_latency": sum(self.performance_metrics["latency"]) + / len(self.performance_metrics["latency"]), + "p90_latency": sorted(self.performance_metrics["latency"])[ + int(0.9 * len(self.performance_metrics["latency"])) + ], + "median_latency": sorted(self.performance_metrics["latency"])[ + len(self.performance_metrics["latency"]) // 2 + ], + "throughput": self.pbar.total / (self.finished_time - self.start_time), + }, + } + print("All requests completed") + print("Performance metrics summary:") + print( + f" Total requests: {performance_data['summary']['total_requests']} at {performance_data['summary']['request_rate']} requests per second" + ) + print(f" Average TTFT: {performance_data['summary']['average_ttft']:.2f}") + print(f" P90 TTFT: {performance_data['summary']['p90_ttft']:.2f}") + print(f" Median TTFT: {performance_data['summary']['median_ttft']:.2f}") + print( + f" Average latency: {performance_data['summary']['average_latency']:.2f}" + ) + print(f" P90 latency: {performance_data['summary']['p90_latency']:.2f}") + print(f" Median latency: {performance_data['summary']['median_latency']:.2f}") + print( + f" Throughput: {performance_data['summary']['throughput']:.2f} requests per second" + ) + log_to_jsonl_file(performance_data, args.log_file) + + +if __name__ == "__main__": + args = parse_args() + flush_cache_url = f"http://{args.host}:{args.port}/flush_cache" + + for request_rate in [16, 14, 12, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1]: + args.request_rate = request_rate + requests.post(flush_cache_url) + time.sleep(1) + WorkloadGenerator(args).run() diff --git a/sglang/benchmark/hicache/bench_serving.py b/sglang/benchmark/hicache/bench_serving.py new file mode 100644 index 0000000..90f4c6c --- /dev/null +++ b/sglang/benchmark/hicache/bench_serving.py @@ -0,0 +1,1029 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/backend_request_func.py +# Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/benchmark_serving.py + +""" +Benchmark online serving with dynamic requests. + +Usage: +python3 -m sglang.bench_serving --backend sglang --num-prompt 10 + +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 3000 --random-input 1024 --random-output 1024 --random-range-ratio 0.5 +python3 -m sglang.bench_serving --backend sglang --dataset-name random --request-rate-range 1,2,4,8,16,32 --random-input 4096 --random-output 1024 --random-range-ratio 0.125 --multi +""" + +import argparse +import asyncio +import json +import os +import random +import sys +import time +import traceback +import warnings +from argparse import ArgumentParser +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple + +import aiohttp +import numpy as np +import requests +from data_processing import MsgContent, SampleOutput, get_dataset +from tqdm.asyncio import tqdm +from transformers import PreTrainedTokenizerBase + +from sglang.bench_serving import get_tokenizer, remove_prefix, set_ulimit + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=20 * 60 * 60) + +global args + + +@dataclass +class RequestFuncInput: + prompts: List[Tuple[MsgContent, int, int]] + api_url: str + model: str + lora_name: str + extra_request_body: Dict[str, Any] + + # For multiturn chat, store the context + prev_messages: List = field(default_factory=list) + finished_prompts: int = 0 + + +@dataclass +class RequestFuncOutput: + generated_text: List[str] = field(default_factory=list) + prompt_len: List[int] = field(default_factory=list) + output_len: List[int] = field(default_factory=list) + latency: List[float] = field(default_factory=list) + ttft: List[float] = field(default_factory=list) + itl: List[float] = field(default_factory=list) # List of inter-token latencies + + success: bool = False + error: str = "" + + +# set ignore_eos True by default +async def async_request_openai_completions( + request_func_input: RequestFuncInput, + queue: asyncio.Queue, + tokenizer: PreTrainedTokenizerBase, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith( + "completions" + ), "OpenAI Completions API URL must end with 'completions'." + + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + payload = { + "model": request_func_input.model, + "temperature": 0.0, + "best_of": 1, + "stream": not args.disable_stream, + "stream_options": {"include_usage": True}, + "ignore_eos": not args.disable_ignore_eos, + **request_func_input.extra_request_body, + } + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", + } + + output = RequestFuncOutput() + + prompt_idx = request_func_input.finished_prompts + messages = request_func_input.prev_messages + prompt, input_len, max_tokens = request_func_input.prompts[prompt_idx] + prompt_len = sum( + prompt[1] + prompt[2] # input_len + output_len + for prompt in request_func_input.prompts[:prompt_idx] + ) + prompt_len += input_len + + # Messages + messages.append( + { + "role": "user", + "content": prompt, + } + ) + payload["messages"] = messages + payload["max_tokens"] = max_tokens + + # output.prompt_len = request_func_input.prompt_len + # print(payload) + + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: + if response.status == 200: + actual_prompt_len = prompt_len - 1 + actual_output_len = 0 + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + timestamp = time.perf_counter() + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if data["usage"] is not None and len(data["usage"]) > 0: + actual_prompt_len = data["usage"]["prompt_tokens"] + actual_output_len = data["usage"]["completion_tokens"] + continue + delta = data["choices"][0]["delta"] + + if delta.get("content", None): + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft.append(ttft) + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + generated_text += delta["content"] + most_recent_timestamp = timestamp + + output.prompt_len.append(actual_prompt_len) # truncate + output.output_len.append(actual_output_len) + output.generated_text.append(generated_text) + output.success = True + output.latency.append(latency) + + # Prepare for the new request + request_func_input.prompts[prompt_idx] = ( + prompt, + input_len, + actual_output_len, # changes from max_tokens to output_len + ) + prompt_idx += 1 + messages.append( + { + "role": "assistant", + "content": generated_text, + } + ) + + # Move the new request to the end of the queue + if prompt_idx < len(request_func_input.prompts): + request_func_input.finished_prompts = prompt_idx + request_func_input.prev_messages = messages + await queue.put(request_func_input) + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_profile(api_url: str) -> RequestFuncOutput: + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + output = RequestFuncOutput() + try: + async with session.post(url=api_url) as response: + if response.status == 200: + output.success = True + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + return output + + +ASYNC_REQUEST_FUNCS = { + "sglang": async_request_openai_completions, + "vllm": async_request_openai_completions, + "lmdeploy": async_request_openai_completions, +} + + +@dataclass +class BenchmarkMetrics: + completed: int + total_input: int + total_output: int + total_output_retokenized: int + request_throughput: float + input_throughput: float + output_throughput: float + output_throughput_retokenized: float + total_throughput: float + total_throughput_retokenized: float + mean_ttft_ms: float + median_ttft_ms: float + std_ttft_ms: float + p90_ttft_ms: float + p99_ttft_ms: float + mean_tpot_ms: float + median_tpot_ms: float + std_tpot_ms: float + p90_tpot_ms: float + p99_tpot_ms: float + mean_itl_ms: float + median_itl_ms: float + std_itl_ms: float + p90_itl_ms: float + p99_itl_ms: float + mean_e2e_latency_ms: float + median_e2e_latency_ms: float + std_e2e_latency_ms: float + p99_e2e_latency_ms: float + concurrency: float + + +async def get_requests( + input_requests_queue: asyncio.Queue, + request_rate: float, + num_actual_requests: int, +) -> AsyncGenerator[RequestFuncInput, None]: + for _ in range(num_actual_requests): + try: + request = await asyncio.wait_for( + input_requests_queue.get(), timeout=300 + ) # Wait for 5 minites then abort + except Exception as e: + print(f"exception: {e}") + break + + yield request + + if request_rate == float("inf"): + continue + + interval = np.random.exponential(1.0 / request_rate) + await asyncio.sleep(interval) + + +def calculate_metrics( + outputs: List[RequestFuncOutput], + dur_s: float, + tokenizer: PreTrainedTokenizerBase, + backend: str, +) -> Tuple[BenchmarkMetrics, List[int]]: + output_lens: List[int] = [] + retokenized_output_lens: List[int] = [] + total_input = 0 + completed = 0 + itls: List[float] = [] + tpots: List[float] = [] + ttfts: List[float] = [] + e2e_latencies: List[float] = [] + output_success = 0 + for i in range(len(outputs)): + if outputs[i].success: + output_success += 1 + assert len(outputs[i].generated_text) == len(outputs[i].latency) + assert len(outputs[i].generated_text) == len(outputs[i].ttft) + for j in range(len(outputs[i].generated_text)): + output_len = outputs[i].output_len[j] + output_lens.append(output_len) + retokenized_output_len = len( + tokenizer.encode( + outputs[i].generated_text[j], add_special_tokens=False + ) + ) + retokenized_output_lens.append(retokenized_output_len) + total_input += outputs[i].prompt_len[j] + if output_len > 1: + tpots.append( + (outputs[i].latency[j] - outputs[i].ttft[j]) / (output_len - 1) + ) + + completed += 1 + itls += outputs[i].itl + ttfts += outputs[i].ttft + e2e_latencies += outputs[i].latency + + else: + output_lens.append(0) + retokenized_output_lens.append(0) + + if completed == 0: + warnings.warn( + "All requests failed. This is likely due to a misconfiguration " + "on the benchmark arguments.", + stacklevel=2, + ) + metrics = BenchmarkMetrics( + completed=completed, + total_input=total_input, + total_output=sum(output_lens), + total_output_retokenized=sum(retokenized_output_lens), + request_throughput=completed / dur_s, + input_throughput=total_input / dur_s, + output_throughput=sum(output_lens) / dur_s, + output_throughput_retokenized=sum(retokenized_output_lens) / dur_s, + total_throughput=(total_input + sum(output_lens)) / dur_s, + total_throughput_retokenized=(total_input + sum(retokenized_output_lens)) + / dur_s, + mean_ttft_ms=np.mean(ttfts or 0) + * 1000, # ttfts is empty if streaming is not supported by backend + median_ttft_ms=np.median(ttfts or 0) * 1000, + std_ttft_ms=np.std(ttfts or 0) * 1000, + p90_ttft_ms=np.percentile(ttfts or 0, 90) * 1000, + p99_ttft_ms=np.percentile(ttfts or 0, 99) * 1000, + mean_tpot_ms=np.mean(tpots or 0) * 1000, + median_tpot_ms=np.median(tpots or 0) * 1000, + std_tpot_ms=np.std(tpots or 0) * 1000, + p90_tpot_ms=np.percentile(tpots or 0, 90) * 1000, + p99_tpot_ms=np.percentile(tpots or 0, 99) * 1000, + mean_itl_ms=np.mean(itls or 0) * 1000, + median_itl_ms=np.median(itls or 0) * 1000, + std_itl_ms=np.std(itls or 0) * 1000, + p90_itl_ms=np.percentile(itls or 0, 90) * 1000, + p99_itl_ms=np.percentile(itls or 0, 99) * 1000, + mean_e2e_latency_ms=np.mean(e2e_latencies) * 1000, + median_e2e_latency_ms=np.median(e2e_latencies) * 1000, + std_e2e_latency_ms=np.std(e2e_latencies) * 1000, + p99_e2e_latency_ms=np.percentile(e2e_latencies, 99) * 1000, + concurrency=np.sum(e2e_latencies) / dur_s, + ) + + return metrics, output_lens + + +async def benchmark( + backend: str, + api_url: str, + base_url: str, + model_id: str, + tokenizer: PreTrainedTokenizerBase, + input_requests: SampleOutput, + request_rate: float, + max_concurrency: Optional[int], + disable_tqdm: bool, + lora_name: str, + extra_request_body: Dict[str, Any], + profile: bool, + enable_shared_prefix: bool, +): + if backend in ASYNC_REQUEST_FUNCS: + request_func = ASYNC_REQUEST_FUNCS[backend] + else: + raise ValueError(f"Unknown backend: {backend}") + + # Limit concurrency + # From https://github.com/vllm-project/vllm/pull/9390 + semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None + + async def limited_request_func(request_func_input, queue, tokenizer, pbar): + if semaphore is None: + return await request_func( + request_func_input=request_func_input, + queue=queue, + tokenizer=tokenizer, + pbar=pbar, + ) + async with semaphore: + return await request_func( + request_func_input=request_func_input, + queue=queue, + tokenizer=tokenizer, + pbar=pbar, + ) + + num_actual_requests = sum(len(r) for r in input_requests) + print(f"Num of shared prefixes or conversations: {len(input_requests)}") + print(f"Num of total requests: {num_actual_requests}") + + # flatten the requests for shared prefix + if enable_shared_prefix: + input_requests = [[r] for requests in input_requests for r in requests] + inputs_requests_queue = asyncio.Queue(maxsize=len(input_requests)) + print("Starting initial single prompt test run...") + # NOTE: Just use the first request of the first conversation for warmup + test_input = RequestFuncInput( + model=model_id, + prompts=input_requests[0][:1], + api_url=api_url, + lora_name=lora_name, + extra_request_body=extra_request_body, + ) + test_output = await request_func( + request_func_input=test_input, queue=inputs_requests_queue, tokenizer=tokenizer + ) + if not test_output.success: + raise ValueError( + "Initial test run failed - Please make sure benchmark arguments " + f"are correctly specified. Error: {test_output.error}" + ) + else: + print("Initial test run completed. Starting main benchmark run...") + + # Check the states + assert inputs_requests_queue.empty() + + # Flush cache + if "sglang" in backend: + requests.post(base_url + "/flush_cache") + + time.sleep(1.0) + + # Start profiler + if profile: + print("Starting profiler...") + profile_output = await async_request_profile( + api_url=base_url + "/start_profile" + ) + if profile_output.success: + print("Profiler started") + + for request in input_requests: + request_func_input = RequestFuncInput( + model=model_id, + prompts=request, + api_url=api_url, + lora_name=lora_name, + extra_request_body=extra_request_body, + ) + inputs_requests_queue.put_nowait(request_func_input) + if ( + not args.enable_multiturn + and not args.enable_shared_prefix + and not args.dataset_name == "generated-shared-prefix" + ): + assert len(input_requests) == num_actual_requests + + pbar = None if disable_tqdm else tqdm(total=num_actual_requests) + + benchmark_start_time = time.perf_counter() + tasks: List[asyncio.Task] = [] + async for request in get_requests( + inputs_requests_queue, request_rate, num_actual_requests + ): + tasks.append( + asyncio.create_task( + limited_request_func( + request_func_input=request, + queue=inputs_requests_queue, + tokenizer=tokenizer, + pbar=pbar, + ) + ) + ) + outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) + + # Stop profiler + if profile: + print("Stopping profiler...") + profile_output = await async_request_profile(api_url=base_url + "/stop_profile") + if profile_output.success: + print("Profiler stopped") + + if pbar is not None: + pbar.close() + + # Compute metrics and print results + benchmark_duration = time.perf_counter() - benchmark_start_time + metrics, output_lens = calculate_metrics( + outputs=outputs, + dur_s=benchmark_duration, + tokenizer=tokenizer, + backend=backend, + ) + + print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) + print("{:<40} {:<10}".format("Backend:", backend)) + print("{:<40} {:<10}".format("Traffic request rate:", request_rate)) + print( + "{:<40} {:<10}".format( + "Max reqeuest concurrency:", + max_concurrency if max_concurrency else "not set", + ) + ) + print("{:<40} {:<10}".format("Successful requests:", metrics.completed)) + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration)) + print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) + print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output)) + print( + "{:<40} {:<10}".format( + "Total generated tokens (retokenized):", metrics.total_output_retokenized + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Request throughput (req/s):", metrics.request_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Input token throughput (tok/s):", metrics.input_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Output token throughput (tok/s):", metrics.output_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Total token throughput (tok/s):", metrics.total_throughput + ) + ) + print("{:<40} {:<10.2f}".format("Concurrency:", metrics.concurrency)) + print("{s:{c}^{n}}".format(s="End-to-End Latency", n=50, c="-")) + print( + "{:<40} {:<10.2f}".format("Mean E2E Latency (ms):", metrics.mean_e2e_latency_ms) + ) + print( + "{:<40} {:<10.2f}".format( + "Median E2E Latency (ms):", metrics.median_e2e_latency_ms + ) + ) + print("{s:{c}^{n}}".format(s="Time to First Token", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean TTFT (ms):", metrics.mean_ttft_ms)) + print("{:<40} {:<10.2f}".format("Median TTFT (ms):", metrics.median_ttft_ms)) + print("{:<40} {:<10.2f}".format("P90 TTFT (ms):", metrics.p90_ttft_ms)) + print("{:<40} {:<10.2f}".format("P99 TTFT (ms):", metrics.p99_ttft_ms)) + print( + "{s:{c}^{n}}".format(s="Time per Output Token (excl. 1st token)", n=50, c="-") + ) + print("{:<40} {:<10.2f}".format("Mean TPOT (ms):", metrics.mean_tpot_ms)) + print("{:<40} {:<10.2f}".format("Median TPOT (ms):", metrics.median_tpot_ms)) + print("{:<40} {:<10.2f}".format("P90 TPOT (ms):", metrics.p90_tpot_ms)) + print("{:<40} {:<10.2f}".format("P99 TPOT (ms):", metrics.p99_tpot_ms)) + print("{s:{c}^{n}}".format(s="Inter-token Latency", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean ITL (ms):", metrics.mean_itl_ms)) + print("{:<40} {:<10.2f}".format("Median ITL (ms):", metrics.median_itl_ms)) + print("{:<40} {:<10.2f}".format("P90 ITL (ms):", metrics.p90_itl_ms)) + print("{:<40} {:<10.2f}".format("P99 ITL (ms):", metrics.p99_itl_ms)) + print("=" * 50) + + if ( + metrics.median_ttft_ms is not None + and metrics.mean_itl_ms is not None + and metrics.output_throughput is not None + ): + result = { + # Arguments + "backend": args.backend, + "dataset_name": args.dataset_name, + "request_rate": request_rate, + "max_concurrency": max_concurrency, + "fixed_output_len": args.fixed_output_len, + "random_input_len": args.random_input_len, + "random_output_len": args.random_output_len, + "random_range_ratio": args.random_range_ratio, + # Results + "duration": benchmark_duration, + "completed": metrics.completed, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "total_output_tokens_retokenized": metrics.total_output_retokenized, + "request_throughput": metrics.request_throughput, + "input_throughput": metrics.input_throughput, + "output_throughput": metrics.output_throughput, + "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, + "median_e2e_latency_ms": metrics.median_e2e_latency_ms, + "std_e2e_latency_ms": metrics.std_e2e_latency_ms, + "p99_e2e_latency_ms": metrics.p99_e2e_latency_ms, + "mean_ttft_ms": metrics.mean_ttft_ms, + "median_ttft_ms": metrics.median_ttft_ms, + "std_ttft_ms": metrics.std_ttft_ms, + "p99_ttft_ms": metrics.p99_ttft_ms, + "mean_tpot_ms": metrics.mean_tpot_ms, + "median_tpot_ms": metrics.median_tpot_ms, + "std_tpot_ms": metrics.std_tpot_ms, + "p99_tpot_ms": metrics.p99_tpot_ms, + "mean_itl_ms": metrics.mean_itl_ms, + "median_itl_ms": metrics.median_itl_ms, + "std_itl_ms": metrics.std_itl_ms, + "p99_itl_ms": metrics.p99_itl_ms, + "concurrency": metrics.concurrency, + "input_throughput": metrics.input_throughput, + "output_throughput": metrics.output_throughput, + "fixed_output_len": args.fixed_output_len, + "random_input_len": args.random_input_len, + "random_output_len": args.random_output_len, + "random_range_ratio": args.random_range_ratio, + "duration": benchmark_duration, + "completed": metrics.completed, + } + else: + print(f"Error running benchmark for request rate: {request_rate}") + print("-" * 30) + + # Determine output file name + if args.output_file: + output_file_name = args.output_file + else: + now = datetime.now().strftime("%m%d") + if args.dataset_name == "random": + output_file_name = f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_{args.random_output_len}.jsonl" + else: + output_file_name = ( + f"{args.backend}_{now}_{args.num_prompts}_{args.dataset_name}.jsonl" + ) + + # Append results to a JSONL file + with open(output_file_name, "a") as file: + file.write(json.dumps(result) + "\n") + + result = { + "duration": benchmark_duration, + "completed": metrics.completed, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "total_output_tokens_retokenized": metrics.total_output_retokenized, + "request_throughput": metrics.request_throughput, + "input_throughput": metrics.input_throughput, + "output_throughput": metrics.output_throughput, + "mean_ttft_ms": metrics.mean_ttft_ms, + "median_ttft_ms": metrics.median_ttft_ms, + "std_ttft_ms": metrics.std_ttft_ms, + "p90_ttft_ms": metrics.p90_ttft_ms, + "p99_ttft_ms": metrics.p99_ttft_ms, + "mean_tpot_ms": metrics.mean_tpot_ms, + "median_tpot_ms": metrics.median_tpot_ms, + "std_tpot_ms": metrics.std_tpot_ms, + "p90_tpot_ms": metrics.p90_tpot_ms, + "p99_tpot_ms": metrics.p99_tpot_ms, + "mean_itl_ms": metrics.mean_itl_ms, + "median_itl_ms": metrics.median_itl_ms, + "std_itl_ms": metrics.std_itl_ms, + "p90_itl_ms": metrics.p90_itl_ms, + "p99_itl_ms": metrics.p99_itl_ms, + "input_lens": [output.prompt_len for output in outputs], + "output_lens": output_lens, + "ttfts": [output.ttft for output in outputs], + "itls": [output.itl for output in outputs], + "generated_texts": [output.generated_text for output in outputs], + "errors": [output.error for output in outputs], + "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, + "median_e2e_latency_ms": metrics.median_e2e_latency_ms, + } + return result + + +def run_benchmark(args_: argparse.Namespace): + global args + args = args_ + + # Set default value for max_concurrency if not present + if not hasattr(args, "max_concurrency"): + args.max_concurrency = None + + # Set global environments + set_ulimit() + random.seed(args.seed) + np.random.seed(args.seed) + + extra_request_body = {} + if args.extra_request_body: + extra_request_body = json.loads(args.extra_request_body) + + # Set url + if args.port is None: + args.port = { + "sglang": 30000, + "lmdeploy": 23333, + "vllm": 8000, + }.get(args.backend, 30000) + + model_url = ( + f"{args.base_url}/v1/models" + if args.base_url + else f"http://{args.host}:{args.port}/v1/models" + ) + + if args.backend in ["sglang", "vllm", "lmdeploy"]: + api_url = ( + f"{args.base_url}/v1/chat/completions" + if args.base_url + else f"http://{args.host}:{args.port}/v1/chat/completions" + ) + base_url = ( + f"http://{args.host}:{args.port}" if args.base_url is None else args.base_url + ) + + # Get model name + if args.model is None: + if args.backend == "truss": + print( + "Please provide a model with `--model` when using truss backend. e.g. --model meta-llama/Llama-3.1-8B-Instruct" + ) + sys.exit(1) + try: + response = requests.get(model_url) + model_list = response.json().get("data", []) + args.model = model_list[0]["id"] if model_list else None + except Exception as e: + print(f"Failed to fetch model from {model_url}. Error: {e}") + print( + "Please specify the correct host and port using `--host` and `--port`." + ) + sys.exit(1) + + if args.model is None: + print("No model specified or found. Please provide a model using `--model`.") + sys.exit(1) + + # Dataset compatibility check + if args.enable_multiturn: + # TODO: Support multiturn for random + if args.dataset_name not in ["sharegpt", "ultrachat", "loogle", "nextqa"]: + print( + "Multiturn conversation is only supported for sharegpt, ultrachat, loogle, and nextqa datasets." + ) + sys.exit(1) + + if args.enable_shared_prefix: + if args.dataset_name not in ["loogle", "nextqa"]: + print("Shared prefix is only supported for loogle and nextqa datasets.") + sys.exit(1) + + print(f"{args}\n") + + # Read dataset + backend = args.backend + model_id = args.model + tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model + + tokenizer = get_tokenizer(tokenizer_id) + + input_requests = get_dataset(args, tokenizer) + + return asyncio.run( + benchmark( + backend=backend, + api_url=api_url, + base_url=base_url, + model_id=model_id, + tokenizer=tokenizer, + input_requests=input_requests, + request_rate=args.request_rate, + max_concurrency=args.max_concurrency, + disable_tqdm=args.disable_tqdm, + lora_name=args.lora_name, + extra_request_body=extra_request_body, + profile=args.profile, + enable_shared_prefix=args.enable_shared_prefix, + ) + ) + + +if __name__ == "__main__": + parser = ArgumentParser(description="Benchmark the online serving throughput.") + parser.add_argument( + "--backend", + type=str, + choices=list(ASYNC_REQUEST_FUNCS.keys()), + default="sglang", + help="Must specify a backend, depending on the LLM Inference Engine.", + ) + parser.add_argument( + "--base-url", + type=str, + default=None, + help="Server or API base url if not using http host and port.", + ) + parser.add_argument( + "--host", type=str, default="0.0.0.0", help="Default host is 0.0.0.0." + ) + parser.add_argument( + "--port", + type=int, + help="If not set, the default port is configured according to its default value for different LLM Inference Engines.", + ) + parser.add_argument( + "--dataset-name", + type=str, + default="sharegpt", + choices=[ + "sharegpt", + "random", + "generated-shared-prefix", + "ultrachat", + "loogle", + "nextqa", + ], + help="Name of the dataset to benchmark on.", + ) + parser.add_argument( + "--dataset-path", type=str, default="", help="Path to the dataset." + ) + parser.add_argument( + "--model", + type=str, + help="Name or path of the model. If not set, the default model will request /v1/models for conf.", + ) + parser.add_argument( + "--tokenizer", + type=str, + help="Name or path of the tokenizer. If not set, using the model conf.", + ) + parser.add_argument( + "--chat-template", + type=str, + help="The buliltin chat template name or the path of the chat template file. This is only used for OpenAI-compatible API server.", + ) + parser.add_argument( + "--num-prompts", + type=int, + default=1000, + help="Number of prompts to process. Default is 1000.", + ) + parser.add_argument( + "--fixed-output-len", + type=int, + default=None, + help="Output length for each request. Overrides the output length from the dataset.", + ) + parser.add_argument( + "--sharegpt-context-len", + type=int, + default=None, + help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.", + ) + parser.add_argument( + "--random-input-len", + type=int, + default=1024, + help="Number of input tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-output-len", + default=1024, + type=int, + help="Number of output tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-range-ratio", + type=float, + default=0.0, + help="Range of sampled ratio of input/output length, " + "used only for random dataset.", + ) + parser.add_argument( + "--request-rate", + type=float, + default=float("inf"), + help="Number of requests per second. If this is inf, then all the requests are sent at time 0. " + "Otherwise, we use Poisson process to synthesize the request arrival times. Default is inf.", + ) + parser.add_argument( + "--max-concurrency", + type=int, + default=None, + help="Maximum number of concurrent requests. This can be used " + "to help simulate an environment where a higher level component " + "is enforcing a maximum number of concurrent requests. While the " + "--request-rate argument controls the rate at which requests are " + "initiated, this argument will control how many are actually allowed " + "to execute at a time. This means that when used in combination, the " + "actual request rate may be lower than specified with --request-rate, " + "if the server is not processing requests fast enough to keep up.", + ) + parser.add_argument( + "--multi", + action="store_true", + help="Use request rate range rather than single value.", + ) + parser.add_argument( + "--request-rate-range", + type=str, + default="2,34,2", + help="Range of request rates in the format start,stop,step. Default is 2,34,2. It also supports a list of request rates, requiring the parameters to not equal three.", + ) + parser.add_argument("--output-file", type=str, help="Output JSONL file name.") + parser.add_argument( + "--enable-multiturn", + action="store_true", + help="Enable multiturn chat for online serving benchmarking. " + "This option is effective on the following datasets: " + "sharegpt, ultrachat, loogle, nextqa", + ) + parser.add_argument( + "--enable-shared-prefix", + action="store_true", + help="Enable shared prefix for online serving benchmarking. " + "This option is effective on the following datasets: " + "loogle, nextqa", + ) + + parser.add_argument( + "--disable-shuffle", + action="store_true", + help="Disable shuffling datasets. This is useful to generate stable output " + "in benchmarking", + ) + parser.add_argument( + "--disable-tqdm", + action="store_true", + help="Specify to disable tqdm progress bar.", + ) + parser.add_argument( + "--disable-stream", + action="store_true", + help="Disable streaming mode.", + ) + parser.add_argument( + "--return-logprob", + action="store_true", + help="Return logprob.", + ) + parser.add_argument("--seed", type=int, default=1, help="The random seed.") + parser.add_argument( + "--disable-ignore-eos", + action="store_true", + help="Disable ignoring EOS.", + ) + parser.add_argument( + "--extra-request-body", + metavar='{"key1": "value1", "key2": "value2"}', + type=str, + help="Append given JSON object to the request payload. You can use this to specify" + "additional generate params like sampling params.", + ) + parser.add_argument( + "--apply-chat-template", + action="store_true", + help="Apply chat template", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Use Torch Profiler. The endpoint must be launched with " + "SGLANG_TORCH_PROFILER_DIR to enable profiler.", + ) + parser.add_argument( + "--lora-name", + type=str, + default=None, + help="The name of LoRA adapter", + ) + + group = parser.add_argument_group("generated-shared-prefix dataset arguments") + group.add_argument( + "--gsp-num-groups", + type=int, + default=64, + help="Number of system prompt groups for generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-prompts-per-group", + type=int, + default=16, + help="Number of prompts per system prompt group for generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-system-prompt-len", + type=int, + default=2048, + help="Target length in tokens for system prompts in generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-question-len", + type=int, + default=128, + help="Target length in tokens for questions in generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-output-len", + type=int, + default=256, + help="Target length in tokens for outputs in generated-shared-prefix dataset", + ) + # videos specific + parser.add_argument( + "--max-frames", + type=int, + default=sys.maxsize, + help="The maximum number of frames to extract from each video. " + "This option is specific to the nextqa dataset (video benchmark). ", + ) + args = parser.parse_args() + + if args.enable_multiturn and args.enable_shared_prefix: + parser.error( + "--enable-multiturn and --enable-shared-prefix cannot be set at the same time." + ) + + run_benchmark(args) diff --git a/sglang/benchmark/hicache/data_processing.py b/sglang/benchmark/hicache/data_processing.py new file mode 100644 index 0000000..e2951a3 --- /dev/null +++ b/sglang/benchmark/hicache/data_processing.py @@ -0,0 +1,589 @@ +import json +import os +import pickle +import random +from pathlib import Path +from typing import List, Optional, Tuple, Union + +import numpy as np +from nextqa import NExTQALoader + +# from nextqa.video import , VideoPrompt +from tqdm.asyncio import tqdm +from transformers import PreTrainedTokenizerBase + +SHAREGPT_URL = "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json" + +from sglang.bench_serving import ( + download_and_cache_file, + gen_prompt, + get_gen_prefix_cache_path, +) +from sglang.lang.chat_template import get_chat_template, get_chat_template_by_model_path +from sglang.srt.openai_api.protocol import ChatCompletionMessageContentPart +from sglang.utils import encode_video_base64 + +# type of content fields, can be only prompts or with images/videos +MsgContent = Union[str, List[ChatCompletionMessageContentPart]] + +# A list of all the conversations. Each conversation is a list of +# tuples. If multiturn is not enabled, the length of list is 1, +# containing only the first Q&A pair. +# For the shared prefix workload (synthetic, loogle, nextqa), it +# is a list of conversations sharing the same prefix (synthetic, +# doc, video) +SampleOutput = List[List[Tuple[MsgContent, int, int]]] + + +def common_filter_chat( + num_requests: int, + new_dataset: List, + tokenizer: PreTrainedTokenizerBase, + min_prompt_len: Optional[int], + min_output_len: Optional[int], + max_prompt_len: Optional[int], + max_output_len: Optional[int], + fixed_output_len: Optional[int], +) -> SampleOutput: + # Filter out sequences that are too long or too short + filtered_dataset: SampleOutput = [] + l = 0 + input_tokens = 0 + output_tokens = 0 + while l < num_requests: + for i in range(len(new_dataset)): + if l == num_requests: + break + processed = [] + for j in new_dataset[i]: + # Tokenize the prompts and completions. + prompt = j[0] + prompt_token_ids = tokenizer.encode(prompt) + prompt_len = len(prompt_token_ids) + + completion = j[1] + completion_token_ids = tokenizer.encode(completion) + output_len = ( + len(completion_token_ids) + if fixed_output_len is None + else fixed_output_len + ) + if ( + min_prompt_len is not None + and prompt_len < min_prompt_len + or min_output_len is not None + and output_len < min_output_len + or max_prompt_len is not None + and prompt_len > max_prompt_len + or max_output_len is not None + and output_len > max_output_len + ): + # Prune too short sequences. + continue + input_tokens += prompt_len + output_tokens += output_len + processed.append((prompt, prompt_len, output_len)) + filtered_dataset.append(processed) + l += 1 + + print(f"#Input tokens: {input_tokens}") + print(f"#Output tokens: {output_tokens}") + return filtered_dataset + + +def sample_sharegpt_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + disable_shuffle: bool = False, + enable_multiturn: bool = True, + fixed_output_len: Optional[int] = None, +) -> SampleOutput: + if fixed_output_len is not None and fixed_output_len < 4: + raise ValueError("output_len too small") + + # Download sharegpt if necessary + if not os.path.isfile(dataset_path): + dataset_path = download_and_cache_file(SHAREGPT_URL) + + # Load the dataset. + with open(dataset_path) as f: + dataset = json.load(f) + # Filter out the conversations with less than 2 turns. + dataset = [data for data in dataset if len(data["conversations"]) >= 2] + + # Keep one conversation in one list + new_dataset = [] + for data in dataset: + if len(data["conversations"]) % 2 != 0: + continue + if data["conversations"][0]["from"] != "human": + continue + chat = [] + total_len = 2 + if enable_multiturn: + total_len = len(data["conversations"]) + for i in range(0, total_len, 2): + # One user One Assistant + chat.append( + ( + data["conversations"][i]["value"], + data["conversations"][i + 1]["value"], + ) + ) + new_dataset.append(chat) + + if not disable_shuffle: + # Shuffle the dataset. + random.shuffle(new_dataset) + + # Filter out sequences that are too long or too short + filtered_dataset: SampleOutput = common_filter_chat( + num_requests, new_dataset, tokenizer, 4, 4, None, None, fixed_output_len + ) + return filtered_dataset + + +def sample_ultrachat_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + disable_shuffle: bool = False, + enable_multiturn: bool = True, + fixed_output_len: Optional[int] = None, +) -> SampleOutput: + if fixed_output_len is not None and fixed_output_len < 4: + raise ValueError("output_len too small") + + # Load the dataset + dataset = [] + with open(dataset_path) as f: + while True: + line = f.readline() + if not line: + break + dataset.append(json.loads(line)) + + # Filter out the conversations with less than 2 turns. + dataset = [data for data in dataset if len(data["data"]) >= 2] + + # Keep one conversation in one list + new_dataset = [] + for data in dataset: + if len(data["data"]) % 2 != 0: + continue + chat = [] + total_len = 2 + if enable_multiturn: + total_len = len(data["data"]) + for i in range(0, total_len, 2): + # One user One Assistant + chat.append((data["data"][i], data["data"][i + 1])) + new_dataset.append(chat) + + # Shuffle the dataset. + if not disable_shuffle: + random.shuffle(new_dataset) + + # Filter out sequences that are too long or too short + filtered_dataset: SampleOutput = common_filter_chat( + num_requests, new_dataset, tokenizer, 4, 4, None, None, fixed_output_len + ) + return filtered_dataset + + +def sample_loogle_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + disable_shuffle: bool = False, + enable_multiturn: bool = True, + enable_shared_prefix: bool = False, + fixed_output_len: Optional[int] = None, +) -> SampleOutput: + if fixed_output_len is not None and fixed_output_len < 4: + raise ValueError("output_len too small") + + # Load the dataset + dataset = [] + with open(dataset_path) as f: + while True: + line = f.readline() + if not line: + break + dataset.append(json.loads(line)) + + # Keep one conversation in one list + new_dataset = [] + # TODO: Add shared prefix support for loogle + # NOTE: Now we preprocess it only for chat + for data in dataset: + chat = [] + if ( + "qa_pairs" not in data + or data["qa_pairs"] == "none" + or len(data["qa_pairs"]) == 0 + ): + # If Q is none (for summarization), + # We add a question for summarization + # And keep the summary up to 1024 words + chat.append( + ( + "Input: " + + data["input"] + + " Question: " + + "Please summarize the input", + data["input"][:1024], + ) + ) + new_dataset.append(chat) + else: + qa_pairs = eval(data["qa_pairs"]) + for i, qa in enumerate(qa_pairs): + if i == 0 or enable_shared_prefix: + # Combine input with the first Q + chat.append( + ("Input: " + data["input"] + " Question: " + qa["Q"], qa["A"]) + ) + elif enable_multiturn: + chat.append((qa["Q"], qa["A"])) + + new_dataset.append(chat) + + # Shuffle the dataset. + if not disable_shuffle: + random.shuffle(new_dataset) + + # Filter out sequences that are too long or too short + filtered_dataset: SampleOutput = common_filter_chat( + num_requests, new_dataset, tokenizer, 4, None, None, None, fixed_output_len + ) + return filtered_dataset + + +def sample_nextqa_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + max_frames: int, # Specific for video + model_path: str, + disable_shuffle: bool = False, + enable_multiturn: bool = True, # No multiturn support for now + backend: str = "sglang-oai", + chat_template_name: Optional[str] = None, + fixed_output_len: Optional[int] = None, +) -> SampleOutput: + """ + Example of messages: + message = { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": base64_data}}, + {"type": "text", "text": video.prompt}, + ], + } + """ + + if fixed_output_len is None: + fixed_output_len = 4096 + + # TODO: Check for multiturn + dataset = NExTQALoader(video_dir=dataset_path, max_frames=max_frames) + new_dataset = [] + for v in dataset: + new_dataset.append(v) + + if not disable_shuffle: + random.shuffle(new_dataset) + + # TODO: prompt len can get from server side + filtered_dataset = [] + l = 0 + while l < num_requests: + for i in range(len(new_dataset)): + if l == num_requests: + break + + video = new_dataset[i] + + # text prompt + prompt = video.prompt + + # NOTE: Chat Template is a must for video benchmark because we have to + # add special image token for later expansion + if backend == "sglang" or backend == "sglang-native": + if "chat_template" in tokenizer.init_kwargs: + chat_template = get_chat_template(tokenizer.get_chat_template()) + elif chat_template_name is not None: + chat_template = get_chat_template(chat_template_name) + else: + chat_template = get_chat_template_by_model_path(model_path) + prompt = chat_template.image_token + prompt + + prompt_token_ids = tokenizer(prompt).input_ids + prompt_len = len(prompt_token_ids) + output_len = fixed_output_len # max output len, not real output len + + # video input + base64_data = encode_video_base64(video.path, video.num_frames) + + # NOTE: This will be replaced by the expanded length from the server + prompt_len += video.num_frames + + # add to content + content = [ + {"type": "image_url", "image_url": {"url": base64_data}}, + {"type": "text", "text": prompt}, + ] + + filtered_dataset.append([(content, prompt_len, output_len)]) + l += 1 + return filtered_dataset + + +def sample_random_requests( + input_len: int, + output_len: int, + num_prompts: int, + range_ratio: float, + tokenizer: PreTrainedTokenizerBase, + dataset_path: str, + disable_shuffle: bool = False, +) -> SampleOutput: + + input_lens = np.random.randint( + max(int(input_len * range_ratio), 1), + input_len + 1, + size=num_prompts, + ) + output_lens = np.random.randint( + int(output_len * range_ratio), + output_len + 1, + size=num_prompts, + ) + + if True: + # Sample token ids from ShareGPT and repeat/truncate them to satisfy the input_lens + + # Download sharegpt if necessary + if not os.path.isfile(dataset_path): + dataset_path = download_and_cache_file(SHAREGPT_URL) + + # Load the dataset. + with open(dataset_path) as f: + dataset = json.load(f) + # Filter out the conversations with less than 2 turns. + dataset = [data for data in dataset if len(data["conversations"]) >= 2] + # Only keep the first two turns of each conversation. + dataset = [ + (data["conversations"][0]["value"], data["conversations"][1]["value"]) + for data in dataset + ] + + if not disable_shuffle: + # Shuffle the dataset. + random.shuffle(dataset) + + # Filter out sequences that are too long or too short + input_requests: SampleOutput = [] + for data in dataset: + i = len(input_requests) + if i == num_prompts: + break + + # Tokenize the prompts and completions. + prompt = data[0] + prompt_token_ids = tokenizer.encode(prompt) + prompt_len = len(prompt_token_ids) + + # Skip empty prompt + if prompt_len == 0: + continue + + if prompt_len > input_lens[i]: + input_ids = prompt_token_ids[: input_lens[i]] + else: + ratio = (input_lens[i] + prompt_len - 1) // prompt_len + input_ids = (prompt_token_ids * ratio)[: input_lens[i]] + prompt = tokenizer.decode(input_ids) + input_requests.append([(prompt, int(input_lens[i]), int(output_lens[i]))]) + else: + # Sample token ids from random integers. This can cause some NaN issues. + offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts) + input_requests = [] + for i in range(num_prompts): + prompt = tokenizer.decode( + [ + (offsets[i] + i + j) % tokenizer.vocab_size + for j in range(input_lens[i]) + ] + ) + input_requests.append([(prompt, int(input_lens[i]), int(output_lens[i]))]) + + print(f"#Input tokens: {np.sum(input_lens)}") + print(f"#Output tokens: {np.sum(output_lens)}") + return input_requests + + +def gen_prompt(tokenizer, token_num): + """Generate a random prompt of specified token length using tokenizer vocabulary.""" + all_available_tokens = list(tokenizer.get_vocab().values()) + selected_tokens = random.choices(all_available_tokens, k=token_num) + return tokenizer.decode(selected_tokens) + + +def get_gen_prefix_cache_path(args, tokenizer): + """Create cache directory under ~/.cache/sglang/benchmark""" + cache_dir = Path.home() / ".cache" / "sglang" / "benchmark" + + # Create a unique cache filename based on the generation parameters + cache_key = ( + f"gen_prefix_{args.gen_num_groups}_{args.gen_prompts_per_group}_" + f"{args.gen_system_prompt_len}_{args.gen_question_len}_{args.gen_output_len}_" + f"{tokenizer.__class__.__name__}.pkl" + ) + return cache_dir / cache_key + + +def sample_generated_shared_prefix_requests( + num_groups: int, + prompts_per_group: int, + system_prompt_len: int, + question_len: int, + output_len: int, + tokenizer: PreTrainedTokenizerBase, + args, + disable_shuffle: bool = False, +) -> SampleOutput: + """Generate benchmark requests with shared system prompts using random tokens and caching.""" + cache_path = get_gen_prefix_cache_path(args, tokenizer) + + # Try to load from cache first + if cache_path.exists(): + print(f"\nLoading cached generated input data from {cache_path}") + with open(cache_path, "rb") as f: + return pickle.load(f) + + print("\nGenerating new input data...") + + # Generate system prompts for each group + system_prompts = [] + for _ in range(num_groups): + system_prompt = gen_prompt(tokenizer, system_prompt_len) + system_prompts.append(system_prompt) + + # Generate questions + questions = [] + for _ in range(num_groups * prompts_per_group): + question = gen_prompt(tokenizer, question_len) + questions.append(question) + + # Combine system prompts with questions + input_requests = [] + total_input_tokens = 0 + total_output_tokens = 0 + + for group_idx in tqdm(range(num_groups), desc="Generating system prompt"): + system_prompt = system_prompts[group_idx] + input_requests.append([]) + for prompt_idx in tqdm( + range(prompts_per_group), desc="Generating questions", leave=False + ): + question = questions[group_idx * prompts_per_group + prompt_idx] + full_prompt = f"{system_prompt}\n\n{question}" + prompt_len = len(tokenizer.encode(full_prompt)) + input_requests[-1].append((full_prompt, prompt_len, output_len)) + total_input_tokens += prompt_len + total_output_tokens += output_len + + if not disable_shuffle: + # Shuffle questions + random.shuffle(input_requests) + + # Print statistics + print(f"\nGenerated shared prefix dataset statistics:") + print(f"Number of groups: {num_groups}") + print(f"Prompts per group: {prompts_per_group}") + print(f"Total prompts: {len(input_requests) * prompts_per_group}") + print(f"Total input tokens: {total_input_tokens}") + print(f"Total output tokens: {total_output_tokens}") + print( + f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens" + ) + print( + f"Average question length: {sum(len(tokenizer.encode(q)) for q in questions) / len(questions):.1f} tokens\n" + ) + + # Save to cache + cache_path.parent.mkdir(parents=True, exist_ok=True) + print(f"Caching generated input data to {cache_path}") + with open(cache_path, "wb") as f: + pickle.dump(input_requests, f) + + return input_requests + + +def get_dataset(args, tokenizer): + if args.dataset_name == "sharegpt": + input_requests = sample_sharegpt_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + disable_shuffle=args.disable_shuffle, + enable_multiturn=args.enable_multiturn, + fixed_output_len=args.fixed_output_len, + ) + elif args.dataset_name == "ultrachat": + input_requests = sample_ultrachat_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + disable_shuffle=args.disable_shuffle, + enable_multiturn=args.enable_multiturn, + fixed_output_len=args.fixed_output_len, + ) + elif args.dataset_name == "loogle": + input_requests = sample_loogle_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + disable_shuffle=args.disable_shuffle, + enable_multiturn=args.enable_multiturn, + enable_shared_prefix=args.enable_shared_prefix, + fixed_output_len=args.fixed_output_len, + ) + elif args.dataset_name == "nextqa": + input_requests = sample_nextqa_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + max_frames=args.max_frames, + model_path=args.model, + disable_shuffle=args.disable_shuffle, + enable_multiturn=args.enable_multiturn, + backend=args.backend, + chat_template_name=args.chat_template, + fixed_output_len=args.fixed_output_len, + ) + elif args.dataset_name == "random": + input_requests = sample_random_requests( + input_len=args.random_input_len, + output_len=args.random_output_len, + num_prompts=args.num_prompts, + range_ratio=args.random_range_ratio, + tokenizer=tokenizer, + dataset_path=args.dataset_path, + ) + elif args.dataset_name == "generated-shared-prefix": + input_requests = sample_generated_shared_prefix_requests( + num_groups=args.gen_num_groups, + prompts_per_group=args.gen_prompts_per_group, + system_prompt_len=args.gen_system_prompt_len, + question_len=args.gen_question_len, + output_len=args.gen_output_len, + args=args, + tokenizer=tokenizer, + ) + else: + raise ValueError(f"Unknown dataset: {args.dataset_name}") + return input_requests diff --git a/sglang/benchmark/hicache/download.sh b/sglang/benchmark/hicache/download.sh new file mode 100644 index 0000000..340bcc3 --- /dev/null +++ b/sglang/benchmark/hicache/download.sh @@ -0,0 +1,66 @@ +#!/usr/bin/bash + +# The usage function +usage() { + echo "Usage: $0 {sharegpt|ultragpt|loogle|nextqa|all}" + exit 1 +} + +# The download function +download() { + case "$1" in + sharegpt) + echo $1 + wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json + ;; + ultragpt) + echo $1 + # Questions about the world + wget https://cloud.tsinghua.edu.cn/seafhttp/files/be1d7b87-22ca-449e-a6a7-c61d1ea7e010/ultrachat_release_230407.json + # Writing and Creation + wget https://cloud.tsinghua.edu.cn/seafhttp/files/61742d2a-25e2-4d08-b2b9-15f47ae50ace/ultrachat_material_release_230417.json + wget https://cloud.tsinghua.edu.cn/seafhttp/files/f71f6aa6-d346-4b16-85b7-8502efa3d608/ultrachat_material_release_230412.json + # External materials + wget https://cloud.tsinghua.edu.cn/seafhttp/files/42d22e28-e899-4975-a70f-5eda163e265d/ultrachat_existent_material_release_230420.json.gz + gunzip ultrachat_existent_material_release_230420.json.gz + ;; + loogle) + echo $1 + git lfs install + git clone git@hf.co:datasets/bigainlco/LooGLE + unzip LooGLE/data.zip + ;; + nextqa) + echo $1 + git lfs install + git clone https://huggingface.co/datasets/lmms-lab/NExTQA + unzip NExTQA/videos.zip + ;; + *) + usage + exit 1 + ;; + esac +} + +# Arg check +if [ "$#" -ne 1 ]; then + usage +fi + +# Invoke + +case "$1" in + sharegpt|ultragpt|loogle|nextqa) + download "$1" + ;; + all) + download sharegpt + download ultragpt + download loogle + download nextqa + ;; + *) + usage + ;; +esac diff --git a/sglang/benchmark/hicache/nextqa.py b/sglang/benchmark/hicache/nextqa.py new file mode 100644 index 0000000..4db6caa --- /dev/null +++ b/sglang/benchmark/hicache/nextqa.py @@ -0,0 +1,159 @@ +import os +import sys +from typing import List + +import av +from datasets import load_dataset + + +def find_video_files(video_dir) -> List[str]: + if os.path.isfile(video_dir): + return [video_dir] + + video_files = [] + for root, dirs, files in os.walk(video_dir): + for file in files: + if file.endswith((".mp4", ".avi", ".mov")): + video_files.append(os.path.join(root, file)) + # if file is dir + elif os.path.isdir(file): + video_files.extend(find_video_files(file)) + return video_files + + +def video_frames(video_path, max_frames) -> int: + container = av.open(video_path) + total_frames = container.streams.video[0].frames + return min(total_frames, max_frames) + + +class Video: + def __init__(self, video_path, num_frames): + self.path = video_path + self.num_frames = num_frames + + def __str__(self): + return f"Video({self.path}, {self.num_frames})" + + def __iter__(self): + return iter((self.path, self.num_frames)) + + +class VideoPrompt(Video): + def __init__(self, video_path, num_frames, prompt): + super().__init__(video_path, num_frames) + self.prompt = prompt + + def __str__(self): + return f"VideoPrompt({self.path}, {self.num_frames}, {self.prompt})" + + def __iter__(self): + return iter((self.path, self.num_frames, self.prompt)) + + +class VideoLoader: + pass + + +class VideoFileLoader(VideoLoader): + """ + Load all the videos in a directory + """ + + def __init__(self, video_dir, batch_size=1, max_frames=sys.maxsize): + super().__init__() + self.video_dir = video_dir + self.video_files = find_video_files(video_dir) + self.batch_size = batch_size + self.max_frames = max_frames + print(f"batch_size: {batch_size}, max_frames: {max_frames}") + + def __iter__(self): # (file, number of frames) + if self.batch_size == 1: + for video_file in self.video_files: + yield Video(video_file, video_frames(video_file, self.max_frames)) + else: + batch = [] + for video_file in self.video_files: + video = Video(video_file, video_frames(video_file, self.max_frames)) + batch.append(video) + if len(batch) == self.batch_size: + yield batch + batch = [] + + +class NExTQALoader(VideoLoader): + """ + Load vdideos and prompts from NExT dataset + set: train, test or validation + """ + + def __init__( + self, video_dir, batch_size=1, max_frames=sys.maxsize, dset="test", task="OE" + ): + """ + task: 'MV' or 'OE' + """ + super().__init__() + self.task = task + print(f"Loading the {dset} data of {task} from lmms-lab/NExTQA") + self.ds = load_dataset("lmms-lab/NExTQA", task) + self.ds = self.ds[dset] + + # self.n = ds.num_rows + self.video_dir = video_dir + self.video_files = find_video_files(video_dir) + self.video_to_path = dict() + for video_file in self.video_files: + video_id = video_file.split("/")[-1].split(".")[0] + self.video_to_path[video_id] = video_file + + self.batch_size = batch_size + self.max_frames = max_frames + + def get_video_prompt(self, entry, max_frames) -> VideoPrompt: + # Get video + video_id = entry["video"] + video_path = self.video_to_path[video_id] + assert os.path.exists(video_path), f"Video not found: {video_path}" + num_frames = min(entry["frame_count"], max_frames) + video = Video(video_path, num_frames) + prompt = entry["question"] + "?" + if self.task == "MC": # add choices + prompt += f' a0: {entry["a0"]}, a1: {entry["a1"]}, a2: {entry["a2"]}, a3: {entry["a3"]}' + return VideoPrompt(video_path, num_frames, prompt) + + def __iter__(self): + if self.batch_size == 1: + for entry in self.ds: + yield self.get_video_prompt(entry, self.max_frames) + else: + batch = [] + for entry in self.ds: + video = self.get_video_prompt(entry, self.max_frames) + batch.append(video) + if len(batch) == self.batch_size: + yield batch + batch = [] + + +# main +if __name__ == "__main__": + video_dir = "./videos" + # video_loader = VideoFileLoader(video_dir, batch_size=16) + # for batch in video_loader: + # print(f"Number of videos in batch: {len(batch)}") + # for video_file, num_frames in batch: + # print(f"Video: {video_file} number of frames: {num_frames}") + + video_loader = NExTQALoader(video_dir, batch_size=16, dset="test", task="OE") + for batch in video_loader: + print(f"Number of videos in batch: {len(batch)}") + for video_file, num_frames, prompt in batch: + print( + f"Video: {video_file} number of frames: {num_frames}, prompt: {prompt}" + ) + # break + # for video_file, prompt in batch: + # print(f"Video: {video_file} prompt: {prompt}") + # break diff --git a/sglang/benchmark/json_decode_regex/README.md b/sglang/benchmark/json_decode_regex/README.md new file mode 100644 index 0000000..26acaba --- /dev/null +++ b/sglang/benchmark/json_decode_regex/README.md @@ -0,0 +1,60 @@ +## Run benchmark + +### Build dataset +``` +pip install wikipedia +python3 build_dataset.py +``` + +### Dependencies + +``` +llama_cpp_python 0.2.19 +guidance 0.1.10 +vllm 0.2.5 +outlines 0.0.22 +``` + +### Benchmark sglang + +Run Llama-7B + +``` +python3 -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +Run Mixtral-8x7B + +``` +python3 -m sglang.launch_server --model-path mistralai/Mixtral-8x7B-Instruct-v0.1 --port 30000 --tp-size 8 +``` + +Benchmark + +``` +python3 bench_sglang.py --num-questions 10 +``` + + +### Benchmark Outlines + vLLM + +Run Llama-7B + +``` +python3 -m outlines.serve.serve --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +Benchmark + +``` +python3 bench_other.py --backend outlines --num-questions 10 +``` + + +### Benchmark guidance + +Run Llama-7B and benchmark + +``` +python3 bench_other.py --backend guidance --num-questions 10 --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` diff --git a/sglang/benchmark/json_decode_regex/bench_other.py b/sglang/benchmark/json_decode_regex/bench_other.py new file mode 100644 index 0000000..d80ea1d --- /dev/null +++ b/sglang/benchmark/json_decode_regex/bench_other.py @@ -0,0 +1,98 @@ +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +from tqdm import tqdm + +from sglang.lang.ir import REGEX_FLOAT, REGEX_INT, REGEX_STR +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +REGEX_LIST = r"\[(" + REGEX_STR + ", )*" + REGEX_STR + r"\]" + + +# fmt: off +def json_decode(document, generate): + s = "Please extract the information of a city from the following wikipedia page.\n" + s += "Page begin.\n" + document + "Page end.\n" + s += "Here is the name, country, and symbol of the city in JSON format.\n" + s += "{\n" + s += ' "name": ' + s += generate(s, max_tokens=8, regex=REGEX_STR + ",") + "\n" + s += ' "country": ' + s += generate(s, max_tokens=8, regex=REGEX_STR + ",") + "\n" + s += ' "latitude": ' + s += generate(s, max_tokens=8, regex=REGEX_FLOAT + ",") + "\n" + s += ' "population": ' + s += generate(s, max_tokens=8, regex=REGEX_INT + ",") + "\n" + s += ' "top 3 landmarks": ' + s += generate(s, max_tokens=24, regex=REGEX_LIST) + "\n" + s += "}\n" + + return s +# fmt: on + + +def main(args): + lines = read_jsonl(args.data_path) + arguments = [] + for i in range(len(lines[: args.num_questions])): + arguments.append( + { + "document": lines[i]["document"], + } + ) + states = [None] * len(arguments) + + # Select backend + call_generate = partial(get_call_generate(args), temperature=0) + + # Run requests + def get_one_answer(i): + states[i] = json_decode(generate=call_generate, **arguments[i]) + + tic = time.time() + if args.parallel == 1: + for i in tqdm(range(len(arguments))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + rets = list( + tqdm( + executor.map(get_one_answer, list(range(len(arguments)))), + total=len(arguments), + ) + ) + for _ in rets: + pass + + latency = time.time() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "json_decode_regex", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="questions.jsonl") + parser.add_argument("--num-questions", type=int, default=20) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/json_decode_regex/bench_sglang.py b/sglang/benchmark/json_decode_regex/bench_sglang.py new file mode 100644 index 0000000..4139ebf --- /dev/null +++ b/sglang/benchmark/json_decode_regex/bench_sglang.py @@ -0,0 +1,101 @@ +import argparse +import json +import time + +import sglang as sgl +from sglang.lang.ir import REGEX_FLOAT, REGEX_INT, REGEX_STR +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + +REGEX_LIST = r"\[(" + REGEX_STR + ", )*" + REGEX_STR + r"\]" + +# fmt: off +@sgl.function +def json_warm_up(s): + s += "The information about Hogwarts is in the following JSON format.\n" + with s.var_scope("json_output"): + s += "{\n" + s += ' "name": ' + sgl.gen("name", max_tokens=8, regex=REGEX_STR + ",") + "\n" + s += ' "country": ' + sgl.gen("country", max_tokens=8, regex=REGEX_STR + ",") + "\n" + s += ' "latitude": ' + sgl.gen("latitude", max_tokens=8, regex=REGEX_FLOAT + ",") + "\n" + s += ' "population": ' + sgl.gen("population", max_tokens=8, regex=REGEX_INT + ",") + "\n" + s += ' "top 3 landmarks": ' + sgl.gen( "landmarks", max_tokens=24, regex=REGEX_LIST) + "\n" + s += "}\n" + print(f'The warmp up json result is:\n{s["json_output"]}') +# fmt: on + +# fmt: off +@sgl.function +def json_decode(s, document): + s += "Please extract the information of a city from the following wikipedia page.\n" + s += "Page begin.\n" + document + "Page end.\n" + s += "Here is the name, country, and symbol of the city in JSON format.\n" + with s.var_scope("json_output"): + s += "{\n" + s += ' "name": ' + sgl.gen("name", max_tokens=8, regex=REGEX_STR + ",") + "\n" + s += ' "country": ' + sgl.gen("country", max_tokens=8, regex=REGEX_STR + ",") + "\n" + s += ' "latitude": ' + sgl.gen("latitude", max_tokens=8, regex=REGEX_FLOAT + ",") + "\n" + s += ' "population": ' + sgl.gen("population", max_tokens=8, regex=REGEX_INT + ",") + "\n" + s += ' "top 3 landmarks": ' + sgl.gen( "landmarks", max_tokens=24, regex=REGEX_LIST) + "\n" + s += "}\n" +# fmt: on + + +def main(args): + lines = read_jsonl(args.data_path) + lines = list(lines) + arguments = [] + for i in range(len(lines[: args.num_questions])): + arguments.append( + { + "document": lines[i]["document"], + } + ) + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Warm up + json_warm_up.run().sync() + + # Run requests + tic = time.time() + states = json_decode.run_batch( + arguments, temperature=0, num_threads=args.parallel, progress_bar=True + ) + latency = time.time() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(f"tmp_{args.backend}_json_results.txt", "w") as fout: + for state in states: + fout.write(state["json_output"] + "\n") + + with open(args.result_file, "a") as fout: + value = { + "task": "json_decode_regex", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="questions.jsonl") + parser.add_argument("--num-questions", type=int, default=20) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/json_decode_regex/build_dataset.py b/sglang/benchmark/json_decode_regex/build_dataset.py new file mode 100644 index 0000000..1396e5e --- /dev/null +++ b/sglang/benchmark/json_decode_regex/build_dataset.py @@ -0,0 +1,58 @@ +import json + +import transformers +import wikipedia + +model_path = "meta-llama/Llama-2-7b-chat-hf" +t = transformers.AutoTokenizer.from_pretrained(model_path) +city_names = [ + "los angles", + "london", + "tokyo", + "beijing", + "singapore", + "paris", + "dubai", + "sydney", + "moscow", + "rome", + "toronto", + "rio de janeiro", + "istanbul", + "berlin", + "auckland", + "buenos aires", + "mexico city", + "mumbai", + "seoul", + "bangkok", + "cairo", + "athens", + "jerusalem", +] + + +def get_content(city_name): + content = str(wikipedia.page(city_name).content) + content = content.replace("\n\n", "\n") + + tokens = t.encode(content) + + expected_tokens = 3000 + truncate_len = int((expected_tokens / len(tokens)) * len(content)) + truncate_content = content[:truncate_len] + truncate_tokens = t.encode(truncate_content) + + # Count token + print( + f"city_name: {city_name}, #tokens: {len(tokens)}, #truncate tokens: {len(truncate_tokens)}" + ) + + return truncate_content + + +if __name__ == "__main__": + with open("questions.jsonl", "w") as fout: + for city_name in city_names: + truncate_content = get_content(city_name) + fout.write(json.dumps({"document": truncate_content}) + "\n") diff --git a/sglang/benchmark/json_jump_forward/README.md b/sglang/benchmark/json_jump_forward/README.md new file mode 100644 index 0000000..38fb67e --- /dev/null +++ b/sglang/benchmark/json_jump_forward/README.md @@ -0,0 +1,88 @@ +## Run benchmark + +### Dependencies + +``` +llama_cpp_python 0.2.38 +guidance 0.1.10 +vllm 0.2.7 +outlines 0.0.25 +``` + +### Build dataset + +When benchmarking long document information retrieval, run the following command to build the dataset: + +```bash +pip install wikipedia +python3 build_dataset.py +``` + +### Benchmark sglang + +Run Llama-7B + +```bash +python3 -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +Benchmark Character Generation + +```bash +python3 bench_sglang.py --mode character +``` + +Benchmark City Information Retrieval + +```bash +python3 bench_sglang.py --mode city +``` + + +### Benchmark Outlines + vLLM + +Run Llama-7B + +```bash +python3 -m outlines.serve.serve --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +Benchmark Character Generation + +```bash +python3 bench_other.py --mode character --backend outlines +``` + +Benchmark City Information Retrieval + +```bash +python3 bench_other.py --mode city --backend outlines +``` + +### Benchmark guidance + +Run Llama-7B and benchmark character generation + +```bash +python3 bench_other.py --mode character --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +Run Llama-7B and benchmark city information retrieval + +```bash +python3 bench_other.py --mode city --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +### Benchmark lmql + +Run Llama-7B and benchmark character generation + +``` +python3 bench_other.py --mode character --backend lmql --parallel 1 +``` + +Run Llama-7B and benchmark city information retrieval + +``` +python3 bench_other.py --mode city --backend lmql --parallel 1 +``` diff --git a/sglang/benchmark/json_jump_forward/bench_other.py b/sglang/benchmark/json_jump_forward/bench_other.py new file mode 100644 index 0000000..9eb5c58 --- /dev/null +++ b/sglang/benchmark/json_jump_forward/bench_other.py @@ -0,0 +1,288 @@ +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +import guidance +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +# there are some FSM bugs with json regex converted from pydantic model +# here use a string regex instead +# regex_string = build_regex_from_object(HarryPoterRole) +character_regex = ( + r"""\{\n""" + + r""" "name": "[\w\d\s]{1,16}",\n""" + + r""" "house": "(Gryffindor|Slytherin|Ravenclaw|Hufflepuff)",\n""" + + r""" "blood status": "(Pure-blood|Half-blood|Muggle-born)",\n""" + + r""" "occupation": "(student|teacher|auror|ministry of magic|death eater|order of the phoenix)",\n""" + + r""" "wand": \{\n""" + + r""" "wood": "[\w\d\s]{1,16}",\n""" + + r""" "core": "[\w\d\s]{1,16}",\n""" + + r""" "length": [0-9]{1,2}\.[0-9]{0,2}\n""" + + r""" \},\n""" + + r""" "alive": "(Alive|Deceased)",\n""" + + r""" "patronus": "[\w\d\s]{1,16}",\n""" + + r""" "bogart": "[\w\d\s]{1,16}"\n""" + + r"""\}""" +) + +city_regex = ( + r"""\{\n""" + + r""" "name": "[\w\d\s]{1,16}",\n""" + + r""" "country": "[\w\d\s]{1,16}",\n""" + + r""" "latitude": [-+]?[0-9]*\.?[0-9]{0,2},\n""" + + r""" "population": [-+]?[0-9]{1,9},\n""" + + r""" "top 3 landmarks": \["[\w\d\s]{1,16}", "[\w\d\s]{1,16}", "[\w\d\s]{1,16}"\]\n""" + + r"""\}""" +) + +# fmt: off +def character_gen(name, generate): + s = name + " is a character in Harry Potter. Please fill in the following information about this character.\n" + s += generate(s, max_tokens=256, regex=character_regex) + return s +# fmt: on + +# fmt: off +def city_gen(document, generate): + s = "Please extract the information of a city from the following wikipedia page.\n" + s += "Page begin.\n" + document + "Page end.\n" + s += "Here is the name, country, and symbol of the city in JSON format.\n" + s += generate(s, max_tokens=256, regex=city_regex) + return s +# fmt: on + + +@guidance +def character_maker(lm, name): + regex_str_no_quote = r"[\w\d\s]+" + regex_float = r"[0-9]+\.[0-9]+" + lm += f"""\ + {name} is a character in Harry Potter. Please fill in the following information about this character. + {{ + "name": "{guidance.gen("name", max_tokens=16, regex=regex_str_no_quote)}", + "house": "{guidance.select(options=['Gryffindor', 'Slytherin', 'Ravenclaw', 'Hufflepuff'], name='house')}", + "blood status": "{guidance.select(options=['Pure-blood', 'Half-blood', 'Muggle-born'], name='blood status')}", + "occupation": "{guidance.select(options=['student', 'teacher', 'auror', 'ministry of magic', 'death eater', 'order of the phoenix'], name='occupation')}", + "wand": {{ + "wood": "{guidance.gen("wood", max_tokens=16, regex=regex_str_no_quote)}", + "core": "{guidance.gen('core', max_tokens=16, regex=regex_str_no_quote)}", + "length": {guidance.gen('length', max_tokens=10, regex=regex_float)} + }}, + "alive": "{guidance.select(options=['Alive', 'Deceased'], name='alive')}", + "patronus": "{guidance.gen('patronus', max_tokens=16, regex=regex_str_no_quote)}", + "bogart": "{guidance.gen('bogart', max_tokens=16, regex=regex_str_no_quote)}" + }} + """ + + return lm + + +async def call_generate_lmql( + prompt, temperature, max_tokens, regex, max_len=4096, model=None, **kwargs +): + assert model is not None + import lmql + + @lmql.query(model=model) + async def program(question, max_tokens, regex): + '''lmql + """{question}[ANSWER]""" where len(TOKENS(ANSWER)) < max_tokens and REGEX(ANSWER, regex) + return ANSWER + ''' + + return await program( + question=prompt, + temperature=temperature, + max_tokens=max_tokens, + max_len=max_len, + regex=regex, + **kwargs, + ) + + +@guidance +def city_maker(lm, document): + regex_str_no_quote = r"[\w\d\s]+" + regex_float = r"[0-9]+\.[0-9]+" + lm += f"""\ + Please extract the information of a city from the following wikipedia page. + Page begin. + {document} + Page end. + Here is the name, country, and symbol of the city in JSON format. + {{ + "name": "{guidance.gen("name", max_tokens=16, regex=regex_str_no_quote)}", + "country": "{guidance.gen("country", max_tokens=16, regex=regex_str_no_quote)}", + "latitude": {guidance.gen("latitude", max_tokens=10, regex=regex_float)}, + "population": {guidance.gen("population", max_tokens=10, regex=r"[0-9]+")}, + "top 3 landmarks": [ + "{guidance.gen("landmark1", max_tokens=16, regex=regex_str_no_quote)}", "{guidance.gen("landmark2", max_tokens=16, regex=regex_str_no_quote)}", "{guidance.gen("landmark3", max_tokens=16, regex=regex_str_no_quote)}" + ] + }} + """ + + return lm + + +def bench_character(args): + arguments = [] + with open(args.data_path, "r") as f: + for line in f: + arguments.append({"name": line.strip()}) + arguments = arguments[: args.num_jsons] + + states = [None] * len(arguments) + + # Select backend + if args.backend == "outlines": + call_generate = partial(get_call_generate(args), temperature=0) + + def get_one_answer(i): + states[i] = character_gen(**arguments[i], generate=call_generate) + + elif args.backend == "guidance": + model = guidance.models.LlamaCpp( + args.model_path, + n_gpu_layers=-1, + n_ctx=args.n_ctx, + ) + + def get_one_answer(i): + lm = model + character_maker(**arguments[i]) + states[i] = lm + + elif args.backend == "lmql": + import asyncio + + import lmql + + model = lmql.model(args.model_path, endpoint=f"{args.host}:{args.port}") + call_generate = partial( + call_generate_lmql, + model=model, + max_tokens=256, + regex=character_regex, + ) + + async def get_one_answer_async(i): + states[i] = await call_generate(prompt=arguments[i]["name"], temperature=0) + + else: + raise ValueError(f"Invalid backend: {args.backend}") + + tic = time.time() + + if args.backend != "lmql": + if args.parallel == 1: + for i in tqdm(range(len(arguments))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + rets = list( + tqdm( + executor.map(get_one_answer, list(range(len(arguments)))), + total=len(arguments), + ) + ) + for _ in rets: + pass + else: + batches = [] + for i in range(0, len(arguments), args.parallel): + batches.append(list(range(i, min(i + args.parallel, len(arguments))))) + loop = asyncio.get_event_loop() + + for bt in tqdm(batches): + loop.run_until_complete( + asyncio.gather(*[get_one_answer_async(i) for i in bt]) + ) + + latency = time.time() - tic + + return states, latency + + +def bench_city_doc(args): + arguments = [] + for line in read_jsonl(args.data_path): + arguments.append({"document": line["document"]}) + arguments = arguments[: args.num_jsons] + + states = [None] * len(arguments) + + # Select backend + if args.backend == "outlines": + call_generate = partial(get_call_generate(args), temperature=0) + + def get_one_answer(i): + states[i] = city_gen(**arguments[i], generate=call_generate) + + elif args.backend == "guidance": + model = guidance.models.LlamaCpp( + args.model_path, + n_gpu_layers=-1, + n_ctx=args.n_ctx, + ) + + def get_one_answer(i): + lm = model + city_maker(**arguments[i]) + states[i] = lm + + else: + raise ValueError(f"Invalid backend: {args.backend}") + + tic = time.time() + if args.parallel == 1: + for i in tqdm(range(len(arguments))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + rets = executor.map(get_one_answer, list(range(len(arguments)))) + for _ in rets: + pass + + latency = time.time() - tic + + return states, latency + + +def main(args): + if args.mode == "character": + args.data_path = "dataset.txt" + states, latency = bench_character(args) + elif args.mode == "city": + args.data_path = "questions.jsonl" + states, latency = bench_city_doc(args) + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}_{args.mode}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "json_jump_forward", + "backend": args.backend, + "latency": round(latency, 3), + "num_jsons": args.num_jsons, + "mode": args.mode, + "parallel": args.parallel, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str) + parser.add_argument("--num-jsons", type=int, default=50) + parser.add_argument( + "--mode", type=str, default="character", choices=["character", "city"] + ) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/json_jump_forward/bench_sglang.py b/sglang/benchmark/json_jump_forward/bench_sglang.py new file mode 100644 index 0000000..10cf269 --- /dev/null +++ b/sglang/benchmark/json_jump_forward/bench_sglang.py @@ -0,0 +1,143 @@ +import argparse +import json +import time + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + +# there are some FSM bugs with json regex converted from pydantic model +# here use a string regex instead +# regex_string = build_regex_from_object(HarryPoterRole) +character_regex = ( + r"""\{\n""" + + r""" "name": "[\w\d\s]{1,16}",\n""" + + r""" "house": "(Gryffindor|Slytherin|Ravenclaw|Hufflepuff)",\n""" + + r""" "blood status": "(Pure-blood|Half-blood|Muggle-born)",\n""" + + r""" "occupation": "(student|teacher|auror|ministry of magic|death eater|order of the phoenix)",\n""" + + r""" "wand": \{\n""" + + r""" "wood": "[\w\d\s]{1,16}",\n""" + + r""" "core": "[\w\d\s]{1,16}",\n""" + + r""" "length": [0-9]{1,2}\.[0-9]{0,2}\n""" + + r""" \},\n""" + + r""" "alive": "(Alive|Deceased)",\n""" + + r""" "patronus": "[\w\d\s]{1,16}",\n""" + + r""" "bogart": "[\w\d\s]{1,16}"\n""" + + r"""\}""" +) + +city_regex = ( + r"""\{\n""" + + r""" "name": "[\w\d\s]{1,16}",\n""" + + r""" "country": "[\w\d\s]{1,16}",\n""" + + r""" "latitude": [-+]?[0-9]*\.?[0-9]{0,2},\n""" + + r""" "population": [-+]?[0-9]{1,9},\n""" + + r""" "top 3 landmarks": \["[\w\d\s]{1,16}", "[\w\d\s]{1,16}", "[\w\d\s]{1,16}"\]\n""" + + r"""\}""" +) + +# fmt: off +@sgl.function +def character_gen(s, name): + s += name + " is a character in Harry Potter. Please fill in the following information about this character.\n" + s += sgl.gen("json_output", max_tokens=256, regex=character_regex) +# fmt: on + +# fmt: off +@sgl.function +def city_gen(s, document): + s += "Please extract the information of a city from the following wikipedia page.\n" + s += "Page begin.\n" + document + "Page end.\n" + s += "Here is the name, country, and symbol of the city in JSON format.\n" + s += sgl.gen("json_output",max_tokens=256, regex=city_regex) +# fmt: on + + +def bench_city_doc(args): + arguments = [] + for line in read_jsonl(args.data_path): + arguments.append({"document": line["document"]}) + arguments = arguments[: args.num_jsons] + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.time() + states = city_gen.run_batch( + arguments, + temperature=0, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.time() - tic + + return states, latency + + +def bench_character(args): + arguments = [] + with open(args.data_path, "r") as f: + for line in f: + arguments.append({"name": line.strip()}) + arguments = arguments[: args.num_jsons] + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.time() + states = character_gen.run_batch( + arguments, + temperature=0, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.time() - tic + + return states, latency + + +def main(args): + if args.mode == "character": + args.data_path = "dataset.txt" + states, latency = bench_character(args) + elif args.mode == "city": + args.data_path = "questions.jsonl" + states, latency = bench_city_doc(args) + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}_{args.mode}.txt", states) + with open(f"{args.backend}_{args.mode}.json", "w") as fout: + for state in states: + fout.write(state["json_output"] + "\n") + + with open(args.result_file, "a") as fout: + value = { + "task": "json_jump_forward", + "backend": args.backend, + "latency": round(latency, 3), + "num_jsons": args.num_jsons, + "mode": args.mode, + "parallel": args.parallel, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str) + parser.add_argument("--num-jsons", type=int, default=50) + parser.add_argument( + "--mode", type=str, default="character", choices=["character", "city"] + ) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/json_jump_forward/build_dataset.py b/sglang/benchmark/json_jump_forward/build_dataset.py new file mode 100644 index 0000000..1396e5e --- /dev/null +++ b/sglang/benchmark/json_jump_forward/build_dataset.py @@ -0,0 +1,58 @@ +import json + +import transformers +import wikipedia + +model_path = "meta-llama/Llama-2-7b-chat-hf" +t = transformers.AutoTokenizer.from_pretrained(model_path) +city_names = [ + "los angles", + "london", + "tokyo", + "beijing", + "singapore", + "paris", + "dubai", + "sydney", + "moscow", + "rome", + "toronto", + "rio de janeiro", + "istanbul", + "berlin", + "auckland", + "buenos aires", + "mexico city", + "mumbai", + "seoul", + "bangkok", + "cairo", + "athens", + "jerusalem", +] + + +def get_content(city_name): + content = str(wikipedia.page(city_name).content) + content = content.replace("\n\n", "\n") + + tokens = t.encode(content) + + expected_tokens = 3000 + truncate_len = int((expected_tokens / len(tokens)) * len(content)) + truncate_content = content[:truncate_len] + truncate_tokens = t.encode(truncate_content) + + # Count token + print( + f"city_name: {city_name}, #tokens: {len(tokens)}, #truncate tokens: {len(truncate_tokens)}" + ) + + return truncate_content + + +if __name__ == "__main__": + with open("questions.jsonl", "w") as fout: + for city_name in city_names: + truncate_content = get_content(city_name) + fout.write(json.dumps({"document": truncate_content}) + "\n") diff --git a/sglang/benchmark/json_jump_forward/dataset.txt b/sglang/benchmark/json_jump_forward/dataset.txt new file mode 100644 index 0000000..c12421e --- /dev/null +++ b/sglang/benchmark/json_jump_forward/dataset.txt @@ -0,0 +1,50 @@ +Harry Potter +Hermione Granger +Ron Weasley +Albus Dumbledore +Severus Snape +Rubeus Hagrid +Draco Malfoy +Ginny Weasley +Fred Weasley +George Weasley +Percy Weasley +Sirius Black +Remus Lupin +Neville Longbottom +Luna Lovegood +Cedric Diggory +Cho Chang +Lord Voldemort +Minerva McGonagall +Filius Flitwick +Dolores Umbridge +Bellatrix Lestrange +Lucius Malfoy +Molly Weasley +Arthur Weasley +Nymphadora Tonks +Dobby +Moaning Myrtle +Peter Pettigrew +Alastor 'Mad-Eye' Moody +Horace Slughorn +Vernon Dursley +Petunia Dursley +Dudley Dursley +Argus Filch +Sybill Trelawney +Gilderoy Lockhart +Fleur Delacour +Viktor Krum +Bill Weasley +Oliver Wood +Cornelius Fudge +Barty Crouch Sr. +Barty Crouch Jr. +Kingsley Shacklebolt +Quirinus Quirrell +Nearly Headless Nick +Aunt Marge +Griphook +Ludo Bagman diff --git a/sglang/benchmark/json_schema/README.md b/sglang/benchmark/json_schema/README.md new file mode 100644 index 0000000..e6790e1 --- /dev/null +++ b/sglang/benchmark/json_schema/README.md @@ -0,0 +1,15 @@ +## Run benchmark + +### Benchmark sglang + +Run Llama-8b + +```bash +python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --port 30000 +``` + +Benchmark + +```bash +python3 bench_sglang.py +``` diff --git a/sglang/benchmark/json_schema/bench_sglang.py b/sglang/benchmark/json_schema/bench_sglang.py new file mode 100644 index 0000000..5a5fd66 --- /dev/null +++ b/sglang/benchmark/json_schema/bench_sglang.py @@ -0,0 +1,146 @@ +import argparse +import json +import time +from typing import List, Tuple + +import jsonschema +from datasets import load_dataset + +import sglang as sgl +from sglang.global_config import global_config +from sglang.srt.hf_transformers_utils import get_tokenizer +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text + + +@sgl.function +def schema_gen(s, message: Tuple[str, str], json_schema: str): + system, user = message + s += sgl.system(system) + s += sgl.user(user) + s += sgl.assistant( + sgl.gen("json_output", temperature=0, max_tokens=256, json_schema=json_schema) + ) + + +def contains_formats(schema, formats: List[str]): + if isinstance(schema, dict): + if schema.get("format", None) in formats: + return True + for value in schema.values(): + if contains_formats(value, formats): + return True + elif isinstance(schema, list): + for item in schema: + if contains_formats(item, formats): + return True + return False + + +def convert_dataset(path: str): + raw_dataset = load_dataset(path) + dataset = [] + for data in raw_dataset["train"]: + messages = data["prompt"] + schema = data["schema"] + obj = json.loads(schema) + + # skip some corrupted examples + if obj.get("type", None) is None: + continue + + # skip schema with format "email" + # which is not supported by outlines for now + if contains_formats(obj, ["email"]): + continue + + system = messages[0] + user = messages[1] + assert system["role"] == "system", "invalid role" + assert user["role"] == "user", "invalid role" + assert len(messages) == 2, "invalid message length" + message = json.dumps(system["content"]), json.dumps(user["content"]) + dataset.append( + { + "message": message, + "json_schema": schema, + } + ) + + return dataset + + +def bench_schema(args): + arguments = convert_dataset(args.data_path) + + if args.num_jsons < 0 or args.num_jsons > len(arguments): + args.num_jsons = len(arguments) + arguments = arguments[: args.num_jsons] + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.time() + states = schema_gen.run_batch( + arguments, + temperature=0, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.time() - tic + + # Check if the outputs are valid + indexs = [] + for i, state in enumerate(states): + try: + schema = json.loads(arguments[i]["json_schema"]) + obj = json.loads(state["json_output"]) + assert jsonschema.validate(obj, schema) is None + except Exception as e: + print(e) + indexs.append(i) + + return states, latency + + +def main(args): + states, latency = bench_schema(args) + + # Compute accuracy + tokenizer = get_tokenizer( + global_config.default_backend.get_server_info()["tokenizer_path"] + ) + output_jsons = [state["json_output"] for state in states] + num_output_tokens = sum(len(tokenizer.encode(x)) for x in output_jsons) + print(f"Latency: {latency:.3f}") + print(f"Output throughput: {num_output_tokens / latency:.3f} token/s") + print(f"#output tokens: {num_output_tokens}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + with open(f"{args.backend}.jsonl", "w") as fout: + for state in states: + fout.write(state["json_output"] + "\n") + + with open(args.result_file, "a") as fout: + value = { + "task": "json_schema", + "backend": args.backend, + "latency": round(latency, 3), + "num_jsons": args.num_jsons, + "parallel": args.parallel, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="NousResearch/json-mode-eval") + parser.add_argument("--num-jsons", type=int, default=-1) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py b/sglang/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py new file mode 100644 index 0000000..b61240b --- /dev/null +++ b/sglang/benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py @@ -0,0 +1,403 @@ +import itertools +import math + +import cudnn +import torch +import torch.utils.benchmark as benchmark +from flashinfer import BatchDecodeWithPagedKVCacheWrapper + +from sglang.srt.layers.attention.triton_ops.decode_attention import decode_attention_fwd +from sglang.srt.utils import should_use_tensor_core + + +def benchmark_forward( + fn, + *inputs, + repeats=10, + amp=False, + amp_dtype=torch.float16, + **kwinputs, +): + def amp_wrapper(*inputs, **kwinputs): + with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=amp): + fn(*inputs, **kwinputs) + + t = benchmark.Timer( + stmt="fn_amp(*inputs, **kwinputs)", + globals={"fn_amp": amp_wrapper, "inputs": inputs, "kwinputs": kwinputs}, + num_threads=torch.get_num_threads(), + ) + m = t.timeit(repeats) + return t, m + + +def time_fwd(func, *args, **kwargs): + time_f = benchmark_forward(func, *args, **kwargs) + return time_f[1].mean * 1e6 + + +def decode_attention_sglang( + q, + kv_data, + batch_size, + kv_len, + head_num_q, + head_num_kv, + head_dim, + num_kv_splits, + warmup=10, +): + + k_buffer = kv_data[0].view(-1, head_num_kv, head_dim) + v_buffer = kv_data[1].view(-1, head_num_kv, head_dim) + o = torch.empty_like(q) + total_tokens = batch_size * kv_len + req_to_token = torch.arange(0, total_tokens).to(0).int().view(batch_size, kv_len) + b_req_idx = torch.arange(0, batch_size).to(0).int() + b_seq_len = torch.full((batch_size,), kv_len, dtype=torch.int32, device="cuda") + max_len_in_batch = kv_len + sm_scale = 1.0 / (head_dim**0.5) + + attn_logits = torch.empty( + (batch_size, head_num_q, num_kv_splits, head_dim + 1), + dtype=torch.float32, + device="cuda", + ) + + for _ in range(warmup): + decode_attention_fwd( + q, + k_buffer, + v_buffer, + o, + req_to_token, + b_req_idx, + b_seq_len, + attn_logits, + num_kv_splits, + sm_scale, + ) + + f = time_fwd( + decode_attention_fwd, + q, + k_buffer, + v_buffer, + o, + req_to_token, + b_req_idx, + b_seq_len, + attn_logits, + num_kv_splits, + sm_scale, + ) + + return f, o + + +def decode_attention_flashinfer(dtype, head_num_q, head_num_kv): + workspace_buffer = torch.empty(128 * 1024 * 1024, dtype=torch.int8, device="cuda") + use_tensor_cores = should_use_tensor_core( + kv_cache_dtype=dtype, + num_attention_heads=head_num_q, + num_kv_heads=head_num_kv, + ) + flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper( + workspace_buffer, "NHD", use_tensor_cores=use_tensor_cores + ) + + class FlashinferAttention(torch.autograd.Function): + @staticmethod + def forward( + ctx, + q, + kv_data, + batch_size, + kv_len, + head_num_q, + head_num_kv, + head_dim, + dtype, + warmup=10, + ): + total_tokens = batch_size * kv_len + kv_indptr = torch.arange(0, batch_size + 1).to(0).int() * kv_len + kv_indices = torch.arange(0, total_tokens).to(0).int() + kv_last_page_len = torch.full( + (batch_size,), 1, dtype=torch.int32, device="cuda" + ) + + flashinfer_decode_wrapper.end_forward() + flashinfer_decode_wrapper.begin_forward( + kv_indptr, + kv_indices, + kv_last_page_len, + head_num_q, + head_num_kv, + head_dim, + 1, + pos_encoding_mode="NONE", + data_type=dtype, + ) + + for _ in range(warmup): + o = flashinfer_decode_wrapper.forward( + q.contiguous().view(-1, head_num_q, head_dim), kv_data + ) + + f = time_fwd( + flashinfer_decode_wrapper.forward, + q.contiguous().view(-1, head_num_q, head_dim), + kv_data, + ) + + return f, o + + return FlashinferAttention + + +def convert_to_cudnn_type(torch_type): + if torch_type == torch.float16: + return cudnn.data_type.HALF + elif torch_type == torch.bfloat16: + return cudnn.data_type.BFLOAT16 + elif torch_type == torch.float32: + return cudnn.data_type.FLOAT + elif torch_type == torch.int32: + return cudnn.data_type.INT32 + elif torch_type == torch.int64: + return cudnn.data_type.INT64 + else: + raise ValueError("Unsupported tensor data type.") + + +def decode_attention_cudnn( + q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype, warmup=10 +): + # Prepare data: continuous q,k,v + dims_q = (batch_size, head_num_q, 1, head_dim) + strides_q = (head_num_q * head_dim, head_dim, head_num_q * head_dim, 1) + q_gpu = q.as_strided(dims_q, strides_q) + o_gpu = ( + torch.empty(batch_size * head_num_q * head_dim) + .half() + .cuda() + .as_strided(dims_q, strides_q) + ) + + dims_kv = (batch_size, head_num_kv, kv_len, head_dim) + strides_kv = ( + kv_len * head_num_kv * head_dim, + head_dim, + head_num_kv * head_dim, + 1, + ) + k_gpu = kv_data[0].as_strided(dims_kv, strides_kv) + v_gpu = kv_data[1].as_strided(dims_kv, strides_kv) + + seq_len_q_gpu = torch.full((batch_size, 1, 1, 1), 1, device="cuda") + seq_len_kv_gpu = torch.full((batch_size, 1, 1, 1), kv_len, device="cuda") + attn_scale = 1.0 / (head_dim**0.5) + + # Prepare data: paged k,v + block_size = 1 + blocks_per_batch = math.ceil(kv_len / block_size) + # [num_blocks, head_num_kv, block_size, head_dim], num_blocks = batch_size * blocks_per_batch + container_k_gpu = torch.cat(k_gpu.chunk(blocks_per_batch, dim=2), dim=0) + container_v_gpu = torch.cat(v_gpu.chunk(blocks_per_batch, dim=2), dim=0) + page_table_k_gpu = ( + torch.linspace( + 0, + batch_size * blocks_per_batch - 1, + batch_size * blocks_per_batch, + device="cuda", + dtype=torch.int32, + ) + .reshape(blocks_per_batch, 1, batch_size, 1) + .transpose(0, 2) + ) + page_table_v_gpu = page_table_k_gpu.clone() + + graph = cudnn.pygraph( + io_data_type=convert_to_cudnn_type(dtype), + intermediate_data_type=cudnn.data_type.FLOAT, + compute_data_type=cudnn.data_type.FLOAT, + ) + + q = graph.tensor_like(q_gpu) + container_k = graph.tensor_like(container_k_gpu) + container_v = graph.tensor_like(container_v_gpu) + page_table_k = graph.tensor_like(page_table_k_gpu) + page_table_v = graph.tensor_like(page_table_v_gpu) + + seq_len_q = graph.tensor_like(seq_len_q_gpu) + seq_len_kv = graph.tensor_like(seq_len_kv_gpu) + + o, _ = graph.sdpa( + name="sdpa", + q=q, + k=container_k, # Container K: non contiguous container with K blocks + v=container_v, # Container V: non contiguous container with V blocks + is_inference=True, + attn_scale=attn_scale, + use_causal_mask=False, + use_padding_mask=True, + seq_len_q=seq_len_q, + seq_len_kv=seq_len_kv, + paged_attention_k_table=page_table_k, # Page Table K: Tensor containing offsets to the container with K blocks + paged_attention_v_table=page_table_v, # Page Table V: Tensor containing offsets to the container with V blocks + paged_attention_max_seq_len_kv=kv_len, # The maximum sequence length for K caches (this is optional, but recommended) + ) + + o.set_output(True).set_dim(dims_q).set_stride(strides_q) + + graph.validate() + graph.build_operation_graph() + graph.create_execution_plans([cudnn.heur_mode.A]) + graph.check_support() + graph.build_plans() + + workspace = torch.empty( + graph.get_workspace_size(), device="cuda", dtype=torch.uint8 + ) + + variant_pack = { + q: q_gpu, + container_k: container_k_gpu, + container_v: container_v_gpu, + page_table_k: page_table_k_gpu, + page_table_v: page_table_v_gpu, + seq_len_q: seq_len_q_gpu, + seq_len_kv: seq_len_kv_gpu, + o: o_gpu, + } + + for _ in range(warmup): + graph.execute(variant_pack, workspace) + + f = time_fwd( + graph.execute, + variant_pack, + workspace, + ) + + return f, o_gpu.squeeze(dim=2) + + +def calculate_diff(): + + dtype = torch.float16 + batch_size = 64 + kv_len = 4096 + head_num_q = 64 + head_num_kv = 8 + head_dim = 128 + + q = torch.randn(batch_size, head_num_q, head_dim, dtype=dtype, device="cuda") + kv_data = ( + torch.randn( + batch_size * kv_len, head_num_kv, head_dim, dtype=dtype, device="cuda" + ), + torch.randn( + batch_size * kv_len, head_num_kv, head_dim, dtype=dtype, device="cuda" + ), + ) + + _, output_sglang = decode_attention_sglang( + q, + kv_data, + batch_size, + kv_len, + head_num_q, + head_num_kv, + head_dim, + num_kv_splits=8, + ) + + attn_flashinfer = decode_attention_flashinfer(dtype, head_num_q, head_num_kv).apply + _, output_flashinfer = attn_flashinfer( + q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype + ) + + _, output_cudnn = decode_attention_cudnn( + q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype + ) + + print(f"SGLang output={output_sglang}") + print(f"FlashInfer output={output_flashinfer}") + print(f"cuDNN output={output_cudnn}") + if torch.allclose(output_sglang, output_flashinfer, atol=1e-2, rtol=1e-2): + print("✅ SGLang[Triton] and FlashInfer match") + else: + print("❌ SGLang[Triton] and FlashInfer differ") + + if torch.allclose(output_sglang, output_cudnn, atol=1e-2, rtol=1e-2): + print("✅ SGLang[Triton] and cuDNN match") + else: + print("❌ SGLang[Triton] and cuDNN differ") + + +if __name__ == "__main__": + calculate_diff() + + head_dim = 128 + dtype = torch.float16 + batch_size_range = [2**i for i in range(0, 8, 2)] + kv_len_range = [2**i for i in range(6, 13, 1)] + configs = list(itertools.product(batch_size_range, kv_len_range)) + + for head_num_q, head_num_kv in [[32, 32], [64, 8], [40, 8]]: + attn_flashinfer = decode_attention_flashinfer( + dtype, head_num_q, head_num_kv + ).apply + for batch_size, kv_len in configs: + q = torch.randn( + batch_size, head_num_q, head_dim, dtype=dtype, device="cuda" + ) + kv_data = ( + torch.randn( + batch_size * kv_len, + head_num_kv, + head_dim, + dtype=dtype, + device="cuda", + ), + torch.randn( + batch_size * kv_len, + head_num_kv, + head_dim, + dtype=dtype, + device="cuda", + ), + ) + us_cudnn, output_cudnn = decode_attention_cudnn( + q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype + ) + us_sglang, output_sglang = decode_attention_sglang( + q, + kv_data, + batch_size, + kv_len, + head_num_q, + head_num_kv, + head_dim, + num_kv_splits=8, + ) + us_flashinfer, _ = attn_flashinfer( + q, kv_data, batch_size, kv_len, head_num_q, head_num_kv, head_dim, dtype + ) + print( + head_num_q, + " ", + head_num_kv, + " ", + batch_size, + " ", + kv_len, + " ", + us_cudnn, + " ", + us_sglang, + " ", + us_flashinfer, + ) diff --git a/sglang/benchmark/kernels/deepseek/README.md b/sglang/benchmark/kernels/deepseek/README.md new file mode 100644 index 0000000..fba775f --- /dev/null +++ b/sglang/benchmark/kernels/deepseek/README.md @@ -0,0 +1,19 @@ +## DeepSeek kernels benchmark + + +### Prerequisites +- You should install [DeepGemm](https://github.com/deepseek-ai/DeepGEMM) from source before run `benchmark_deepgemm_fp8_gemm.py` and `benchmark_deepgemm_fp8_group_gemm.py`. + +### Benchmark +- `benchmark_deepgemm_fp8_gemm.py` + ```bash + python benchmark_deepgemm_fp8_gemm.py --run_correctness --tp_size 1 + ``` + +- `benchmark_deepgemm_fp8_group_gemm.py` + ```bash + python benchmark_deepgemm_fp8_group_gemm.py --run_correctness --tp_size 1 + ``` + + - You can use the `--run_correctness` parameter to verify all kernels results's correctness. + - You can use the `--tp_size` parameter to benchmark all FP8 w8a8 block-wise matrix multiplications involved in DeepSeek V3/R1 under the current tensor parallelism (TP) setting. This benchmark compares DeepSeek's open-source [DeepGemm](https://github.com/deepseek-ai/DeepGEMM) implementation with SGLang's and VLLM Triton implementation. diff --git a/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py b/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py new file mode 100644 index 0000000..7151f86 --- /dev/null +++ b/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py @@ -0,0 +1,398 @@ +from typing import Tuple + +import deep_gemm +import tilelang +import tilelang.language as T +import torch +import triton +from deep_gemm import ceil_div, get_col_major_tma_aligned_tensor +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + w8a8_block_fp8_matmul as vllm_w8a8_block_fp8_matmul, +) + +from sglang.srt.layers.quantization.fp8_kernel import w8a8_block_fp8_matmul + + +# Adapted from https://github.com/tile-ai/tilelang/blob/a8cfdce92795cb861c9033573534653ee040b5ed/examples/deepseek_deepgemm/example_deepgemm_fp8_2xAcc.py#L1 +def tl_gemm( + M, + N, + K, + in_dtype, + out_dtype, + accum_dtype, +): + assert in_dtype in [ + "e4m3_float8", + ], "Currently only e4m3_float8 is supported" + assert out_dtype in [ + "bfloat16", + "float16", + ], "Currently only bfloat16 and float16 are supported" + + TILE_SIZE = (128, 128, 128) + block_M = TILE_SIZE[0] + block_N = TILE_SIZE[1] + block_K = TILE_SIZE[2] + + A_shape = (M, K) + Scales_A_shape = (M, T.ceildiv(K, block_K)) + B_shape = (N, K) + Scales_B_shape = (T.ceildiv(N, block_N), T.ceildiv(K, block_K)) + A_shared_shape = (block_M, block_K) + B_shared_shape = (block_N, block_K) + C_shared_shape = (block_M, block_N) + + @T.prim_func + def main( + A: T.Buffer(A_shape, in_dtype), + scales_a: T.Buffer(Scales_A_shape, "float32"), + B: T.Buffer(B_shape, in_dtype), + scales_b: T.Buffer(Scales_B_shape, "float32"), + C: T.Buffer((M, N), out_dtype), + ): + with T.Kernel(T.ceildiv(N, block_N), T.ceildiv(M, block_M), threads=128) as ( + bx, + by, + ): + + A_shared = T.alloc_shared(A_shared_shape, in_dtype) + B_shared = T.alloc_shared(B_shared_shape, in_dtype) + C_shared = T.alloc_shared(C_shared_shape, out_dtype) + Scale_C_shared = T.alloc_shared((block_M), "float32") + C_local = T.alloc_fragment(C_shared_shape, accum_dtype) + C_local_accum = T.alloc_fragment(C_shared_shape, accum_dtype) + + # Improve L2 Cache + T.use_swizzle(panel_size=10) + + T.clear(C_local) + T.clear(C_local_accum) + K_iters = T.ceildiv(K, block_K) + for k in T.Pipelined(K_iters, num_stages=4): + # Load A into shared memory + T.copy(A[by * block_M, k * block_K], A_shared) + # Load B into shared memory + T.copy(B[bx * block_N, k * block_K], B_shared) + # Load scale into shared memory + Scale_B = scales_b[bx, k] + for i in T.Parallel(block_M): + Scale_C_shared[i] = scales_a[by * block_M + i, k] * Scale_B + + T.gemm(A_shared, B_shared, C_local, transpose_B=True) + # Promote to enable 2xAcc + for i, j in T.Parallel(block_M, block_N): + C_local_accum[i, j] += C_local[i, j] * Scale_C_shared[i] + T.clear(C_local) + # TMA store + T.copy(C_local_accum, C_shared) + T.copy(C_shared, C[by * block_M, bx * block_N]) + + return main + + +def per_token_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 2 and x.size(1) % 128 == 0 + m, n = x.shape + x_view = x.view(m, -1, 128) + x_amax = x_view.abs().float().amax(dim=2).view(m, -1).clamp(1e-4) + return (x_view * (448.0 / x_amax.unsqueeze(2))).to(torch.float8_e4m3fn).view( + m, n + ), (x_amax / 448.0).view(m, -1) + + +def per_block_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + assert x.dim() == 2 + m, n = x.shape + x_padded = torch.zeros( + (ceil_div(m, 128) * 128, ceil_div(n, 128) * 128), dtype=x.dtype, device=x.device + ) + x_padded[:m, :n] = x + x_view = x_padded.view(-1, 128, x_padded.size(1) // 128, 128) + x_amax = x_view.abs().float().amax(dim=(1, 3), keepdim=True).clamp(1e-4) + x_scaled = (x_view * (448.0 / x_amax)).to(torch.float8_e4m3fn) + return x_scaled.view_as(x_padded)[:m, :n].contiguous(), (x_amax / 448.0).view( + x_view.size(0), x_view.size(2) + ) + + +def fp8_gemm_deepgemm( + x_fp8: torch.Tensor, + x_scale: torch.Tensor, + y_fp8: torch.Tensor, + y_scale: torch.Tensor, + m: int, + n: int, + k: int, +): + """DeepGEMM implementation of FP8 GEMM""" + out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16) + + # Run DeepGEMM kernel + deep_gemm.gemm_fp8_fp8_bf16_nt((x_fp8, x_scale), (y_fp8, y_scale), out) + return out + + +def fp8_gemm_sglang( + x_fp8: torch.Tensor, + x_scale: torch.Tensor, + y_fp8: torch.Tensor, + y_scale: torch.Tensor, + m: int, + n: int, + k: int, +): + """SGLang implementation of FP8 GEMM""" + block_size = [128, 128] # Matches the block size in per_block_cast_to_fp8 + + # Run SGLang kernel + out = w8a8_block_fp8_matmul( + x_fp8, y_fp8, x_scale, y_scale, block_size, torch.bfloat16 + ) + return out + + +def fp8_gemm_vllm( + x_fp8: torch.Tensor, + x_scale: torch.Tensor, + y_fp8: torch.Tensor, + y_scale: torch.Tensor, + m: int, + n: int, + k: int, +): + """vLLM implementation of FP8 GEMM""" + block_size = [128, 128] # Matches the block size in per_block_cast_to_fp8 + + # Run vLLM kernel + out = vllm_w8a8_block_fp8_matmul( + x_fp8, y_fp8, x_scale, y_scale, block_size, torch.bfloat16 + ) + return out + + +def calculate_diff(m: int, n: int, k: int): + x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + + x_fp8, x_scale = per_token_cast_to_fp8(x.clone()) + y_fp8, y_scale = per_block_cast_to_fp8(y.clone()) + x_scale_col_major = get_col_major_tma_aligned_tensor(x_scale.clone()) + + out_deepgemm = fp8_gemm_deepgemm( + x_fp8.clone(), + x_scale_col_major.clone(), + y_fp8.clone(), + y_scale.clone(), + m, + n, + k, + ) + out_sglang = fp8_gemm_sglang( + x_fp8.clone(), x_scale.clone(), y_fp8.clone(), y_scale.clone(), m, n, k + ) + + tilelang_func = tl_gemm(m, n, k, "e4m3_float8", "bfloat16", "float32") + tilelang_kernel = tilelang.compile(tilelang_func, out_idx=[-1]) + out_tilelang = tilelang_kernel( + x_fp8.clone(), x_scale.clone(), y_fp8.clone(), y_scale.clone() + ) + + diff_sglang_deepgemm = torch.abs(out_deepgemm - out_sglang).mean().item() + diff_tilelang_deepgemm = torch.abs(out_deepgemm - out_tilelang).mean().item() + diff_tilelang_sglang = torch.abs(out_tilelang - out_sglang).mean().item() + + print(f"Shape m={m}, n={n}, k={k}:") + print(f"DeepGEMM output: {out_deepgemm[0, 0:5]}") + print(f"SGLang output: {out_sglang[0, 0:5]}") + print(f"TileLang output: {out_tilelang[0, 0:5]}") + print(f"Mean absolute difference (SGLang-DeepGEMM): {diff_sglang_deepgemm}") + print(f"Mean absolute difference (TileLang-DeepGEMM): {diff_tilelang_deepgemm}") + print(f"Mean absolute difference (TileLang-SGLang): {diff_tilelang_sglang}") + + sglang_deepgemm_match = torch.allclose( + out_deepgemm, out_sglang, atol=1e-2, rtol=1e-2 + ) + tilelang_deepgemm_match = torch.allclose( + out_deepgemm, out_tilelang, atol=1e-2, rtol=1e-2 + ) + tilelang_sglang_match = torch.allclose( + out_tilelang, out_sglang, atol=1e-2, rtol=1e-2 + ) + + if sglang_deepgemm_match and tilelang_deepgemm_match and tilelang_sglang_match: + print("✅ All implementations match\n") + else: + print("❌ Some implementations differ:") + print(f" - SGLang vs DeepGEMM: {'✅' if sglang_deepgemm_match else '❌'}") + print(f" - TileLang vs DeepGEMM: {'✅' if tilelang_deepgemm_match else '❌'}") + print(f" - TileLang vs SGLang: {'✅' if tilelang_sglang_match else '❌'}\n") + + +def get_weight_shapes(tp_size): + # cannot TP + total = [ + (512 + 64, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (7168, 16384), + (7168, 18432), + ] + # N can TP + n_tp = [ + (18432 * 2, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (24576, 1536), + (4096, 7168), + ] + # K can TP + k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] + + weight_shapes = [] + for t in total: + weight_shapes.append(t) + for n_t in n_tp: + new_t = (n_t[0] // tp_size, n_t[1]) + weight_shapes.append(new_t) + for k_t in k_tp: + new_t = (k_t[0], k_t[1] // tp_size) + weight_shapes.append(new_t) + + return weight_shapes + + +def create_benchmark_configs(tp_size): + configs = [] + weight_shapes = get_weight_shapes(tp_size) + batch_sizes = [8, 16, 32, 64, 128, 256, 1024, 2048, 4096] + + for n, k in weight_shapes: + for m in batch_sizes: + configs.append((m, n, k, tp_size)) + + return configs + + +def get_benchmark(tp_size): + all_configs = create_benchmark_configs(tp_size) + + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["m", "n", "k", "tp_size"], + x_vals=[list(config) for config in all_configs], + line_arg="provider", + line_vals=["deepgemm", "sglang", "tilelang"], + line_names=["DeepGEMM", "SGLang", "TileLang"], + styles=[("blue", "-"), ("red", "-"), ("green", "-")], + ylabel="ms", + plot_name=f"fp8-gemm-performance-comparison-tp{tp_size}", + args={}, + ) + ) + def benchmark(m, n, k, tp_size, provider): + print(f"Shape (m={m}, n={n}, k={k}, tp={tp_size}), Provider: {provider}") + x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + + # Preprocess data before benchmarking + x_fp8, x_scale = per_token_cast_to_fp8(x) + y_fp8, y_scale = per_block_cast_to_fp8(y) + x_scale_col_major = get_col_major_tma_aligned_tensor(x_scale.clone()) + + quantiles = [0.5, 0.2, 0.8] + + if provider == "deepgemm": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: fp8_gemm_deepgemm( + x_fp8.clone(), + x_scale_col_major.clone(), + y_fp8.clone(), + y_scale.clone(), + m, + n, + k, + ), + quantiles=quantiles, + ) + elif provider == "sglang": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: fp8_gemm_sglang( + x_fp8.clone(), + x_scale.clone(), + y_fp8.clone(), + y_scale.clone(), + m, + n, + k, + ), + quantiles=quantiles, + ) + else: # tilelang + tilelang_func = tl_gemm(m, n, k, "e4m3_float8", "bfloat16", "float32") + tilelang_kernel = tilelang.compile(tilelang_func, out_idx=[-1]) + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: tilelang_kernel( + x_fp8.clone(), + x_scale.clone(), + y_fp8.clone(), + y_scale.clone(), + ), + quantiles=quantiles, + ) + + # Calculate TFLOPS + flops = 2 * m * n * k # multiply-adds + tflops = flops / (ms * 1e-3) / 1e12 + + # Print shape-specific results with TFLOPS + print(f"Time: {ms*1000:.2f} ms, TFLOPS: {tflops:.2f}") + return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms + + return benchmark + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--save_path", + type=str, + default="./configs/benchmark_ops/fp8_gemm/", + help="Path to save fp8 gemm benchmark results", + ) + parser.add_argument( + "--run_correctness", + action="store_true", + default=True, + help="Whether to run correctness test", + ) + parser.add_argument( + "--tp_size", + type=int, + default=1, + help="Tensor parallelism size to benchmark (default: 1)", + ) + args = parser.parse_args() + + # Set random seed for reproducibility + torch.manual_seed(0) + torch.cuda.manual_seed(0) + + # Enable TF32, adapted from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_core.py#L148 + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + + # Run correctness tests on a few examples + if args.run_correctness: + print("Running correctness tests...") + calculate_diff(64, 512, 7168) # Small test + calculate_diff(64, 7168, 16384) # Medium test + calculate_diff(64, 18432, 7168) # Large test + + # Get the benchmark function with the specified tp_size + benchmark = get_benchmark(args.tp_size) + + print(f"Running performance benchmark for TP size = {args.tp_size}...") + benchmark.run(print_data=True, save_path=args.save_path) diff --git a/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py b/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py new file mode 100644 index 0000000..2c3e8df --- /dev/null +++ b/sglang/benchmark/kernels/deepseek/benchmark_deepgemm_fp8_group_gemm.py @@ -0,0 +1,486 @@ +from typing import Tuple + +import deep_gemm +import torch +import triton +import triton.language as tl +from deep_gemm import calc_diff, get_col_major_tma_aligned_tensor + +# Import shared functionality from the regular GEMM benchmark +from sglang.benchmark.kernels.deepseek.benchmark_deepgemm_fp8_gemm import ( + per_block_cast_to_fp8, + per_token_cast_to_fp8, +) + + +def construct_grouped_and_flat_fp8( + x: torch.Tensor, y: torch.Tensor, num_groups: int, is_masked: bool +) -> Tuple[ + Tuple[torch.Tensor, torch.Tensor], # grouped x_fp8 + Tuple[torch.Tensor, torch.Tensor], # grouped y_fp8 + Tuple[torch.Tensor, torch.Tensor], # flat x_fp8 + Tuple[torch.Tensor, torch.Tensor], # flat y_fp8 + torch.Tensor, # output + torch.Tensor, # reference output +]: + # Verify input shapes + m, k = x.shape + n, k_y = y.shape + assert k == k_y, f"Incompatible shapes: x({m}, {k}), y({n}, {k_y})" + assert m % num_groups == 0, f"m({m}) must be divisible by num_groups({num_groups})" + assert m % 4 == 0, f"TMA alignment error: {m}" + + # Reshape inputs for grouped processing + m_per_group = m // num_groups + x_grouped = x.view(num_groups, m_per_group, k) + y_grouped = y.unsqueeze(0).expand(num_groups, n, k) + + # Initialize output tensors + out = torch.empty((num_groups, m_per_group, n), device="cuda", dtype=torch.bfloat16) + ref_out = torch.einsum("gmk,gnk->gmn", x_grouped, y_grouped) + + # Quantize grouped tensors + x_fp8_grouped = ( + torch.empty_like(x_grouped, dtype=torch.float8_e4m3fn), + torch.empty( + (num_groups, m_per_group, k // 128), device="cuda", dtype=torch.float + ), + ) + y_fp8_grouped = ( + torch.empty_like(y_grouped, dtype=torch.float8_e4m3fn), + torch.empty( + (num_groups, (n + 127) // 128, k // 128), device="cuda", dtype=torch.float + ), + ) + for i in range(num_groups): + x_fp8_grouped[0][i], x_fp8_grouped[1][i] = per_token_cast_to_fp8(x_grouped[i]) + y_fp8_grouped[0][i], y_fp8_grouped[1][i] = per_block_cast_to_fp8(y_grouped[i]) + + # Quantize flat tensors + x_fp8_flat = per_token_cast_to_fp8(x) + y_fp8_flat = per_block_cast_to_fp8(y) + + # For non-masked input, merge the group and M dims in output + if not is_masked: + x_fp8_grouped = ( + x_fp8_grouped[0].view(-1, k), + per_token_cast_to_fp8(x_grouped.view(-1, k))[1], + ) + out, ref_out = out.view(-1, n), ref_out.view(-1, n) + + # Transpose earlier for testing + x_fp8_grouped = ( + x_fp8_grouped[0], + get_col_major_tma_aligned_tensor(x_fp8_grouped[1]), + ) + x_fp8_flat = (x_fp8_flat[0], get_col_major_tma_aligned_tensor(x_fp8_flat[1])) + + return x_fp8_grouped, y_fp8_grouped, x_fp8_flat, y_fp8_flat, out, ref_out + + +# Since we don't have a group gemm kernel in SGLang/vLLM, we implemented a +# custom kernel based on the Triton tutorial. +# https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html +@triton.jit +def fp8_gemm_group_triton_kernel( + # Pointers to matrices + a_ptr, + b_ptr, + c_ptr, + # Pointers to scaling factors + a_scale_ptr, + b_scale_ptr, + # Matrix dimensions + M, + N, + K, + # The stride variables represent how much to increase the ptr by when moving by 1 + # element in a particular dimension. + stride_am, + stride_ak, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + # Strides for scaling factors + stride_a_scale_m, + stride_a_scale_k, + stride_b_scale_n, + stride_b_scale_k, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + """Kernel for computing the matmul C = A x B with FP8 inputs and scaling factors. + A has shape (M, K), B has shape (K, N) and C has shape (M, N) + + Note: Block sizes must be multiples of 32 for optimal TMA performance. + """ + # Map program ids to the block of C it should compute + pid_group = tl.program_id(axis=0) # Group ID + pid_n = tl.program_id(axis=1) # N dimension ID + + # Compute the M block ID within this group + group_size_m = min(M - pid_group * GROUP_SIZE_M, GROUP_SIZE_M) + pid_m_within_group = tl.program_id(axis=2) % group_size_m + pid_m = pid_group * GROUP_SIZE_M + pid_m_within_group + + # Create pointers for the first blocks of A and B + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + # Initialize accumulator + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + + # Main loop + for k_block in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + k_offset = k_block * BLOCK_SIZE_K + + # Load the next block of A and B, with masks + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k_offset, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k_offset, other=0.0) + + # Calculate indices for scaling factors for this K block + a_scale_ptrs = a_scale_ptr + ( + offs_am * stride_a_scale_m + k_block * stride_a_scale_k + ) + b_scale_ptrs = b_scale_ptr + ( + pid_n * stride_b_scale_n + k_block * stride_b_scale_k + ) + + # Perform matrix multiplication in FP8 + res = tl.dot(a, b) + + # Load scaling factors for the current block + a_scale = tl.load(a_scale_ptrs)[:, None] # [BLOCK_SIZE_M, 1] + b_scale = tl.load(b_scale_ptrs) + + # Apply scaling factors to the accumulated result + accumulator += res * a_scale * b_scale + + # Advance pointers + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + # Convert to bfloat16 for output + c = accumulator.to(tl.bfloat16) + + # Write back the result + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + +def fp8_gemm_group_triton(a_tuple, b_tuple, c, num_groups): + """ + Perform matrix multiplication with FP8 inputs and proper scaling. + + Args: + a_tuple: Tuple of (quantized_tensor, scale_factors) for input A + b_tuple: Tuple of (quantized_tensor, scale_factors) for input B + c: Output tensor in BF16 format + num_groups: Number of groups for grouped GEMM + + Returns: + Result tensor in BF16 format + """ + # Unpack the tuples + a, a_scale = a_tuple + b, b_scale = b_tuple + + M, K = a.shape + _, N = b.shape + + # Configure block sizes - must be multiples of 32 for TMA alignment + BLOCK_SIZE_M = 128 + BLOCK_SIZE_N = 128 + BLOCK_SIZE_K = 128 + + # Calculate grid dimensions + num_pid_m = triton.cdiv(M, BLOCK_SIZE_M) + num_pid_n = triton.cdiv(N, BLOCK_SIZE_N) + num_groups_grid = triton.cdiv(num_pid_m, num_groups) + + # 3D grid launch - (group, n_blocks, m_blocks_per_group) + grid = (num_groups_grid, num_pid_n, min(num_groups, num_pid_m)) + + fp8_gemm_group_triton_kernel[grid]( + a, + b, + c, + a_scale, + b_scale, + M, + N, + K, + a.stride(0), + a.stride(1), + b.stride(0), + b.stride(1), + c.stride(0), + c.stride(1), + a_scale.stride(0), + 1, # Stride in the K dimension may be 1 + b_scale.stride(0), + 1 if b_scale.dim() > 1 else 0, + BLOCK_SIZE_M=BLOCK_SIZE_M, + BLOCK_SIZE_N=BLOCK_SIZE_N, + BLOCK_SIZE_K=BLOCK_SIZE_K, + GROUP_SIZE_M=num_groups, + ) + + return c + + +def fp8_gemm_group_deepgemm(x_fp8_grouped, y_fp8_grouped, out, m_indices): + deep_gemm.m_grouped_gemm_fp8_fp8_bf16_nt_contiguous( + x_fp8_grouped, + y_fp8_grouped, + out, + m_indices, + ) + return out + + +def calculate_diff(m: int, n: int, k: int, num_groups: int): + print(f"Shape (m={m}, n={n}, k={k}") + x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + x_fp8_grouped, y_fp8_grouped, x_fp8_flat, y_fp8_flat, out, out_torch = ( + construct_grouped_and_flat_fp8(x, y, num_groups, is_masked=False) + ) + m_per_group = m // num_groups + out_deepgemm = out.clone() + m_indices = torch.arange(0, num_groups, device="cuda", dtype=torch.int) + m_indices = ( + m_indices.unsqueeze(-1).expand(num_groups, m_per_group).contiguous().view(-1) + ) + + fp8_gemm_group_deepgemm( + x_fp8_grouped, + y_fp8_grouped, + out_deepgemm, + m_indices, + ) + torch.cuda.synchronize() + + # Prepare inputs for Triton + a, a_scale = x_fp8_flat + b, b_scale = y_fp8_flat + b = b.T.contiguous() + # Ensure scales are in the right format and contiguous + a_scale, b_scale = a_scale.contiguous(), b_scale.contiguous() + M, _ = a.shape + _, N = b.shape + c = torch.empty((M, N), device=a.device, dtype=torch.bfloat16) + out_triton = fp8_gemm_group_triton((a, a_scale), (b, b_scale), c, num_groups) + torch.cuda.synchronize() + + diff_torch_deepgemm = torch.abs(out_torch - out_deepgemm).mean().item() + diff_torch_triton = torch.abs(out_torch - out_triton).mean().item() + diff_deepgemm_triton = torch.abs(out_deepgemm - out_triton).mean().item() + + print(f"Shape m={m}, n={n}, k={k}:") + print(f"Torch output: {out_torch[0, 0:5]}") + print(f"DeepGEMM output: {out_deepgemm[0, 0:5]}") + print(f"Triton output: {out_triton[0, 0:5]}") + print(f"Mean absolute difference (Torch-DeepGEMM): {diff_torch_deepgemm}") + print(f"Mean absolute difference (Torch-Triton): {diff_torch_triton}") + print(f"Mean absolute difference (DeepGEMM-Triton): {diff_deepgemm_triton}") + + deepgemm_torch_diff = calc_diff(out_deepgemm, out_torch) + triton_torch_diff = calc_diff(out_triton, out_torch) + deepgemm_triton_diff = calc_diff(out_deepgemm, out_triton) + + DIFF_THRESHOLD = 0.001 + all_match = ( + deepgemm_torch_diff < DIFF_THRESHOLD + and triton_torch_diff < DIFF_THRESHOLD + and deepgemm_triton_diff < DIFF_THRESHOLD + ) + if all_match: + print("✅ All implementations match\n") + else: + print("❌ Some implementations differ:") + print( + f" - Torch vs DeepGEMM: {'✅' if deepgemm_torch_diff < DIFF_THRESHOLD else '❌'}" + f" - Torch vs Triton: {'✅' if triton_torch_diff < DIFF_THRESHOLD else '❌'}" + f" - DeepGEMM vs Triton: {'✅' if deepgemm_triton_diff < DIFF_THRESHOLD else '❌'}" + ) + + +def get_weight_shapes(tp_size): + # cannot TP + total = [ + (512 + 64, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (7168, 16384), + (7168, 18432), + ] + # N can TP + n_tp = [ + (18432 * 2, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (24576, 1536), + (4096, 7168), + ] + # K can TP + k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] + + weight_shapes = [] + for t in total: + weight_shapes.append(t) + for n_t in n_tp: + new_t = (n_t[0] // tp_size, n_t[1]) + weight_shapes.append(new_t) + for k_t in k_tp: + new_t = (k_t[0], k_t[1] // tp_size) + weight_shapes.append(new_t) + + return weight_shapes + + +def create_benchmark_configs(tp_size): + configs = [] + weight_shapes = get_weight_shapes(tp_size) + batch_sizes = [2048, 4096] + group_sizes = [4, 8] + for n, k in weight_shapes: + for m in batch_sizes: + for num_groups in group_sizes: + configs.append((m, n, k, num_groups, tp_size)) + + return configs + + +def get_benchmark(tp_size): + all_configs = create_benchmark_configs(tp_size) + + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["m", "n", "k", "num_groups", "tp_size"], + x_vals=[config for config in all_configs], + line_arg="provider", + line_vals=["deepgemm", "triton"], + line_names=["DeepGEMM", "Triton"], + styles=[("blue", "-"), ("red", "-")], + ylabel="ms", + plot_name=f"fp8-group-gemm-performance-comparison-tp{tp_size}", + args={}, + ) + ) + def benchmark(m, n, k, num_groups, tp_size, provider): + print( + f"Shape (m={m}, n={n}, k={k}, tp={tp_size}, num_groups={num_groups}, Provider: {provider}" + ) + x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16) + y = torch.randn((n, k), device="cuda", dtype=torch.bfloat16) + x_fp8_grouped, y_fp8_grouped, x_fp8_flat, y_fp8_flat, out, out_torch = ( + construct_grouped_and_flat_fp8(x, y, num_groups, is_masked=False) + ) + m_per_group = m // num_groups + m_indices = torch.arange(0, num_groups, device="cuda", dtype=torch.int) + m_indices = ( + m_indices.unsqueeze(-1) + .expand(num_groups, m_per_group) + .contiguous() + .view(-1) + ) + + quantiles = [0.5, 0.2, 0.8] + + if provider == "deepgemm": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: fp8_gemm_group_deepgemm( + x_fp8_grouped, + y_fp8_grouped, + out, + m_indices, + ), + quantiles=quantiles, + ) + elif provider == "triton": + # Prepare inputs for Triton + # We did it outside of the lambda function to make it fair comparison like deepgemm + a, a_scale = x_fp8_flat + b, b_scale = y_fp8_flat + b = b.T.contiguous() + # Ensure scales are in the right format and contiguous + a_scale, b_scale = a_scale.contiguous(), b_scale.contiguous() + M, _ = a.shape + _, N = b.shape + c = torch.empty((M, N), device=a.device, dtype=torch.bfloat16) + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: fp8_gemm_group_triton( + (a, a_scale), + (b, b_scale), + c, + num_groups, + ), + quantiles=quantiles, + ) + + # Calculate TFLOPS + flops = 2 * m * n * k # multiply-adds + tflops = flops / (ms * 1e-3) / 1e12 + + print(f"Time: {ms*1000:.2f} ms, TFLOPS: {tflops:.2f}") + return ms * 1000, max_ms * 1000, min_ms * 1000 # convert to ms + + return benchmark + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--save_path", + type=str, + default="./configs/benchmark_ops/fp8_group_gemm/", + help="Path to save deepgemm fp8 group gemm benchmark results", + ) + parser.add_argument( + "--run_correctness", + action="store_true", + help="Whether to run correctness test", + ) + parser.add_argument( + "--tp_size", + type=int, + default=1, + help="Tensor parallelism size to benchmark (default: 1)", + ) + args = parser.parse_args() + + # Set random seed for reproducibility + torch.manual_seed(0) + torch.cuda.manual_seed(0) + + # Enable TF32, adapted from https://github.com/deepseek-ai/DeepGEMM/blob/main/tests/test_core.py#L148 + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + + # Run correctness tests on a few examples + if args.run_correctness: + print("Running correctness tests...") + calculate_diff(8192, 7168, 4096, 4) + calculate_diff(8192, 2048, 7168, 4) + calculate_diff(4096, 7168, 4096, 8) + calculate_diff(4096, 2048, 7168, 8) + calculate_diff(4096, 576, 7168, 8) + + # Get the benchmark function with the specified tp_size + benchmark = get_benchmark(args.tp_size) + + print(f"Running performance benchmark for TP size = {args.tp_size}...") + benchmark.run(print_data=True, save_path=args.save_path) diff --git a/sglang/benchmark/kernels/fused_moe_triton/README.md b/sglang/benchmark/kernels/fused_moe_triton/README.md new file mode 100644 index 0000000..2a3e37f --- /dev/null +++ b/sglang/benchmark/kernels/fused_moe_triton/README.md @@ -0,0 +1,49 @@ +## Benchmark Kernels + +This directory contains benchmarking tools for MoE (Mixture of Experts) kernels. + +### Tuning Tool + +- `tuning_fused_moe_triton.py`: A tool for tuning the `fused_moe_triton` kernel. Adapted from [vllm's benchmark_moe.py](https://github.com/vllm-project/vllm/blob/main/benchmarks/kernels/benchmark_moe.py), with added support for various model architectures. + +Example usage: +```bash +# Tune Qwen2-57B with FP8 and TP=4 +python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ + --model Qwen/Qwen2-57B-A14B-Instruct \ + --tp-size 4 \ + --dtype fp8_w8a8 \ + --tune + +# Tune Mixtral-8x7B with default settings +python benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py \ + --model mistralai/Mixtral-8x7B-Instruct-v0.1 \ + --tune +``` + +After tuning, a configuration file (e.g., `E=64,N=640,device_name=NVIDIA_GeForce_RTX_4090,dtype=fp8_w8a8.json`) will be generated in the current directory. You can move this file to `sglang/srt/layers/fused_moe_triton/configs/` to use it in `sglang`. + +### Performance Comparison Tool + +- `benchmark_vllm_vs_sglang_fused_moe_triton.py`: A tool for comparing the performance of fused MoE kernels between vllm and sglang implementations. Supports various model architectures and data types. + +Example usage: +```bash +# Compare with default settings (Mixtral model) +python benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py + +# Compare with FP8 mode for Qwen2-57B +python benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py \ + --model Qwen/Qwen2-57B-A14B-Instruct \ + --use-fp8 + +# Compare with custom TP size +python benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py \ + --tp-size 4 +``` + +The benchmark results will be saved as plots and data files in the specified output directory (default: `./configs/benchmark_ops/vllm_sglang_fused_moe/`). + +- `benchmark_torch_compile_fused_moe.py`: A tool for benchmarking the performance of the fused MoE kernel with `torch.compile` and original fused MoE kernel. + +Usage is the same as `benchmark_vllm_vs_sglang_fused_moe_triton.py`. diff --git a/sglang/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py b/sglang/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py new file mode 100644 index 0000000..8206684 --- /dev/null +++ b/sglang/benchmark/kernels/fused_moe_triton/benchmark_torch_compile_fused_moe.py @@ -0,0 +1,294 @@ +import argparse + +import torch +import triton +from torch.nn import functional as F +from transformers import AutoConfig + +from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( + fused_moe as fused_moe_triton, +) +from sglang.srt.model_executor.cuda_graph_runner import set_torch_compile_config + + +def get_model_config(model_name: str, tp_size: int): + """Get model configuration parameters""" + config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + + if config.architectures[0] == "DbrxForCausalLM": + E = config.ffn_config.moe_num_experts + topk = config.ffn_config.moe_top_k + intermediate_size = config.ffn_config.ffn_hidden_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] == "JambaForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] == "Qwen2MoeForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] in ["DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM"]: + E = config.n_routed_experts + topk = config.num_experts_per_tok + intermediate_size = config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] in [ + "Grok1ForCausalLM", + "Grok1ImgGen", + "Grok1AForCausalLM", + ]: + E = config.num_local_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + else: + # Default: Mixtral + E = config.num_local_experts + topk = config.num_experts_per_tok + intermediate_size = config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + + shape_configs = { + "num_experts": E, + "topk": topk, + "hidden_size": config.hidden_size, + "shard_intermediate_size": shard_intermediate_size, + "dtype": config.torch_dtype, + } + print(f"{shape_configs=}") + return shape_configs + + +def fused_topk_native( + hidden_states: torch.Tensor, + gating_output: torch.Tensor, + topk: int, + renormalize: bool, +): + assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch" + M, _ = hidden_states.shape + topk_weights = torch.empty( + M, topk, dtype=torch.float32, device=hidden_states.device + ) + topk_ids = torch.empty(M, topk, dtype=torch.int32, device=hidden_states.device) + topk_weights = F.softmax(gating_output.float(), dim=-1) + topk_weights, topk_ids = torch.topk(topk_weights, topk, dim=-1) + if renormalize: + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + return topk_weights, topk_ids + + +@torch.compile(dynamic=False) +def fused_moe_torch( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=False, + w1_scale=None, + w2_scale=None, + a1_scale=None, + a2_scale=None, +) -> torch.Tensor: + assert not use_fp8_w8a8, "Not supported" + + topk_weights, topk_ids = fused_topk_native( + hidden_states=x, + gating_output=input_gating, + topk=topk, + renormalize=True, + ) + w13_weights = w1[topk_ids] + w1_weights, w3_weights = torch.chunk(w13_weights, 2, dim=2) + w2_weights = w2[topk_ids] + x1 = torch.einsum("ti,taoi -> tao", x, w1_weights) + x1 = F.silu(x1) + x3 = torch.einsum("ti, taoi -> tao", x, w3_weights) + expert_outs = torch.einsum("tao, taio -> tai", (x1 * x3), w2_weights) + return torch.einsum("tai,ta -> ti", expert_outs, topk_weights.to(expert_outs.dtype)) + + +def fused_moe_torch_compile( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=False, + w1_scale=None, + w2_scale=None, + a1_scale=None, + a2_scale=None, +): + return fused_moe_torch( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + ) + + +def fused_moe_sglang_api( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=False, + w1_scale=None, + w2_scale=None, + a1_scale=None, + a2_scale=None, +): + return fused_moe_triton( + x, + w1, + w2, + input_gating, + topk, + renormalize=True, + inplace=True, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + ) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size"], + x_vals=list(range(1, 5)), + line_arg="provider", + line_vals=[ + "fused_moe_triton", + "fused_moe_torch_compile", + ], + line_names=[ + "fused_moe_triton", + "fused_moe_torch_compile", + ], + styles=[ + ("blue", "-"), + ("green", "-"), + ], + ylabel="Time (ms)", + plot_name="fused-moe-performance", + args={}, + ) +) +def benchmark(batch_size, provider, model_config, use_fp8=False): + print(f"benchmark {provider} with batch_size={batch_size}") + torch.set_default_device("cuda") + torch.cuda.manual_seed_all(0) + set_torch_compile_config() + + num_tokens = batch_size + num_experts = model_config["num_experts"] + hidden_size = model_config["hidden_size"] + shard_intermediate_size = model_config["shard_intermediate_size"] + topk = model_config["topk"] + dtype = model_config["dtype"] + + x = torch.randn(num_tokens, hidden_size, dtype=dtype) + + if use_fp8: + init_dtype = dtype + w1 = torch.randn( + num_experts, shard_intermediate_size, hidden_size, dtype=init_dtype + ) + w2 = torch.randn( + num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype + ) + w1 = w1.to(torch.float8_e4m3fn) + w2 = w2.to(torch.float8_e4m3fn) + w1_scale = torch.randn(num_experts, dtype=torch.float32) + w2_scale = torch.randn(num_experts, dtype=torch.float32) + a1_scale = torch.randn(1, dtype=torch.float32) + a2_scale = torch.randn(1, dtype=torch.float32) + else: + w1 = torch.randn(num_experts, shard_intermediate_size, hidden_size, dtype=dtype) + w2 = torch.randn( + num_experts, hidden_size, shard_intermediate_size // 2, dtype=dtype + ) + w1_scale = w2_scale = a1_scale = a2_scale = None + + input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32) + + # Warmup + api_func = ( + fused_moe_torch_compile + if provider == "fused_moe_torch_compile" + else fused_moe_sglang_api + ) + for _ in range(10): + y = api_func( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=use_fp8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + ) + torch.cuda.synchronize() + + quantiles = [0.5, 0.2, 0.8] + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: api_func( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=use_fp8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + )[0], + quantiles=quantiles, + ) + return ms, min_ms, max_ms + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1" + ) + parser.add_argument("--tp-size", type=int, default=2) + parser.add_argument("--use-fp8", action="store_true") + parser.add_argument( + "--save-path", + type=str, + default="./configs/benchmark_ops/fused_moe_torch_compile/", + ) + args = parser.parse_args() + + model_config = get_model_config(args.model, args.tp_size) + benchmark.run( + show_plots=True, + print_data=True, + save_path=args.save_path, + model_config=model_config, + use_fp8=args.use_fp8, + ) + + +if __name__ == "__main__": + main() diff --git a/sglang/benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py b/sglang/benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py new file mode 100644 index 0000000..0760116 --- /dev/null +++ b/sglang/benchmark/kernels/fused_moe_triton/benchmark_vllm_vs_sglang_fused_moe_triton.py @@ -0,0 +1,299 @@ +import argparse + +import torch +import triton +import vllm +from transformers import AutoConfig +from vllm.model_executor.layers.fused_moe.fused_moe import fused_moe as fused_moe_vllm + +from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( + fused_moe as fused_moe_sglang, +) + + +def get_model_config(model_name: str, tp_size: int): + """Get model configuration parameters""" + config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + + if config.architectures[0] == "DbrxForCausalLM": + E = config.ffn_config.moe_num_experts + topk = config.ffn_config.moe_top_k + intermediate_size = config.ffn_config.ffn_hidden_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] == "JambaForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] == "Qwen2MoeForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] in ["DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM"]: + E = config.n_routed_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + elif config.architectures[0] in [ + "Grok1ForCausalLM", + "Grok1ImgGen", + "Grok1AForCausalLM", + ]: + E = config.num_local_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + else: + # Default: Mixtral + E = config.num_local_experts + topk = config.num_experts_per_tok + intermediate_size = config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // tp_size + + vllm_version_num = ( + vllm.__version_tuple__[0] * 100 + + vllm.__version_tuple__[1] * 10 + + vllm.__version_tuple__[2] + ) + block_shape = None + if ( + hasattr(config, "quantization_config") + and "weight_block_size" in config.quantization_config + ): + block_shape = config.quantization_config["weight_block_size"] + assert len(block_shape) == 2 + assert ( + vllm_version_num >= 66 + ), "Block-wise quantized fp8 fused_moe is only supported for VLLM>=0.6.6.post1" + + shape_configs = { + "num_experts": E, + "topk": topk, + "hidden_size": config.hidden_size, + "shard_intermediate_size": shard_intermediate_size, + "dtype": config.torch_dtype, + "block_shape": block_shape, + } + print(f"{shape_configs=}") + return shape_configs + + +def fused_moe_vllm_api( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=False, + w1_scale=None, + w2_scale=None, + a1_scale=None, + a2_scale=None, + block_shape=None, +): + if block_shape is not None: + return fused_moe_vllm( + x, + w1, + w2, + input_gating, + topk, + renormalize=True, + inplace=True, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + ) + else: + return fused_moe_vllm( + x, + w1, + w2, + input_gating, + topk, + renormalize=True, + inplace=True, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + ) + + +def fused_moe_sglang_api( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=False, + w1_scale=None, + w2_scale=None, + a1_scale=None, + a2_scale=None, + block_shape=None, +): + return fused_moe_sglang( + x, + w1, + w2, + input_gating, + topk, + renormalize=True, + inplace=True, + use_fp8_w8a8=use_fp8_w8a8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + ) + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size"], + x_vals=list(range(1, 513)), + line_arg="provider", + line_vals=[ + "vllm_fused_moe_triton", + "sglang_fused_moe_triton", + ], + line_names=[ + "vllm_fused_moe_triton", + "sglang_fused_moe_triton", + ], + styles=[ + ("blue", "-"), + ("green", "-"), + ], + ylabel="Time (ms)", + plot_name="fused-moe-performance", + args={}, + ) +) +def benchmark(batch_size, provider, model_config, use_fp8=False): + print(f"benchmark {provider} with batch_size={batch_size}") + torch.set_default_device("cuda") + torch.cuda.manual_seed_all(0) + + num_tokens = batch_size + num_experts = model_config["num_experts"] + hidden_size = model_config["hidden_size"] + shard_intermediate_size = model_config["shard_intermediate_size"] + topk = model_config["topk"] + dtype = model_config["dtype"] + block_shape = getattr(model_config, "block_shape", None) + + x = torch.randn(num_tokens, hidden_size, dtype=dtype) + w1_scale = w2_scale = a1_scale = a2_scale = None + + if use_fp8: + init_dtype = dtype + w1 = torch.randn( + num_experts, shard_intermediate_size, hidden_size, dtype=init_dtype + ) + w2 = torch.randn( + num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype + ) + w1 = w1.to(torch.float8_e4m3fn) + w2 = w2.to(torch.float8_e4m3fn) + + if block_shape is None: + w1_scale = torch.randn(num_experts, dtype=torch.float32) + w2_scale = torch.randn(num_experts, dtype=torch.float32) + a1_scale = torch.randn(1, dtype=torch.float32) + a2_scale = torch.randn(1, dtype=torch.float32) + else: + block_n, block_k = block_shape[0], block_shape[1] + n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n + n_tiles_w2 = (hidden_size + block_n - 1) // block_n + k_tiles_w1 = (hidden_size + block_k - 1) // block_k + k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k + w1_scale = torch.rand( + (num_experts, n_tiles_w1, k_tiles_w1), dtype=torch.float32 + ) + w2_scale = torch.rand( + (num_experts, n_tiles_w2, k_tiles_w2), dtype=torch.float32 + ) + else: + w1 = torch.randn(num_experts, shard_intermediate_size, hidden_size, dtype=dtype) + w2 = torch.randn( + num_experts, hidden_size, shard_intermediate_size // 2, dtype=dtype + ) + + input_gating = torch.randn(num_tokens, num_experts, dtype=torch.float32) + + # Warmup + api_func = ( + fused_moe_vllm_api + if provider == "vllm_fused_moe_triton" + else fused_moe_sglang_api + ) + for _ in range(10): + y = api_func( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=use_fp8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + ) + torch.cuda.synchronize() + + quantiles = [0.5, 0.2, 0.8] + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: api_func( + x, + w1, + w2, + input_gating, + topk, + use_fp8_w8a8=use_fp8, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + )[0], + quantiles=quantiles, + ) + return ms, min_ms, max_ms + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1" + ) + parser.add_argument("--tp-size", type=int, default=2) + parser.add_argument("--use-fp8", action="store_true") + parser.add_argument( + "--save-path", + type=str, + default="./configs/benchmark_ops/vllm_sglang_fused_moe/", + ) + args = parser.parse_args() + + model_config = get_model_config(args.model, args.tp_size) + benchmark.run( + show_plots=True, + print_data=True, + save_path=args.save_path, + model_config=model_config, + use_fp8=args.use_fp8, + ) + + +if __name__ == "__main__": + main() diff --git a/sglang/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py b/sglang/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py new file mode 100644 index 0000000..89309ac --- /dev/null +++ b/sglang/benchmark/kernels/fused_moe_triton/tuning_fused_moe_triton.py @@ -0,0 +1,564 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/main/benchmarks/kernels/benchmark_moe.py +import argparse +import json +import time +from datetime import datetime +from typing import Any, Dict, List, Tuple, TypedDict + +import ray +import torch +import triton +from ray.experimental.tqdm_ray import tqdm +from transformers import AutoConfig + +from sglang.srt.layers.moe.fused_moe_triton.fused_moe import ( + fused_moe, + get_config_dtype_str, + get_config_file_name, + get_default_config, + get_moe_configs, +) +from sglang.srt.utils import is_hip + +_is_hip_ = is_hip() + + +class BenchmarkConfig(TypedDict): + BLOCK_SIZE_M: int + BLOCK_SIZE_N: int + BLOCK_SIZE_K: int + GROUP_SIZE_M: int + num_warps: int + num_stages: int + + +def benchmark_config( + config: BenchmarkConfig, + num_tokens: int, + num_experts: int, + shard_intermediate_size: int, + hidden_size: int, + topk: int, + dtype: torch.dtype, + use_fp8_w8a8: bool, + use_int8_w8a8: bool, + use_int8_w8a16: bool, + block_shape: List[int] = None, + num_iters: int = 100, +) -> float: + init_dtype = torch.float16 if use_fp8_w8a8 else dtype + x = torch.randn(num_tokens, hidden_size, dtype=dtype) + if use_int8_w8a16 or use_int8_w8a8: + w1 = torch.randint( + -127, + 127, + ( + num_experts, + shard_intermediate_size, + hidden_size, + ), + dtype=torch.int8, + ) + w2 = torch.randint( + -127, + 127, + ( + num_experts, + hidden_size, + shard_intermediate_size // 2, + ), + dtype=torch.int8, + ) + else: + w1 = torch.randn( + num_experts, shard_intermediate_size, hidden_size, dtype=init_dtype + ) + w2 = torch.randn( + num_experts, hidden_size, shard_intermediate_size // 2, dtype=init_dtype + ) + gating_output = torch.randn(num_iters, num_tokens, num_experts, dtype=torch.float32) + + w1_scale = None + w2_scale = None + a1_scale = None + a2_scale = None + if use_int8_w8a16: + w1_scale = torch.randn( + (num_experts, 2 * shard_intermediate_size), dtype=torch.float32 + ) + w2_scale = torch.randn((hidden_size, num_experts), dtype=torch.float32) + if use_fp8_w8a8 or use_int8_w8a8: + if use_int8_w8a8 and block_shape is None: + w1_scale = torch.randn( + num_experts, shard_intermediate_size, dtype=torch.float32 + ) + w2_scale = torch.randn(num_experts, hidden_size, dtype=torch.float32) + elif block_shape is None: + w1_scale = torch.randn(num_experts, dtype=torch.float32) + w2_scale = torch.randn(num_experts, dtype=torch.float32) + a1_scale = torch.randn(1, dtype=torch.float32) + a2_scale = torch.randn(1, dtype=torch.float32) + else: + block_n, block_k = block_shape[0], block_shape[1] + n_tiles_w1 = (shard_intermediate_size + block_n - 1) // block_n + n_tiles_w2 = (hidden_size + block_n - 1) // block_n + k_tiles_w1 = (hidden_size + block_k - 1) // block_k + k_tiles_w2 = (shard_intermediate_size // 2 + block_k - 1) // block_k + w1_scale = torch.rand( + (num_experts, n_tiles_w1, k_tiles_w1), dtype=torch.float32 + ) + w2_scale = torch.rand( + (num_experts, n_tiles_w2, k_tiles_w2), dtype=torch.float32 + ) + + if use_fp8_w8a8: + w1 = w1.to(torch.float8_e4m3fnuz if _is_hip_ else torch.float8_e4m3fn) + w2 = w2.to(torch.float8_e4m3fnuz if _is_hip_ else torch.float8_e4m3fn) + + input_gating = torch.empty(num_tokens, num_experts, dtype=torch.float32) + + def prepare(i: int): + input_gating.copy_(gating_output[i]) + + def run(): + from sglang.srt.layers.moe.fused_moe_triton import override_config + + with override_config(config): + fused_moe( + x, + w1, + w2, + input_gating, + topk, + renormalize=True, + inplace=True, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=use_int8_w8a8, + use_int8_w8a16=use_int8_w8a16, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + block_shape=block_shape, + ) + + # JIT compilation & warmup + run() + torch.cuda.synchronize() + + # Capture 10 invocations with CUDA graph + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + for _ in range(10): + run() + torch.cuda.synchronize() + + # Warmup + for _ in range(5): + graph.replay() + torch.cuda.synchronize() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + + latencies: List[float] = [] + for i in range(num_iters): + prepare(i) + torch.cuda.synchronize() + + start_event.record() + graph.replay() + end_event.record() + end_event.synchronize() + latencies.append(start_event.elapsed_time(end_event)) + avg = sum(latencies) / (num_iters * 10) * 1000 # us + graph.reset() + return avg + + +def get_rocm_configs_compute_bound() -> List[Dict[str, int]]: + configs: List[BenchmarkConfig] = [] + waves_per_eu_range = 0 + for num_stages in [2]: + for block_m in [32, 64, 128, 256]: + for block_k in [32, 64, 128, 256]: + for block_n in [16, 32, 64, 128, 256]: + for num_warps in [1, 2, 4, 8]: + for group_size in [1, 4, 8, 16, 32]: + configs.append( + { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": group_size, + "num_warps": num_warps, + "num_stages": num_stages, + "waves_per_eu": waves_per_eu_range, + } + ) + return configs + + +def get_configs_compute_bound() -> List[Dict[str, int]]: + # Reduced search space for faster tuning. + # TODO(woosuk): Increase the search space and use a performance model to + # prune the search space. + configs: List[BenchmarkConfig] = [] + if _is_hip_: + configs = get_rocm_configs_compute_bound() + else: + for num_stages in [2, 3, 4, 5]: + for block_m in [16, 32, 64, 128, 256]: + for block_k in [64, 128, 256]: + for block_n in [32, 64, 128, 256]: + for num_warps in [4, 8]: + for group_size in [1, 16, 32, 64]: + configs.append( + { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": group_size, + "num_warps": num_warps, + "num_stages": num_stages, + } + ) + return configs + + +@ray.remote(num_gpus=1) +class BenchmarkWorker: + + def __init__(self, seed: int) -> None: + torch.set_default_device("cuda") + torch.cuda.manual_seed_all(0) + self.seed = seed + + def benchmark( + self, + num_tokens: int, + num_experts: int, + shard_intermediate_size: int, + hidden_size: int, + topk: int, + dtype: torch.dtype, + use_fp8_w8a8: bool, + use_int8_w8a8: bool, + use_int8_w8a16: bool, + block_shape: List[int], + ) -> Tuple[Dict[str, int], float]: + torch.cuda.manual_seed_all(0) + dtype_str = get_config_dtype_str( + dtype, use_int8_w8a16=use_int8_w8a16, use_fp8_w8a8=use_fp8_w8a8 + ) + # NOTE(woosuk): The current naming convention uses w2.shape[2], which + # is the intermediate size after silu_and_mul. + block_n = block_shape[0] if block_shape else 0 + block_k = block_shape[1] if block_shape else 0 + op_config = get_moe_configs( + num_experts, shard_intermediate_size // 2, dtype_str, block_n, block_k + ) + if op_config is None: + config = get_default_config( + num_tokens, + num_experts, + shard_intermediate_size, + hidden_size, + topk, + dtype_str, + False, + block_shape, + ) + else: + config = op_config[min(op_config.keys(), key=lambda x: abs(x - num_tokens))] + kernel_time = benchmark_config( + config, + num_tokens, + num_experts, + shard_intermediate_size, + hidden_size, + topk, + dtype, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + block_shape, + ) + return config, kernel_time + + def tune( + self, + num_tokens: int, + num_experts: int, + shard_intermediate_size: int, + hidden_size: int, + topk: int, + dtype: torch.dtype, + use_fp8_w8a8: bool, + use_int8_w8a8: bool, + use_int8_w8a16: bool, + block_shape: List[int], + search_space: List[Dict[str, int]], + ) -> Dict[str, int]: + best_config = None + best_time = float("inf") + for config in tqdm(search_space): + try: + kernel_time = benchmark_config( + config, + num_tokens, + num_experts, + shard_intermediate_size, + hidden_size, + topk, + dtype, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + block_shape, + num_iters=10, + ) + except triton.runtime.autotuner.OutOfResources: + # Some configurations may be invalid and fail to compile. + continue + + if kernel_time < best_time: + best_time = kernel_time + best_config = config + now = datetime.now() + print(f"{now.ctime()}] Completed tuning for batch_size={num_tokens}") + assert best_config is not None + return best_config + + +def sort_config(config: BenchmarkConfig) -> BenchmarkConfig: + return { + "BLOCK_SIZE_M": config["BLOCK_SIZE_M"], + "BLOCK_SIZE_N": config["BLOCK_SIZE_N"], + "BLOCK_SIZE_K": config["BLOCK_SIZE_K"], + "GROUP_SIZE_M": config["GROUP_SIZE_M"], + "num_warps": config["num_warps"], + "num_stages": config["num_stages"], + **( + {"waves_per_eu": config["waves_per_eu"]} if "waves_per_eu" in config else {} + ), + } + + +def save_configs( + configs: Dict[int, BenchmarkConfig], + num_experts: int, + shard_intermediate_size: int, + hidden_size: int, + topk: int, + dtype: torch.dtype, + use_fp8_w8a8: bool, + use_int8_w8a8: bool, + use_int8_w8a16: bool, + block_shape: List[int], +) -> None: + dtype_str = get_config_dtype_str( + dtype, + use_int8_w8a16=use_int8_w8a16, + use_fp8_w8a8=use_fp8_w8a8, + use_int8_w8a8=use_int8_w8a8, + ) + + # NOTE(woosuk): The current naming convention uses w2.shape[2], which + # is the intermediate size after silu_and_mul. + filename = get_config_file_name( + num_experts, + shard_intermediate_size // 2, + dtype_str, + block_shape, + ) + + print(f"Writing best config to {filename}...") + with open(filename, "w") as f: + json.dump(configs, f, indent=4) + f.write("\n") + + +def main(args: argparse.Namespace): + print(args) + + config = AutoConfig.from_pretrained(args.model, trust_remote_code=True) + if config.architectures[0] == "DbrxForCausalLM": + E = config.ffn_config.moe_num_experts + topk = config.ffn_config.moe_top_k + intermediate_size = config.ffn_config.ffn_hidden_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size + elif config.architectures[0] == "JambaForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size + elif config.architectures[0] == "Qwen2MoeForCausalLM": + E = config.num_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size + elif config.architectures[0] in ["DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM"]: + E = config.n_routed_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size + elif config.architectures[0] in [ + "Grok1ForCausalLM", + "Grok1ImgGen", + "Grok1AForCausalLM", + ]: + E = config.num_local_experts + topk = config.num_experts_per_tok + intermediate_size = config.moe_intermediate_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size + else: + # Default: Mixtral + E = config.num_local_experts + topk = config.num_experts_per_tok + intermediate_size = config.intermediate_size + shard_intermediate_size = 2 * intermediate_size // args.tp_size + + hidden_size = config.hidden_size + dtype = config.torch_dtype + use_fp8_w8a8 = args.dtype == "fp8_w8a8" + use_int8_w8a8 = args.dtype == "int8_w8a8" + use_int8_w8a16 = args.dtype == "int8_w8a16" + block_shape = None + if ( + hasattr(config, "quantization_config") + and "weight_block_size" in config.quantization_config + ): + block_shape = config.quantization_config["weight_block_size"] + assert len(block_shape) == 2 + + if args.batch_size is None: + batch_sizes = [ + 1, + 2, + 4, + 8, + 16, + 24, + 32, + 48, + 64, + 96, + 128, + 256, + 512, + 1024, + 1536, + 2048, + 3072, + 4096, + ] + else: + batch_sizes = [args.batch_size] + + ray.init() + num_gpus = int(ray.available_resources()["GPU"]) + workers = [BenchmarkWorker.remote(args.seed) for _ in range(num_gpus)] + + def _distribute(method: str, inputs: List[Any]) -> List[Any]: + outputs = [] + worker_idx = 0 + for input_args in inputs: + worker = workers[worker_idx] + worker_method = getattr(worker, method) + output = worker_method.remote(*input_args) + outputs.append(output) + worker_idx = (worker_idx + 1) % num_gpus + return ray.get(outputs) + + if args.tune: + search_space = get_configs_compute_bound() + if block_shape is not None: + block_n, block_k = block_shape[0], block_shape[1] + search_space = [ + config + for config in search_space + if block_k % config["BLOCK_SIZE_K"] == 0 + ] + print(f"Start tuning over {len(search_space)} configurations...") + + start = time.time() + configs = _distribute( + "tune", + [ + ( + batch_size, + E, + shard_intermediate_size, + hidden_size, + topk, + dtype, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + block_shape, + search_space, + ) + for batch_size in batch_sizes + ], + ) + best_configs = { + M: sort_config(config) for M, config in zip(batch_sizes, configs) + } + save_configs( + best_configs, + E, + shard_intermediate_size, + hidden_size, + topk, + dtype, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + block_shape, + ) + end = time.time() + print(f"Tuning took {end - start:.2f} seconds") + else: + outputs = _distribute( + "benchmark", + [ + ( + batch_size, + E, + shard_intermediate_size, + hidden_size, + topk, + dtype, + use_fp8_w8a8, + use_int8_w8a8, + use_int8_w8a16, + block_shape, + ) + for batch_size in batch_sizes + ], + ) + + for batch_size, (config, kernel_time) in zip(batch_sizes, outputs): + print(f"Batch size: {batch_size}, config: {config}") + print(f"Kernel time: {kernel_time:.2f} us") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--model", type=str, default="mistralai/Mixtral-8x7B-Instruct-v0.1" + ) + parser.add_argument("--tp-size", "-tp", type=int, default=2) + parser.add_argument( + "--dtype", + type=str, + choices=["auto", "fp8_w8a8", "int8_w8a16", "int8_w8a8"], + default="auto", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--batch-size", type=int, required=False) + parser.add_argument("--tune", action="store_true") + args = parser.parse_args() + + main(args) diff --git a/sglang/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_decode.py b/sglang/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_decode.py new file mode 100644 index 0000000..78d8149 --- /dev/null +++ b/sglang/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_decode.py @@ -0,0 +1,576 @@ +import itertools +import math +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import rearrange +from sgl_kernel import lightning_attention_decode as sgl_lightning_attention_decode + + +@triton.jit +def _decode_kernel( + Q, + K, + V, + KV, + Out, + S, + b: tl.constexpr, + h: tl.constexpr, + n: tl.constexpr, + d: tl.constexpr, + d_original: tl.constexpr, + e: tl.constexpr, + e_original: tl.constexpr, +): + off_bh = tl.program_id(0) + off_h = off_bh % h + + qk_offset = off_bh * n * d + v_offset = off_bh * n * e + o_offset = off_bh * n * e + kv_offset = off_bh * d * e + + s = tl.load(S + off_h) + ratio = tl.exp(-s) + + d_idx = tl.arange(0, d) + e_idx = tl.arange(0, e) + + # Create masks for original dimensions + d_mask = d_idx < d_original + e_mask = e_idx < e_original + + # Load with masking + q = tl.load(Q + qk_offset + d_idx, mask=d_mask, other=0.0) + k = tl.load(K + qk_offset + d_idx, mask=d_mask, other=0.0) + v = tl.load(V + v_offset + e_idx, mask=e_mask, other=0.0) + + # Load KV with 2D masking + kv = tl.load( + KV + kv_offset + d_idx[:, None] * e + e_idx[None, :], + mask=(d_mask[:, None] & e_mask[None, :]), + other=0.0, + ) + + # Compute outer product using element-wise operations + k_v_prod = k[:, None] * v[None, :] + kv = ratio * kv + k_v_prod + + # Store KV with 2D masking + tl.store( + KV + kv_offset + d_idx[:, None] * e + e_idx[None, :], + kv.to(KV.dtype.element_ty), + mask=(d_mask[:, None] & e_mask[None, :]), + ) + + # Compute matrix-vector multiplication using element-wise operations and reduction + o = tl.sum(q[:, None] * kv, axis=0) + + # Store output with masking + tl.store(Out + o_offset + e_idx, o.to(Out.dtype.element_ty), mask=e_mask) + + +def lightning_attn_decode(q, k, v, kv, s): + """Triton implementation of Lightning Attention decode operation""" + b, h, n, d = q.shape + e = v.shape[-1] + assert n == 1, "Sequence length must be 1 in decode mode" + + # Get padded dimensions (power of 2) + d_padded = next_power_of_2(d) + e_padded = next_power_of_2(e) + + # Create output tensor (padded) + o_padded = torch.empty(b, h, n, e_padded, dtype=v.dtype, device=v.device) + + # Create padded tensors without actually padding the data + q_padded = torch.empty(b, h, n, d_padded, dtype=q.dtype, device=q.device) + k_padded = torch.empty(b, h, n, d_padded, dtype=k.dtype, device=k.device) + v_padded = torch.empty(b, h, n, e_padded, dtype=v.dtype, device=v.device) + kv_padded = torch.empty( + b, h, d_padded, e_padded, dtype=torch.float32, device=kv.device + ) + + # Copy data to padded tensors + q_padded[..., :d] = q + k_padded[..., :d] = k + v_padded[..., :e] = v + kv_padded[..., :d, :e] = kv + + # Launch kernel + grid = (b * h, 1) + _decode_kernel[grid]( + q_padded, + k_padded, + v_padded, + kv_padded, + o_padded, + s, + b=b, + h=h, + n=n, + d=d_padded, + d_original=d, + e=e_padded, + e_original=e, + ) + + # Get unpadded outputs + o = o_padded[..., :e] + kv_out = kv_padded[..., :d, :e] + + return o, kv_out + + +def next_power_of_2(n): + return 2 ** (int(math.ceil(math.log(n, 2)))) + + +class MiniMaxText01LightningAttention(nn.Module): + def __init__(self, config=None, layer_idx: Optional[int] = None, **kwargs): + super().__init__() + if config is None: + config = type("Config", (), kwargs) + + bias = False + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads) + + self.out_proj = nn.Linear( + self.head_dim * self.num_heads, self.hidden_size, bias=bias + ) + self.act = get_activation_fn(config.hidden_act) + self.norm = MiniMaxText01RMSNorm(self.head_dim * self.num_heads) + + self.qkv_proj = nn.Linear( + self.hidden_size, 3 * self.head_dim * self.num_heads, bias=bias + ) + self.output_gate = nn.Linear( + self.hidden_size, self.head_dim * self.num_heads, bias=bias + ) + + # for inference only + self.offset = 0 + self.layer_idx = layer_idx + + def forward( + self, + hidden_states, + attn_mask: Optional[torch.Tensor] = None, # (b, h, n, m) + output_attentions: bool = False, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + use_cache: bool = False, + slope_rate: Optional[torch.Tensor] = None, + **kwargs, + ): + if (not self.training) and (not do_eval): + return self.inference( + hidden_states, + attn_mask, + output_attentions, + past_key_value, + use_cache, + slope_rate, + ) + + def inference( + self, + x, + attn_mask: Optional[torch.Tensor] = None, # (b, n) + output_attentions: bool = False, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + use_cache: bool = False, + slope_rate: Optional[torch.Tensor] = None, # (h, 1, 1) + ): + # x: b n d + b, n, d = x.shape + # linear map + qkv = self.act(self.qkv_proj(x)) + new_shape = qkv.size()[:-1] + (self.num_heads, -1) + qkv = qkv.view(*new_shape) + q, k, v = torch.split(qkv, [self.head_dim] * 3, dim=3) + q = q.transpose(1, 2) # [b, n, h, d] -> [b, h, n, d] + k = k.transpose(1, 2) # [b, n, h, d] -> [b, h, n, d] + v = v.transpose(1, 2) # [b, n, h, d] -> [b, h, n, e] + + self.offset += 1 + ratio = torch.exp(-slope_rate) # [h, 1, 1] + + # decode mode + kv = past_key_value # [b, h, d, e] + output = [] + for i in range(n): + # kv: [b, h, d, e] + # ratio: [h, 1, 1] + # k: [b, h, n, d] + # v: [b, h, n, e] + # k[:, :, i : i + 1]: [b, h, 1, d] + # v[:, :, i : i + 1]: [b, h, 1, e] + # ratio * kv: [b, h, d, e] + # torch.einsum( + # "... n d, ... n e -> ... d e", + # k[:, :, i : i + 1], + # v[:, :, i : i + 1], + # ) + # [b, h, d, e] + [b, h, d, e] -> [b, h, d, e] + kv = ratio * kv + torch.einsum( + "... n d, ... n e -> ... d e", + k[:, :, i : i + 1], + v[:, :, i : i + 1], + ) + # q[:, :, i : i + 1]: [b, h, 1, d] + # kv.to(q.dtype): [b, h, d, e] + # torch.einsum( + # "... n e, ... e d -> ... n d", q[:, :, i : i + 1], kv.to(q.dtype) + # ) + # [b, h, 1, d] * [b, h, d, e] -> [b, h, 1, e] + qkv = torch.einsum( + "... n e, ... e d -> ... n d", q[:, :, i : i + 1], kv.to(q.dtype) + ) + output.append(qkv) + output = torch.cat(output, dim=-2) + + # reshape + output = rearrange(output, "b h n d -> b n (h d)") + # normalize + output = self.norm(output) + # gate + output = F.sigmoid(self.output_gate(x)) * output + # outproj + output = self.out_proj(output) + + attn_weights = None + + return output, attn_weights, kv + + +def get_activation_fn(activation): + if activation == "gelu": + return F.gelu + elif activation == "relu": + return F.relu + elif activation == "elu": + return F.elu + elif activation == "sigmoid": + return F.sigmoid + elif activation == "exp": + + def f(x): + with torch.no_grad(): + x_max = torch.max(x, dim=-1, keepdims=True).values + y = torch.exp(x - x_max) + return y + + return f + elif activation == "leak": + return F.leaky_relu + elif activation == "1+elu": + + def f(x): + return 1 + F.elu(x) + + return f + elif activation == "2+elu": + + def f(x): + return 2 + F.elu(x) + + return f + elif activation == "silu" or activation == "swish": + return F.silu + elif activation == "sine": + return torch.sin + else: + return lambda x: x + + +class MiniMaxText01RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + MiniMaxText01RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +def test_lightning_attention_implementations(model_params): + torch.manual_seed(42) + + batch_size = 64 + seq_len = 1 + dtype = torch.bfloat16 + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + hidden_states = torch.randn( + batch_size, seq_len, model_params["hidden_size"], dtype=dtype, device=device + ) + + attention_mask = torch.ones(batch_size, seq_len, dtype=dtype, device=device) + + slope_rate = _build_slope_tensor(model_params["num_attention_heads"]).to(device) + + model_attn = MiniMaxText01LightningAttention(**model_params).to(dtype).to(device) + model_attn.eval() + + d = model_params["head_dim"] + past_kv = torch.randn( + batch_size, + model_params["num_attention_heads"], + d, + d, + device=device, + ) + with torch.no_grad(): + model_output, _, new_kv = model_attn.inference( + hidden_states, + attn_mask=attention_mask, + slope_rate=slope_rate, + past_key_value=past_kv, + ) + + qkv = model_attn.act(model_attn.qkv_proj(hidden_states)) + new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1) + qkv = qkv.view(*new_shape) + q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + past_kv = past_kv.contiguous() + slope_rate = slope_rate.contiguous() + + # Test Triton implementation + triton_output, triton_new_kv = lightning_attn_decode(q, k, v, past_kv, slope_rate) + triton_output = triton_output.transpose(1, 2).contiguous() + triton_output = triton_output.view(batch_size, seq_len, -1) + triton_output = model_attn.norm(triton_output) + triton_output = torch.sigmoid(model_attn.output_gate(hidden_states)) * triton_output + triton_output = model_attn.out_proj(triton_output) + + # Test SGL implementation + sgl_output = torch.empty_like(v) + sgl_new_kv = torch.empty_like(past_kv) + sgl_lightning_attention_decode(q, k, v, past_kv, slope_rate, sgl_output, sgl_new_kv) + + sgl_output = sgl_output.transpose(1, 2).contiguous() + sgl_output = sgl_output.view(batch_size, seq_len, -1) + sgl_output = model_attn.norm(sgl_output) + sgl_output = torch.sigmoid(model_attn.output_gate(hidden_states)) * sgl_output + sgl_output = model_attn.out_proj(sgl_output) + + # Verify Triton implementation results + torch.testing.assert_close( + model_output, + triton_output, + rtol=1e-3, + atol=1e-2, + msg="Triton lightning attention implementation produces different output results", + ) + torch.testing.assert_close( + new_kv, + triton_new_kv, + rtol=1e-3, + atol=1e-2, + msg="Triton lightning attention implementation produces different kv results", + ) + + # Verify SGL implementation results + torch.testing.assert_close( + model_output, + sgl_output, + rtol=1e-3, + atol=1e-2, + msg="SGL lightning attention implementation produces different output results", + ) + torch.testing.assert_close( + new_kv, + sgl_new_kv, + rtol=1e-3, + atol=1e-2, + msg="SGL lightning attention implementation produces different kv results", + ) + + print("✅ All implementations match") + + +def _build_slope_tensor(n_attention_heads: int): + def get_slopes(n): + def get_slopes_power_of_2(n): + start = 2 ** (-(2 ** -(math.log2(n) - 3))) + ratio = start + return [start * ratio**i for i in range(n)] + + if math.log2(n).is_integer(): + return get_slopes_power_of_2(n) + else: + closest_power_of_2 = 2 ** math.floor(math.log2(n)) + return ( + get_slopes_power_of_2(closest_power_of_2) + + get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2] + ) + + slopes = torch.tensor(get_slopes(n_attention_heads)).reshape( + n_attention_heads, 1, 1 + ) + return slopes + + +def get_benchmark(): + batch_size_range = [i for i in range(1, 33)] # max 32 + seq_length_range = [1] # decode mode sequence length is fixed to 1 + configs = list(itertools.product(batch_size_range, seq_length_range)) + + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "seq_len"], + x_vals=[list(_) for _ in configs], + line_arg="provider", + line_vals=["Original", "Triton", "SGL"], + line_names=[ + "Original PyTorch Implementation", + "Triton Implementation", + "SGL Implementation", + ], + styles=[("blue", "-"), ("green", "-"), ("red", "-")], + ylabel="us", + plot_name="lightning-attention-decode-performance", + args={}, + ) + ) + def benchmark(batch_size, seq_len, provider): + dtype = torch.bfloat16 + device = torch.device("cuda") + + params = { + "hidden_size": 6144, + "num_attention_heads": 64, + "head_dim": 96, + "hidden_act": "gelu", + } + + hidden_states = torch.randn( + batch_size, seq_len, params["hidden_size"], dtype=dtype, device=device + ) + + attention_mask = torch.ones(batch_size, seq_len, dtype=dtype, device=device) + + slope_rate = _build_slope_tensor(params["num_attention_heads"]).to(device) + model_attn = MiniMaxText01LightningAttention(**params).to(dtype).to(device) + model_attn.eval() + + d = params["head_dim"] + past_kv = torch.randn( + batch_size, + params["num_attention_heads"], + d, + d, + device=device, + ) + + quantiles = [0.5, 0.2, 0.8] + if provider == "Original": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: model_attn.inference( + hidden_states, + attn_mask=attention_mask, + slope_rate=slope_rate, + past_key_value=past_kv, + ), + quantiles=quantiles, + ) + elif provider == "Triton": + + def run_triton(): + qkv = model_attn.act(model_attn.qkv_proj(hidden_states)) + new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1) + qkv = qkv.view(*new_shape) + q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + output, new_kv = lightning_attn_decode(q, k, v, past_kv, slope_rate) + output = output.transpose(1, 2).contiguous() + output = output.view(batch_size, seq_len, -1) + output = model_attn.norm(output) + output = torch.sigmoid(model_attn.output_gate(hidden_states)) * output + return model_attn.out_proj(output) + + ms, min_ms, max_ms = triton.testing.do_bench( + run_triton, + quantiles=quantiles, + ) + else: # SGL + + def run_sgl(): + qkv = model_attn.act(model_attn.qkv_proj(hidden_states)) + new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1) + qkv = qkv.view(*new_shape) + q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1) + q = q.transpose(1, 2).contiguous() + k = k.transpose(1, 2).contiguous() + v = v.transpose(1, 2).contiguous() + + output = torch.empty_like(v) + new_kv = torch.empty_like(past_kv) + sgl_lightning_attention_decode( + q, k, v, past_kv, slope_rate, output, new_kv + ) + + output = output.transpose(1, 2).contiguous() + output = output.view(batch_size, seq_len, -1) + output = model_attn.norm(output) + output = torch.sigmoid(model_attn.output_gate(hidden_states)) * output + return model_attn.out_proj(output) + + ms, min_ms, max_ms = triton.testing.do_bench( + run_sgl, + quantiles=quantiles, + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + return benchmark + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--save_path", + type=str, + default="./configs/benchmark_ops/lightning_attention_decode/", + help="Path to save lightning attention decode benchmark results", + ) + args = parser.parse_args() + + params = { + "hidden_size": 6144, + "num_attention_heads": 64, + "head_dim": 96, + "hidden_act": "silu", + } + # Run correctness test first + # Adapted from https://huggingface.co/MiniMaxAI/MiniMax-Text-01/blob/main/config.json + test_lightning_attention_implementations(params) + + # Run performance benchmark + benchmark = get_benchmark() + benchmark.run(print_data=True, save_path=args.save_path) diff --git a/sglang/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_prefill.py b/sglang/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_prefill.py new file mode 100644 index 0000000..3bf9054 --- /dev/null +++ b/sglang/benchmark/kernels/minmax-text-01-lightning_attention/benchmark_lightning_attention_prefill.py @@ -0,0 +1,603 @@ +import itertools +import math +import os +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +import triton +import triton.language as tl +from einops import rearrange + + +# Adapted from https://github.com/OpenNLPLab/lightning-attention/blob/main/lightning_attn/ops/triton/lightning_attn2.py +@triton.jit +def _fwd_kernel( + Q, + K, + V, + Out, + S, # log lambda + b: tl.constexpr, + h: tl.constexpr, + n: tl.constexpr, + d: tl.constexpr, + e: tl.constexpr, + BLOCK: tl.constexpr, + NUM_BLOCK: tl.constexpr, + BLOCK_MODEL: tl.constexpr, +): + ##### get offset + off_bh = tl.program_id(0) + off_h = off_bh % h + off_e = tl.program_id(1) + qk_offset = off_bh * n * d + v_offset = off_bh * n * e + o_offset = off_bh * n * e + # channel offset + e_offset = off_e * BLOCK_MODEL + + ##### get block ptr + Q_block_ptr = Q + qk_offset + tl.arange(0, d)[None, :] + K_trans_block_ptr = K + qk_offset + tl.arange(0, d)[:, None] + V_block_ptr = V + v_offset + e_offset + tl.arange(0, BLOCK_MODEL)[None, :] + O_block_ptr = Out + o_offset + e_offset + tl.arange(0, BLOCK_MODEL)[None, :] + S_block_ptr = S + off_h + + ##### init diag decay(Lambda); q, k decay; kv + s = tl.load(S_block_ptr) + # q, k decay + off_block = tl.arange( + 0, BLOCK + ) # Not bug, this is a bit different from algorithm 1, but is mathematically equivalent + q_decay = tl.exp(-s.to(tl.float32) * off_block[:, None]) + k_trans_decay = tl.exp(-s.to(tl.float32) * (BLOCK - off_block[None, :])) + block_decay = tl.exp(-s.to(tl.float32) * BLOCK) + # diag decay + index = off_block[:, None] - off_block[None, :] + s_index = s * index + s_index = tl.where(index >= 0, -s_index, float("-inf")) + diag_decay = tl.exp(s_index) + kv = tl.zeros([d, BLOCK_MODEL], dtype=tl.float32) + + ##### compute + for i in range(NUM_BLOCK): + # load + q = tl.load( + Q_block_ptr + off_block[:, None] * d, mask=off_block[:, None] < n, other=0.0 + ).to(tl.float32) + k_trans = tl.load( + K_trans_block_ptr + off_block[None, :] * d, + mask=off_block[None, :] < n, + other=0.0, + ).to(tl.float32) + v = tl.load( + V_block_ptr + off_block[:, None] * e, mask=off_block[:, None] < n, other=0.0 + ).to(tl.float32) + + # compute + qk = tl.dot(q, k_trans) * diag_decay + o_intra = tl.dot(qk, v) + o_inter = tl.dot(q, kv) * q_decay + o = o_intra + o_inter + + # save and update + tl.store( + O_block_ptr + off_block[:, None] * e, + o.to(O_block_ptr.dtype.element_ty), + mask=off_block[:, None] < n, + ) + kv = block_decay * kv + tl.dot(k_trans * k_trans_decay, v) + off_block += BLOCK + + +def lightning_attn2(q, k, v, s): + q = q.contiguous() + k = k.contiguous() + v = v.contiguous() + s = s.contiguous() + + b, h, n, d = q.shape + e = v.shape[-1] + + # Pad d to next power of 2 + d_padded = next_power_of_2(d) + if d_padded != d: + q_padded = F.pad(q, (0, d_padded - d)) + k_padded = F.pad(k, (0, d_padded - d)) + else: + q_padded = q + k_padded = k + + # Pad e to next power of 2 + e_padded = next_power_of_2(e) + if e_padded != e: + v_padded = F.pad(v, (0, e_padded - e)) + else: + v_padded = v + + o_padded = torch.empty((b, h, n, e_padded), dtype=q.dtype, device=q.device) + + BLOCK = 64 + NUM_BLOCK = triton.cdiv(q.shape[2], BLOCK) + # parallel over channel + BLOCK_MODEL = min(triton.next_power_of_2(e_padded), 32) + grid = (b * h, triton.cdiv(e_padded, BLOCK_MODEL)) + + _fwd_kernel[grid]( + q_padded, + k_padded, + v_padded, + o_padded, + s, + b, + h, + n, + d_padded, + e_padded, + BLOCK=BLOCK, + NUM_BLOCK=NUM_BLOCK, + BLOCK_MODEL=BLOCK_MODEL, + ) + + # Remove padding from output + if e_padded != e: + o = o_padded[..., :e] + else: + o = o_padded + + return o + + +def is_support(dim): + return 16 % dim + + +def next_power_of_2(n): + return 2 ** (int(math.ceil(math.log(n, 2)))) + + +def lightning_attn_func(q, k, v, s): + b, h, n, d = q.shape + e = v.shape[-1] + assert is_support(d) and is_support(e) + + # pad v's feature dim to power of 2 + e_pad = next_power_of_2(e) + need_pad = e_pad != e + if need_pad: + v = F.pad(v, (0, e_pad - e)) + + if d > 128: + # split over head + if 64 % d: + m = 64 + elif 32 % d: + m = 32 + elif 16 % d: + m = 16 + arr = [m * i for i in range(d // m + 1)] + if arr[-1] != d: + arr.append(d) + n = len(arr) + o = 0 + for i in range(n - 1): + start = arr[i] + end = arr[i + 1] + q1 = q[..., start:end] + k1 = k[..., start:end] + o += lightning_attn2(q1, k1, v, s) + else: + o = lightning_attn2(q, k, v, s) + + if need_pad: + o = o[:, :, :, :e] + + return o + + +debug = eval(os.environ.get("debug", default="False")) + +BLOCK = 256 + + +# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->MiniMaxText01 +class MiniMaxText01RMSNorm(nn.Module): + def __init__(self, hidden_size, eps=1e-6): + """ + MiniMaxText01RMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +# Copied from https://huggingface.co/MiniMaxAI/MiniMax-Text-01/blob/main/modeling_minimax_text_01.py +def get_activation_fn(activation): + if debug: + logger.info(f"activation: {activation}") + if activation == "gelu": + return F.gelu + elif activation == "relu": + return F.relu + elif activation == "elu": + return F.elu + elif activation == "sigmoid": + return F.sigmoid + elif activation == "exp": + + def f(x): + with torch.no_grad(): + x_max = torch.max(x, dim=-1, keepdims=True).values + y = torch.exp(x - x_max) + + return y + + return f + elif activation == "leak": + return F.leaky_relu + elif activation == "1+elu": + + def f(x): + return 1 + F.elu(x) + + return f + elif activation == "2+elu": + + def f(x): + return 2 + F.elu(x) + + return f + elif activation == "silu" or activation == "swish": + return F.silu + elif activation == "sine": + return torch.sin + else: + logger.info(f"activation: does not support {activation}, use Identity!!!") + return lambda x: x + + +# Copied from https://huggingface.co/MiniMaxAI/MiniMax-Text-01/blob/main/modeling_minimax_text_01.py +class MiniMaxText01LightningAttention(nn.Module): + def __init__(self, config=None, layer_idx: Optional[int] = None, **kwargs): + super().__init__() + if config is None: + config = type("Config", (), kwargs) + + bias = False + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads) + + self.out_proj = nn.Linear( + self.head_dim * self.num_heads, self.hidden_size, bias=bias + ) + self.act = get_activation_fn(config.hidden_act) + self.norm = MiniMaxText01RMSNorm(self.head_dim * self.num_heads) + + self.qkv_proj = nn.Linear( + self.hidden_size, 3 * self.head_dim * self.num_heads, bias=bias + ) + self.output_gate = nn.Linear( + self.hidden_size, self.head_dim * self.num_heads, bias=bias + ) + + # for inference only + self.offset = 0 + self.layer_idx = layer_idx + + def forward( + self, + hidden_states, + attn_mask: Optional[torch.Tensor] = None, # (b, h, n, m) + output_attentions: bool = False, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + use_cache: bool = False, + slope_rate: Optional[torch.Tensor] = None, + **kwargs, + ): + if (not self.training) and (not do_eval): + return self.inference( + hidden_states, + attn_mask, + output_attentions, + past_key_value, + use_cache, + slope_rate, + ) + + def inference( + self, + x, + attn_mask: Optional[torch.Tensor] = None, # (b, n) + output_attentions: bool = False, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + use_cache: bool = False, + slope_rate: Optional[torch.Tensor] = None, # (h, 1, 1) + ): + # x: b n d + b, n, d = x.shape + # linear map + qkv = self.act(self.qkv_proj(x)) + new_shape = qkv.size()[:-1] + (self.num_heads, -1) + qkv = qkv.view(*new_shape) + q, k, v = torch.split(qkv, [self.head_dim] * 3, dim=3) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + if past_key_value is None: + self.offset = q.shape[-2] + else: + self.offset += 1 + + # for align with metaseq + ratio = torch.exp(-slope_rate) + + # only use for the first time + if past_key_value is None: + slope_rate = slope_rate.to(torch.float32) + if attn_mask is not None: + v = v.masked_fill( + (1 - attn_mask).unsqueeze(1).unsqueeze(-1).to(torch.bool), 0 + ) + NUM_BLOCK = (n + BLOCK - 1) // BLOCK + b, h, n, d = q.shape + e = v.shape[-1] + # other + array = torch.arange(BLOCK).to(q) + 1 + q_decay = torch.exp(-slope_rate * array.reshape(-1, 1)) + k_decay = torch.exp(-slope_rate * (BLOCK - array.reshape(-1, 1))) + index = array[:, None] - array[None, :] + s_index = ( + slope_rate + * index[ + None, + None, + ] + ) + s_index = torch.where(index >= 0, -s_index, float("-inf")) + diag_decay = torch.exp(s_index) + + kv = torch.zeros(b, h, d, e).to(torch.float32).to(q.device) + output = torch.empty((b, h, n, e), dtype=q.dtype, device=q.device) + for i in range(NUM_BLOCK): + si = i * BLOCK + ei = min(si + BLOCK, n) + m = ei - si + qi = q[:, :, si:ei].contiguous() + ki = k[:, :, si:ei].contiguous() + vi = v[:, :, si:ei].contiguous() + qkv_none_diag = torch.matmul(qi * q_decay[:, :m], kv).to(torch.float32) + + # diag + qk = ( + torch.matmul(qi, ki.transpose(-1, -2)).to(torch.float32) + * diag_decay[:, :, :m, :m] + ) + qkv_diag = torch.matmul(qk, vi.to(torch.float32)) + block_decay = torch.exp(-slope_rate * m) + output[:, :, si:ei] = qkv_none_diag + qkv_diag + kv = block_decay * kv + torch.matmul( + (ki * k_decay[:, -m:]).transpose(-1, -2).to(vi.dtype), vi + ) + + else: + kv = past_key_value + output = [] + for i in range(n): + kv = ratio * kv + torch.einsum( + "... n d, ... n e -> ... d e", + k[:, :, i : i + 1], + v[:, :, i : i + 1], + ) + qkv = torch.einsum( + "... n e, ... e d -> ... n d", q[:, :, i : i + 1], kv.to(q.dtype) + ) + output.append(qkv) + output = torch.cat(output, dim=-2) + # reshape + output = rearrange(output, "b h n d -> b n (h d)") + # normalize + output = self.norm(output) + # gate + output = F.sigmoid(self.output_gate(x)) * output + # outproj + output = self.out_proj(output) + + attn_weights = None + + return output, attn_weights, kv + + +def _build_slope_tensor(n_attention_heads: int): + def get_slopes(n): + def get_slopes_power_of_2(n): + start = 2 ** (-(2 ** -(math.log2(n) - 3))) + ratio = start + return [start * ratio**i for i in range(n)] + + if math.log2(n).is_integer(): + return get_slopes_power_of_2( + n + ) # In the paper, we only train models that have 2^a heads for some a. This function has + else: # some good properties that only occur when the input is a power of 2. To maintain that even + closest_power_of_2 = 2 ** math.floor( + math.log2(n) + ) # when the number of heads is not a power of 2, we use this workaround. + return ( + get_slopes_power_of_2(closest_power_of_2) + + get_slopes(2 * closest_power_of_2)[0::2][: n - closest_power_of_2] + ) + + # h, 1, 1 + slopes = torch.tensor(get_slopes(n_attention_heads)).reshape( + n_attention_heads, 1, 1 + ) + + return slopes + + +def test_lightning_attention_implementations(model_params): + torch.manual_seed(42) + + batch_size = 2 + seq_len = 1024 + dtype = torch.bfloat16 + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + + hidden_states = torch.randn( + batch_size, seq_len, model_params["hidden_size"], dtype=dtype, device=device + ) + + attention_mask = torch.ones(batch_size, seq_len, dtype=dtype, device=device) + + slope_rate = _build_slope_tensor(model_params["num_attention_heads"]).to(device) + + model_attn = MiniMaxText01LightningAttention(**model_params).to(dtype).to(device) + model_attn.eval() + + with torch.no_grad(): + model_output, _, _ = model_attn.inference( + hidden_states, attn_mask=attention_mask, slope_rate=slope_rate + ) + + qkv = model_attn.act(model_attn.qkv_proj(hidden_states)) + new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1) + qkv = qkv.view(*new_shape) + q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + lib_output = lightning_attn_func(q, k, v, slope_rate) + lib_output = lib_output.transpose(1, 2).contiguous() + lib_output = lib_output.view(batch_size, seq_len, -1) + lib_output = model_attn.norm(lib_output) + lib_output = torch.sigmoid(model_attn.output_gate(hidden_states)) * lib_output + lib_output = model_attn.out_proj(lib_output) + + torch.testing.assert_close( + model_output, + lib_output, + rtol=1e-3, + atol=1e-2, + msg="Lightning attention implementations produce different results", + ) + + print("✅ Two implementations match") + + +def get_benchmark(): + batch_size_range = [2**i for i in range(0, 7)] # max 64 + seq_length_range = [256, 512, 1024, 2048, 4096] # max 4096 + configs = list(itertools.product(batch_size_range, seq_length_range)) + + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "seq_len"], + x_vals=[list(_) for _ in configs], + line_arg="provider", + line_vals=["MiniMax-Text-01", "OpenNLPLab"], + line_names=[ + "MiniMax-Text-01 Model Implementation", + "OpenNLPLab Library Implementation", + ], + styles=[("blue", "-"), ("green", "-")], + ylabel="us", + plot_name="lightning-attention-prefill-performance", + args={}, + ) + ) + def benchmark(batch_size, seq_len, provider): + dtype = torch.bfloat16 + device = torch.device("cuda") + + params = { + "hidden_size": 6144, + "num_attention_heads": 64, + "head_dim": 96, + "hidden_act": "gelu", + } + + hidden_states = torch.randn( + batch_size, seq_len, params["hidden_size"], dtype=dtype, device=device + ) + + attention_mask = torch.ones(batch_size, seq_len, dtype=dtype, device=device) + + slope_rate = _build_slope_tensor(params["num_attention_heads"]).to(device) + model_attn = MiniMaxText01LightningAttention(**params).to(dtype).to(device) + model_attn.eval() + + quantiles = [0.5, 0.2, 0.8] + if provider == "MiniMax-Text-01": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: model_attn.inference( + hidden_states, attn_mask=attention_mask, slope_rate=slope_rate + ), + quantiles=quantiles, + ) + else: + + def run_lib(): + qkv = model_attn.act(model_attn.qkv_proj(hidden_states)) + new_shape = qkv.size()[:-1] + (model_attn.num_heads, -1) + qkv = qkv.view(*new_shape) + q, k, v = torch.split(qkv, [model_attn.head_dim] * 3, dim=-1) + q = q.transpose(1, 2) + k = k.transpose(1, 2) + v = v.transpose(1, 2) + + lib_output = lightning_attn_func(q, k, v, slope_rate) + lib_output = lib_output.transpose(1, 2).contiguous() + lib_output = lib_output.view(batch_size, seq_len, -1) + lib_output = model_attn.norm(lib_output) + lib_output = ( + torch.sigmoid(model_attn.output_gate(hidden_states)) * lib_output + ) + return model_attn.out_proj(lib_output) + + ms, min_ms, max_ms = triton.testing.do_bench( + run_lib, + quantiles=quantiles, + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + return benchmark + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--save_path", + type=str, + default="./configs/benchmark_ops/lightning_attention_prefill/", + help="Path to save lightning attention prefill benchmark results", + ) + args = parser.parse_args() + + # Run correctness test first + # Adapted from https://huggingface.co/MiniMaxAI/MiniMax-Text-01/blob/main/config.json + params = { + "hidden_size": 6144, + "num_attention_heads": 64, + "head_dim": 96, + "hidden_act": "silu", + } + test_lightning_attention_implementations(params) + + # Run performance benchmark + benchmark = get_benchmark() + benchmark.run(print_data=True, save_path=args.save_path) diff --git a/sglang/benchmark/kernels/quantization/bench_int8_quant.py b/sglang/benchmark/kernels/quantization/bench_int8_quant.py new file mode 100644 index 0000000..94b7956 --- /dev/null +++ b/sglang/benchmark/kernels/quantization/bench_int8_quant.py @@ -0,0 +1,94 @@ +import argparse + +import torch +import triton +from vllm._custom_ops import scaled_int8_quant as vllm_scaled_int8_quant + +from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8 + + +@torch.compile(backend="inductor") +def torch_int8_quant(x): + int8_max = torch.iinfo(torch.int8).max + + abs_max = x.abs().max(dim=-1, keepdim=True).values + scales = abs_max.to(torch.float32) / float(int8_max) + + q_x = (x / scales).round().to(torch.int8) + + return q_x, scales + + +def _test_accuracy_once(M, K, input_dtype, device): + x = torch.randn(M, K, dtype=input_dtype, device=device) * 5000 + out, scales, _ = vllm_scaled_int8_quant(x, symmetric=True) + out1, scales1 = per_token_quant_int8(x) + out2, scales2 = torch_int8_quant(x) + torch.testing.assert_close(out, out2, atol=1, rtol=0) + torch.testing.assert_close(out, out1, atol=1, rtol=0) + torch.testing.assert_close(scales, scales2) + torch.testing.assert_close(scales1, scales2) + print(f"M: {M}, K: {K}, type: {input_dtype} OK") + + +def test_accuracy(): + Ms = [1, 13, 128, 1024, 2048, 4096] + Ks = [512, 1024, 2048, 8192] + input_dtypes = [torch.float16, torch.bfloat16] + for M in Ms: + for K in Ks: + for input_dtype in input_dtypes: + _test_accuracy_once(M, K, input_dtype, "cuda") + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size"], + x_vals=[1, 16, 32, 64, 128, 256, 512, 1024, 2048], + x_log=False, + line_arg="provider", + line_vals=["vllm op", "triton", "torch.compile"], + line_names=["vllm op", "triton", "torch.compile"], + styles=[("blue", "-"), ("orange", "-"), ("red", "-")], + ylabel="ms", + plot_name="int8 per token quant", + args={}, + ) +) +def benchmark(batch_size, provider): + M, K = batch_size, 16384 + x = torch.randn(M, K, dtype=torch.float16, device="cuda") * 1000 + + quantiles = [0.5, 0.2, 0.8] + if provider == "vllm op": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: vllm_scaled_int8_quant(x, symmetric=True), + quantiles=quantiles, + ) + if provider == "triton": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: per_token_quant_int8(x), + quantiles=quantiles, + ) + if provider == "torch.compile": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: torch_int8_quant(x), + quantiles=quantiles, + ) + + return ms, min_ms, max_ms + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--save_path", + type=str, + default="./bench_int8_quant_res", + help="Path to save int8 quant benchmark results", + ) + args = parser.parse_args() + + test_accuracy() + + benchmark.run(print_data=True, show_plots=True, save_path=args.save_path) diff --git a/sglang/benchmark/kernels/quantization/tuning_block_wise_kernel.py b/sglang/benchmark/kernels/quantization/tuning_block_wise_kernel.py new file mode 100644 index 0000000..197939f --- /dev/null +++ b/sglang/benchmark/kernels/quantization/tuning_block_wise_kernel.py @@ -0,0 +1,474 @@ +# Copyright 2025 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import argparse +import json +import multiprocessing as mp +import os +import time +from datetime import datetime +from typing import Any, Dict, List + +import torch +import triton +from tqdm import tqdm + +mp.set_start_method("spawn", force=True) + +from sglang.srt.layers.quantization.fp8_kernel import ( + _w8a8_block_fp8_matmul, + _w8a8_block_fp8_matmul_unrolledx4, +) +from sglang.srt.layers.quantization.int8_kernel import _w8a8_block_int8_matmul +from sglang.srt.utils import get_device_core_count, get_device_name, is_hip + +is_hip_ = is_hip() + +DTYPE_MAP = { + "float32": torch.float32, + "float16": torch.float16, + "half": torch.half, + "bfloat16": torch.bfloat16, +} + + +def w8a8_block_matmul( + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + block_size: List[int], + config: Dict[str, Any], + output_dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + """This function performs matrix multiplication with block-wise quantization. + + It takes two input tensors `A` and `B` with scales `As` and `Bs`. + The output is returned in the specified `output_dtype`. + + Args: + A: The input tensor, e.g., activation. + B: The input tensor, e.g., weight. + As: The per-token-group quantization scale for `A`. + Bs: The per-block quantization scale for `B`. + block_size: The block size for per-block quantization. It should be 2-dim, e.g., [128, 128]. + output_dytpe: The dtype of the returned tensor. + + Returns: + torch.Tensor: The result of matmul. + """ + assert len(block_size) == 2 + block_n, block_k = block_size[0], block_size[1] + + assert A.shape[-1] == B.shape[-1] + assert A.shape[:-1] == As.shape[:-1] and A.is_contiguous() + assert triton.cdiv(A.shape[-1], block_k) == As.shape[-1] + M = A.numel() // A.shape[-1] + + assert B.ndim == 2 and B.is_contiguous() and Bs.ndim == 2 + N, K = B.shape + assert triton.cdiv(N, block_n) == Bs.shape[0] + assert triton.cdiv(K, block_k) == Bs.shape[1] + + C_shape = A.shape[:-1] + (N,) + C = A.new_empty(C_shape, dtype=output_dtype) + + def grid(META): + return ( + triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), + ) + + # Use manually unrolledx4 kernel on AMD GPU when the grid size is small. + # Empirical testing shows the sweet spot lies when it's less than the # of + # compute units available on the device. + num_workgroups = triton.cdiv(M, config["BLOCK_SIZE_M"]) * triton.cdiv( + N, config["BLOCK_SIZE_N"] + ) + + if A.dtype == torch.float8_e4m3fnuz or A.dtype == torch.float8_e4m3fn: + kernel = ( + _w8a8_block_fp8_matmul_unrolledx4 + if (is_hip_ == True and num_workgroups <= get_device_core_count()) + else _w8a8_block_fp8_matmul + ) + else: + kernel = _w8a8_block_int8_matmul + + kernel[grid]( + A, + B, + C, + As, + Bs, + M, + N, + K, + block_n, + block_k, + A.stride(-2), + A.stride(-1), + B.stride(1), + B.stride(0), + C.stride(-2), + C.stride(-1), + As.stride(-2), + As.stride(-1), + Bs.stride(1), + Bs.stride(0), + **config, + ) + + return C + + +def get_rocm_configs_compute_bound(): + configs = [] + waves_per_eu_range = 0 + for num_stages in [2]: + for block_m in [32, 64, 128, 256]: + for block_k in [32, 64, 128, 256]: + for block_n in [16, 32, 64, 128, 256]: + for num_warps in [4, 8]: + for group_size in [1, 4, 8, 16, 32]: + configs.append( + { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": group_size, + "num_warps": num_warps, + "num_stages": num_stages, + "waves_per_eu": waves_per_eu_range, + } + ) + return configs + + +def get_configs_compute_bound(): + configs = [] + if is_hip_: + configs = get_rocm_configs_compute_bound() + else: + for num_stages in [2, 3, 4, 5]: + for block_m in [16, 32, 64, 128, 256]: + for block_k in [64, 128]: + for block_n in [32, 64, 128, 256]: + for num_warps in [4, 8]: + for group_size in [1, 16, 32, 64]: + configs.append( + { + "BLOCK_SIZE_M": block_m, + "BLOCK_SIZE_N": block_n, + "BLOCK_SIZE_K": block_k, + "GROUP_SIZE_M": group_size, + "num_warps": num_warps, + "num_stages": num_stages, + } + ) + return configs + + +def get_weight_shapes(tp_size): + # NOTE(HandH1998): The weight shapes only works for DeepSeek-V3. Modify them, if you tune for another different model. + # cannot TP + total = [ + (512 + 64, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (7168, 16384), + (7168, 18432), + ] + # N can TP + n_tp = [ + (18432 * 2, 7168), + ((128 + 64) * 128, 7168), + (128 * (128 + 128), 512), + (24576, 1536), + (4096, 7168), + ] + # K can TP + k_tp = [(7168, 18432), (7168, 16384), (7168, 2048)] + + weight_shapes = [] + for t in total: + weight_shapes.append(t) + for n_t in n_tp: + new_t = (n_t[0] // tp_size, n_t[1]) + weight_shapes.append(new_t) + for k_t in k_tp: + new_t = (k_t[0], k_t[1] // tp_size) + weight_shapes.append(new_t) + return weight_shapes + + +def benchmark_config( + A, B, As, Bs, block_size, config, out_dtype=torch.float16, num_iters=10 +): + def run(): + w8a8_block_matmul(A, B, As, Bs, block_size, config, out_dtype) + + torch.cuda.synchronize() + # JIT complication & warmup + for _ in range(5): + run() + torch.cuda.synchronize() + + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + + latencies: List[float] = [] + for i in range(num_iters): + torch.cuda.synchronize() + start_event.record() + run() + end_event.record() + end_event.synchronize() + latencies.append(start_event.elapsed_time(end_event)) + avg = sum(latencies) / (num_iters * 10) * 1000 # us + return avg + + +def tune(M, N, K, block_size, out_dtype, search_space, input_type): + factor_for_scale = 1e-2 + + if input_type == "fp8": + fp8_info = torch.finfo( + torch.float8_e4m3fnuz if is_hip_ else torch.float8_e4m3fn + ) + fp8_max, fp8_min = fp8_info.max, fp8_info.min + + A_fp32 = ( + (torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max + ) + A = A_fp32.clamp(min=fp8_min, max=fp8_max).to( + torch.float8_e4m3fnuz if is_hip_ else torch.float8_e4m3fn + ) + + B_fp32 = ( + (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max + ) + B = B_fp32.clamp(min=fp8_min, max=fp8_max).to( + torch.float8_e4m3fnuz if is_hip_ else torch.float8_e4m3fn + ) + else: + int8_info = torch.iinfo(torch.int8) + int8_max, int8_min = int8_info.max, int8_info.min + + A_fp32 = ( + (torch.rand(M, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * int8_max + ) + A = A_fp32.clamp(min=int8_min, max=int8_max).to(torch.int8) + + B_fp32 = ( + (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * int8_max + ) + B = B_fp32.clamp(min=int8_min, max=int8_max).to(torch.int8) + + block_n, block_k = block_size[0], block_size[1] + n_tiles = (N + block_n - 1) // block_n + k_tiles = (K + block_k - 1) // block_k + + As = torch.rand(M, k_tiles, dtype=torch.float32, device="cuda") * factor_for_scale + Bs = ( + torch.rand(n_tiles, k_tiles, dtype=torch.float32, device="cuda") + * factor_for_scale + ) + + best_config = None + best_time = float("inf") + for config in tqdm(search_space): + try: + kernel_time = benchmark_config( + A, + B, + As, + Bs, + block_size, + config, + out_dtype, + num_iters=10, + ) + except triton.runtime.autotuner.OutOfResources: + # Some configurations may be invalid and fail to compile. + continue + + if kernel_time < best_time: + best_time = kernel_time + best_config = config + now = datetime.now() + print(f"{now.ctime()}] Completed tuning for batch_size={M}") + assert best_config is not None + return best_config + + +def save_configs( + N, + K, + block_n, + block_k, + configs, + save_path, + input_type="fp8", +) -> None: + os.makedirs(save_path, exist_ok=True) + device_name = get_device_name().replace(" ", "_") + json_file_name = f"N={N},K={K},device_name={device_name},dtype={input_type}_w8a8,block_shape=[{block_n}, {block_k}].json" + + config_file_path = os.path.join(save_path, json_file_name) + print(f"Writing best config to {config_file_path}...") + + with open(config_file_path, "w") as f: + json.dump(configs, f, indent=4) + f.write("\n") + + +def get_available_gpu_count(): + """Get the number of available GPUs.""" + return torch.cuda.device_count() + + +def tune_on_gpu(args_dict): + """Run tuning on a specific GPU.""" + gpu_id = args_dict["gpu_id"] + batch_sizes = args_dict["batch_sizes"] + weight_shapes = args_dict["weight_shapes"] + args = args_dict["args"] + + torch.cuda.set_device(gpu_id) + print(f"Starting tuning on GPU {gpu_id} with batch sizes {batch_sizes}") + + block_n = args.block_n + block_k = args.block_k + out_dtype = DTYPE_MAP[args.out_dtype] + save_path = args.save_path + input_type = args.input_type + + search_space = get_configs_compute_bound() + search_space = [ + config for config in search_space if block_k % config["BLOCK_SIZE_K"] == 0 + ] + + start = time.time() + results = {} + for shape in tqdm(weight_shapes, desc=f"GPU {gpu_id} - Shapes"): + N, K = shape[0], shape[1] + print(f"[GPU {gpu_id}] Tune for weight shape of `N: {N}, K: {K}`") + benchmark_results = [ + tune( + batch_size, + N, + K, + [block_n, block_k], + out_dtype, + search_space, + input_type, + ) + for batch_size in tqdm(batch_sizes, desc=f"GPU {gpu_id} - Batch sizes") + ] + best_configs = {M: config for M, config in zip(batch_sizes, benchmark_results)} + save_configs(N, K, block_n, block_k, best_configs, save_path, input_type) + + end = time.time() + print(f"Tuning on GPU {gpu_id} took {end - start:.2f} seconds") + + +def distribute_batch_sizes(batch_sizes, num_gpus): + """Distribute batch sizes across available GPUs.""" + batches_per_gpu = [] + for i in range(num_gpus): + start_idx = i * len(batch_sizes) // num_gpus + end_idx = (i + 1) * len(batch_sizes) // num_gpus + batches_per_gpu.append(batch_sizes[start_idx:end_idx]) + return batches_per_gpu + + +def main(args): + print(args) + + num_gpus = get_available_gpu_count() + if num_gpus == 0: + raise RuntimeError("No GPU available for tuning") + print(f"Found {num_gpus} GPUs for parallel tuning") + + torch.cuda.init() + + if args.batch_size is None: + batch_sizes = [ + 1, + 2, + 4, + 8, + 16, + 24, + 32, + 48, + 64, + 96, + 128, + 256, + 512, + 1024, + 1536, + 2048, + 3072, + 4096, + ] + else: + batch_sizes = [args.batch_size] + num_gpus = 1 # If only one batch size, use only one GPU + + weight_shapes = get_weight_shapes(args.tp_size) + + batches_per_gpu = distribute_batch_sizes(batch_sizes, num_gpus) + + process_args = [] + for gpu_id in range(num_gpus): + process_args.append( + { + "gpu_id": gpu_id, + "batch_sizes": batches_per_gpu[gpu_id], + "weight_shapes": weight_shapes, # Each GPU processes all weight shapes + "args": args, + } + ) + + ctx = mp.get_context("spawn") + with ctx.Pool(num_gpus) as pool: + pool.map(tune_on_gpu, process_args) + + print("Multi-GPU tuning completed") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument("--tp-size", "-tp", type=int, default=8) + parser.add_argument( + "--input-type", type=str, choices=["fp8", "int8"], default="fp8" + ) + parser.add_argument( + "--out-dtype", + type=str, + choices=["float32", "float16", "bfloat16", "half"], + default="float16", + ) + parser.add_argument("--block-n", type=int, default=128) + parser.add_argument("--block-k", type=int, default=128) + parser.add_argument("--batch-size", type=int, required=False) + parser.add_argument( + "--save-path", type=str, default="python/sglang/srt/layers/quantization/configs" + ) + args = parser.parse_args() + + main(args) diff --git a/sglang/benchmark/kernels/rmsnorm/benchmark_rmsnorm.py b/sglang/benchmark/kernels/rmsnorm/benchmark_rmsnorm.py new file mode 100644 index 0000000..aeeea62 --- /dev/null +++ b/sglang/benchmark/kernels/rmsnorm/benchmark_rmsnorm.py @@ -0,0 +1,230 @@ +import itertools +from typing import Optional, Tuple, Union + +import torch +import triton +from flashinfer.norm import fused_add_rmsnorm, rmsnorm +from torch import nn +from vllm import _custom_ops as vllm_ops + + +class HuggingFaceRMSNorm(nn.Module): + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward( + self, + x: torch.Tensor, + residual: Optional[torch.Tensor] = None, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]: + orig_dtype = x.dtype + x = x.to(torch.float32) + if residual is not None: + x = x + residual.to(torch.float32) + residual = x.to(orig_dtype) + + variance = x.pow(2).mean(dim=-1, keepdim=True) + x = x * torch.rsqrt(variance + self.variance_epsilon) + x = x.to(orig_dtype) * self.weight + if residual is None: + return x + else: + return x, residual + + +def rmsnorm_naive( + x: torch.Tensor, + weight: torch.Tensor, + residual: Optional[torch.Tensor] = None, + eps: float = 1e-6, +): + naive_norm = HuggingFaceRMSNorm(x.shape[-1], eps=eps) + naive_norm.weight = nn.Parameter(weight) + naive_norm = naive_norm.to(x.device) + + orig_shape = x.shape + x = x.view(-1, x.shape[-1]) + if residual is not None: + residual = residual.view(-1, residual.shape[-1]) + + output = naive_norm(x, residual) + + if isinstance(output, tuple): + output = (output[0].view(orig_shape), output[1].view(orig_shape)) + else: + output = output.view(orig_shape) + return output + + +def rmsnorm_flashinfer( + x: torch.Tensor, + weight: torch.Tensor, + residual: Optional[torch.Tensor] = None, + eps: float = 1e-6, +): + orig_shape = x.shape + x = x.view(-1, x.shape[-1]) + if residual is not None: + residual = residual.view(-1, residual.shape[-1]) + + if residual is not None: + fused_add_rmsnorm(x, residual, weight, eps) + output = (x, residual) + else: + output = rmsnorm(x, weight, eps) + + if isinstance(output, tuple): + output = (output[0].view(orig_shape), output[1].view(orig_shape)) + else: + output = output.view(orig_shape) + return output + + +def rmsnorm_vllm( + x: torch.Tensor, + weight: torch.Tensor, + residual: Optional[torch.Tensor] = None, + eps: float = 1e-6, +): + orig_shape = x.shape + x = x.view(-1, x.shape[-1]) + if residual is not None: + residual = residual.view(-1, residual.shape[-1]) + + if residual is not None: + vllm_ops.fused_add_rms_norm(x, residual, weight, eps) + output = (x, residual) + else: + out = torch.empty_like(x) + vllm_ops.rms_norm(out, x, weight, eps) + output = out + + if isinstance(output, tuple): + output = (output[0].view(orig_shape), output[1].view(orig_shape)) + else: + output = output.view(orig_shape) + return output + + +def calculate_diff(batch_size, seq_len, hidden_size, use_residual=True): + dtype = torch.bfloat16 + x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda") + weight = torch.ones(hidden_size, dtype=dtype, device="cuda") + residual = torch.randn_like(x) if use_residual else None + + output_naive = rmsnorm_naive( + x.clone(), weight, residual.clone() if residual is not None else None + ) + output_flashinfer = rmsnorm_flashinfer( + x.clone(), weight, residual.clone() if residual is not None else None + ) + output_vllm = rmsnorm_vllm( + x.clone(), weight, residual.clone() if residual is not None else None + ) + + if use_residual: + output_naive = output_naive[0] + output_flashinfer = output_flashinfer[0] + output_vllm = output_vllm[0] + + print(f"Naive output={output_naive}") + print(f"FlashInfer output={output_flashinfer}") + print(f"VLLM output={output_vllm}") + + if torch.allclose( + output_naive, output_flashinfer, atol=1e-2, rtol=1e-2 + ) and torch.allclose(output_naive, output_vllm, atol=1e-2, rtol=1e-2): + print("✅ All implementations match") + else: + print("❌ Implementations differ") + + +batch_size_range = [2**i for i in range(0, 7, 2)] +seq_length_range = [2**i for i in range(6, 11, 1)] +head_num_range = [32, 48] +configs = list(itertools.product(head_num_range, batch_size_range, seq_length_range)) + + +def get_benchmark(use_residual): + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["head_num", "batch_size", "seq_len"], + x_vals=[list(_) for _ in configs], + line_arg="provider", + line_vals=["huggingface", "flashinfer", "vllm"], + line_names=["HuggingFace", "FlashInfer", "vLLM"], + styles=[("blue", "-"), ("green", "-"), ("red", "-")], + ylabel="us", + plot_name=f"rmsnorm-performance-{'with' if use_residual else 'without'}-residual", + args={}, + ) + ) + def benchmark(head_num, batch_size, seq_len, provider): + dtype = torch.bfloat16 + hidden_size = head_num * 128 # assuming head_dim = 128 + + x = torch.randn(batch_size, seq_len, hidden_size, dtype=dtype, device="cuda") + weight = torch.ones(hidden_size, dtype=dtype, device="cuda") + residual = torch.randn_like(x) if use_residual else None + + quantiles = [0.5, 0.2, 0.8] + + if provider == "huggingface": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: rmsnorm_naive( + x.clone(), + weight, + residual.clone() if residual is not None else None, + ), + quantiles=quantiles, + ) + elif provider == "flashinfer": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: rmsnorm_flashinfer( + x.clone(), + weight, + residual.clone() if residual is not None else None, + ), + quantiles=quantiles, + ) + else: + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: rmsnorm_vllm( + x.clone(), + weight, + residual.clone() if residual is not None else None, + ), + quantiles=quantiles, + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + return benchmark + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--use_residual", action="store_true", help="Whether to use residual connection" + ) + parser.add_argument( + "--save_path", + type=str, + default="./configs/benchmark_ops/rmsnorm/", + help="Path to save rmsnorm benchmark results", + ) + args = parser.parse_args() + + # Run correctness test + calculate_diff( + batch_size=4, seq_len=128, hidden_size=4096, use_residual=args.use_residual + ) + + # Get the benchmark function with proper use_residual setting + benchmark = get_benchmark(args.use_residual) + # Run performance benchmark + benchmark.run(print_data=True, save_path=args.save_path) diff --git a/sglang/benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py b/sglang/benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py new file mode 100644 index 0000000..1ce43c8 --- /dev/null +++ b/sglang/benchmark/kernels/scheduler_batch/benchmark_write_req_to_token_pool_triton.py @@ -0,0 +1,342 @@ +import itertools +import os + +import torch +import triton +import triton.language as tl + + +@triton.jit +def write_req_to_token_pool_triton( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + req_to_token_ptr_stride: tl.constexpr, +): + BLOCK_SIZE: tl.constexpr = 512 + pid = tl.program_id(0) + + req_pool_index = tl.load(req_pool_indices + pid) + pre_len = tl.load(pre_lens + pid) + seq_len = tl.load(seq_lens + pid) + + # TODO: optimize this? + cumsum_start = 0 + for i in range(pid): + cumsum_start += tl.load(extend_lens + i) + + num_loop = tl.cdiv(seq_len - pre_len, BLOCK_SIZE) + for i in range(num_loop): + offset = tl.arange(0, BLOCK_SIZE) + i * BLOCK_SIZE + mask = offset < (seq_len - pre_len) + value = tl.load(out_cache_loc + cumsum_start + offset, mask=mask) + tl.store( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + offset + + pre_len, + value, + mask=mask, + ) + + +@triton.jit +def write_req_to_token_pool_triton_optimize( + req_to_token_ptr, # [max_batch, max_context_len] + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + req_to_token_ptr_stride: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + pid_batch = tl.program_id(0) + pid_token = tl.program_id(1) + + req_pool_index = tl.load(req_pool_indices + pid_batch) + pre_len = tl.load(pre_lens + pid_batch) + seq_len = tl.load(seq_lens + pid_batch) + extend_len = seq_len - pre_len + + cumsum_start = 0 + for i in range(pid_batch): + cumsum_start += tl.load(extend_lens + i) + + token_start = pid_token * BLOCK_SIZE + + offset = tl.arange(0, BLOCK_SIZE) + actual_offset = token_start + offset + mask = actual_offset < extend_len + + src_ptr = out_cache_loc + cumsum_start + actual_offset + src_ptr = tl.max_contiguous(tl.multiple_of(src_ptr, BLOCK_SIZE), BLOCK_SIZE) + value = tl.load(src_ptr, mask=mask) + dst_ptr = ( + req_to_token_ptr + + req_pool_index * req_to_token_ptr_stride + + actual_offset + + pre_len + ) + dst_ptr = tl.max_contiguous(tl.multiple_of(dst_ptr, BLOCK_SIZE), BLOCK_SIZE) + + tl.store(dst_ptr, value, mask=mask) + + +def write_req_to_token_pool_reference( + req_to_token: torch.Tensor, + req_pool_indices: torch.Tensor, + pre_lens: torch.Tensor, + seq_lens: torch.Tensor, + extend_lens: torch.Tensor, + out_cache_loc: torch.Tensor, +) -> None: + """Reference implementation using PyTorch""" + for i in range(len(req_pool_indices)): + req_pool_idx = req_pool_indices[i].item() + pre_len = pre_lens[i].item() + seq_len = seq_lens[i].item() + extend_len = extend_lens[i].item() + + cumsum_start = sum(extend_lens[:i].tolist()) + + # Copy values from out_cache_loc to req_to_token + req_to_token[req_pool_idx, pre_len:seq_len] = out_cache_loc[ + cumsum_start : cumsum_start + extend_len + ] + + +def test_write_req_to_token_pool(): + max_batch = 4097 + max_context_len = 6148 + batch_size = 1 + extend_len = 14 + + # Initialize input tensors + req_to_token = torch.zeros( + (max_batch, max_context_len), dtype=torch.int32, device="cuda" + ) + req_pool_indices = torch.tensor([42], dtype=torch.int32, device="cuda") + pre_lens = torch.tensor([8], dtype=torch.int32, device="cuda") + seq_lens = torch.tensor([22], dtype=torch.int32, device="cuda") + extend_lens = torch.tensor([extend_len], dtype=torch.int32, device="cuda") + out_cache_loc = torch.arange(extend_len, dtype=torch.int32, device="cuda") + + # Create copies for reference implementation + req_to_token_ref = req_to_token.clone() + req_to_token_opt = req_to_token.clone() + + # Run original triton kernel + write_req_to_token_pool_triton[(batch_size,)]( + req_to_token, + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + max_context_len, + ) + + # Run optimized triton kernel + def grid(batch_size, extend_len): + num_token_blocks = triton.cdiv(extend_len, 512) + return (batch_size, num_token_blocks) + + write_req_to_token_pool_triton_optimize[grid(batch_size, extend_len)]( + req_to_token_opt, + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + max_context_len, + BLOCK_SIZE=512, + ) + + # Run reference implementation + write_req_to_token_pool_reference( + req_to_token_ref, + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + ) + + # Compare results + torch.testing.assert_close(req_to_token, req_to_token_ref) + torch.testing.assert_close(req_to_token_opt, req_to_token_ref) + + # Test case 2: batch size > 1 + batch_size = 3 + extend_lens_list = [14, 20, 30] + total_extend_len = sum(extend_lens_list) + + req_to_token = torch.zeros( + (max_batch, max_context_len), dtype=torch.int32, device="cuda" + ) + req_pool_indices = torch.tensor([42, 100, 200], dtype=torch.int32, device="cuda") + pre_lens = torch.tensor([8, 10, 15], dtype=torch.int32, device="cuda") + seq_lens = torch.tensor([22, 30, 45], dtype=torch.int32, device="cuda") + extend_lens = torch.tensor(extend_lens_list, dtype=torch.int32, device="cuda") + out_cache_loc = torch.arange(total_extend_len, dtype=torch.int32, device="cuda") + + req_to_token_ref = req_to_token.clone() + req_to_token_opt = req_to_token.clone() + + # Run original triton kernel + write_req_to_token_pool_triton[(batch_size,)]( + req_to_token, + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + max_context_len, + ) + + # Run optimized triton kernel + max_extend_len = max(extend_lens_list) + write_req_to_token_pool_triton_optimize[grid(batch_size, max_extend_len)]( + req_to_token_opt, + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + max_context_len, + BLOCK_SIZE=512, + ) + + # Run reference implementation + write_req_to_token_pool_reference( + req_to_token_ref, + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + ) + + # Compare results + torch.testing.assert_close(req_to_token, req_to_token_ref) + torch.testing.assert_close(req_to_token_opt, req_to_token_ref) + + +def get_benchmark(): + batch_sizes = [1, 2, 4, 8, 16, 32, 64, 128] + extend_lens = [32, 64, 128, 256, 512, 1024, 2048, 4096, 8192] + configs = list(itertools.product(batch_sizes, extend_lens)) + + @triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["batch_size", "extend_len"], + x_vals=configs, + line_arg="provider", + line_vals=["reference", "triton", "triton_optimize"], + line_names=["PyTorch", "Triton", "Triton Optimized"], + styles=[("blue", "-"), ("green", "-"), ("red", "-")], + ylabel="us", + plot_name="write-req-to-token-pool-performance", + args={}, + ) + ) + def benchmark(batch_size, extend_len, provider): + max_batch = 256 + max_context_len = 16384 + + extend_lens_list = [extend_len] * batch_size + total_extend_len = sum(extend_lens_list) + + req_to_token = torch.zeros( + (max_batch, max_context_len), dtype=torch.int32, device="cuda" + ) + req_pool_indices = torch.arange(batch_size, dtype=torch.int32, device="cuda") + pre_lens = torch.ones(batch_size, dtype=torch.int32, device="cuda") * 8 + seq_lens = pre_lens + extend_len + extend_lens = torch.tensor(extend_lens_list, dtype=torch.int32, device="cuda") + out_cache_loc = torch.arange(total_extend_len, dtype=torch.int32, device="cuda") + + quantiles = [0.5, 0.2, 0.8] + + if provider == "reference": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: write_req_to_token_pool_reference( + req_to_token.clone(), + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + ), + quantiles=quantiles, + ) + elif provider == "triton": + ms, min_ms, max_ms = triton.testing.do_bench( + lambda: write_req_to_token_pool_triton[(batch_size,)]( + req_to_token.clone(), + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + max_context_len, + ), + quantiles=quantiles, + ) + else: + + def run_optimized(): + block_size = 128 if extend_len <= 1024 else 512 + grid_config = (batch_size, triton.cdiv(extend_len, block_size)) + write_req_to_token_pool_triton_optimize[grid_config]( + req_to_token.clone(), + req_pool_indices, + pre_lens, + seq_lens, + extend_lens, + out_cache_loc, + max_context_len, + BLOCK_SIZE=block_size, + ) + + ms, min_ms, max_ms = triton.testing.do_bench( + run_optimized, quantiles=quantiles + ) + + return 1000 * ms, 1000 * max_ms, 1000 * min_ms + + return benchmark + + +def run_benchmark(save_path: str = "./configs/benchmark_ops/write_req_to_token_pool/"): + """Run benchmark and save results""" + + # Ensure save path exists + os.makedirs(save_path, exist_ok=True) + + # Run correctness test + test_write_req_to_token_pool() + print("Correctness test passed!") + + # Run performance test + benchmark = get_benchmark() + benchmark.run(print_data=True, save_path=save_path) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument( + "--save_path", + type=str, + default="./configs/benchmark_ops/write_req_to_token_pool/", + help="Path to save benchmark results", + ) + args = parser.parse_args() + + run_benchmark(args.save_path) diff --git a/sglang/benchmark/line_retrieval/README.md b/sglang/benchmark/line_retrieval/README.md new file mode 100644 index 0000000..3c81a63 --- /dev/null +++ b/sglang/benchmark/line_retrieval/README.md @@ -0,0 +1,37 @@ +## Download data + +``` +wget https://raw.githubusercontent.com/merrymercy/merrymercy.github.io/master/files/random_words.json +python3 gen_data.py --number 1000 +``` + +## Run benchmark + +### Benchmark sglang +``` +python3 -m sglang.launch_server --model-path codellama/CodeLlama-7b-hf --port 30000 +``` + +``` +python3 bench_sglang.py --src-index 600 --num-q 50 --parallel 1 +``` + + +### + +``` +# original +Accuracy: 0.940, latency: 332.83 s + +# parallel encoding (no_adjust, offset = 1000) +Accuracy: 0.760, latency: 238.46 s + +# parallel encoding (no_adjust, offset = 3000) +Accuracy: 0.760, latency: 238.46 s + +# parallel encoding (no_adjust, offset = 0) +Accuracy: 0.520, latency: 238.46 s + +# parallel encoding (adjust_cache) +Accuracy: 0.460, latency: 257.66 s +``` diff --git a/sglang/benchmark/line_retrieval/bench_sglang.py b/sglang/benchmark/line_retrieval/bench_sglang.py new file mode 100644 index 0000000..922d500 --- /dev/null +++ b/sglang/benchmark/line_retrieval/bench_sglang.py @@ -0,0 +1,149 @@ +import argparse +import json +import re +import time + +import numpy as np + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text + + +@sgl.function +def line_retrieval(s, prefix, suffix, body_0, body_1, body_2, body_3): + s += prefix + "\n" + + contexts = [body_0, body_1, body_2, body_3] + position_ids_offset = [i * 1000 for i in range(len(contexts))] + forks = s.fork(len(contexts), position_ids_offset) + forks += lambda i: contexts[i] + "\n" + forks.join(mode="concate_and_append") + + s += "\n" + suffix + s += sgl.gen("answer", max_tokens=16) + + +def eval_model(args, line_obj, num_hoops, src_indices, dst_percents): + arguments = [] + labels = [] + sum_src_indices = [] + sum_dst_indices = [] + + for i in range(len(src_indices)): + for j in range(len(dst_percents)): + src_index = src_indices[i] + dst_percent = dst_percents[j] + + query_indices = line_obj["group_by_num_hoops"][str(num_hoops)] + query_indices = [ + q + for q in query_indices + if all(l <= src_index for l in line_obj["links"][q]) and q < src_index + ] + dst_index = query_indices[ + min(int(len(query_indices) * dst_percent), len(query_indices) - 1) + ] + label = line_obj["values"][dst_index] + + body = line_obj["lines"][: src_index + 1] + suffix = line_obj["suffix"].replace("???", line_obj["indices"][dst_index]) + body_part_len = len(body) // 4 + + arguments.append( + { + "prefix": line_obj["prefix"], + "body_0": "\n".join(body[:body_part_len]), + "body_1": "\n".join(body[body_part_len : 2 * body_part_len]), + "body_2": "\n".join(body[2 * body_part_len : 3 * body_part_len]), + "body_3": "\n".join(body[3 * body_part_len :]), + "suffix": suffix, + } + ) + labels.append(label) + sum_src_indices.append(src_index) + sum_dst_indices.append(dst_index) + + # Select backend + backend = select_sglang_backend(args) + + tic = time.time() + states = line_retrieval.run_batch( + arguments, + temperature=0, + backend=backend, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.time() - tic + + corrects = [] + for i in range(len(arguments)): + output = states[i]["answer"] + prompt_len = states[i].get_meta_info("answer").get("prompt_length", -1) + label = labels[i] + + # Try all numbers + findall = re.findall("\d+", output) + if not findall: + response_number = output + else: + for response_number in findall: + if response_number == label: + break + + correct = response_number == label + corrects.append(correct) + + # Log results + summary = ( + f"Line index: {sum_src_indices[i]} -> {sum_dst_indices[i]}, " + f"Prompt len: {prompt_len}, " + f"Correct: {correct}, " + f"Label: {label}, Predicted: {response_number}, " + ) + print(summary) + + accuracy = np.mean(corrects) + print(f"Accuracy: {accuracy:.3f}, latency: {latency:.2f} s") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "line_retrieval", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": len(arguments), + "other": { + "num_questions": len(arguments), + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +def main(args): + line_obj = json.load(open(args.data_path, "r")) + + num_hoops = args.num_hoops + for src_index in args.src_index: + src_indices = [src_index] + num_queries = args.num_queries_per_src + dst_percents = [i * (1 / (num_queries)) for i in range(num_queries)] + eval_model(args, line_obj, num_hoops, src_indices, dst_percents) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="lines_1000_0.0.json") + parser.add_argument("--src-index", type=int, nargs="+", default=[100]) + parser.add_argument("--num-queries-per-src", type=int, default=10) + parser.add_argument("--num-hoops", type=int, default=1) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/line_retrieval/gen_data.py b/sglang/benchmark/line_retrieval/gen_data.py new file mode 100644 index 0000000..c88ecba --- /dev/null +++ b/sglang/benchmark/line_retrieval/gen_data.py @@ -0,0 +1,139 @@ +""" +Generate line data for line retrieval task. + +Usage: +python3 gen_data.py --number 1000 +""" + +import argparse +import json +from collections import defaultdict + +import numpy as np +from tqdm import tqdm + + +def generate_lines(random_words, num_lines, redirect_ratio): + prefix = "Here is a list of lines, each with its corresponding REGISTER_CONTENT value. Please memorize them. Be prepared to provide the REGISTER_CONTENT value for a specific line index when I ask." + suffix = "The list has ended. Please give the final REGISTER_CONTENT value for a specific line after resovling the redirections and references. For example, the REGISTER_CONTENT of Line __idx0__ is __val0__. The REGISTER_CONTENT of Line __idx1__ is __val1__. The REGISTER_CONTENT of Line __idx2__ is __val2__. The REGISTER_CONTENT of Line ??? is" + + # Raw lines + visited_indices = set([None]) + visited_values = set([None]) + + lines = [] + redirects = [] + indices = [] + values = [] + for i in tqdm(range(num_lines)): + line_index = None + while line_index in visited_indices: + line_index = "-".join(np.random.choice(random_words, size=(2,))) + visited_indices.add(line_index) + + line_value = np.random.randint(low=0, high=999999) + line_value = f"{line_value:06}" + + line = f"Line {line_index}: The REGISTER_CONTENT is {line_value}." + lines.append(line) + redirects.append(None) + indices.append(line_index) + values.append(line_value) + + # Add redirect + if redirect_ratio > 0: + num_redirect_lines = int(len(lines) * redirect_ratio) + redirect_indices = np.random.choice( + np.arange(len(lines)), size=(num_redirect_lines,), replace=False + ) + for i in redirect_indices: + target_idx = np.random.choice(min(i * 2 + 100, num_lines)) + lines[i] = ( + f"Line {indices[i]}: The REGISTER_CONTENT is the same as Line {indices[target_idx]}." + ) + redirects[i] = target_idx + + # Build links and find sources + links = [[] for _ in range(num_lines)] + contains_ring = set() + for i in range(num_lines): + if redirects[i] is None: + continue + + tmp_link = [] + cur = i + visited = set() + while redirects[cur] is not None: + visited.add(cur) + tmp_link.append(redirects[cur]) + cur = redirects[cur] + + if cur in visited: + contains_ring.add(i) + tmp_link = None + break + values[i] = values[cur] + links[i] = tmp_link + + # Group by num_links + group_by_num_hoops = defaultdict(list) + for i in range(num_lines): + if i in contains_ring: + continue + group_by_num_hoops[len(links[i]) + 1].append(i) + + keys = sorted(list(group_by_num_hoops.keys())) + for num_links in keys: + print(f"#links: {num_links}, #lines: {len(group_by_num_hoops[num_links])}") + + # Append few-shot examples + hoop1_candidates = list(group_by_num_hoops[1]) + hoop1_candidate_keys = {c: max([c] + links[c]) for c in hoop1_candidates} + hoop1_candidates.sort(key=lambda c: hoop1_candidate_keys[c]) + hoop2_candidates = list(group_by_num_hoops[2]) + hoop2_candidate_keys = {c: max([c] + links[c]) for c in hoop2_candidates} + hoop2_candidates.sort(key=lambda c: hoop2_candidate_keys[c]) + + i = hoop1_candidates[5] + suffix = suffix.replace("__idx0__", indices[i]).replace("__val0__", values[i]) + if len(hoop2_candidates): + i = hoop2_candidates[0] + suffix = suffix.replace("__idx1__", indices[i]).replace("__val1__", values[i]) + i = hoop2_candidates[1] + suffix = suffix.replace("__idx2__", indices[i]).replace("__val2__", values[i]) + else: + i = hoop1_candidates[1] + suffix = suffix.replace("__idx1__", indices[i]).replace("__val1__", values[i]) + i = hoop1_candidates[10] + suffix = suffix.replace("__idx2__", indices[i]).replace("__val2__", values[i]) + + obj = { + "prefix": prefix, + "suffix": suffix, + "lines": lines, + "indices": indices, + "values": values, + "links": links, + "group_by_num_hoops": group_by_num_hoops, + "contains_ring": sorted(list(contains_ring)), + } + return obj + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--number", type=int) + parser.add_argument("--redirect-ratio", type=float, default=0.0) + args = parser.parse_args() + + num_lines = args.number + + random_words_filename = "random_words.json" + random_words = json.load(open(random_words_filename, "r")) + + np.random.seed(42) + obj = generate_lines(random_words, num_lines, args.redirect_ratio) + + fout = f"lines_{num_lines}_{args.redirect_ratio:.1f}.json" + with open(fout, "w") as fout: + json.dump(obj, fout, indent=2) diff --git a/sglang/benchmark/llava_bench/README.md b/sglang/benchmark/llava_bench/README.md new file mode 100644 index 0000000..a71f339 --- /dev/null +++ b/sglang/benchmark/llava_bench/README.md @@ -0,0 +1,61 @@ +## Download benchmark images + +``` +python3 download_images.py +``` + +image benchmark source: https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild + +### Other Dependency +``` +pip3 install "sglang[all]" +pip3 install "torch>=2.1.2" "transformers>=4.36" pillow +``` + +## Run benchmark + +### Benchmark sglang +Launch a server +``` +python3 -m sglang.launch_server --model-path liuhaotian/llava-v1.6-vicuna-7b --tokenizer-path llava-hf/llava-1.5-7b-hf --port 30000 +``` + +Run benchmark +``` +# Run with local models +python3 bench_sglang.py --num-questions 60 + +# Run with OpenAI models +python3 bench_sglang.py --num-questions 60 --backend gpt-4-vision-preview +``` + +### Bench LLaVA original code +``` +git clone git@github.com:haotian-liu/LLaVA.git +cd LLaVA +git reset --hard 9a26bd1435b4ac42c282757f2c16d34226575e96 +pip3 install -e . + +cd ~/sglang/benchmark/llava_bench +CUDA_VISIBLE_DEVICES=0 bash bench_hf_llava_bench.sh +``` + + +### Benchmark llama.cpp + +``` +# Install +CMAKE_ARGS="-DLLAMA_CUBLAS=on" pip install llama-cpp-python +pip install sse_starlette starlette_context pydantic_settings + +# Download weights +mkdir -p ~/model_weights/llava-v1.5-7b/ +wget https://huggingface.co/mys/ggml_llava-v1.5-7b/resolve/main/ggml-model-f16.gguf -O ~/model_weights/llava-v1.5-7b/ggml-model-f16.gguf +wget https://huggingface.co/mys/ggml_llava-v1.5-7b/resolve/main/mmproj-model-f16.gguf -O ~/model_weights/llava-v1.5-7b/mmproj-model-f16.gguf +``` + +``` +python3 -m llama_cpp.server --model ~/model_weights/llava-v1.5-7b/ggml-model-f16.gguf --clip_model_path ~/model_weights/llava-v1.5-7b/mmproj-model-f16.gguf --chat_format llava-1-5 --port 23000 + +OPENAI_BASE_URL=http://localhost:23000/v1 python3 bench_sglang.py --backend gpt-4-vision-preview --num-q 1 +``` diff --git a/sglang/benchmark/llava_bench/bench_hf_llava_bench.sh b/sglang/benchmark/llava_bench/bench_hf_llava_bench.sh new file mode 100644 index 0000000..a51a715 --- /dev/null +++ b/sglang/benchmark/llava_bench/bench_hf_llava_bench.sh @@ -0,0 +1,9 @@ +#!/bin/bash + +python -m llava.eval.model_vqa \ + --model-path liuhaotian/llava-v1.5-7b \ + --question-file ./questions.jsonl \ + --image-folder ./images \ + --answers-file ./answers_hf.jsonl \ + --temperature 0 \ + --conv-mode vicuna_v1 diff --git a/sglang/benchmark/llava_bench/bench_hf_mme.sh b/sglang/benchmark/llava_bench/bench_hf_mme.sh new file mode 100644 index 0000000..6ed332f --- /dev/null +++ b/sglang/benchmark/llava_bench/bench_hf_mme.sh @@ -0,0 +1,9 @@ +#!/bin/bash + +python -m llava.eval.model_vqa_loader \ + --model-path liuhaotian/llava-v1.5-7b \ + --question-file ./mme_pack/llava_mme_bench_replace.jsonl \ + --image-folder ./mme_pack/MME_Benchmark_release_version \ + --answers-file ./answers_hf_mme.jsonl \ + --temperature 0 \ + --conv-mode vicuna_v1 diff --git a/sglang/benchmark/llava_bench/bench_sglang.py b/sglang/benchmark/llava_bench/bench_sglang.py new file mode 100644 index 0000000..f84c8a9 --- /dev/null +++ b/sglang/benchmark/llava_bench/bench_sglang.py @@ -0,0 +1,96 @@ +import argparse +import json +import os +import time + +import tqdm + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + + +@sgl.function +def image_qa(s, image_file, question): + s += sgl.user(sgl.image(image_file) + question) + s += sgl.assistant(sgl.gen("answer", max_tokens=args.max_tokens)) + + +def main(args): + lines = list(read_jsonl(args.question_file))[: args.num_questions] + arguments = [ + { + "image_file": os.path.abspath(args.image_folder + "/" + l["image"]), + "question": l["text"], + } + for l in lines + ] + # arguments = [ + # {"image_file": + # Image.open(os.path.abspath(args.image_folder + "/" + l["image"])), + # "question": l["text"]} for l in lines + # ] + + states = [None] * len(lines) + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.time() + if args.parallel == 1: + for i in tqdm.tqdm(range(len(lines))): + image_file = arguments[i]["image_file"] + question = arguments[i]["question"] + ret = image_qa.run(image_file=image_file, question=question, temperature=0) + states[i] = ret + else: + states = image_qa.run_batch( + arguments, temperature=0, num_threads=args.parallel, progress_bar=True + ) + latency = time.time() - tic + + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + print(f"Write output to {args.answer_file}") + with open(args.answer_file, "w") as fout: + for i in range(len(lines)): + value = { + "question_id": lines[i]["question_id"], + "prompt": lines[i]["text"], + "text": states[i]["answer"].strip(), + "model_id": backend.model_info["model_path"], + "answer_id": i, + "metadata": {}, + } + fout.write(json.dumps(value) + "\n") + + with open(args.result_file, "a") as fout: + value = { + "task": "llava_bench", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": len(lines), + "parallel": args.parallel, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--question-file", type=str, default="questions.jsonl") + parser.add_argument("--answer-file", type=str, default="answers.jsonl") + parser.add_argument("--image-folder", type=str, default="./images") + parser.add_argument("--temperature", type=float, default=0.0) + parser.add_argument("--num-questions", type=int, default=None) + parser.add_argument("--max-tokens", type=int, default=768) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/llava_bench/bench_sglang_mme.sh b/sglang/benchmark/llava_bench/bench_sglang_mme.sh new file mode 100644 index 0000000..8c8c09c --- /dev/null +++ b/sglang/benchmark/llava_bench/bench_sglang_mme.sh @@ -0,0 +1,2 @@ +MME_FOLDER=./mme_pack +python3 bench_sglang.py --num-questions 5000 --question-file $MME_FOLDER/llava_mme_bench_replace.jsonl --answer-file answer_mme.jsonl --image-folder $MME_FOLDER/MME_Benchmark_release_version --max-tokens 4 diff --git a/sglang/benchmark/llava_bench/download_images.py b/sglang/benchmark/llava_bench/download_images.py new file mode 100644 index 0000000..701190a --- /dev/null +++ b/sglang/benchmark/llava_bench/download_images.py @@ -0,0 +1,20 @@ +import os + +# Create the 'images' directory if it doesn't exist +if not os.path.exists("images"): + os.makedirs("images") + +# Base URL +base_url = "https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild/resolve/main/images/" + +# Loop through image numbers +for i in range(1, 25): + # Format the image number with leading zeros + image_number = str(i).zfill(3) + image_url = base_url + image_number + ".jpg" + image_path = "images/" + image_number + ".jpg" + + # Download the image using wget + os.system(f"wget -O {image_path} {image_url}") + +print("Download complete.") diff --git a/sglang/benchmark/llm_judge/README.md b/sglang/benchmark/llm_judge/README.md new file mode 100644 index 0000000..164b19a --- /dev/null +++ b/sglang/benchmark/llm_judge/README.md @@ -0,0 +1,33 @@ +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 25 --parallel 8 +python3 bench_sglang.py --num-questions 16 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --backend vllm --num-questions 25 +``` + + +### Benchmark guidance +``` +python3 bench_other.py --backend guidance --num-questions 25 --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +### Benchmark lmql + +``` +python3 bench_other.py --backend lmql --num-questions 25 --parallel 1 +``` diff --git a/sglang/benchmark/llm_judge/bench_other.py b/sglang/benchmark/llm_judge/bench_other.py new file mode 100644 index 0000000..2231bcd --- /dev/null +++ b/sglang/benchmark/llm_judge/bench_other.py @@ -0,0 +1,151 @@ +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +system_prompt = "Please serve as an impartial judge and rigorously evaluate the quality of the following article. Apply the most stringent standards possible, showing no leniency." + +dimension_prompts = [ + "Content: This refers to the essences of the essay. The substance should be well researched, accurate, relevant to the topic and should show a thorough understanding of the subject. The essay should also reflect a clear goal or purpose.", + "Organization and Structure: An essay needs to be properly structured with a clear introduction, body, and conclusion. The essay should flow naturally, with one paragraph leading seamlessly into the next.", + "Argument and Analysis: The argument made in the essay should be logical, coherent and clearly articulated. Each point made should be backed up by solid evidence and thorough analysis.", + "Clarity and Precision: The essay should be written in a clear and concise manner. The points made should be easily understood by the reader. The language used should also be precise and unambiguous.", + "Grammar and Punctuation: Proper use of grammar and punctuation is vital in an academic essay. Errors in grammar and punctuation not only distract the reader but can also negatively impact the meaning and interpretation of the content.", + "Referencing and Citation: An essay should contain proper citations and references for all sources used. This not only prevents accusations of plagiarism but also gives credit to the authors of the works that have contributed to the essay. The citation should adhere to a specific format as required by the academic institution or specified by the professor.", +] + + +def multi_dimension_judge(article, generate): + s = system_prompt + s += "\n```\n" + article + "\n```\n\n" + + judges = [] + for i in range(len(dimension_prompts)): + comp = generate( + s + + "USER: Please judge the quality based on the following metric. " + + dimension_prompts[i] + + " Please provide a single-paragraph judgement. " + + "Focus on the provided metric and do not say other things. " + 'End your judgement paragraph with the word "END"\nJUDGE:', + max_tokens=256, + stop="END", + ) + judges.append(comp) + + s += "I will judge the quality based on the following metrics.\n" + for i in range(len(dimension_prompts)): + s += dimension_prompts[i].split(":")[0] + ": " + judges[i].strip() + "\n" + + s += "In summary, on a scale of 1 to 10, I would give the article a score of" + s += generate(s, max_tokens=2, stop=None) + + return s + + +async def multi_dimension_judge_async(article, generate): + s = system_prompt + s += "\n```\n" + article + "\n```\n\n" + + judges = [] + for i in range(len(dimension_prompts)): + comp = await generate( + s + + "USER: Please judge the quality based on the following metric. " + + dimension_prompts[i] + + " Please provide a single-paragraph judgement. " + + "Focus on the provided metric and do not say other things. " + 'End your judgement paragraph with the word "END"\nJUDGE:', + max_tokens=256, + stop="END", + ) + judges.append(comp) + + s += "I will judge the quality based on the following metrics.\n" + for i in range(len(dimension_prompts)): + s += dimension_prompts[i].split(":")[0] + ": " + judges[i].strip() + "\n" + + s += "In summary, on a scale of 1 to 10, I would give the article a score of" + s += await generate(s, max_tokens=2, stop=None) + + return s + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_questions] + states = [None] * len(lines) + + # Select backend + call_generate = partial(get_call_generate(args), temperature=0) + + # Run requests + tic = time.time() + + if args.backend != "lmql": + + def get_one_answer(i): + states[i] = multi_dimension_judge(lines[i], call_generate) + + if args.parallel == 1: + for i in tqdm(range(len(lines))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(lines)))), + total=len(lines), + ) + ) + + else: + import asyncio + + async def get_one_answer_async(i): + states[i] = await multi_dimension_judge_async(lines[i], call_generate) + + batches = [] + for i in range(0, len(lines), args.parallel): + batches.append(list(range(i, min(i + args.parallel, len(lines))))) + + loop = asyncio.get_event_loop() + for bt in tqdm(batches): + loop.run_until_complete( + asyncio.gather(*[get_one_answer_async(i) for i in bt]) + ) + + latency = time.time() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "llm_judge", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="articles.jsonl") + parser.add_argument("--num-questions", type=int, default=20) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/llm_judge/bench_sglang.py b/sglang/benchmark/llm_judge/bench_sglang.py new file mode 100644 index 0000000..38c9597 --- /dev/null +++ b/sglang/benchmark/llm_judge/bench_sglang.py @@ -0,0 +1,97 @@ +import argparse +import json +import time + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + +system_prompt = "Please serve as an impartial judge and rigorously evaluate the quality of the following article. Apply the most stringent standards possible, showing no leniency." + +dimension_prompts = [ + "Content: This refers to the essences of the essay. The substance should be well researched, accurate, relevant to the topic and should show a thorough understanding of the subject. The essay should also reflect a clear goal or purpose.", + "Organization and Structure: An essay needs to be properly structured with a clear introduction, body, and conclusion. The essay should flow naturally, with one paragraph leading seamlessly into the next.", + "Argument and Analysis: The argument made in the essay should be logical, coherent and clearly articulated. Each point made should be backed up by solid evidence and thorough analysis.", + "Clarity and Precision: The essay should be written in a clear and concise manner. The points made should be easily understood by the reader. The language used should also be precise and unambiguous.", + "Grammar and Punctuation: Proper use of grammar and punctuation is vital in an academic essay. Errors in grammar and punctuation not only distract the reader but can also negatively impact the meaning and interpretation of the content.", + "Referencing and Citation: An essay should contain proper citations and references for all sources used. This not only prevents accusations of plagiarism but also gives credit to the authors of the works that have contributed to the essay. The citation should adhere to a specific format as required by the academic institution or specified by the professor.", +] + + +@sgl.function +def multi_dimension_judge(s, article): + s += system_prompt + s += "\n```\n" + article + "\n```\n\n" + + forks = s.fork(len(dimension_prompts)) + for i in range(len(dimension_prompts)): + forks[i] += ( + "USER: Please judge the quality based on the following metric. " + + dimension_prompts[i] + + " Please provide a single-paragraph judgement. " + + "Focus on the provided metric and do not say other things. " + 'End your judgement paragraph with the word "END"\nJUDGE:' + ) + forks[i] += sgl.gen("judgement", max_tokens=256, stop="END") + forks.join() + + s += "I will judge the quality based on the following metrics.\n" + for i in range(len(dimension_prompts)): + s += ( + dimension_prompts[i].split(":")[0] + + ": " + + forks[i]["judgement"].strip() + + "\n" + ) + + s += "In summary, on a scale of 1 to 10, I would give the article a score of" + s += sgl.gen("score", max_tokens=2) + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_questions] + arguments = [{"article": l} for l in lines] + + # Select backend + backend = select_sglang_backend(args) + + # Run requests + tic = time.time() + states = multi_dimension_judge.run_batch( + arguments, + temperature=0, + backend=backend, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.time() - tic + + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "llm_judge", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="articles.jsonl") + parser.add_argument("--num-questions", type=int, default=20) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/long_json_decode/README.md b/sglang/benchmark/long_json_decode/README.md new file mode 100644 index 0000000..37fceee --- /dev/null +++ b/sglang/benchmark/long_json_decode/README.md @@ -0,0 +1,33 @@ +## Run benchmark + +### Benchmark sglang +``` +python3 -m sglang.launch_server --model-path codellama/CodeLlama-7b-instruct-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 5 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model codellama/CodeLlama-7b-instruct-hf --disable-log-requests --port 21000 --gpu 0.97 +``` + +``` +python3 bench_other.py --backend vllm --num-questions 5 +``` + + +### Benchmark guidance +``` +python3 bench_other.py --backend guidance --num-questions 5 --parallel 1 --n-ctx 11000 --model-path path/to/code-llama/gguf +``` + + +### Build dataset +``` +pip install wikipedia +python3 build_dataset.py +``` diff --git a/sglang/benchmark/long_json_decode/bench_other.py b/sglang/benchmark/long_json_decode/bench_other.py new file mode 100644 index 0000000..a83c797 --- /dev/null +++ b/sglang/benchmark/long_json_decode/bench_other.py @@ -0,0 +1,89 @@ +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + + +def json_decode(document, generate): + s = "Please extract the information of a city from the following wikipedia page.\n" + s += "Page begin.\n" + document + "Page end.\n" + s += "Here is the name, country, and symbol of the city in JSON format.\n" + s += "{\n" + s += ' "name": "' + s += generate(s, max_tokens=8, stop='"') + '",\n' + s += ' "country": "' + s += generate(s, max_tokens=8, stop='"') + '",\n' + s += ' "air port code": "' + s += generate(s, max_tokens=8, stop='"') + '",\n' + s += ' "top 3 landmarks": "' + s += generate(s, max_tokens=24, stop='"') + '",\n' + s += "}\n" + return s + + +def main(args): + lines = read_jsonl(args.data_path) + arguments = [] + for i in range(len(lines[: args.num_questions])): + arguments.append( + { + "document": lines[i]["document"], + } + ) + states = [None] * len(arguments) + + # Select backend + call_generate = partial(get_call_generate(args), temperature=0) + + # Run requests + def get_one_answer(i): + states[i] = json_decode(generate=call_generate, **arguments[i]) + + tic = time.time() + if args.parallel == 1: + for i in tqdm(range(len(arguments))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(arguments)))), + total=len(arguments), + ) + ) + + latency = time.time() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "long_json_decode", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="questions.jsonl") + parser.add_argument("--num-questions", type=int, default=100) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/long_json_decode/bench_sglang.py b/sglang/benchmark/long_json_decode/bench_sglang.py new file mode 100644 index 0000000..6e19a73 --- /dev/null +++ b/sglang/benchmark/long_json_decode/bench_sglang.py @@ -0,0 +1,81 @@ +import argparse +import json +import time + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + + +@sgl.function +def json_decode(s, document): + s += "Please extract the information of a city from the following wikipedia page.\n" + s += "Page begin.\n" + document + "Page end.\n" + s += "Here is the name, country, and symbol of the city in JSON format.\n" + s += "{\n" + s += ' "name": "' + sgl.gen("name", max_tokens=8, stop='"') + '",\n' + s += ' "country": "' + sgl.gen("country", max_tokens=8, stop='"') + '",\n' + s += ( + ' "air port code": "' + + sgl.gen("air port code", max_tokens=8, stop='"') + + '",\n' + ) + s += ( + ' "top 3 landmarks": "' + + sgl.gen("landmarks", max_tokens=24, stop='"') + + '",\n' + ) + s += "}\n" + + +def main(args): + lines = read_jsonl(args.data_path) + arguments = [] + for i in range(len(lines[: args.num_questions])): + arguments.append( + { + "document": lines[i]["document"], + } + ) + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.time() + states = json_decode.run_batch( + arguments, temperature=0, num_threads=args.parallel, progress_bar=True + ) + latency = time.time() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "long_json_decode", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="questions.jsonl") + parser.add_argument("--num-questions", type=int, default=10) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/long_json_decode/build_dataset.py b/sglang/benchmark/long_json_decode/build_dataset.py new file mode 100644 index 0000000..7b15d39 --- /dev/null +++ b/sglang/benchmark/long_json_decode/build_dataset.py @@ -0,0 +1,27 @@ +import json + +import transformers +import wikipedia + +name = "meta-llama/Llama-2-7b-chat-hf" +t = transformers.AutoTokenizer.from_pretrained(name) +city_names = ["los angles", "london", "tokyo", "beijing", "singapore"] + + +for city_name in city_names: + content = str(wikipedia.page(city_name).content) + content = content.replace("\n\n", "\n") + + tokens = t.encode(content) + + truncate_len = int((10000 / len(tokens)) * len(content)) + truncate_content = content[:truncate_len] + truncate_tokens = t.encode(truncate_content) + + # Count token + print( + f"city_name: {city_name}, #tokens: {len(tokens)}, #truncate tokens: {len(truncate_tokens)}" + ) + + with open("questions.jsonl", "a") as fout: + fout.write(json.dumps({"document": truncate_content}) + "\n") diff --git a/sglang/benchmark/lora/launch_server.py b/sglang/benchmark/lora/launch_server.py new file mode 100644 index 0000000..33ad848 --- /dev/null +++ b/sglang/benchmark/lora/launch_server.py @@ -0,0 +1,68 @@ +import argparse +import os + +NUM_LORAS = 4 +LORA_PATH = { + "base": "meta-llama/Llama-2-7b-hf", + "lora": "winddude/wizardLM-LlaMA-LoRA-7B", +} + + +def launch_server(args): + base_path = LORA_PATH["base"] + lora_path = LORA_PATH["lora"] + + if args.base_only: + cmd = f"python3 -m sglang.launch_server --model {base_path} " + else: + cmd = f"python3 -m sglang.launch_server --model {base_path} --lora-paths " + for i in range(NUM_LORAS): + lora_name = f"lora{i}" + cmd += f"{lora_name}={lora_path} " + cmd += f"--disable-radix --disable-cuda-graph " + cmd += f"--max-loras-per-batch {args.max_loras_per_batch} " + cmd += f"--max-running-requests {args.max_running_requests} " + cmd += f"--lora-backend {args.lora_backend} " + cmd += f"--tp-size {args.tp_size} " + if args.disable_custom_all_reduce: + cmd += "--disable-custom-all-reduce" + print(cmd) + os.system(cmd) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--base-only", + action="store_true", + ) + parser.add_argument( + "--max-loras-per-batch", + type=int, + default=8, + ) + parser.add_argument( + "--max-running-requests", + type=int, + default=8, + ) + parser.add_argument( + "--lora-backend", + type=str, + default="triton", + ) + parser.add_argument( + "--tp-size", + type=int, + default=1, + help="Tensor parallel size for distributed inference", + ) + # disable_custom_all_reduce + parser.add_argument( + "--disable-custom-all-reduce", + action="store_true", + help="Disable custom all reduce when device does not support p2p communication", + ) + args = parser.parse_args() + + launch_server(args) diff --git a/sglang/benchmark/lora/lora_bench.py b/sglang/benchmark/lora/lora_bench.py new file mode 100644 index 0000000..1d93942 --- /dev/null +++ b/sglang/benchmark/lora/lora_bench.py @@ -0,0 +1,475 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import argparse +import asyncio +import json +import random +import resource +import sys +import time +import traceback +from argparse import ArgumentParser +from datetime import datetime +from typing import Any, Dict, List, Optional, Tuple + +import aiohttp +import numpy as np +from launch_server import LORA_PATH, NUM_LORAS +from tqdm.asyncio import tqdm +from transformers import PreTrainedTokenizerBase + +from sglang.bench_serving import ( + AIOHTTP_TIMEOUT, + RequestFuncInput, + RequestFuncOutput, + calculate_metrics, + get_request, + get_tokenizer, + remove_prefix, + sample_random_requests, +) + +global args + + +# set ignore_eos True by default +async def async_request_openai_completions( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + # assert api_url.endswith( + # "completions" + # ), "OpenAI Completions API URL must end with 'completions'." + + prompt = request_func_input.prompt + + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + # payload = { + # "model": request_func_input.model, + # "prompt": prompt, + # "temperature": 0.0, + # "best_of": 1, + # "max_tokens": request_func_input.output_len, + # "stream": not args.disable_stream, + # "ignore_eos": not args.disable_ignore_eos, + # **request_func_input.extra_request_body, + # } + # headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"} + if args.base_only: + payload = { + "text": prompt, + "sampling_params": {"max_new_tokens": request_func_input.output_len}, + } + else: + payload = { + "text": prompt, + "sampling_params": {"max_new_tokens": request_func_input.output_len}, + "lora_path": f"lora{random.randint(0, NUM_LORAS - 1)}", + } + headers = {"Authorization": ""} + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if data["text"]: + # if data["choices"][0]["text"]: + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + # generated_text += data["choices"][0]["text"] + generated_text += data["text"] + + output.generated_text = generated_text + output.success = True + output.latency = latency + output.output_len = request_func_input.output_len + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +ASYNC_REQUEST_FUNCS = { + "sglang": async_request_openai_completions, +} + + +async def benchmark( + backend: str, + api_url: str, + model_id: str, + tokenizer: PreTrainedTokenizerBase, + input_requests: List[Tuple[str, int, int]], + request_rate: float, + disable_tqdm: bool, + extra_request_body: Dict[str, Any], +): + if backend in ASYNC_REQUEST_FUNCS: + request_func = ASYNC_REQUEST_FUNCS[backend] + else: + raise ValueError(f"Unknown backend: {backend}") + + print("Starting initial single prompt test run...") + test_prompt, test_prompt_len, test_output_len = input_requests[0] + test_input = RequestFuncInput( + model=model_id, + prompt=test_prompt, + api_url=api_url, + prompt_len=test_prompt_len, + output_len=test_output_len, + lora_name="dummy", # the lora_name argument will not be used + extra_request_body=extra_request_body, + ) + test_output = await request_func(request_func_input=test_input) + if not test_output.success: + raise ValueError( + "Initial test run failed - Please make sure benchmark arguments " + f"are correctly specified. Error: {test_output.error}" + ) + else: + print("Initial test run completed. Starting main benchmark run...") + + pbar = None if disable_tqdm else tqdm(total=len(input_requests)) + + benchmark_start_time = time.perf_counter() + tasks: List[asyncio.Task] = [] + async for request in get_request(input_requests, request_rate): + prompt, prompt_len, output_len = request + request_func_input = RequestFuncInput( + model=model_id, + prompt=prompt, + api_url=api_url, + prompt_len=prompt_len, + output_len=output_len, + lora_name="dummy", + extra_request_body=extra_request_body, + ) + tasks.append( + asyncio.create_task( + request_func(request_func_input=request_func_input, pbar=pbar) + ) + ) + outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) + + if pbar is not None: + pbar.close() + + benchmark_duration = time.perf_counter() - benchmark_start_time + + metrics, output_lens = calculate_metrics( + input_requests=input_requests, + outputs=outputs, + dur_s=benchmark_duration, + tokenizer=tokenizer, + backend=backend, + ) + + print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) + print("{:<40} {:<10}".format("Backend:", backend)) + print("{:<40} {:<10}".format("Traffic request rate:", request_rate)) + print("{:<40} {:<10}".format("Successful requests:", metrics.completed)) + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration)) + print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) + print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output)) + print( + "{:<40} {:<10}".format( + "Total generated tokens (retokenized):", metrics.total_output_retokenized + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Request throughput (req/s):", metrics.request_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Input token throughput (tok/s):", metrics.input_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Output token throughput (tok/s):", metrics.output_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format("Total throughput (tok/s):", metrics.total_throughput) + ) + print("{s:{c}^{n}}".format(s="End-to-End Latency", n=50, c="-")) + print( + "{:<40} {:<10.2f}".format("Mean E2E Latency (ms):", metrics.mean_e2e_latency_ms) + ) + print( + "{:<40} {:<10.2f}".format( + "Median E2E Latency (ms):", metrics.median_e2e_latency_ms + ) + ) + print("{s:{c}^{n}}".format(s="Time to First Token", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean TTFT (ms):", metrics.mean_ttft_ms)) + print("{:<40} {:<10.2f}".format("Median TTFT (ms):", metrics.median_ttft_ms)) + print("{:<40} {:<10.2f}".format("P99 TTFT (ms):", metrics.p99_ttft_ms)) + print( + "{s:{c}^{n}}".format(s="Time per Output Token (excl. 1st token)", n=50, c="-") + ) + print("{:<40} {:<10.2f}".format("Mean TPOT (ms):", metrics.mean_tpot_ms)) + print("{:<40} {:<10.2f}".format("Median TPOT (ms):", metrics.median_tpot_ms)) + print("{:<40} {:<10.2f}".format("P99 TPOT (ms):", metrics.p99_tpot_ms)) + print("{s:{c}^{n}}".format(s="Inter-token Latency", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean ITL (ms):", metrics.mean_itl_ms)) + print("{:<40} {:<10.2f}".format("Median ITL (ms):", metrics.median_itl_ms)) + print("{:<40} {:<10.2f}".format("P99 ITL (ms):", metrics.p99_itl_ms)) + print("=" * 50) + + if ( + metrics.median_ttft_ms is not None + and metrics.mean_itl_ms is not None + and metrics.output_throughput is not None + ): + result = { + "backend": args.backend, + "request_rate": request_rate, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "total_output_tokens_retokenized": metrics.total_output_retokenized, + "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, + "median_e2e_latency_ms": metrics.median_e2e_latency_ms, + "median_ttft_ms": metrics.median_ttft_ms, + "median_itl_ms": metrics.median_itl_ms, + "output_throughput": metrics.output_throughput, + "random_input_len": args.random_input_len, + "random_output_len": args.random_output_len, + "random_range_ratio": args.random_range_ratio, + "duration": benchmark_duration, + "completed": metrics.completed, + } + else: + print(f"Error running benchmark for request rate: {request_rate}") + print("-" * 30) + + # Determine output file name + if args.output_file: + output_file_name = args.output_file + else: + now = datetime.now().strftime("%m%d") + output_file_name = f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_{args.random_output_len}.jsonl" + + # Append results to a JSONL file + with open(output_file_name, "a") as file: + file.write(json.dumps(result) + "\n") + + result = { + "duration": benchmark_duration, + "completed": metrics.completed, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "total_output_tokens_retokenized": metrics.total_output_retokenized, + "request_throughput": metrics.request_throughput, + "input_throughput": metrics.input_throughput, + "output_throughput": metrics.output_throughput, + "mean_ttft_ms": metrics.mean_ttft_ms, + "median_ttft_ms": metrics.median_ttft_ms, + "std_ttft_ms": metrics.std_ttft_ms, + "p99_ttft_ms": metrics.p99_ttft_ms, + "mean_tpot_ms": metrics.mean_tpot_ms, + "median_tpot_ms": metrics.median_tpot_ms, + "std_tpot_ms": metrics.std_tpot_ms, + "p99_tpot_ms": metrics.p99_tpot_ms, + "mean_itl_ms": metrics.mean_itl_ms, + "median_itl_ms": metrics.median_itl_ms, + "std_itl_ms": metrics.std_itl_ms, + "p99_itl_ms": metrics.p99_itl_ms, + "input_lens": [output.prompt_len for output in outputs], + "output_lens": output_lens, + "ttfts": [output.ttft for output in outputs], + "itls": [output.itl for output in outputs], + "generated_texts": [output.generated_text for output in outputs], + "errors": [output.error for output in outputs], + "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, + "median_e2e_latency_ms": metrics.median_e2e_latency_ms, + } + return result + + +def run_benchmark(args_: argparse.Namespace): + global args + args = args_ + + # Set global environments + set_ulimit() + random.seed(args.seed) + np.random.seed(args.seed) + + # Set url + if args.port is None: + args.port = { + "sglang": 30000, + }.get(args.backend, 30000) + + # api_url = ( + # f"{args.base_url}/v1/completions" + # if args.base_url + # else f"http://{args.host}:{args.port}/v1/completions" + # ) + api_url = ( + f"{args.base_url}/generate" + if args.base_url + else f"http://{args.host}:{args.port}/generate" + ) + + print(f"{args}\n") + + # Read dataset + backend = args.backend + model_id = args.model = LORA_PATH["base"] + tokenizer_id = args.model + + tokenizer = get_tokenizer(tokenizer_id) + + input_requests = sample_random_requests( + input_len=args.random_input_len, + output_len=args.random_output_len, + num_prompts=args.num_prompts, + range_ratio=args.random_range_ratio, + tokenizer=tokenizer, + dataset_path="", + ) + + return asyncio.run( + benchmark( + backend=backend, + api_url=api_url, + model_id=model_id, + tokenizer=tokenizer, + input_requests=input_requests, + request_rate=args.request_rate, + disable_tqdm=False, + extra_request_body={}, + ) + ) + + +def set_ulimit(target_soft_limit=65535): + resource_type = resource.RLIMIT_NOFILE + current_soft, current_hard = resource.getrlimit(resource_type) + + if current_soft < target_soft_limit: + try: + resource.setrlimit(resource_type, (target_soft_limit, current_hard)) + except ValueError as e: + print(f"Fail to set RLIMIT_NOFILE: {e}") + + +if __name__ == "__main__": + parser = ArgumentParser(description="Benchmark the online lora serving throughput.") + parser.add_argument( + "--backend", + type=str, + choices=list(ASYNC_REQUEST_FUNCS.keys()), + default="sglang", + help="Must specify a backend, depending on the LLM Inference Engine.", + ) + parser.add_argument( + "--base-url", + type=str, + default=None, + help="Server or API base url if not using http host and port.", + ) + parser.add_argument( + "--host", type=str, default="0.0.0.0", help="Default host is 0.0.0.0." + ) + parser.add_argument( + "--port", + type=int, + help="If not set, the default port is configured according to its default value for different LLM Inference Engines.", + ) + parser.add_argument( + "--num-prompts", + type=int, + default=50, + help="Number of prompts to process. Default is 1000.", + ) + parser.add_argument( + "--random-input-len", + type=int, + default=1024, + help="Number of input tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-output-len", + type=int, + default=128, + help="Number of output tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-range-ratio", + type=float, + default=0.0, + help="Range of sampled ratio of input/output length, " + "used only for random dataset.", + ) + parser.add_argument( + "--request-rate", + type=float, + default=float("inf"), + help="Number of requests per second. If this is inf, then all the requests are sent at time 0. " + "Otherwise, we use Poisson process to synthesize the request arrival times. Default is inf.", + ) + parser.add_argument( + "--base-only", + action="store_true", + ) + parser.add_argument("--output-file", type=str, help="Output JSONL file name.") + parser.add_argument("--seed", type=int, default=1, help="The random seed.") + args = parser.parse_args() + run_benchmark(args) diff --git a/sglang/benchmark/mmlu/README.md b/sglang/benchmark/mmlu/README.md new file mode 100644 index 0000000..16de20c --- /dev/null +++ b/sglang/benchmark/mmlu/README.md @@ -0,0 +1,59 @@ +## Download data +``` +bash download_data.sh +``` + +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --nsub 10 +``` + +``` +# OpenAI models +python3 bench_sglang.py --backend gpt-3.5-turbo --parallel 8 +``` + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --nsub 10 --backend vllm +``` + + +### Benchmark lightllm +``` +# A10G +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 + +# V100 +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 4500 --port 22000 +``` + +``` +python3 bench_other.py --nsub 10 --backend lightllm +``` + + +### Benchmark guidance +``` +python3 bench_other.py --nsub 10 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + + +### Benchmark lmql +``` +CUDA_VISIBLE_DEVICES=0,1 lmql serve-model meta-llama/Llama-2-7b-chat-hf --cuda --port 23000 +``` + +``` +python3 bench_other.py --nsub 10 --backend lmql --parallel 2 +``` diff --git a/sglang/benchmark/mmlu/bench_other.py b/sglang/benchmark/mmlu/bench_other.py new file mode 100644 index 0000000..c5d48da --- /dev/null +++ b/sglang/benchmark/mmlu/bench_other.py @@ -0,0 +1,173 @@ +import argparse +import asyncio +import json +import os +import time +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +import pandas as pd +import tiktoken +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate + +choices = ["A", "B", "C", "D"] + +tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo") + + +def format_subject(subject): + l = subject.split("_") + s = "" + for entry in l: + s += " " + entry + return s + + +def format_example(df, idx, include_answer=True): + prompt = df.iloc[idx, 0] + k = df.shape[1] - 2 + for j in range(k): + prompt += "\n{}. {}".format(choices[j], df.iloc[idx, j + 1]) + prompt += "\nAnswer:" + if include_answer: + prompt += " {}\n\n".format(df.iloc[idx, k + 1]) + return prompt + + +def gen_prompt(train_df, subject, k=-1): + prompt = "The following are multiple choice questions (with answers) about{}.\n\n".format( + format_subject(subject) + ) + if k == -1: + k = train_df.shape[0] + for i in range(k): + prompt += format_example(train_df, i) + return prompt + + +def evaluate(args, subject, dev_df, test_df, call_generate): + prompts = [] + labels = [] + + # Construct prompts + k = args.ntrain + train_prompt = gen_prompt(dev_df, subject, k) + while len(tokenizer.encode(train_prompt)) > 1536: + k -= 1 + train_prompt = gen_prompt(dev_df, subject, k) + + for i in range(test_df.shape[0]): + prompt_end = format_example(test_df, i, include_answer=False) + prompt = train_prompt + prompt_end + prompts.append(prompt) + + label = test_df.iloc[i, test_df.shape[1] - 1] + labels.append(label) + + preds = [None] * len(prompts) + max_tokens = 1 + + # Run requests + if args.backend != "lmql": + # Use thread pool + def get_one_answer(i): + pred = call_generate(prompts[i], temperature=0, max_tokens=max_tokens) + preds[i] = pred.strip()[0] + + tic = time.time() + if args.parallel == 1: + for i in range(len(prompts)): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + executor.map(get_one_answer, list(range(len(prompts)))) + else: + # Use asyncio + async def batched_call(batch_size): + for i in range(0, len(prompts), batch_size): + tasks = [] + for p in prompts[i : i + batch_size]: + tasks.append(call_generate(p, temperature=0, max_tokens=max_tokens)) + rets = await asyncio.gather(*tasks) + for j in range(len(rets)): + preds[i + j] = rets[j].strip()[0] + + tic = time.time() + asyncio.run(batched_call(batch_size=args.parallel)) + latency = time.time() - tic + + # Compute accuracy + cors = [pred == label for pred, label in zip(preds, labels)] + acc = np.mean(cors) + cors = np.array(cors) + + print( + "Average accuracy {:.3f}, latency {:.2f}, #q: {} - {}".format( + acc, latency, len(prompts), subject + ) + ) + + return cors, acc, latency + + +def main(args): + subjects = sorted( + [ + f.split("_test.csv")[0] + for f in os.listdir(os.path.join(args.data_dir, "test")) + if "_test.csv" in f + ] + ) + + all_cors = [] + all_latencies = [] + num_requests = 0 + + # Select backend + call_generate = get_call_generate(args) + + for subject in tqdm(subjects[: args.nsub]): + dev_df = pd.read_csv( + os.path.join(args.data_dir, "dev", subject + "_dev.csv"), header=None + )[: args.ntrain] + test_df = pd.read_csv( + os.path.join(args.data_dir, "test", subject + "_test.csv"), header=None + ) + + cors, acc, latency = evaluate(args, subject, dev_df, test_df, call_generate) + all_cors.append(cors) + all_latencies.append(latency) + num_requests += len(test_df) + + total_latency = np.sum(all_latencies) + print("Total latency: {:.3f}".format(total_latency)) + + weighted_acc = np.mean(np.concatenate(all_cors)) + print("Average accuracy: {:.3f}".format(weighted_acc)) + + # Write results + with open(args.result_file, "a") as fout: + value = { + "task": "mmlu", + "backend": args.backend, + "num_gpus": 1, + "latency": round(total_latency, 3), + "accuracy": round(weighted_acc, 3), + "num_requests": num_requests, + "other": { + "nsub": args.nsub, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--ntrain", type=int, default=5) + parser.add_argument("--data_dir", type=str, default="data") + parser.add_argument("--nsub", type=int, default=60) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/mmlu/bench_sglang.py b/sglang/benchmark/mmlu/bench_sglang.py new file mode 100644 index 0000000..210b611 --- /dev/null +++ b/sglang/benchmark/mmlu/bench_sglang.py @@ -0,0 +1,173 @@ +import argparse +import json +import os +import time + +import numpy as np +import pandas as pd +import tiktoken + +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) + +choices = ["A", "B", "C", "D"] + +tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo") + + +def format_subject(subject): + l = subject.split("_") + s = "" + for entry in l: + s += " " + entry + return s + + +def format_example(df, idx, include_answer=True): + prompt = df.iloc[idx, 0] + k = df.shape[1] - 2 + for j in range(k): + prompt += "\n{}. {}".format(choices[j], df.iloc[idx, j + 1]) + prompt += "\nAnswer:" + if include_answer: + prompt += " {}\n\n".format(df.iloc[idx, k + 1]) + return prompt + + +def gen_prompt(train_df, subject, k=-1): + prompt = "The following are multiple choice questions (with answers) about{}.\n\n".format( + format_subject(subject) + ) + if k == -1: + k = train_df.shape[0] + for i in range(k): + prompt += format_example(train_df, i) + return prompt + + +def main(args): + subjects = sorted( + [ + f.split("_test.csv")[0] + for f in os.listdir(os.path.join(args.data_dir, "test")) + if "_test.csv" in f + ] + ) + + # Build prompts + arguments = [] + labels = [] + num_questions = [] + + for subject in subjects[: args.nsub]: + dev_df = pd.read_csv( + os.path.join(args.data_dir, "dev", subject + "_dev.csv"), header=None + )[: args.ntrain] + test_df = pd.read_csv( + os.path.join(args.data_dir, "test", subject + "_test.csv"), header=None + ) + num_questions.append(test_df.shape[0]) + + k = args.ntrain + few_shot_examples = gen_prompt(dev_df, subject, k) + while len(tokenizer.encode(few_shot_examples)) > 1536: + k -= 1 + few_shot_examples = gen_prompt(dev_df, subject, k) + + for i in range(test_df.shape[0]): + prompt_end = format_example(test_df, i, include_answer=False) + + arguments.append( + { + "examples": few_shot_examples, + "question": prompt_end, + } + ) + + label = test_df.iloc[i, test_df.shape[1] - 1] + labels.append(label) + + ##################################### + ######### SGL Program Begin ######### + ##################################### + + import sglang as sgl + + if args.backend.startswith("gpt-"): + + @sgl.function + def few_shot_mmlu(s, examples, question): + s += sgl.user(examples + question) + s += sgl.assistant(sgl.gen("answer")) + + else: + + @sgl.function + def few_shot_mmlu(s, examples, question): + s += examples + question + sgl.gen("answer") + + ##################################### + ########## SGL Program End ########## + ##################################### + + # Select backend + backend = select_sglang_backend(args) + + # Run + tic = time.time() + states = few_shot_mmlu.run_batch( + arguments, + temperature=0, + max_new_tokens=1, + backend=backend, + num_threads=args.parallel, + progress_bar=True, + ) + preds = [ + s["answer"].strip()[0] if len(s["answer"].strip()) > 0 else "" for s in states + ] + latency = time.time() - tic + + # Compute accuracy + cors = [pred == label for pred, label in zip(preds, labels)] + + pt = 0 + for subject, num_qs in zip(subjects[: args.nsub], num_questions): + print( + f"subject: {subject}, #q:{num_qs}, acc: {np.mean(cors[pt: pt + num_qs]):.3f}" + ) + pt += num_qs + assert pt == len(cors) + weighted_acc = np.mean(cors) + + # Print results + print("Total latency: {:.3f}".format(latency)) + print("Average accuracy: {:.3f}".format(weighted_acc)) + + # Write results + with open(args.result_file, "a") as fout: + value = { + "task": "mmlu", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(weighted_acc, 3), + "num_requests": len(arguments), + "other": { + "nsub": args.nsub, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--ntrain", "-k", type=int, default=5) + parser.add_argument("--data_dir", "-d", type=str, default="data") + parser.add_argument("--save_dir", "-s", type=str, default="results") + parser.add_argument("--nsub", type=int, default=60) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/mmlu/download_data.sh b/sglang/benchmark/mmlu/download_data.sh new file mode 100644 index 0000000..fb22416 --- /dev/null +++ b/sglang/benchmark/mmlu/download_data.sh @@ -0,0 +1,2 @@ +wget https://people.eecs.berkeley.edu/~hendrycks/data.tar +tar xf data.tar diff --git a/sglang/benchmark/mmmu/README.md b/sglang/benchmark/mmmu/README.md new file mode 100644 index 0000000..00cdf16 --- /dev/null +++ b/sglang/benchmark/mmmu/README.md @@ -0,0 +1,30 @@ +## Run evaluation + +### Evaluate sglang + +Host the VLM: + +``` +python -m sglang.launch_server --model-path Qwen/Qwen2-VL-7B-Instruct --chat-template qwen2-vl --port 30000 +``` + +Benchmark: + +``` +python benchmark/mmmu/bench_sglang.py --port 30000 +``` + +It's recommended to reduce the memory usage by appending something ike `--mem-fraction-static 0.6` to the command above. + +### Evaluate hf + +``` +python benchmark/mmmu/bench_hf.py --model-path Qwen/Qwen2-VL-7B-Instruct +``` + +Some popular model results: + +1. Qwen/Qwen2-VL-2B-Instruct: 0.241 +2. Qwen/Qwen2-VL-7B-Instruct: 0.255 +3. Qwen/Qwen2.5-VL-3B-Instruct: 0.245 +4. Qwen/Qwen2.5-VL-7B-Instruct: 0.242 diff --git a/sglang/benchmark/mmmu/bench_hf.py b/sglang/benchmark/mmmu/bench_hf.py new file mode 100644 index 0000000..c6588c7 --- /dev/null +++ b/sglang/benchmark/mmmu/bench_hf.py @@ -0,0 +1,105 @@ +import argparse + +import torch +from data_utils import save_json +from eval_utils import ( + EvalArgs, + eval_result, + get_sampling_params, + prepare_samples, + process_result, +) +from tqdm import tqdm +from transformers import AutoModel, AutoProcessor, GenerationConfig + + +@torch.no_grad() +def eval_mmmu(args): + eval_args = EvalArgs.from_cli_args(args) + try: + from transformers import AutoModelForImageTextToText + + model = AutoModelForImageTextToText.from_pretrained( + args.model_path, + torch_dtype="auto", + trust_remote_code=True, + ) + except Exception as first_exception: + try: + model = AutoModel.from_pretrained( + args.model_path, + torch_dtype="auto", + trust_remote_code=True, + init_tts=False, + ) + except Exception as second_exception: + raise RuntimeError( + f"Failed to load model: First attempt failed with {first_exception}, " + f"second attempt failed with {second_exception}" + ) from second_exception + + model = model.eval().cuda() + + processor = AutoProcessor.from_pretrained( + args.model_path, torch_dtype="auto", device_map="auto", trust_remote_code=True + ) + + samples = prepare_samples(eval_args) + out_samples = dict() + + sampling_params = get_sampling_params(eval_args) + generation_config = GenerationConfig( + max_new_tokens=sampling_params["max_new_tokens"], + do_sample=False, + ) + + answer_dict = {} + for sample in tqdm(samples): + prompt = sample["final_input_prompt"] + image = sample["image"] + prefix = prompt.split("<")[0] + suffix = prompt.split(">")[1] + assert image is not None + contents = [] + if prefix: + contents += [{"type": "text", "text": prefix}] + contents += [ + { + "type": "image", + "image": sample["image_path"], + } + ] + if suffix: + contents += [{"type": "text", "text": suffix}] + messages = [{"role": "user", "content": contents}] + model_inputs = processor.apply_chat_template( + messages, + tokenize=True, + return_dict=True, + add_generation_prompt=True, + return_tensors="pt", + ).to(model.device) + input_len = model_inputs["input_ids"].shape[-1] + generation = model.generate(**model_inputs, generation_config=generation_config) + generation = generation[0][input_len:] + response = processor.decode(generation, skip_special_tokens=True) + print(f"response: {response}") + process_result(response, sample, answer_dict, out_samples) + + args.output_path = f"{args.model_path}_val_hf.json" + save_json(args.output_path, out_samples) + eval_result(model_answer_path=args.output_path, answer_dict=answer_dict) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--model-path", + type=str, + help="The path of the model weights. This can be a local folder or a Hugging Face repo ID.", + required=True, + ) + EvalArgs.add_cli_args(parser) + args = parser.parse_args() + + eval_mmmu(args) diff --git a/sglang/benchmark/mmmu/bench_sglang.py b/sglang/benchmark/mmmu/bench_sglang.py new file mode 100644 index 0000000..19f60d6 --- /dev/null +++ b/sglang/benchmark/mmmu/bench_sglang.py @@ -0,0 +1,93 @@ +""" +Bench the sglang-hosted vLM with benchmark MMMU + +Usage: + Host the VLM: python -m sglang.launch_server --model-path Qwen/Qwen2-VL-7B-Instruct --chat-template qwen2-vl --port 30000 + + Benchmark: python benchmark/mmmu/bench_sglang.py --port 30000 + +The eval output will be logged +""" + +import argparse +import time + +import openai +from data_utils import save_json +from eval_utils import ( + EvalArgs, + eval_result, + get_sampling_params, + prepare_samples, + process_result, +) +from tqdm import tqdm + +from sglang.test.test_utils import add_common_sglang_args_and_parse + + +def eval_mmmu(args): + eval_args = EvalArgs.from_cli_args(args) + + out_samples = dict() + + sampling_params = get_sampling_params(eval_args) + + samples = prepare_samples(eval_args) + + answer_dict = {} + + # had to use an openai server, since SglImage doesn't support image data + client = openai.Client(api_key="sk", base_url=f"http://127.0.0.1:{args.port}/v1") + + start = time.time() + for i, sample in enumerate(tqdm(samples)): + prompt = sample["final_input_prompt"] + prefix = prompt.split("<")[0] + suffix = prompt.split(">")[1] + image = sample["image"] + assert image is not None + image_path = sample["image_path"] + # TODO: batch + response = client.chat.completions.create( + model="default", + messages=[ + { + "role": "user", + "content": [ + { + "type": "text", + "text": prefix, + }, + { + "type": "image_url", + "image_url": {"url": image_path}, + }, + { + "type": "text", + "text": suffix, + }, + ], + } + ], + temperature=0, + max_completion_tokens=sampling_params["max_new_tokens"], + max_tokens=sampling_params["max_new_tokens"], + ) + response = response.choices[0].message.content + process_result(response, sample, answer_dict, out_samples) + + print(f"Benchmark time: {time.time() - start}") + + args.output_path = f"./val_sglang.json" + save_json(args.output_path, out_samples) + eval_result(model_answer_path=args.output_path, answer_dict=answer_dict) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + args = add_common_sglang_args_and_parse(parser) + EvalArgs.add_cli_args(parser) + args = parser.parse_args() + + eval_mmmu(args) diff --git a/sglang/benchmark/mmmu/data_utils.py b/sglang/benchmark/mmmu/data_utils.py new file mode 100644 index 0000000..cf89169 --- /dev/null +++ b/sglang/benchmark/mmmu/data_utils.py @@ -0,0 +1,221 @@ +"""Utils for data load, save, and process (e.g., prompt construction)""" + +import json +import os +import re + +import yaml + +DOMAIN_CAT2SUB_CAT = { + "Art and Design": ["Art", "Art_Theory", "Design", "Music"], + "Business": ["Accounting", "Economics", "Finance", "Manage", "Marketing"], + "Science": [ + "Biology", + "Chemistry", + "Geography", + "Math", + "Physics", + ], + "Health and Medicine": [ + "Basic_Medical_Science", + "Clinical_Medicine", + "Diagnostics_and_Laboratory_Medicine", + "Pharmacy", + "Public_Health", + ], + "Humanities and Social Science": [ + "History", + "Literature", + "Sociology", + "Psychology", + ], + "Tech and Engineering": [ + "Agriculture", + "Architecture_and_Engineering", + "Computer_Science", + "Electronics", + "Energy_and_Power", + "Materials", + "Mechanical_Engineering", + ], +} + + +CAT_SHORT2LONG = { + "acc": "Accounting", + "agri": "Agriculture", + "arch": "Architecture_and_Engineering", + "art": "Art", + "art_theory": "Art_Theory", + "bas_med": "Basic_Medical_Science", + "bio": "Biology", + "chem": "Chemistry", + "cli_med": "Clinical_Medicine", + "cs": "Computer_Science", + "design": "Design", + "diag_med": "Diagnostics_and_Laboratory_Medicine", + "econ": "Economics", + "elec": "Electronics", + "ep": "Energy_and_Power", + "fin": "Finance", + "geo": "Geography", + "his": "History", + "liter": "Literature", + "manage": "Manage", + "mark": "Marketing", + "mate": "Materials", + "math": "Math", + "mech": "Mechanical_Engineering", + "music": "Music", + "phar": "Pharmacy", + "phys": "Physics", + "psy": "Psychology", + "pub_health": "Public_Health", + "socio": "Sociology", +} + + +# DATA SAVING +def save_json(filename, ds): + with open(filename, "w") as f: + json.dump(ds, f, indent=4) + + +def get_multi_choice_info(options): + """ + Given the list of options for multiple choice question + Return the index2ans and all_choices + """ + + start_chr = "A" + all_choices = [] + index2ans = {} + for i, option in enumerate(options): + index2ans[chr(ord(start_chr) + i)] = option + all_choices.append(chr(ord(start_chr) + i)) + + return index2ans, all_choices + + +def load_yaml(file_path): + with open(file_path, "r") as stream: + try: + yaml_dict = yaml.safe_load(stream) + except yaml.YAMLError as exc: + print(exc) + + return yaml_dict + + +def parse_img_path(text): + matches = re.findall("", text) + return matches + + +def process_single_sample(data): + question = data["question"] + o_imgs_paths = [] + for option in data["options"]: + current_o_imgs_paths = parse_img_path(option) + for img_path in current_o_imgs_paths: + o_imgs_paths.append(img_path) + + if len(o_imgs_paths) > 1: # multiple images in options, used for random selection + return { + "id": data["id"], + "question": question, + "options": data["options"], + "answer": data["answer"], + "image": None, + "question_type": data["question_type"], + } + else: + return { + "id": data["id"], + "question": question, + "options": data["options"], + "answer": data["answer"], + "image": data["image_1"], + "question_type": data["question_type"], + } + + +# DATA SAVING +def save_json(filename, ds): + print(f"answers saved to: {filename}") + os.makedirs(os.path.dirname(filename), exist_ok=True) + with open(filename, "w") as f: + json.dump(ds, f, indent=4) + + +def save_jsonl(filename, data): + """ + Save a dictionary of data to a JSON Lines file with the filename as key and caption as value. + + Args: + filename (str): The path to the file where the data should be saved. + data (dict): The dictionary containing the data to save where key is the image path and value is the caption. + """ + with open(filename, "w", encoding="utf-8") as f: + for img_path, caption in data.items(): + # Extract the base filename without the extension + base_filename = os.path.basename(img_path) + # Create a JSON object with the filename as the key and caption as the value + json_record = json.dumps({base_filename: caption}, ensure_ascii=False) + # Write the JSON object to the file, one per line + f.write(json_record + "\n") + + +def save_args(args, path_dir): + argsDict = args.__dict__ + with open(path_dir + "setting.txt", "w") as f: + f.writelines("------------------ start ------------------" + "\n") + for eachArg, value in argsDict.items(): + f.writelines(eachArg + " : " + str(value) + "\n") + f.writelines("------------------- end -------------------") + + +# DATA PROCESSING +def construct_prompt(sample, config): + question = sample["question"] + options = eval(sample["options"]) + example = "" + if sample["question_type"] == "multiple-choice": + start_chr = "A" + prediction_range = [] + index2ans = {} + for option in options: + prediction_range.append(start_chr) + example += f"({start_chr}) {option}\n" + index2ans[start_chr] = option + start_chr = chr(ord(start_chr) + 1) + empty_prompt_sample_structure = config["multi_choice_example_format"] + empty_prompt = empty_prompt_sample_structure.format(question, example) + res_dict = {} + res_dict["index2ans"] = index2ans + res_dict["correct_choice"] = sample["answer"] + res_dict["all_choices"] = prediction_range + res_dict["empty_prompt"] = empty_prompt + if config["task_instructions"]: + res_dict["final_input_prompt"] = ( + config["task_instructions"].strip() + "\n\n" + empty_prompt + ) + else: + res_dict["final_input_prompt"] = empty_prompt + + res_dict["gt_content"] = options[ord(sample["answer"].upper()) - ord("A")] + else: + empty_prompt_sample_structure = config["short_ans_example_format"] + empty_prompt = empty_prompt_sample_structure.format(question) + res_dict = {} + res_dict["empty_prompt"] = empty_prompt + if config["task_instructions"]: + res_dict["final_input_prompt"] = ( + config["task_instructions"].strip() + "\n\n" + empty_prompt + ) + else: + res_dict["final_input_prompt"] = empty_prompt + res_dict["gt_content"] = sample["answer"] + + res_dict.update(sample) + return res_dict diff --git a/sglang/benchmark/mmmu/eval_utils.py b/sglang/benchmark/mmmu/eval_utils.py new file mode 100644 index 0000000..046bed6 --- /dev/null +++ b/sglang/benchmark/mmmu/eval_utils.py @@ -0,0 +1,561 @@ +"""Response Parsing and Evaluation for various models""" + +import argparse +import dataclasses +import json +import os +import pprint +import random +import re +from typing import Dict, Optional + +import numpy as np +import torch +from data_utils import ( + CAT_SHORT2LONG, + DOMAIN_CAT2SUB_CAT, + construct_prompt, + load_yaml, + process_single_sample, +) +from datasets import concatenate_datasets, load_dataset +from tqdm import tqdm + + +@dataclasses.dataclass +class EvalArgs: + seed: int = 42 + split: str = "validation" + # Default setting to make the benchmark available on A100 for most 7B models + image_pixels_limit: int = 4300000 + result_filename: str = "" + prompt_format_file: str = "prompt_format.yaml" + dataset_path: str = "MMMU/MMMU" + extra_request_body: Optional[str] = None + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument( + "--result-filename", type=str, default=EvalArgs.result_filename + ) + + parser.add_argument( + "--image-pixels-limit", type=int, default=EvalArgs.image_pixels_limit + ) + + parser.add_argument( + "--dataset-path", + type=str, + default=EvalArgs.dataset_path, + help="path to the dataset", + ) + parser.add_argument("--seed", type=int, default=1, help="The random seed.") + parser.add_argument( + "--prompt-format-file", + type=str, + help="The path to the prompt format of mmmu. If not, a default format llava_config.yaml will be used", + ) + parser.add_argument("--split", type=str, default=EvalArgs.split) + parser.add_argument( + "--extra-request-body", + metavar='{"key1": "value1", "key2": "value2"}', + type=str, + default=EvalArgs.extra_request_body, + help="Append given JSON object to the request payload. You can use this to specify" + "additional generate params like sampling params.", + ) + + @classmethod + def from_cli_args(cls, args: argparse.Namespace): + attrs = [attr.name for attr in dataclasses.fields(cls)] + return cls(**{attr: getattr(args, attr) for attr in attrs}) + + +def set_seed(seed_value): + """ + Set the seed for PyTorch (both CPU and CUDA), Python, and NumPy for reproducible results. + + :param seed_value: An integer value to be used as the seed. + """ + torch.manual_seed(seed_value) + if torch.cuda.is_available(): + torch.cuda.manual_seed(seed_value) + torch.cuda.manual_seed_all(seed_value) # For multi-GPU setups + random.seed(seed_value) + np.random.seed(seed_value) + torch.backends.cudnn.deterministic = True + torch.backends.cudnn.benchmark = False + + +def prepare_samples(eval_args: EvalArgs): + print("preparing samples...") + # Build prompts + set_seed(eval_args.seed) + + prompt_format_file = ( + eval_args.prompt_format_file + if eval_args.prompt_format_file is not None + else os.path.join(os.path.dirname(__file__), "prompt_format.yaml") + ) + # load config and process to one value + eval_args.config = load_yaml(prompt_format_file) + for key, value in eval_args.config.items(): + if key != "eval_params" and type(value) == list: + assert len(value) == 1, "key {} has more than one value".format(key) + eval_args.config[key] = value[0] + + # run for each subject + sub_dataset_list = [] + + for subject in tqdm(CAT_SHORT2LONG.values()): + sub_dataset = load_dataset( + eval_args.dataset_path, subject, split=eval_args.split + ) + sub_dataset_list.append(sub_dataset) + # break + + # merge all dataset + dataset = concatenate_datasets(sub_dataset_list) + + ## prepare images + samples = [] + skip_count = 0 + + # use image file as input to ensure the consistency between sglang and hf + images_path = os.path.expanduser("~/.cache/mmmu/images") + os.makedirs(images_path, exist_ok=True) + print(f"Saving images to: {images_path}") + + for i, sample in enumerate(tqdm(dataset)): + sample = process_single_sample(sample) + sample = construct_prompt(sample, eval_args.config) + image = sample["image"] + + width, height = image.size + if width * height >= eval_args.image_pixels_limit: + skip_count += 1 + continue + image_path = f"{images_path}/image_{i}.png" + if not os.path.exists(image_path): + image.save(image_path) + sample["image_path"] = image_path + samples.append(sample) + + print( + f"skipping {skip_count} samples with large images, {round((float(skip_count) / len(dataset)) * 100, 2)}% of dataset" + ) + print("samples have been prepared") + return samples + + +def get_sampling_params(eval_args): + max_new_tokens = 30 + temperature = 0.001 + + extra_request_body = {} + if eval_args.extra_request_body: + extra_request_body = json.loads(eval_args.extra_request_body) + + return { + "temperature": temperature, + "max_new_tokens": max_new_tokens, + **extra_request_body, + } + + +# ----------- Process Multi-choice ------------- +def parse_multi_choice_response(response, all_choices, index2ans): + """ + Parse the prediction from the generated response. + Return the predicted index e.g., A, B, C, D. + """ + for char in [",", ".", "!", "?", ";", ":", "'"]: + response = response.strip(char) + response = " " + response + " " # add space to avoid partial match + + index_ans = True + ans_with_brack = False + candidates = [] + for choice in all_choices: # e.g., (A) (B) (C) (D) + if f"({choice})" in response: + candidates.append(choice) + ans_with_brack = True + + if len(candidates) == 0: + for choice in all_choices: # e.g., A B C D + if f" {choice} " in response: + candidates.append(choice) + + # if all above doesn't get candidates, check if the content is larger than 5 tokens and try to parse the example + if len(candidates) == 0 and len(response.split()) > 5: + for index, ans in index2ans.items(): + if ans.lower() in response.lower(): + candidates.append(index) + index_ans = False # it's content ans. + + if len(candidates) == 0: # still not get answer, randomly choose one. + pred_index = random.choice(all_choices) + elif len(candidates) > 1: + start_indexes = [] + if index_ans: + if ans_with_brack: + for can in candidates: + index = response.rfind(f"({can})") + start_indexes.append(index) # -1 will be ignored anyway + # start_indexes = [generated_response.index(f'({can})') for can in candidates] + else: + for can in candidates: + index = response.rfind(f" {can} ") + start_indexes.append(index) + else: + for can in candidates: + index = response.lower().rfind(index2ans[can].lower()) + start_indexes.append(index) + # get the last one + pred_index = candidates[np.argmax(start_indexes)] + else: # if only one candidate, use it. + pred_index = candidates[0] + + return pred_index + + +# ----------- Process Open ------------- +def check_is_number(string): + """ + Check if the given string a number. + """ + try: + float(string.replace(",", "")) + return True + except ValueError: + # check if there's comma inside + return False + + +def normalize_str(string): + """ + Normalize the str to lower case and make them float numbers if possible. + """ + # check if characters in the string + + # if number, numerize it. + string = string.strip() + + is_number = check_is_number(string) + + if is_number: + string = string.replace(",", "") + string = float(string) + # leave 2 decimal + string = round(string, 2) + return [string] + else: # it's likely to be a string + # lower it + string = string.lower() + if len(string) == 1: + return [" " + string, string + " "] # avoid trivial matches + return [string] + + +def extract_numbers(string): + """ + Exact all forms of numbers from a string with regex. + """ + # Pattern for numbers with commas + pattern_commas = r"-?\b\d{1,3}(?:,\d{3})+\b" + # Pattern for scientific notation + pattern_scientific = r"-?\d+(?:\.\d+)?[eE][+-]?\d+" + # Pattern for simple numbers without commas + pattern_simple = r"-?(?:\d+\.\d+|\.\d+|\d+\b)(?![eE][+-]?\d+)(?![,\d])" + + # Extract numbers with commas + numbers_with_commas = re.findall(pattern_commas, string) + # Extract numbers in scientific notation + numbers_scientific = re.findall(pattern_scientific, string) + # Extract simple numbers without commas + numbers_simple = re.findall(pattern_simple, string) + + # Combine all extracted numbers + all_numbers = numbers_with_commas + numbers_scientific + numbers_simple + return all_numbers + + +def parse_open_response(response): + """ + Parse the prediction from the generated response. + Return a list of predicted strings or numbers. + """ + + # content = content.strip("\n").strip(".").strip(" ") + def get_key_subresponses(response): + key_responses = [] + response = response.strip().strip(".").lower() + sub_responses = re.split(r"\.\s(?=[A-Z])|\n", response) + indicators_of_keys = [ + "could be ", + "so ", + "is ", + "thus ", + "therefore ", + "final ", + "answer ", + "result ", + ] + key_responses = [] + for index, resp in enumerate(sub_responses): + # if last one, accept it's an equation (the entire response can be just one sentence with equation) + if index == len(sub_responses) - 1: + indicators_of_keys.extend(["="]) + shortest_key_response = None # the shortest response that may contain the answer (tail part of the response) + for indicator in indicators_of_keys: + if indicator in resp: + if not shortest_key_response: + shortest_key_response = resp.split(indicator)[-1].strip() + else: + if len(resp.split(indicator)[-1].strip()) < len( + shortest_key_response + ): + shortest_key_response = resp.split(indicator)[-1].strip() + # key_responses.append(resp.split(indicator)[1].strip()) + + if shortest_key_response: + # and it's not trivial + if shortest_key_response.strip() not in [ + ":", + ",", + ".", + "!", + "?", + ";", + ":", + "'", + ]: + key_responses.append(shortest_key_response) + if len(key_responses) == 0: # did not found any + return [response] + return key_responses + + # pdb.set_trace() + key_responses = get_key_subresponses(response) + + pred_list = key_responses.copy() # keep the original string response + for resp in key_responses: + pred_list.extend(extract_numbers(resp)) + + tmp_pred_list = [] + for i in range(len(pred_list)): + tmp_pred_list.extend(normalize_str(pred_list[i])) + pred_list = tmp_pred_list + + # remove duplicates + pred_list = list(set(pred_list)) + + return pred_list + + +# ----------- Evaluation ------------- + + +def eval_multi_choice(gold_i, pred_i): + """ + Evaluate a multiple choice instance. + """ + correct = False + # only they are exactly the same, we consider it as correct + if isinstance(gold_i, list): + for answer in gold_i: + if answer == pred_i: + correct = True + break + else: # gold_i is a string + if gold_i == pred_i: + correct = True + return correct + + +def eval_open(gold_i, pred_i): + """ + Evaluate an open question instance + """ + correct = False + if isinstance(gold_i, list): + # use float to avoid trivial matches + norm_answers = [] + for answer in gold_i: + norm_answers.extend(normalize_str(answer)) + else: + norm_answers = normalize_str(gold_i) + for pred in pred_i: # pred is already normalized in parse response phase + if isinstance(pred, str): # if it's a string, then find if ans in the pred_i + for norm_ans in norm_answers: + # only see if the string answer in the string pred + if isinstance(norm_ans, str) and norm_ans in pred: + if not correct: + correct = True + break + else: # it's a float number + if pred in norm_answers: + if not correct: + correct = True + break + return correct + + +# ----------- Batch Evaluation ------------- +def evaluate(samples): + """ + Batch evaluation for multiple choice and open questions. + """ + pred_correct = 0 + judge_dict = dict() + for sample in samples: + gold_i = sample["answer"] + pred_i = sample["parsed_pred"] + if sample["question_type"] == "multiple-choice": + correct = eval_multi_choice(gold_i, pred_i) + else: # open question + correct = eval_open(gold_i, pred_i) + + if correct: + judge_dict[sample["id"]] = "Correct" + pred_correct += 1 + else: + # print(f"Wrong! expected {pred_i}, answered with {gold_i}") + judge_dict[sample["id"]] = "Wrong" + + if len(samples) == 0: + return {"acc": 0} + return judge_dict, {"acc": pred_correct / len(samples)} + + +# ----------- Calculate Accuracy ------------- +def calculate_ins_level_acc(results: Dict): + """Calculate the instruction level accuracy for given Subject results""" + acc = 0 + ins_num = 0 + for cat_results in results.values(): + acc += cat_results["acc"] * cat_results["num_example"] + ins_num += cat_results["num_example"] + if ins_num == 0: + return 0 + return acc / ins_num + + +def process_result(response, sample, answer_dict, out_samples): + if sample["question_type"] == "multiple-choice": + pred_ans = parse_multi_choice_response( + response, sample["all_choices"], sample["index2ans"] + ) + else: # open question + pred_ans = response + + out_samples[sample["id"]] = pred_ans + + # set ground truth answer + answer_dict[sample["id"]] = { + "question_type": sample["question_type"], + "ground_truth": sample["answer"], + } + + +def eval_result(model_answer_path, answer_dict): + print("Evaluating...") + output_dict = json.load(open(model_answer_path)) + # answer_dict = json.load(open(answer_path)) + + # group by category + output_dict_w_cat = {} + for data_id, parsed_pred in output_dict.items(): + category = "_".join(data_id.split("_")[1:-1]) + if category not in output_dict_w_cat: + output_dict_w_cat.update({category: {}}) + output_dict_w_cat[category].update({data_id: parsed_pred}) + + # group by category + answer_dict_w_cat = {} + for data_id, parsed_pred in answer_dict.items(): + category = "_".join(data_id.split("_")[1:-1]) + if category not in answer_dict_w_cat: + answer_dict_w_cat.update({category: {}}) + answer_dict_w_cat[category].update({data_id: parsed_pred}) + + evaluation_result = {} + + for category in CAT_SHORT2LONG.values(): + # print("Evaluating: {}".format(category)) + # get cat_outputs and cat_answers + try: + cat_outputs = output_dict_w_cat[category] + cat_answers = answer_dict_w_cat[category] + except KeyError: + # print("Skipping {} for not found".format(category)) + continue + + exampels_to_eval = [] + for data_id, parsed_pred in cat_outputs.items(): + question_type = cat_answers[data_id]["question_type"] + if question_type != "multiple-choice": + parsed_pred = parse_open_response( + parsed_pred + ) # mainly for type consistency (make it number, etc.) + else: + parsed_pred = parsed_pred + + exampels_to_eval.append( + { + "id": data_id, + "question_type": question_type, + "answer": cat_answers[data_id]["ground_truth"], + "parsed_pred": parsed_pred, + } + ) + + judge_dict, metric_dict = evaluate(exampels_to_eval) + metric_dict.update({"num_example": len(exampels_to_eval)}) + + evaluation_result[category] = metric_dict + + printable_results = {} + # pdb.set_trace() + # add domain Subject + for domain, in_domain_cats in DOMAIN_CAT2SUB_CAT.items(): + in_domain_cat_results = {} + for cat_name in in_domain_cats: # use the order in DOMAIN_CAT2SUB_CAT + if cat_name in evaluation_result.keys(): + in_domain_cat_results[cat_name] = evaluation_result[cat_name] + else: + pass + in_domain_ins_acc = calculate_ins_level_acc(in_domain_cat_results) + in_domain_data_num = sum( + [ + cat_results["num_example"] + for cat_results in in_domain_cat_results.values() + ] + ) + printable_results["Overall-" + domain] = { + "num": int(in_domain_data_num), + "acc": round(in_domain_ins_acc, 3), + } + # add sub category + for cat_name, cat_results in in_domain_cat_results.items(): + printable_results[cat_name] = { + "num": int(cat_results["num_example"]), + "acc": round(cat_results["acc"], 3), + } + + # table.append(["-----------------------------", "-----", "----"]) + all_ins_acc = calculate_ins_level_acc(evaluation_result) + overall_acc = round(all_ins_acc, 3) + printable_results["Overall"] = { + "num": sum( + [cat_results["num_example"] for cat_results in evaluation_result.values()] + ), + "acc": overall_acc, + } + pprint.pprint(printable_results) + out = model_answer_path + with open(out, "w", encoding="utf-8") as outfile: + json.dump(printable_results, outfile) + print(f"eval out saved to {out}") + + print(f"Overall accuracy: {overall_acc}") diff --git a/sglang/benchmark/mmmu/prompt_format.yaml b/sglang/benchmark/mmmu/prompt_format.yaml new file mode 100644 index 0000000..1a0f721 --- /dev/null +++ b/sglang/benchmark/mmmu/prompt_format.yaml @@ -0,0 +1,15 @@ +task_instructions: +- "" +multi_choice_example_format: +- "{} + +{} + +Answer with the option's letter from the given choices directly." + +short_ans_example_format: +- "{} + +Answer the question using a single word or phrase." +temperature: +- 0 diff --git a/sglang/benchmark/mtbench/README.md b/sglang/benchmark/mtbench/README.md new file mode 100644 index 0000000..e6babf9 --- /dev/null +++ b/sglang/benchmark/mtbench/README.md @@ -0,0 +1,48 @@ +## Download Dataset + +```sh +wget -O question.jsonl https://raw.githubusercontent.com/lm-sys/FastChat/main/fastchat/llm_judge/data/mt_bench/question.jsonl +``` + +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 80 +``` + +### Benchmark sglang EAGLE +``` +python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3-8B-Instruct --speculative-algo EAGLE \ + --speculative-draft lmsys/sglang-EAGLE-LLaMA3-Instruct-8B --speculative-num-steps 5 \ + --speculative-eagle-topk 8 --speculative-num-draft-tokens 64 --dtype float16 --port 30000 +``` + +``` +python3 bench_sglang_eagle.py --num-questions 80 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-questions 80 --backend vllm +``` + + +### Benchmark lightllm +``` +# A10G +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 +``` + +``` +python3 bench_other.py --num-questions 80 --backend lightllm +``` diff --git a/sglang/benchmark/mtbench/bench_other.py b/sglang/benchmark/mtbench/bench_other.py new file mode 100644 index 0000000..2c321e8 --- /dev/null +++ b/sglang/benchmark/mtbench/bench_other.py @@ -0,0 +1,111 @@ +import argparse +import json +import os +import time +import uuid +from concurrent.futures import ThreadPoolExecutor + +from fastchat.model import get_conversation_template +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate + + +def load_questions(filename): + questions = [] + with open(filename, "r") as fin: + for line in fin: + obj = json.loads(line) + questions.append(obj) + return questions + + +def write_answers(filename, model_id, questions, answers): + with open(os.path.expanduser(filename), "w") as fout: + for i in range(len(answers)): + ans_json = { + "question_id": questions[i]["question_id"], + "answer_id": uuid.uuid4().hex, + "model_id": model_id, + "choices": { + "index": 0, + "turns": [answers[i][0], answers[i][1]], + }, + "tstamp": time.time(), + } + fout.write(json.dumps(ans_json) + "\n") + + +def main(args): + questions = load_questions(args.question_file) + questions = (questions * 10)[: args.num_questions] + max_tokens = 256 + model_id = "llama-2-chat" + + conv_main = get_conversation_template(model_id) + + # Select backend + call_generate = get_call_generate(args) + + answers = [None] * len(questions) + + def get_answer(i): + conv = conv_main.copy() + cur_answers = [] + for j in range(2): + q = questions[i]["turns"][j] + conv.append_message(conv.roles[0], q) + conv.append_message(conv.roles[1], None) + + prompt = conv.get_prompt() + output = call_generate(prompt, temperature=0, max_tokens=max_tokens).strip() + + cur_answers.append(output) + conv.update_last_message(output) + + answers[i] = cur_answers + + # Run requests + tic = time.time() + if args.parallel == 1: + for i in tqdm(range(len(questions))): + get_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_answer, list(range(len(questions)))), + total=len(questions), + ) + ) + + latency = time.time() - tic + + print(f"#questions: {len(questions)}, Latency: {latency:.2f}") + + # Write results + answer_file = args.answer_file or f"tmp_output_{args.backend}.txt" + write_answers(answer_file, model_id, questions, answers) + + with open(args.result_file, "a") as fout: + value = { + "task": "mtbench", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--question-file", type=str, default="question.jsonl") + parser.add_argument("--answer-file", type=str, default=None) + parser.add_argument("--num-questions", type=int, default=80) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/mtbench/bench_sglang.py b/sglang/benchmark/mtbench/bench_sglang.py new file mode 100644 index 0000000..b57d164 --- /dev/null +++ b/sglang/benchmark/mtbench/bench_sglang.py @@ -0,0 +1,99 @@ +import argparse +import json +import os +import time +import uuid + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) + + +def load_questions(filename): + questions = [] + with open(filename, "r") as fin: + for line in fin: + obj = json.loads(line) + questions.append(obj) + return questions + + +def write_answers(filename, model_id, questions, answers): + with open(os.path.expanduser(filename), "w") as fout: + for i in range(len(answers)): + ans_json = { + "question_id": questions[i]["question_id"], + "answer_id": uuid.uuid4().hex, + "model_id": model_id, + "choices": { + "index": 0, + "turns": [answers[i][0], answers[i][1]], + }, + "tstamp": time.time(), + } + fout.write(json.dumps(ans_json) + "\n") + + +@sgl.function +def answer_mt_bench(s, question_1, question_2): + s += sgl.system() + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1")) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2")) + + +def main(args): + # Construct prompts + questions = load_questions(args.question_file)[: args.num_questions] + arguments = [ + {"question_1": q["turns"][0], "question_2": q["turns"][1]} for q in questions + ] + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.time() + rets = answer_mt_bench.run_batch( + arguments, + temperature=0, + max_new_tokens=256, + num_threads=args.parallel, + progress_bar=True, + ) + answers = [[s["answer_1"], s["answer_2"]] for s in rets] + latency = time.time() - tic + + print(f"#questions: {len(questions)}, Latency: {latency:.2f}") + + # Write results + model_id = backend.model_info["model_path"] + answer_file = args.answer_file or f"tmp_output_{args.backend}.txt" + write_answers(answer_file, model_id, questions, answers) + + with open(args.result_file, "a") as fout: + value = { + "task": "mtbench", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--question-file", type=str, default="question.jsonl") + parser.add_argument("--answer-file", type=str, default=None) + parser.add_argument("--num-questions", type=int, default=80) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/mtbench/bench_sglang_eagle.py b/sglang/benchmark/mtbench/bench_sglang_eagle.py new file mode 100644 index 0000000..e1207af --- /dev/null +++ b/sglang/benchmark/mtbench/bench_sglang_eagle.py @@ -0,0 +1,137 @@ +""" +Adapted from https://github.com/chromecast56/sglang/blob/6f145d2eadb93a116134f703358ce76f15381045/benchmark/mtbench/bench_sglang.py + +Benchmark SGLang EAGLE/EAGLE3 Speculative Decoding + +Usage: +python3 benchmark/mtbench/bench_sglang_eagle.py --num-questions 80 --parallel 1 +""" + +import argparse +import json +import os +import time +import uuid + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) + + +def load_questions(filename): + questions = [] + with open(filename, "r") as fin: + for line in fin: + obj = json.loads(line) + questions.append(obj) + return questions + + +def write_answers(filename, model_id, questions, answers): + with open(os.path.expanduser(filename), "w") as fout: + for i in range(len(answers)): + ans_json = { + "question_id": questions[i]["question_id"], + "answer_id": uuid.uuid4().hex, + "model_id": model_id, + "choices": { + "index": 0, + "turns": [answers[i][0], answers[i][1]], + }, + "tstamp": time.time(), + } + fout.write(json.dumps(ans_json) + "\n") + + +@sgl.function +def answer_mt_bench(s, question_1, question_2): + s += sgl.system( + "You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.\n\nIf a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information." + ) + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1")) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2")) + + +def main(args): + # Construct prompts + questions = load_questions(args.question_file)[: args.num_questions] + arguments = [ + {"question_1": q["turns"][0], "question_2": q["turns"][1]} for q in questions + ] + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.time() + rets = answer_mt_bench.run_batch( + arguments, + temperature=0, + max_new_tokens=2048, + num_threads=args.parallel, + progress_bar=True, + ) + answers = [[s["answer_1"], s["answer_2"]] for s in rets] + + latency = time.time() - tic + num_output_tokens = sum( + s.get_meta_info("answer_1")["completion_tokens"] + + s.get_meta_info("answer_2")["completion_tokens"] + for s in rets + ) + + # NOTE: acceptance length is just completion_tokens / spec_verify_ct + # {'id': '3bb9c5ead109488d8ed5ee9cbecaec29', 'finish_reason': {'type': 'length', 'length': 256}, 'prompt_tokens': 37, 'spec_verify_ct': 101, 'completion_tokens': 256, 'cached_tokens': 0} + + output_throughput = num_output_tokens / latency + + has_verify = "spec_verify_ct" in rets[0].get_meta_info("answer_1") + if has_verify: + num_verify_tokens = sum( + s.get_meta_info("answer_1")["spec_verify_ct"] + + s.get_meta_info("answer_2")["spec_verify_ct"] + for s in rets + ) + + accept_length = num_output_tokens / num_verify_tokens + else: + accept_length = 1.0 + + print( + f"#questions: {len(questions)}, Throughput: {output_throughput:.2f} token/s, Acceptance length: {accept_length:.2f}" + ) + + # Write results + model_id = backend.model_info["model_path"] + answer_file = args.answer_file or f"tmp_output_{args.backend}.txt" + write_answers(answer_file, model_id, questions, answers) + + with open(args.result_file, "a") as fout: + value = { + "task": "mtbench", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "throughput": round(output_throughput, 3), + "accept_length": round(accept_length, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--question-file", type=str, default="question.jsonl") + parser.add_argument("--answer-file", type=str, default=None) + parser.add_argument("--num-questions", type=int, default=80) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/multi_chain_reasoning/README.md b/sglang/benchmark/multi_chain_reasoning/README.md new file mode 100644 index 0000000..4c9f740 --- /dev/null +++ b/sglang/benchmark/multi_chain_reasoning/README.md @@ -0,0 +1,49 @@ +## Download data +``` +wget https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl +``` + +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --schedule-conservativeness 1.3 +``` + +``` +python3 bench_sglang.py --num-questions 64 +python3 bench_sglang.py --num-questions 32 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-questions 64 --backend vllm +``` + + +### Benchmark lightllm +``` +# A10G +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 +``` + +``` +python3 bench_other.py --num-questions 64 --backend lightllm +``` + + +### Benchmark guidance +``` +python3 bench_other.py --num-questions 8 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +### Benchmark lmql + +``` +python3 bench_other.py --num-questions 64 --backend lmql --parallel 1 +``` diff --git a/sglang/benchmark/multi_chain_reasoning/bench_other.py b/sglang/benchmark/multi_chain_reasoning/bench_other.py new file mode 100644 index 0000000..e0ff2be --- /dev/null +++ b/sglang/benchmark/multi_chain_reasoning/bench_other.py @@ -0,0 +1,186 @@ +import argparse +import ast +import asyncio +import json +import re +import time +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +prompt_lib = [ + "Let us think step by step.", + "Approach this methodically. Let's dissect the problem into smaller, more manageable parts.", + "It's important to proceed step by step, ensuring accuracy at each stage.", + "Take a deep breath and break this down.", + "A little bit of arithmetic and a logical approach will help us quickly arrive at the solution to this problem.", + "I am extremely good at math.", +] + + +def multi_chain_gsm8k(question, num_chains, call_generate): + s = "Question: " + question + "\n" + # s += call_generate(s + "Answer: " + prompt_lib[0], max_tokens=256, + # stop="Question", temperature=0) + # return s + + comps = [] + for i in range(num_chains): + comps.append( + call_generate( + s + "Answer: " + prompt_lib[i % num_chains], + max_tokens=256, + temperature=0.3, + stop="Question", + ) + ) + + s += "Answer: To answer this question, here are some possible solutions. " + s += "After considering all of them, I will do a majority vote.\n\n" + for i in range(num_chains): + s += f"Solution {i+1}: " + comps[i].strip() + "\n\n" + s += "\nBy considering the above solutions and doing a majority vote, I think the final answer (a single integer number) is " + s += call_generate(s, max_tokens=16, temperature=0, stop=None) + return s + + +async def multi_chain_gsm8k_async(question, num_chains, call_generate): + s = "Question: " + question + "\n" + # s += call_generate(s + "Answer: " + prompt_lib[0], max_tokens=256, + # stop="Question", temperature=0) + # return s + + comps = [] + for i in range(num_chains): + comps.append( + await call_generate( + s + "Answer: " + prompt_lib[i % num_chains], + max_tokens=256, + temperature=0.3, + stop="Question", + ) + ) + + s += "Answer: To answer this question, here are some possible solutions. " + s += "After considering all of them, I will do a majority vote.\n\n" + for i in range(num_chains): + s += f"Solution {i+1}: " + comps[i].strip() + "\n\n" + s += "\nBy considering the above solutions and doing a majority vote, I think the final answer (a single integer number) is " + s += await call_generate(s, max_tokens=16, temperature=0, stop=None) + return s + + +def main(args): + lines = read_jsonl(args.data_path) + + # Construct prompts + k = args.num_shot + + questions = [] + labels = [] + for i in range(len(lines[: args.num_questions])): + questions.append(lines[i]["question"]) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + + states = [None] * len(labels) + + # Select backend + call_generate = get_call_generate(args) + + # Run requests + if args.backend != "lmql": + # Use thread pool + def get_one_answer(i): + answer = multi_chain_gsm8k(questions[i], args.num_chains, call_generate) + states[i] = answer + + tic = time.time() + if args.parallel == 1: + for i in tqdm(range(len(questions))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(questions)))), + total=len(questions), + ) + ) + + else: + # Use asyncio + async def get_one_answer_asyncio(i): + answer = await multi_chain_gsm8k_async( + questions[i], args.num_chains, call_generate + ) + states[i] = answer + + tic = time.time() + loop = asyncio.get_event_loop() + batches = [ + list(range(i, min(i + args.parallel, len(questions)))) + for i in range(0, len(questions), args.parallel) + ] + for bt in tqdm(batches): + tasks = [get_one_answer_asyncio(k) for k in bt] + loop.run_until_complete(asyncio.gather(*tasks)) + + latency = time.time() - tic + + preds = [] + for i in range(len(states)): + preds.append(get_answer_value(states[i])) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + print(f"Latency: {latency:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "multi_chain_gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-shot", type=int, default=0) + parser.add_argument("--num-chains", type=int, default=5) + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=50) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/multi_chain_reasoning/bench_sglang.py b/sglang/benchmark/multi_chain_reasoning/bench_sglang.py new file mode 100644 index 0000000..98a6b51 --- /dev/null +++ b/sglang/benchmark/multi_chain_reasoning/bench_sglang.py @@ -0,0 +1,140 @@ +import argparse +import ast +import json +import re +import time + +import numpy as np + +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +prompt_lib = [ + "Let us think step by step.", + "Approach this methodically. Let's dissect the problem into smaller, more manageable parts.", + "It's important to proceed step by step, ensuring accuracy at each stage.", + "Take a deep breath and break this down.", + "A little bit of arithmetic and a logical approach will help us quickly arrive at the solution to this problem.", + "I am extremely good at math.", +] + + +def main(args): + lines = read_jsonl(args.data_path) + + # Construct prompts + # k = args.num_shot + # few_shot_examples = get_few_shot_examples(lines, k) + + questions = [] + labels = [] + for i in range(len(lines[: args.num_questions])): + questions.append(lines[i]["question"]) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q} for q in questions] + + num_chains = args.num_chains + + ##################################### + ######### SGL Program Begin ######### + ##################################### + + import sglang as sgl + + @sgl.function + def multi_chain_gsm8k(s, question): + s += "Question: " + question + "\n" + # s += "Answer: " + prompt_lib[0] + sgl.gen("answer", max_tokens=256, stop="Question", + # temperature=0) + # return + + forks = s.fork(num_chains) + for i in range(num_chains): + forks[i] += ( + "Answer: " + + prompt_lib[i % num_chains] + + sgl.gen("chain", max_tokens=256, temperature=0.3, stop="Question") + ) + forks.join() + + s += "Answer: To answer this question, here are some possible solutions. " + s += "After considering all of them, I will do a majority vote.\n\n" + for i in range(num_chains): + s += f"Solution {i+1}: " + forks[i]["chain"].strip() + "\n\n" + s += "\nBy considering the above solutions and doing a majority vote, I think the final answer (a single integer number) is " + s += sgl.gen("answer", max_tokens=16) + + ##################################### + ########## SGL Program End ########## + ##################################### + + # Select backend + backend = select_sglang_backend(args) + + # Run requests + tic = time.time() + states = multi_chain_gsm8k.run_batch( + arguments, + temperature=0, + backend=backend, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.time() - tic + + preds = [] + for i in range(len(states)): + preds.append(get_answer_value(states[i]["answer"])) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + print(f"Latency: {latency:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "multi_chain_gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--num-shot", type=int, default=0) + parser.add_argument("--num-chains", type=int, default=5) + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=50) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/multi_document_qa/README.md b/sglang/benchmark/multi_document_qa/README.md new file mode 100644 index 0000000..09f49c7 --- /dev/null +++ b/sglang/benchmark/multi_document_qa/README.md @@ -0,0 +1,47 @@ +## Run benchmark + +### Benchmark sglang +``` +python3 -m sglang.launch_server --model-path codellama/CodeLlama-7b-instruct-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 10 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model codellama/CodeLlama-7b-instruct-hf --disable-log-requests --port 21000 --gpu 0.97 +``` + +``` +python3 bench_other.py --backend vllm --num-questions 64 +``` + + +### Benchmark guidance +``` +python3 bench_other.py --backend guidance --num-questions 32 --parallel 1 --n-ctx 11000 --model-path path/to/code-llama/gguf +``` + + + +### Build dataset + +``` +pip install PyPDF2 +python3 build_dataset.py +``` + +```python +import PyPDF2 + +with open('llama2.pdf', 'rb') as file: + reader = PyPDF2.PdfReader(file) + text = '' + for page_num in range(len(reader.pages)): + text += reader.pages[page_num].extract_text() + with open('output.txt', 'w') as text_file: + text_file.write(text) +``` diff --git a/sglang/benchmark/multi_document_qa/bench_other.py b/sglang/benchmark/multi_document_qa/bench_other.py new file mode 100644 index 0000000..97ff416 --- /dev/null +++ b/sglang/benchmark/multi_document_qa/bench_other.py @@ -0,0 +1,114 @@ +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +USER_PREFIX = "[INST] " +USER_SUFFIX = " [/INST]" +ASSISTANT_PREFIX = "" +ASSISTANT_SUFFIX = " " + + +def multi_document_qa(docs, question, generate): + s = USER_PREFIX + s += "Pleaes answer a question according to given documents.\n" + s += "Question:" + question + "Documents begin.\n" + + s += "".join(docs) + + s += "\nDocuments end." + s += ( + "\n\nBased on the above documents, please answer this question:\n" + + question + + "\nAnswer in three words or fewer." + ) + s += USER_SUFFIX + s += ASSISTANT_PREFIX + answer = generate(s, max_tokens=16, stop=None) + return answer + + +def main(args): + lines = read_jsonl(args.data_path) + l = lines[0] + arguments = [] + labels = [] + + num_docs = 10 + if args.backend == "guidance": + num_docs = 7 # due to OOM + + for i in range(len(l["questions"][: args.num_questions])): + arguments.append( + { + "docs": l["documents"][:num_docs], + "question": l["questions"][i], + } + ) + labels.append(l["answers"][i]) + states = [None] * len(arguments) + + # Select backend + call_generate = partial(get_call_generate(args), temperature=0) + + # Run requests + def get_one_answer(i): + states[i] = multi_document_qa(generate=call_generate, **arguments[i]) + + tic = time.time() + if args.parallel == 1: + for i in tqdm(range(len(labels))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(labels)))), + total=len(labels), + ) + ) + + latency = time.time() - tic + + # Compute accuracy + print(states) + correct = 0 + for s, label in zip(states, labels): + answer = s.lower() + if all(x in answer for x in label.lower().split(" ")): + correct += 1 + accuracy = correct / len(labels) + print(f"Accuracy: {accuracy:.3f}") + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "multi_document_qa", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "accuracy": accuracy, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="questions.jsonl") + parser.add_argument("--num-questions", type=int, default=100) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/multi_document_qa/bench_sglang.py b/sglang/benchmark/multi_document_qa/bench_sglang.py new file mode 100644 index 0000000..2c2db4f --- /dev/null +++ b/sglang/benchmark/multi_document_qa/bench_sglang.py @@ -0,0 +1,93 @@ +import argparse +import json +import time + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + + +@sgl.function +def multi_document_qa(s, docs, question): + s += sgl.user_begin() + s += "Pleaes answer a question according to given documents.\n" + s += "Question:" + question + "Documents begin.\n" + + forks = s.fork(len(docs)) + forks += lambda i: docs[i] + forks.join("concate_and_append") + + s += "\nDocuments end." + s += ( + "\n\nBased on the above documents, please answer this question:\n" + + question + + "\nAnswer in three words or fewer." + ) + s += sgl.user_end() + s += sgl.assistant(sgl.gen("answer", max_tokens=16)) + + +def main(args): + lines = read_jsonl(args.data_path) + l = lines[0] + arguments = [] + labels = [] + for i in range(len(l["questions"][: args.num_questions])): + arguments.append( + { + "docs": l["documents"][:10], + "question": l["questions"][i], + } + ) + labels.append(l["answers"][i]) + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + # Run requests + tic = time.time() + states = multi_document_qa.run_batch( + arguments, temperature=0, num_threads=args.parallel, progress_bar=True + ) + latency = time.time() - tic + + # Compute accuracy + print([s["answer"] for s in states]) + correct = 0 + for s, label in zip(states, labels): + answer = s["answer"].lower() + if all(x in answer for x in label.lower().split(" ")): + correct += 1 + accuracy = correct / len(labels) + print(f"Accuracy: {accuracy:.3f}") + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "multi_document_qa", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "accuracy": accuracy, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="questions.jsonl") + parser.add_argument("--num-questions", type=int, default=100) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/multi_document_qa/build_dataset.py b/sglang/benchmark/multi_document_qa/build_dataset.py new file mode 100644 index 0000000..27df947 --- /dev/null +++ b/sglang/benchmark/multi_document_qa/build_dataset.py @@ -0,0 +1,70 @@ +import json + +import transformers + +content = "\n".join( + open("llama2.txt", "r", encoding="utf-8", errors="ignore").readlines() +) +content = content.replace("\n\n", "\n") + +# Count token +name = "meta-llama/Llama-2-7b-chat-hf" +t = transformers.AutoTokenizer.from_pretrained(name) +print(f"num tokens: {len(t.encode(content))}") + +# Segment +SEP = "\n\n" +parts = content.split(SEP) +print(f"num segments: {len(parts)}") + +segment_len = 1100 + +segments = [] +tmp = [] +tmp_len = 0 +for i in range(len(parts)): + tmp.append(parts[i]) + tmp_len += len(t.encode(parts[i])) + + if tmp_len > segment_len: + segments.append(SEP.join(tmp)) + tmp = [] + tmp_len = 0 + +for i, s in enumerate(segments): + print(i, len(t.encode(segments[i]))) + +# Dump +with open("questions.jsonl", "w") as fout: + fout.write( + json.dumps( + { + "documents": segments[:30], + "questions": [ + "What is the name of the fine-tuned LLMs?", + "Which figure shows the helpfulness human evaluation results for Llama 2-Chat?", + "What is the number of parameters in the largest Llama 2 model?", + "What is the batch size of fine-tuning?", + "Where can we find the details of potential data contamination?", + "What is the full name of MPT?", + "What is the power consumption of RSC in Watt?", + "How many tokens of data do they train on?", + "Which model's release is delayed due to a lack of time to sufficiently red team?", + "Which activation function is used in Llama?", + ], + "answers": [ + "Llama 2 Chat", + "1", + "70 B", + "64", + "A 6", + "MosaicML", + "400", + "2 trillion", + "34 B", + "SwiGLU", + ], + } + ) + + "\n" + ) diff --git a/sglang/benchmark/multi_turn_chat/README.md b/sglang/benchmark/multi_turn_chat/README.md new file mode 100644 index 0000000..0fb5b21 --- /dev/null +++ b/sglang/benchmark/multi_turn_chat/README.md @@ -0,0 +1,66 @@ +### Benchmark sglang + +Run Llama-7B + +``` +python3 -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +Run Mixtral-8x7B +(When there is a CUDA out-of-memory error, try to reduce the `--mem-fraction-static`) + +``` +python3 -m sglang.launch_server --model-path mistralai/Mixtral-8x7B-Instruct-v0.1 --port 30000 --tp-size 8 +``` + +Benchmark(short output) + +``` +python3 bench_sglang.py --tokenizer meta-llama/Llama-2-7b-chat-hf +``` + +Benchmark(long output) + +``` +python3 bench_sglang.py --tokenizer meta-llama/Llama-2-7b-chat-hf --long +``` + +### Benchmark vLLM + +Run Llama-7B + +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +Run Mixtral-8x7B + +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model mistralai/Mixtral-8x7B-Instruct-v0.1 --disable-log-requests --port 21000 --tensor-parallel-size 8 +``` + +Benchmark(short output) + +``` +python3 bench_other.py --tokenizer meta-llama/Llama-2-7b-chat-hf --backend vllm +``` + +Benchmark(long output) + +``` +python3 bench_other.py --tokenizer meta-llama/Llama-2-7b-chat-hf --backend vllm --long +``` + +### Benchmark guidance + +Benchmark Llama-7B (short output) + +``` +python3 bench_other.py --tokenizer meta-llama/Llama-2-7b-chat-hf --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +Benchmark Llama-7B (long output) + +``` +python3 bench_other.py --tokenizer meta-llama/Llama-2-7b-chat-hf --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf --long +``` diff --git a/sglang/benchmark/multi_turn_chat/bench_other.py b/sglang/benchmark/multi_turn_chat/bench_other.py new file mode 100644 index 0000000..81d67ab --- /dev/null +++ b/sglang/benchmark/multi_turn_chat/bench_other.py @@ -0,0 +1,93 @@ +import json +import time +from argparse import ArgumentParser +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +from data_gen import gen_arguments +from tqdm import tqdm +from vllm.transformers_utils.tokenizer import get_tokenizer + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text + + +def multi_turns(generate, qas): + s = "" + for qa in qas: + s += qa["prompt"] + s += generate(s, max_tokens=qa["new_tokens"]) + + return s + + +def main(args): + print(args) + + tokenizer = get_tokenizer(args.tokenizer, trust_remote_code=args.trust_remote_code) + + multi_qas = gen_arguments(args, tokenizer) + + states = [None] * args.num_qa + + call_generate = partial(get_call_generate(args), temperature=0) + + def get_one_answer(i): + states[i] = multi_turns(generate=call_generate, **multi_qas[i]) + + tic = time.time() + if args.parallel == 1: + for i in tqdm(range(len(multi_qas))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + rets = list( + tqdm( + executor.map(get_one_answer, list(range(len(multi_qas)))), + total=len(multi_qas), + ) + ) + for _ in rets: + pass + + latency = time.time() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "multi_turn_chat", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_qa, + "num_turns": args.turns, + "other": { + "parallel": args.parallel, + "output_mode": "long" if args.long else "short", + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--turns", type=int, default=4) + parser.add_argument("--num-qa", type=int, default=20) + parser.add_argument("--min-len-q", type=int, default=256) + parser.add_argument("--max-len-q", type=int, default=512) + parser.add_argument("--min-len-a", type=int, default=4) + parser.add_argument("--max-len-a", type=int, default=8) + parser.add_argument("--tokenizer", type=str, required=True) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--long", action="store_true") + args = add_common_other_args_and_parse(parser) + + if args.long: + args.min_len_a = 256 + args.max_len_a = 512 + args.num_qa = 20 + main(args) diff --git a/sglang/benchmark/multi_turn_chat/bench_sglang.py b/sglang/benchmark/multi_turn_chat/bench_sglang.py new file mode 100644 index 0000000..7feaced --- /dev/null +++ b/sglang/benchmark/multi_turn_chat/bench_sglang.py @@ -0,0 +1,79 @@ +import json +import time +from argparse import ArgumentParser + +from data_gen import gen_arguments +from vllm.transformers_utils.tokenizer import get_tokenizer + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text + + +@sgl.function +def multi_turns(s, qas): + for qa in qas: + s += qa["prompt"] + s += sgl.gen(max_tokens=qa["new_tokens"], ignore_eos=True) + + +def main(args): + tokenizer = get_tokenizer(args.tokenizer, trust_remote_code=args.trust_remote_code) + + multi_qas = gen_arguments(args, tokenizer) + + backend = select_sglang_backend(args) + + tic = time.time() + states = multi_turns.run_batch( + multi_qas, + temperature=0, + backend=backend, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.time() - tic + + print(f"Latency: {latency:.3f}") + + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "multi_turn_chat", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_qa, + "num_turns": args.turns, + "other": { + "parallel": args.parallel, + "output_mode": "long" if args.long else "short", + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--turns", type=int, default=4) + parser.add_argument("--num-qa", type=int, default=20) + parser.add_argument("--min-len-q", type=int, default=256) + parser.add_argument("--max-len-q", type=int, default=512) + parser.add_argument("--min-len-a", type=int, default=4) + parser.add_argument("--max-len-a", type=int, default=8) + parser.add_argument("--tokenizer", type=str, required=True) + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--long", action="store_true") + args = add_common_sglang_args_and_parse(parser) + + if args.long: + args.min_len_a = 256 + args.max_len_a = 512 + args.num_qa = 20 + + print(args) + main(args) diff --git a/sglang/benchmark/multi_turn_chat/data_gen.py b/sglang/benchmark/multi_turn_chat/data_gen.py new file mode 100644 index 0000000..043c07a --- /dev/null +++ b/sglang/benchmark/multi_turn_chat/data_gen.py @@ -0,0 +1,29 @@ +import random +import string + +random.seed(42) + + +def gen_prompt(tokenizer, token_num): + cha_set = string.ascii_letters + string.digits + ret = "".join(random.choices(cha_set, k=token_num)) + while len(tokenizer(ret).input_ids) < token_num: + ret += random.choice(cha_set) + return ret + + +def gen_arguments(args, tokenizer): + multi_qas = [{"qas": []} for _ in range(args.num_qa)] + for i in range(args.num_qa): + qas = multi_qas[i]["qas"] + for _ in range(args.turns): + prompt_len = random.randint(args.min_len_q, args.max_len_q) + new_tokens = random.randint(args.min_len_a, args.max_len_a) + qas.append( + { + "prompt": gen_prompt(tokenizer, prompt_len), + "new_tokens": new_tokens, + } + ) + + return multi_qas diff --git a/sglang/benchmark/multi_turn_chat/long_prompt_multi_turn.py b/sglang/benchmark/multi_turn_chat/long_prompt_multi_turn.py new file mode 100644 index 0000000..20f6dd5 --- /dev/null +++ b/sglang/benchmark/multi_turn_chat/long_prompt_multi_turn.py @@ -0,0 +1,129 @@ +import json +import random +import time +from argparse import ArgumentParser +from pathlib import Path + +from tqdm import tqdm + +import sglang as sgl +from sglang.srt.hf_transformers_utils import get_tokenizer +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text + + +def gen_prompt(tokenizer, token_num): + all_available_tokens = list(tokenizer.get_vocab().values()) + selected_tokens = random.choices(all_available_tokens, k=token_num) + ret = tokenizer.decode(selected_tokens) + return ret + + +def get_cache_path(args): + # Create cache directory under ~/.cache/sglang + cache_dir = Path.home() / ".cache" / "sglang" + + # Create a unique cache filename based on the arguments that affect generation + cache_key = f"qa_{args.num_qa}_{args.turns}_{args.system_prompt_len}_{args.len_q}_{args.len_a}_{args.tokenizer.replace('/', '_')}.json" + return cache_dir / cache_key + + +def gen_arguments(args, tokenizer): + cache_path = get_cache_path(args) + + # Try to load from cache first + if cache_path.exists(): + print(f"Loading cached arguments from {cache_path}") + with open(cache_path, "r") as f: + return json.load(f) + + print("Generating new arguments...") + # First progress bar for system prompts + multi_qas = [] + for _ in tqdm(range(args.num_qa), desc="Generating system prompts"): + multi_qas.append( + {"system_prompt": gen_prompt(tokenizer, args.system_prompt_len), "qas": []} + ) + + # Nested progress bars for QA pairs + for i in tqdm(range(args.num_qa), desc="Generating QA pairs"): + qas = multi_qas[i]["qas"] + for j in range(args.turns): + qas.append( + { + "prompt": gen_prompt(tokenizer, args.len_q), + "new_tokens": args.len_a, + } + ) + + # Save to cache + cache_path.parent.mkdir(parents=True, exist_ok=True) + with open(cache_path, "w") as f: + json.dump(multi_qas, f) + print(f"Cached arguments saved to {cache_path}") + + return multi_qas + + +@sgl.function +def multi_turns(s, system_prompt, qas): + s += system_prompt + + for i, qa in enumerate(qas): + s += qa["prompt"] + s += sgl.gen(max_tokens=qa["new_tokens"], ignore_eos=True) + + +def main(args): + tokenizer = get_tokenizer(args.tokenizer, trust_remote_code=args.trust_remote_code) + + multi_qas = gen_arguments(args, tokenizer) + + backend = select_sglang_backend(args) + + tic = time.time() + states = multi_turns.run_batch( + multi_qas, + temperature=0, + backend=backend, + num_threads="auto", + progress_bar=True, + ) + latency = time.time() - tic + + print(f"Latency: {latency:.3f}") + + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "multi_turn_system_prompt_chat", + "backend": args.backend, + "latency": round(latency, 3), + "num_requests": args.num_qa, + "num_turns": args.turns, + "other": { + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = ArgumentParser() + parser.add_argument("--turns", type=int, default=8) + parser.add_argument("--num-qa", type=int, default=128) + parser.add_argument("--system-prompt-len", type=int, default=2048) + parser.add_argument("--len-q", type=int, default=32) + parser.add_argument("--len-a", type=int, default=128) + parser.add_argument( + "--tokenizer", type=str, default="meta-llama/Meta-Llama-3-8B-Instruct" + ) + parser.add_argument("--trust-remote-code", action="store_true") + args = add_common_sglang_args_and_parse(parser) + + print(args) + main(args) diff --git a/sglang/benchmark/react/README.md b/sglang/benchmark/react/README.md new file mode 100644 index 0000000..51bcaa4 --- /dev/null +++ b/sglang/benchmark/react/README.md @@ -0,0 +1,34 @@ +## Run benchmark + +NOTE: This is an implementation for replaying a given trace for throughput/latency benchmark purposes. It is not an actual ReAct agent implementation. + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 100 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-questions 100 --backend vllm +``` + + +### Benchmark guidance +``` +python3 bench_other.py --num-questions 100 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +### Benchmark lmql + +``` +python3 bench_other.py --num-questions 100 --backend lmql --parallel 1 +``` diff --git a/sglang/benchmark/react/bench_other.py b/sglang/benchmark/react/bench_other.py new file mode 100644 index 0000000..91c5546 --- /dev/null +++ b/sglang/benchmark/react/bench_other.py @@ -0,0 +1,202 @@ +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor + +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + + +def get_prompt(question): + prompt = ( + """Solve a question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types: +(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search. +(2) Lookup[keyword], which returns the next sentence containing keyword in the current passage. +(3) Finish[answer], which returns the answer and finishes the task. +Here are some examples. +Question: What is the elevation range for the area that the eastern sector of the Colorado orogeny extends into? +Thought 1: I need to search Colorado orogeny, find the area that the eastern sector of the Colorado orogeny extends into, then find the elevation range of the area. +Action 1: Search[Colorado orogeny] +Observation 1: The Colorado orogeny was an episode of mountain building (an orogeny) in Colorado and surrounding areas. +Thought 2: It does not mention the eastern sector. So I need to look up eastern sector. +Action 2: Lookup[eastern sector] +Observation 2: (Result 1 / 1) The eastern sector extends into the High Plains and is called the Central Plains orogeny. +Thought 3: The eastern sector of Colorado orogeny extends into the High Plains. So I need to search High Plains and find its elevation range. +Action 3: Search[High Plains] +Observation 3: High Plains refers to one of two distinct land regions: +Thought 4: I need to instead search High Plains (United States). +Action 4: Search[High Plains (United States)] +Observation 4: The High Plains are a subregion of the Great Plains. From east to west, the High Plains rise in elevation from around 1,800 to 7,000 ft (550 to 2,130 m).[3] +Thought 5: High Plains rise in elevation from around 1,800 to 7,000 ft, so the answer is 1,800 to 7,000 ft. +Action 5: Finish[1,800 to 7,000 ft] +Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who? +Thought 1: The question simplifies to "The Simpsons" character Milhouse is named after who. I only need to search Milhouse and find who it is named after. +Action 1: Search[Milhouse] +Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening. +Thought 2: The paragraph does not tell who Milhouse is named after, maybe I can look up "named after". +Action 2: Lookup[named after] +Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous. +Thought 3: Milhouse was named after U.S. president Richard Nixon, so the answer is Richard Nixon. +Action 3: Finish[Richard Nixon] +Question: Which documentary is about Finnish rock groups, Adam Clayton Powell or The Saimaa Gesture? +Thought 1: I need to search Adam Clayton Powell and The Saimaa Gesture, and find which documentary is about Finnish rock groups. +Action 1: Search[Adam Clayton Powell] +Observation 1: Could not find [Adam Clayton Powell]. Similar: ['Adam Clayton Powell III', 'Seventh Avenue (Manhattan)', 'Adam Clayton Powell Jr. State Office Building', 'Isabel Washington Powell', 'Adam Powell', 'Adam Clayton Powell (film)', 'Giancarlo Esposito']. +Thought 2: To find the documentary, I can search Adam Clayton Powell (film). +Action 2: Search[Adam Clayton Powell (film)] +Observation 2: Adam Clayton Powell is a 1989 American documentary film directed by Richard Kilberg. +The film is about the rise and fall of influential African-American politician Adam Clayton Powell Jr.[3][4] It was later aired as part of the PBS series The American Experience. +Thought 3: Adam Clayton Powell (film) is a documentary about an African-American politician, not Finnish rock groups. So the documentary about Finnish rock groups must instead be The Saimaa Gesture. +Action 3: Finish[The Saimaa Gesture] +Question: What profession does Nicholas Ray and Elia Kazan have in common? +Thought 1: I need to search Nicholas Ray and Elia Kazan, find their professions, then find the profession they have in common. +Action 1: Search[Nicholas Ray] +Observation 1: Nicholas Ray (born Raymond Nicholas Kienzle Jr., August 7, 1911 – June 16, 1979) was an American film director, screenwriter, and actor best known for the 1955 film Rebel Without a Cause. +Thought 2: Professions of Nicholas Ray are director, screenwriter, and actor. I need to search Elia Kazan next and find his professions. +Action 2: Search[Elia Kazan] +Observation 2: Elia Kazan was an American film and theatre director, producer, screenwriter and actor. +Thought 3: Professions of Elia Kazan are director, producer, screenwriter, and actor. So profession Nicholas Ray and Elia Kazan have in common is director, screenwriter, and actor. +Action 3: Finish[director, screenwriter, actor] +Question: Which magazine was started first Arthur's Magazine or First for Women? +Thought 1: I need to search Arthur's Magazine and First for Women, and find which was started first. +Action 1: Search[Arthur's Magazine] +Observation 1: Arthur's Magazine (1844-1846) was an American literary periodical published in Philadelphia in the 19th century. +Thought 2: Arthur's Magazine was started in 1844. I need to search First for Women next. +Action 2: Search[First for Women] +Observation 2: First for Women is a woman's magazine published by Bauer Media Group in the USA.[1] The magazine was started in 1989. +Thought 3: First for Women was started in 1989. 1844 (Arthur's Magazine) < 1989 (First for Women), so Arthur's Magazine was started first. +Action 3: Finish[Arthur's Magazine] +Question: Were Pavel Urysohn and Leonid Levin known for the same type of work? +Thought 1: I need to search Pavel Urysohn and Leonid Levin, find their types of work, then find if they are the same. +Action 1: Search[Pavel Urysohn] +Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory. +Thought 2: Pavel Urysohn is a mathematician. I need to search Leonid Levin next and find its type of work. +Action 2: Search[Leonid Levin] +Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist. +Thought 3: Leonid Levin is a mathematician and computer scientist. So Pavel Urysohn and Leonid Levin have the same type of work. +Action 3: Finish[yes] +""" + + question + ) + return prompt + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_questions] + arguments = [{"question": k, "triplets": v} for l in lines for k, v in l.items()] + + states = [] + + # Select backend + call_generate = get_call_generate(args) + + def run_single_agent(argument): + question = argument["question"] + triplets = argument["triplets"] + prompt = get_prompt(question) + for i in range(1, len(triplets) + 2): + prompt += "Thought " + str(i) + ":" + states.append(prompt) + answer = call_generate( + prompt, max_tokens=200, temperature=0, stop="Observation" + ) + if i > len(triplets): + break + prompt += ( + triplets[i - 1]["thought"] + + "\nAction " + + str(i) + + ":" + + triplets[i - 1]["action"] + + "\nObservation " + + str(i) + + ":" + + triplets[i - 1]["observation"] + + "\n" + ) + + states.append(answer) + + async def run_single_agent_async(argument): + question = argument["question"] + triplets = argument["triplets"] + prompt = get_prompt(question) + for i in range(1, len(triplets) + 2): + prompt += "Thought " + str(i) + ":" + states.append(prompt) + answer = await call_generate( + prompt, max_tokens=200, temperature=0, stop="Observation", max_len=4096 + ) + if i > len(triplets): + break + prompt += ( + triplets[i - 1]["thought"] + + "\nAction " + + str(i) + + ":" + + triplets[i - 1]["action"] + + "\nObservation " + + str(i) + + ":" + + triplets[i - 1]["observation"] + + "\n" + ) + + states.append(answer) + + tic = time.time() + + if args.backend != "lmql": + if args.parallel == 1: + for arg in tqdm(arguments): + run_single_agent(arg) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(run_single_agent, arguments), total=len(arguments) + ) + ) + + else: + import asyncio + + loop = asyncio.get_event_loop() + batches = [ + [] for _ in range((len(arguments) + args.parallel - 1) // args.parallel) + ] + for i, arg in enumerate(arguments): + batches[i // args.parallel].append(arg) + for bt in tqdm(batches): + tasks = [run_single_agent_async(arg) for arg in bt] + loop.run_until_complete(asyncio.gather(*tasks)) + + latency = time.time() - tic + + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "ReAct Agents", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": len(arguments), + "other": { + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="hotpotqa_100.jsonl") + parser.add_argument("--num-questions", type=int, default=10) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/react/bench_sglang.py b/sglang/benchmark/react/bench_sglang.py new file mode 100644 index 0000000..b07105e --- /dev/null +++ b/sglang/benchmark/react/bench_sglang.py @@ -0,0 +1,153 @@ +import argparse +import json +import time + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + + +@sgl.function +def webthink(s, question, triplets): + s += ( + """Solve a question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types: +(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search. +(2) Lookup[keyword], which returns the next sentence containing keyword in the current passage. +(3) Finish[answer], which returns the answer and finishes the task. +Here are some examples. +Question: What is the elevation range for the area that the eastern sector of the Colorado orogeny extends into? +Thought 1: I need to search Colorado orogeny, find the area that the eastern sector of the Colorado orogeny extends into, then find the elevation range of the area. +Action 1: Search[Colorado orogeny] +Observation 1: The Colorado orogeny was an episode of mountain building (an orogeny) in Colorado and surrounding areas. +Thought 2: It does not mention the eastern sector. So I need to look up eastern sector. +Action 2: Lookup[eastern sector] +Observation 2: (Result 1 / 1) The eastern sector extends into the High Plains and is called the Central Plains orogeny. +Thought 3: The eastern sector of Colorado orogeny extends into the High Plains. So I need to search High Plains and find its elevation range. +Action 3: Search[High Plains] +Observation 3: High Plains refers to one of two distinct land regions: +Thought 4: I need to instead search High Plains (United States). +Action 4: Search[High Plains (United States)] +Observation 4: The High Plains are a subregion of the Great Plains. From east to west, the High Plains rise in elevation from around 1,800 to 7,000 ft (550 to 2,130 m).[3] +Thought 5: High Plains rise in elevation from around 1,800 to 7,000 ft, so the answer is 1,800 to 7,000 ft. +Action 5: Finish[1,800 to 7,000 ft] +Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who? +Thought 1: The question simplifies to "The Simpsons" character Milhouse is named after who. I only need to search Milhouse and find who it is named after. +Action 1: Search[Milhouse] +Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening. +Thought 2: The paragraph does not tell who Milhouse is named after, maybe I can look up "named after". +Action 2: Lookup[named after] +Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous. +Thought 3: Milhouse was named after U.S. president Richard Nixon, so the answer is Richard Nixon. +Action 3: Finish[Richard Nixon] +Question: Which documentary is about Finnish rock groups, Adam Clayton Powell or The Saimaa Gesture? +Thought 1: I need to search Adam Clayton Powell and The Saimaa Gesture, and find which documentary is about Finnish rock groups. +Action 1: Search[Adam Clayton Powell] +Observation 1: Could not find [Adam Clayton Powell]. Similar: ['Adam Clayton Powell III', 'Seventh Avenue (Manhattan)', 'Adam Clayton Powell Jr. State Office Building', 'Isabel Washington Powell', 'Adam Powell', 'Adam Clayton Powell (film)', 'Giancarlo Esposito']. +Thought 2: To find the documentary, I can search Adam Clayton Powell (film). +Action 2: Search[Adam Clayton Powell (film)] +Observation 2: Adam Clayton Powell is a 1989 American documentary film directed by Richard Kilberg. +The film is about the rise and fall of influential African-American politician Adam Clayton Powell Jr.[3][4] It was later aired as part of the PBS series The American Experience. +Thought 3: Adam Clayton Powell (film) is a documentary about an African-American politician, not Finnish rock groups. So the documentary about Finnish rock groups must instead be The Saimaa Gesture. +Action 3: Finish[The Saimaa Gesture] +Question: What profession does Nicholas Ray and Elia Kazan have in common? +Thought 1: I need to search Nicholas Ray and Elia Kazan, find their professions, then find the profession they have in common. +Action 1: Search[Nicholas Ray] +Observation 1: Nicholas Ray (born Raymond Nicholas Kienzle Jr., August 7, 1911 – June 16, 1979) was an American film director, screenwriter, and actor best known for the 1955 film Rebel Without a Cause. +Thought 2: Professions of Nicholas Ray are director, screenwriter, and actor. I need to search Elia Kazan next and find his professions. +Action 2: Search[Elia Kazan] +Observation 2: Elia Kazan was an American film and theatre director, producer, screenwriter and actor. +Thought 3: Professions of Elia Kazan are director, producer, screenwriter, and actor. So profession Nicholas Ray and Elia Kazan have in common is director, screenwriter, and actor. +Action 3: Finish[director, screenwriter, actor] +Question: Which magazine was started first Arthur's Magazine or First for Women? +Thought 1: I need to search Arthur's Magazine and First for Women, and find which was started first. +Action 1: Search[Arthur's Magazine] +Observation 1: Arthur's Magazine (1844-1846) was an American literary periodical published in Philadelphia in the 19th century. +Thought 2: Arthur's Magazine was started in 1844. I need to search First for Women next. +Action 2: Search[First for Women] +Observation 2: First for Women is a woman's magazine published by Bauer Media Group in the USA.[1] The magazine was started in 1989. +Thought 3: First for Women was started in 1989. 1844 (Arthur's Magazine) < 1989 (First for Women), so Arthur's Magazine was started first. +Action 3: Finish[Arthur's Magazine] +Question: Were Pavel Urysohn and Leonid Levin known for the same type of work? +Thought 1: I need to search Pavel Urysohn and Leonid Levin, find their types of work, then find if they are the same. +Action 1: Search[Pavel Urysohn] +Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory. +Thought 2: Pavel Urysohn is a mathematician. I need to search Leonid Levin next and find its type of work. +Action 2: Search[Leonid Levin] +Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist. +Thought 3: Leonid Levin is a mathematician and computer scientist. So Pavel Urysohn and Leonid Levin have the same type of work. +Action 3: Finish[yes] +""" + + question + ) + for i in range(1, len(triplets) + 2): + s += "Thought " + str(i) + ":" + # NOTE: This is an implementation for replaying a given trace for benchmark purposes. It is not an actual ReAct agent implementation. + ss = s.fork(1) + ss[0] += sgl.gen(name="thought_action", max_tokens=200, stop="Observation") + ss.join() + # to verify the correctness of output, this should be collected + # print(ss[0]["thought_action"]) + if i > len(triplets): + break + s += ( + triplets[i - 1]["thought"] + + "\nAction " + + str(i) + + ":" + + triplets[i - 1]["action"] + + "\nObservation " + + str(i) + + ":" + + triplets[i - 1]["observation"] + + "\n" + ) + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_questions] + arguments = [{"question": k, "triplets": v} for l in lines for k, v in l.items()] + + # Select backend + backend = select_sglang_backend(args) + sgl.set_default_backend(backend) + + states = [] + tic = time.time() + states = webthink.run_batch( + arguments, + temperature=0, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.time() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "ReAct Agents", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": len(arguments), + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="hotpotqa_100.jsonl") + parser.add_argument("--num-questions", type=int, default=10) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/reasoning_benchmark/README.md b/sglang/benchmark/reasoning_benchmark/README.md new file mode 100644 index 0000000..f9d26e0 --- /dev/null +++ b/sglang/benchmark/reasoning_benchmark/README.md @@ -0,0 +1,77 @@ +# Run benchmark + +This benchmark is primarily intended to be used with reasoning models like `DeepSeek-R1` and its distilled models like `DeepSeek-R1-Distill-Qwen-1.5B`. Please use + +```bash +pip install antlr4-python3-runtime +``` + +for `parse_latex` which we use for symbolic equality check. + +## Benchmark sglang + +1. Launch the Server +```bash +python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --port 30000 +``` + +Note that depending on the GPU this benchmark will take quiet some time. To employ data parallelism please use: + +```bash +python3 -m sglang_router.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --port 30000 --dp-size 4 +``` + + +2. Benchmarking + +We use [suggested](https://github.com/deepseek-ai/DeepSeek-R1) parameters of `temperature=0.6`, `top_p=.95`, `max_new_tokens=32768`. The command line argument `num-tries` can be used to evaluate the model multiple times on the same question. We use the suggested `64` from the repo for AIME 2024. For LIMO, we use `8` as the number of tries due to the size of the dataset. + +By default evaluate on LIMO dataset. + +```bash +python3 bench_sglang.py --parallel 256 --num-tries 64 --port 30000 +``` + +Evaluate on AIME 2024 dataset. + +```bash +python3 bench_sglang.py --parallel 256 --port 30000 --data-path Maxwell-Jia/AIME_2024 --question-key Problem --answer-key Answer --num-tries 64 +``` + +Evaluate on [AIME 2025 I dataset](https://huggingface.co/datasets/opencompass/AIME2025). For benchmark result see [here](https://matharena.ai/). + +```bash +python3 bench_sglang.py --parallel 256 --port 30000 --data-path opencompass/AIME2025 --question-key question --answer-key answer --num-tries 64 +``` +## Results + +### Evaluation Results +| Dataset | Num Tries | Accuracy | Reference | Standard Error | +|------------|-----------|----------|-----------|-----------| +| LIMO | 8 | 47.7% | ? | ? | +| AIME 2024 | 64 | 33.2% | 28.9% | 3.4% | +| AIME 2025 I| 64 | 29.9% | 25.0% | ? | + +### Statistic Analysis Results +Set up SGLang engine for statistic analysis, for high efficiency we use `--dp-size 8` for data parallelism: +```bash +python3 -m sglang_router.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B --port 30000 --dp-size 8 +``` +**Experiment 1**: +We fixed the number of attempts (num_tries) and conducted multiple runs to assess the consistency of the model's performance. The results show that all recorded accuracies lie within ± one standard error deviation from the mean. This suggests that **our metric serves as an effective upper bound for the deviation of reported accuracy**. + +To collect the accuracy, run the following command 30 times: +```bash +python3 bench_sglang.py --parallel 64 --port 30000 --data-path Maxwell-Jia/AIME_2024 --question-key Problem --answer-key Answer --num-tries 64 +``` + +![acc_hist](figure/Acc_histplot.png) + + +**Experiment 2**: We explored the relationship between the number of attempts (num_tries) and the standard error (SE) by varying num_tries across a range (e.g., 8, 16, 32, ..., 256) and performing a single run for each value. The results demonstrate that as the number of attempts increases, the standard error decreases, leading to **greater stability in answer accuracy**. + +To reveal the relationship, run the command 6 times and adjust the parameter `--num-tries` for each run: +```bash +python3 bench_sglang.py --parallel 64 --port 30000 --data-path Maxwell-Jia/AIME_2024 --question-key Problem --answer-key Answer --num-tries +``` +![SE_num_tries](figure/SE_numtries.png) diff --git a/sglang/benchmark/reasoning_benchmark/answer_extraction.py b/sglang/benchmark/reasoning_benchmark/answer_extraction.py new file mode 100644 index 0000000..45ce59a --- /dev/null +++ b/sglang/benchmark/reasoning_benchmark/answer_extraction.py @@ -0,0 +1,269 @@ +# Adapted from https://github.com/deepseek-ai/DeepSeek-Math/blob/main/evaluation/data_processing/answer_extraction.py + +import re + +import regex + + +def _fix_fracs(string): + substrs = string.split("\\frac") + new_str = substrs[0] + if len(substrs) > 1: + substrs = substrs[1:] + for substr in substrs: + new_str += "\\frac" + if len(substr) > 0 and substr[0] == "{": + new_str += substr + else: + try: + assert len(substr) >= 2 + except: + return string + a = substr[0] + b = substr[1] + if b != "{": + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}{" + b + "}" + post_substr + else: + new_str += "{" + a + "}{" + b + "}" + else: + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}" + b + post_substr + else: + new_str += "{" + a + "}" + b + string = new_str + return string + + +def _fix_a_slash_b(string): + if len(string.split("/")) != 2: + return string + a = string.split("/")[0] + b = string.split("/")[1] + try: + if "sqrt" not in a: + a = int(a) + if "sqrt" not in b: + b = int(b) + assert string == "{}/{}".format(a, b) + new_string = "\\frac{" + str(a) + "}{" + str(b) + "}" + return new_string + except: + return string + + +def _fix_sqrt(string): + _string = re.sub(r"\\sqrt(-?[0-9.a-zA-Z]+)", r"\\sqrt{\1}", string) + _string = re.sub(r"\\sqrt\s+(\w+)$", r"\\sqrt{\1}", _string) + return _string + + +def _fix_tan(string): + _string = re.sub(r"\\tan(-?[0-9.a-zA-Z]+)", r"\\tan{\1}", string) + _string = re.sub(r"\\tan\s+(\w+)$", r"\\tan{\1}", _string) + return _string + + +def strip_string(string): + string = str(string).strip() + # linebreaks + string = string.replace("\n", "") + + # right "." + string = string.rstrip(".") + + # remove inverse spaces + string = string.replace("\\!", "") + # string = string.replace("\\ ", "") + + # replace \\ with \ + # string = string.replace("\\\\", "\\") + # string = string.replace("\\\\", "\\") + + if string.startswith("\\text{") and string.endswith("}"): + string = string.split("{", 1)[1][:-1] + + # replace tfrac and dfrac with frac + string = string.replace("tfrac", "frac") + string = string.replace("dfrac", "frac") + string = string.replace("cfrac", "frac") + + # remove \left and \right + string = string.replace("\\left", "") + string = string.replace("\\right", "") + + # Remove unit: miles, dollars if after is not none + _string = re.sub(r"\\text{.*?}$", "", string).strip() + if _string != "" and _string != string: + # print("Warning: unit not removed: '{}' -> '{}'".format(string, _string)) + string = _string + + # Remove circ (degrees) + string = string.replace("^{\\circ}", "").strip() + string = string.replace("^\\circ", "").strip() + + string = regex.sub(r"\{(c|m)?m\}(\^(2|3))?", "", string).strip() + string = regex.sub(r"p\.m\.$", "", string).strip() + string = regex.sub(r"(\d)\s*t$", r"\1", string).strip() + + # remove dollar signs + string = string.replace("\\$", "") + string = string.replace("$", "") + + # string = string.replace("\\text", "") + string = string.replace("x\\in", "") + + # remove percentage + string = string.replace("\\%", "%") + string = string.replace("\%", "%") + # string = string.replace("%", "") + + # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string + string = string.replace(" .", " 0.") + string = string.replace("{.", "{0.") + + # cdot + string = string.replace("\\cdot", "") + + # inf + string = string.replace("infinity", "\\infty") + if "\\infty" not in string: + string = string.replace("inf", "\\infty") + string = string.replace("+\\inity", "\\infty") + + # and + # string = string.replace("and", "") + string = string.replace("\\mathbf", "") + string = string.replace("\\mathrm", "") + + # use regex to remove \mbox{...} + string = re.sub(r"\\mbox{.*?}", "", string) + + # quote + string.replace("'", "") + string.replace('"', "") + + # i, j + if "j" in string and "i" not in string: + string = string.replace("j", "i") + + # replace a.000b where b is not number or b is end, with ab, use regex + string = re.sub(r"(\d+)\.0+([^\d])", r"\1\2", string) + string = re.sub(r"(\d+)\.0+$", r"\1", string) + + # if empty, return empty string + if len(string) == 0: + return string + if string[0] == ".": + string = "0" + string + + # to consider: get rid of e.g. "k = " or "q = " at beginning + # if len(string.split("=")) == 2: + # if len(string.split("=")[0]) <= 2: + # string = string.split("=")[1] + + string = _fix_sqrt(string) + string = _fix_tan(string) + string = string.replace(" ", "") + + # \frac1b or \frac12 --> \frac{1}{b} and \frac{1}{2}, etc. Even works with \frac1{72} (but not \frac{72}1). Also does a/b --> \\frac{a}{b} + string = _fix_fracs(string) + + # NOTE: X/Y changed to \frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y + string = _fix_a_slash_b(string) + + string = regex.sub(r"(\\|,|\.)+$", "", string) + + return string + + +def extract_boxed_answers(text): + answers = [] + for piece in text.split("boxed{")[1:]: + n = 0 + for i in range(len(piece)): + if piece[i] == "{": + n += 1 + elif piece[i] == "}": + n -= 1 + if n < 0: + if i + 1 < len(piece) and piece[i + 1] == "%": + answers.append(piece[: i + 1]) + else: + answers.append(piece[:i]) + break + return answers + + +def extract_program_output(pred_str): + """ + extract output between the last ```output\n...\n``` + """ + if "```output" not in pred_str: + return "" + if "```output" in pred_str: + pred_str = pred_str.split("```output")[-1] + if "```" in pred_str: + pred_str = pred_str.split("```")[0] + output = pred_str.strip() + return output + + +def extract_answer(pred_str, exhaust=False): + pred = [] + if "final answer is $" in pred_str and "$. I hope" in pred_str: + tmp = pred_str.split("final answer is $", 1)[1] + pred = [tmp.split("$. I hope", 1)[0].strip()] + elif "boxed" in pred_str: + pred = extract_boxed_answers(pred_str) + elif "he answer is" in pred_str: + pred = [pred_str.split("he answer is")[-1].strip()] + else: + program_output = extract_program_output(pred_str) + if program_output != "": + # fall back to program + pred.append(program_output) + else: # use the last number + pattern = "-?\d*\.?\d+" + ans = re.findall(pattern, pred_str.replace(",", "")) + if len(ans) >= 1: + ans = ans[-1] + else: + ans = "" + if ans: + pred.append(ans) + + # multiple line + _pred = [] + for ans in pred: + ans = ans.strip().split("\n")[0] + ans = ans.lstrip(":") + ans = ans.rstrip(".") + ans = ans.rstrip("/") + ans = strip_string(ans) + _pred.append(ans) + if exhaust: + return _pred + else: + return _pred[-1] if _pred else "" + + +def extract_math_answer(question, reasoning, task): + answer = [] + for ans in extract_answer(reasoning, exhaust=True): + if "separated by commas" in question and all(ch not in ans for ch in "()[]"): + answer.extend([a.strip() for a in ans.split(",")]) + elif regex.search(r"\\text\{\s*and\s*\}", ans): + answer.extend( + [ + a.strip() + for a in regex.sub(r"\\text\{\s*and\s*\}", "[SEP]", ans).split( + "[SEP]" + ) + ] + ) + else: + answer.append(ans.strip()) + return answer diff --git a/sglang/benchmark/reasoning_benchmark/bench_sglang.py b/sglang/benchmark/reasoning_benchmark/bench_sglang.py new file mode 100644 index 0000000..c832049 --- /dev/null +++ b/sglang/benchmark/reasoning_benchmark/bench_sglang.py @@ -0,0 +1,135 @@ +import argparse +import json +import time + +import answer_extraction +import eval_utils +import numpy as np +from datasets import load_dataset + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text + + +@sgl.function +def reasoning_gen(s, question: str): + s += sgl.user( + question + + "\nPlease reason step by step, and put your final answer within \boxed{}." + ) + s += sgl.assistant( + sgl.gen( + "answer", + ) + ) + + +def convert_dataset(path: str, question_key: str, answer_key: str, num_tries: int): + raw_dataset = load_dataset(path) + questions = [] + answers = [] + for data in raw_dataset["train"]: + question = data[question_key] + answer = data[answer_key] + for _ in range(num_tries): + questions.append({"question": question}) + answers.append({"answer": answer}) + return questions, answers + + +def main(args): + # Select backend + sgl.set_default_backend(select_sglang_backend(args)) + + # Get dataset + questions, answers = convert_dataset( + args.data_path, args.question_key, args.answer_key, args.num_tries + ) + + # Run requests + tic = time.time() + states = reasoning_gen.run_batch( + questions, + num_threads=args.parallel, + progress_bar=True, + temperature=0.6, + max_new_tokens=32768, + top_p=0.95, + ) + latency = time.time() - tic + + # Extract results and record outcomes in a list. + outcomes = [] + for i, state in enumerate(states): + try: + pred_answer = answer_extraction.extract_math_answer( + questions[i]["question"], state["answer"], "limo" + ) + gt_answer = str(answers[i]["answer"]) + pred_answer = ( + pred_answer[-1] if isinstance(pred_answer, list) else pred_answer + ) + is_correct = 1 if eval_utils.math_equal(pred_answer, gt_answer) else 0 + except Exception as e: + print(f"Error extracting answer: {e}") + is_correct = 0 + + outcomes.append(is_correct) + + # Calculate overall accuracy using numpy + overall_accuracy = np.mean(outcomes) + print(f"Overall Accuracy: {overall_accuracy}") + + # Calculate mean standard error over questions if num_tries >= 2 + if args.num_tries > 1: + outcomes_np = np.array(outcomes).reshape(-1, args.num_tries) + # Using sample standard deviation with ddof=1 + std_per_question = np.std(outcomes_np, axis=1, ddof=1) + # Compute the standard error for each question: std / sqrt(num_tries) + se_per_question = std_per_question / np.sqrt(args.num_tries) + mean_se = se_per_question.mean() + print(f"Mean Standard Error of Accuracy across questions: {mean_se}") + else: + mean_se = None + print("Not enough samples per question to compute standard error.") + + # Calculate output throughput + num_output_tokens = sum( + s.get_meta_info("answer")["completion_tokens"] for s in states + ) + output_throughput = num_output_tokens / latency + print(f"Output throughput: {output_throughput} token/s") + + # Dump results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + # Write results + with open(args.result_file, "a") as fout: + value = { + "task": "limo", + "backend": args.backend, + "latency": round(latency, 3), + "overall_accuracy": round(overall_accuracy, 3), + "mean_se_accuracy": round(mean_se, 3) if mean_se is not None else None, + "num_requests": len(questions), + "other": { + "num_questions": len(questions), + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="GAIR/LIMO") + parser.add_argument("--question-key", type=str, default="question") + parser.add_argument("--answer-key", type=str, default="answer") + parser.add_argument("--num-tries", type=int, default=1) + add_common_sglang_args_and_parse(parser) + args = parser.parse_args() + main(args) diff --git a/sglang/benchmark/reasoning_benchmark/eval_utils.py b/sglang/benchmark/reasoning_benchmark/eval_utils.py new file mode 100644 index 0000000..ab73695 --- /dev/null +++ b/sglang/benchmark/reasoning_benchmark/eval_utils.py @@ -0,0 +1,206 @@ +# Adapted from https://github.com/deepseek-ai/DeepSeek-Math/blob/main/evaluation/eval/eval_utils.py + +from math import isclose + +import regex +from sympy import N, simplify +from sympy.parsing.latex import parse_latex +from sympy.parsing.sympy_parser import parse_expr + + +def parse_digits(num): + # format: 234.23 || 23% + num = regex.sub(",", "", str(num)) + try: + return float(num) + except: + if num.endswith("%"): + num = num[:-1] + if num.endswith("\\"): + num = num[:-1] + try: + return float(num) / 100 + except: + pass + return None + + +def is_digit(num): + # paired with parse_digits + return parse_digits(num) is not None + + +def symbolic_equal(a, b): + def _parse(s): + for f in [parse_latex, parse_expr]: + try: + return f(s) + except: + pass + return s + + a = _parse(a) + b = _parse(b) + + try: + if simplify(a - b) == 0: + return True + except: + pass + + try: + if isclose(N(a), N(b), abs_tol=1e-3): + return True + except: + pass + return False + + +def math_equal(prediction, reference, include_percentage=True, is_close=True): + """ + Exact match of math if and only if: + 1. numerical equal: both can convert to float and are equal + 2. symbolic equal: both can convert to sympy expression and are equal + """ + if str(prediction) == str(reference): + return True + + try: # 1. numerical equal + if is_digit(prediction) and is_digit(reference): + prediction = parse_digits(prediction) + reference = parse_digits(reference) + # number questions + if include_percentage: + gt_result = [reference / 100, reference, reference * 100] + else: + gt_result = [reference] + for item in gt_result: + try: + if is_close: + if isclose(item, prediction, abs_tol=1e-3): + return True + else: + if item == prediction: + return True + except Exception: + continue + return False + except: + pass + + if not prediction and prediction not in [0, False]: + return False + + # 2. symbolic equal + reference = str(reference).strip() + prediction = str(prediction).strip() + + if ( + regex.match(r"(\(|\[).+(\)|\])", prediction) is not None + and regex.match(r"(\(|\[).+(\)|\])", reference) is not None + ): + pred_parts = prediction[1:-1].split(",") + ref_parts = reference[1:-1].split(",") + if len(pred_parts) == len(ref_parts): + if all( + [ + math_equal( + pred_parts[i], ref_parts[i], include_percentage, is_close + ) + for i in range(len(pred_parts)) + ] + ): + return True + + # Add back matrix comparison + if ( + ( + prediction.startswith("\\begin{pmatrix}") + or prediction.startswith("\\begin{bmatrix}") + ) + and ( + prediction.endswith("\\end{pmatrix}") + or prediction.endswith("\\end{bmatrix}") + ) + and ( + reference.startswith("\\begin{pmatrix}") + or reference.startswith("\\begin{bmatrix}") + ) + and ( + reference.endswith("\\end{pmatrix}") or reference.endswith("\\end{bmatrix}") + ) + ): + pred_lines = [ + line.strip() + for line in prediction[ + len("\\begin{pmatrix}") : -len("\\end{pmatrix}") + ].split("\\\\") + if line.strip() + ] + ref_lines = [ + line.strip() + for line in reference[ + len("\\begin{pmatrix}") : -len("\\end{pmatrix}") + ].split("\\\\") + if line.strip() + ] + matched = True + if len(pred_lines) == len(ref_lines): + for pred_line, ref_line in zip(pred_lines, ref_lines): + pred_parts = pred_line.split("&") + ref_parts = ref_line.split("&") + if len(pred_parts) == len(ref_parts): + if not all( + [ + math_equal( + pred_parts[i], + ref_parts[i], + include_percentage, + is_close, + ) + for i in range(len(pred_parts)) + ] + ): + matched = False + break + else: + matched = False + if not matched: + break + else: + matched = False + if matched: + return True + + # Add back equation comparison + if prediction.count("=") == 1 and reference.count("=") == 1: + pred = prediction.split("=") + pred = f"{pred[0].strip()} - ({pred[1].strip()})" + ref = reference.split("=") + ref = f"{ref[0].strip()} - ({ref[1].strip()})" + if symbolic_equal(pred, ref) or symbolic_equal(f"-({pred})", ref): + return True + elif ( + prediction.count("=") == 1 + and len(prediction.split("=")[0].strip()) <= 2 + and "=" not in reference + ): + if math_equal( + prediction.split("=")[1], reference, include_percentage, is_close + ): + return True + elif ( + reference.count("=") == 1 + and len(reference.split("=")[0].strip()) <= 2 + and "=" not in prediction + ): + if math_equal( + prediction, reference.split("=")[1], include_percentage, is_close + ): + return True + + # symbolic equal with sympy + if symbolic_equal(prediction, reference): + return True + + return False diff --git a/sglang/benchmark/tip_suggestion/.gitignore b/sglang/benchmark/tip_suggestion/.gitignore new file mode 100644 index 0000000..3322de9 --- /dev/null +++ b/sglang/benchmark/tip_suggestion/.gitignore @@ -0,0 +1 @@ +!topic.jsonl diff --git a/sglang/benchmark/tip_suggestion/README.md b/sglang/benchmark/tip_suggestion/README.md new file mode 100644 index 0000000..97ebe1d --- /dev/null +++ b/sglang/benchmark/tip_suggestion/README.md @@ -0,0 +1,33 @@ +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 64 +python3 bench_sglang.py --num-questions 32 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --backend vllm --num-questions 64 +``` + + +### Benchmark guidance +``` +python3 bench_other.py --backend guidance --num-questions 32 --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +### Benchmark lmql + +``` +python3 bench_other.py --backend lmql --num-questions 32 --parallel 1 +``` diff --git a/sglang/benchmark/tip_suggestion/bench_other.py b/sglang/benchmark/tip_suggestion/bench_other.py new file mode 100644 index 0000000..fcc4fd6 --- /dev/null +++ b/sglang/benchmark/tip_suggestion/bench_other.py @@ -0,0 +1,133 @@ +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor +from functools import partial + +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +number = 5 + + +def expand_tip(topic, tip, generate): + s = ( + """Please expand a tip for a topic into a detailed paragraph. + +Topic: staying healthy +Tip: Regular Exercise +Paragraph: Incorporate physical activity into your daily routine. This doesn't necessarily mean intense gym workouts; it can be as simple as walking, cycling, or yoga. Regular exercise helps in maintaining a healthy weight, improves cardiovascular health, boosts mental health, and can enhance cognitive function, which is crucial for fields that require intense intellectual engagement. + +Topic: building a campfire +Tip: Choose the Right Location +Paragraph: Always build your campfire in a safe spot. This means selecting a location that's away from trees, bushes, and other flammable materials. Ideally, use a fire ring if available. If you're building a fire pit, it should be on bare soil or on a bed of stones, not on grass or near roots which can catch fire underground. Make sure the area above is clear of low-hanging branches. + +Topic: writing a blog post +Tip: structure your content effectively +Paragraph: A well-structured post is easier to read and more enjoyable. Start with an engaging introduction that hooks the reader and clearly states the purpose of your post. Use headings and subheadings to break up the text and guide readers through your content. Bullet points and numbered lists can make information more digestible. Ensure each paragraph flows logically into the next, and conclude with a summary or call-to-action that encourages reader engagement. + +Topic: """ + + topic + + "\nTip: " + + tip + + "\nParagraph:" + ) + return generate(s, max_tokens=128, stop=["\n\n"]) + + +def suggest_tips(topic, generate): + s = "Please act as a helpful assistant. Your job is to provide users with useful tips on a specific topic.\n" + s += "USER: Give some tips for " + topic + ".\n" + s += ( + "ASSISTANT: Okay. Here are " + + str(number) + + " concise tips, each under 8 words:\n" + ) + + tips = [] + for i in range(1, 1 + number): + s += f"{i}." + tip = generate(s, max_tokens=24, stop=[".", "\n"]) + s += tip + ".\n" + tips.append(tip) + + paragraphs = [expand_tip(topic, tip, generate=generate) for tip in tips] + + for i in range(1, 1 + number): + s += f"Tip {i}:" + paragraphs[i - 1] + "\n" + return s + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_questions] + states = [None] * len(lines) + + # Select backend + call_generate = partial(get_call_generate(args), temperature=0) + + # Run requests + tic = time.time() + if args.backend != "lmql": + + def get_one_answer(i): + states[i] = suggest_tips(lines[i]["topic"], call_generate) + + if args.parallel == 1: + for i in tqdm(range(len(lines))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(lines)))), + total=len(lines), + ) + ) + + else: + import asyncio + + from lmql_funcs import suggest_tips_async + + async def get_one_answer_async(i): + states[i] = await suggest_tips_async(lines[i]["topic"], call_generate) + + batches = [] + for i in range(0, len(lines), args.parallel): + batches.append(list(range(i, min(i + args.parallel, len(lines))))) + loop = asyncio.get_event_loop() + for batch in tqdm(batches): + loop.run_until_complete( + asyncio.gather(*[get_one_answer_async(i) for i in batch]) + ) + latency = time.time() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "tip_suggestion", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="topic.jsonl") + parser.add_argument("--num-questions", type=int, default=100) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/tip_suggestion/bench_sglang.py b/sglang/benchmark/tip_suggestion/bench_sglang.py new file mode 100644 index 0000000..6d17821 --- /dev/null +++ b/sglang/benchmark/tip_suggestion/bench_sglang.py @@ -0,0 +1,100 @@ +import argparse +import json +import time + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + +number = 5 + + +@sgl.function +def expand_tip(s, topic, tip): + s += ( + """Please expand a tip for a topic into a detailed paragraph. + +Topic: staying healthy +Tip: Regular Exercise +Paragraph: Incorporate physical activity into your daily routine. This doesn't necessarily mean intense gym workouts; it can be as simple as walking, cycling, or yoga. Regular exercise helps in maintaining a healthy weight, improves cardiovascular health, boosts mental health, and can enhance cognitive function, which is crucial for fields that require intense intellectual engagement. + +Topic: building a campfire +Tip: Choose the Right Location +Paragraph: Always build your campfire in a safe spot. This means selecting a location that's away from trees, bushes, and other flammable materials. Ideally, use a fire ring if available. If you're building a fire pit, it should be on bare soil or on a bed of stones, not on grass or near roots which can catch fire underground. Make sure the area above is clear of low-hanging branches. + +Topic: writing a blog post +Tip: structure your content effectively +Paragraph: A well-structured post is easier to read and more enjoyable. Start with an engaging introduction that hooks the reader and clearly states the purpose of your post. Use headings and subheadings to break up the text and guide readers through your content. Bullet points and numbered lists can make information more digestible. Ensure each paragraph flows logically into the next, and conclude with a summary or call-to-action that encourages reader engagement. + +Topic: """ + + topic + + "\nTip: " + + tip + + "\nParagraph:" + ) + s += sgl.gen("paragraph", max_tokens=128, stop=["\n\n"], temperature=0) + + +@sgl.function +def suggest_tips(s, topic): + s += "Please act as a helpful assistant. Your job is to provide users with useful tips on a specific topic.\n" + s += "USER: Give some tips for " + topic + ".\n" + s += ( + "ASSISTANT: Okay. Here are " + + str(number) + + " concise tips, each under 8 words:\n" + ) + + paragraphs = [] + for i in range(1, 1 + number): + s += f"{i}." + sgl.gen(f"tip_{i}", max_tokens=24, stop=[".", "\n"]) + ".\n" + paragraphs.append(expand_tip(topic=topic, tip=s[f"tip_{i}"])) + + for i in range(1, 1 + number): + s += f"Tip {i}:" + paragraphs[i - 1]["paragraph"] + "\n" + + +def main(args): + lines = read_jsonl(args.data_path)[: args.num_questions] + arguments = [{"topic": l["topic"]} for l in lines] + + # Select backend + sgl.set_default_backend(select_sglang_backend(args)) + + # Run requests + tic = time.time() + states = suggest_tips.run_batch( + arguments, temperature=0, num_threads=args.parallel, progress_bar=True + ) + latency = time.time() - tic + + # Compute accuracy + print(f"Latency: {latency:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", states) + + with open(args.result_file, "a") as fout: + value = { + "task": "tip_suggestion", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="topic.jsonl") + parser.add_argument("--num-questions", type=int, default=100) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/tip_suggestion/lmql_funcs.py b/sglang/benchmark/tip_suggestion/lmql_funcs.py new file mode 100644 index 0000000..7790bbe --- /dev/null +++ b/sglang/benchmark/tip_suggestion/lmql_funcs.py @@ -0,0 +1,50 @@ +number = 5 + + +async def expand_tip_async(topic, tip, generate): + s = ( + """Please expand a tip for a topic into a detailed paragraph. + +Topic: staying healthy +Tip: Regular Exercise +Paragraph: Incorporate physical activity into your daily routine. This doesn't necessarily mean intense gym workouts; it can be as simple as walking, cycling, or yoga. Regular exercise helps in maintaining a healthy weight, improves cardiovascular health, boosts mental health, and can enhance cognitive function, which is crucial for fields that require intense intellectual engagement. + +Topic: building a campfire +Tip: Choose the Right Location +Paragraph: Always build your campfire in a safe spot. This means selecting a location that's away from trees, bushes, and other flammable materials. Ideally, use a fire ring if available. If you're building a fire pit, it should be on bare soil or on a bed of stones, not on grass or near roots which can catch fire underground. Make sure the area above is clear of low-hanging branches. + +Topic: writing a blog post +Tip: structure your content effectively +Paragraph: A well-structured post is easier to read and more enjoyable. Start with an engaging introduction that hooks the reader and clearly states the purpose of your post. Use headings and subheadings to break up the text and guide readers through your content. Bullet points and numbered lists can make information more digestible. Ensure each paragraph flows logically into the next, and conclude with a summary or call-to-action that encourages reader engagement. + +Topic: """ + + topic + + "\nTip: " + + tip + + "\nParagraph:" + ) + return await generate(s, max_tokens=128, stop="\n\n") + + +async def suggest_tips_async(topic, generate): + s = "Please act as a helpful assistant. Your job is to provide users with useful tips on a specific topic.\n" + s += "USER: Give some tips for " + topic + ".\n" + s += ( + "ASSISTANT: Okay. Here are " + + str(number) + + " concise tips, each under 8 words:\n" + ) + + tips = [] + for i in range(1, 1 + number): + s += f"{i}." + # NOTE: stop is different due to lmql does not support a list of stop tokens + tip = await generate(s, max_tokens=24, stop=".\n") + s += tip + ".\n" + tips.append(tip) + + paragraphs = [await expand_tip_async(topic, tip, generate=generate) for tip in tips] + + for i in range(1, 1 + number): + s += f"Tip {i}:" + paragraphs[i - 1] + "\n" + return s diff --git a/sglang/benchmark/tip_suggestion/topic.jsonl b/sglang/benchmark/tip_suggestion/topic.jsonl new file mode 100644 index 0000000..c8ac6d1 --- /dev/null +++ b/sglang/benchmark/tip_suggestion/topic.jsonl @@ -0,0 +1,50 @@ +{"topic": "organizing a successful charity event", "number": 6} +{"topic": "improving personal credit scores", "number": 7} +{"topic": "staying motivated during job searches", "number": 5} +{"topic": "maintaining a work-life balance", "number": 9} +{"topic": "reducing carbon footprint at home", "number": 8} +{"topic": "starting a book club", "number": 5} +{"topic": "learning to play a musical instrument", "number": 7} +{"topic": "getting into freelance writing", "number": 6} +{"topic": "beginner yoga poses", "number": 8} +{"topic": "preparing for graduate school exams", "number": 5} +{"topic": "exploring minimalist living", "number": 9} +{"topic": "effective grocery shopping", "number": 7} +{"topic": "winter camping", "number": 5} +{"topic": "starting a podcast on a budget", "number": 8} +{"topic": "creating a capsule wardrobe", "number": 6} +{"topic": "improving your writing skills", "number": 7} +{"topic": "learning a new software quickly", "number": 9} +{"topic": "reducing anxiety before public speaking", "number": 5} +{"topic": "planning a solo travel adventure", "number": 8} +{"topic": "beginner skateboarders", "number": 6} +{"topic": "studying abroad", "number": 7} +{"topic": "planting a vegetable garden", "number": 5} +{"topic": "adopting a shelter pet", "number": 9} +{"topic": "learning to cook ethnic cuisines", "number": 8} +{"topic": "effective conflict resolution", "number": 5} +{"topic": "starting a vlog", "number": 7} +{"topic": "keeping a daily journal", "number": 6} +{"topic": "improving sleep hygiene", "number": 8} +{"topic": "beginner mountain climbers", "number": 5} +{"topic": "creating a mobile app", "number": 9} +{"topic": "maintaining a saltwater aquarium", "number": 7} +{"topic": "preparing for a baby's arrival", "number": 6} +{"topic": "writing a fantasy novel", "number": 5} +{"topic": "effective team leadership", "number": 8} +{"topic": "making a documentary film", "number": 9} +{"topic": "learning about historical events", "number": 7} +{"topic": "baking gluten-free treats", "number": 6} +{"topic": "improving mental arithmetic skills", "number": 5} +{"topic": "building a treehouse", "number": 8} +{"topic": "getting started with watercolor painting", "number": 9} +{"topic": "creating a YouTube tutorial series", "number": 7} +{"topic": "landscape photography", "number": 5} +{"topic": "navigating cultural differences", "number": 6} +{"topic": "preparing for a marathon", "number": 8} +{"topic": "building an online business", "number": 9} +{"topic": "learning to dance at home", "number": 5} +{"topic": "self-publishing a book", "number": 7} +{"topic": "starting an urban farm", "number": 6} +{"topic": "improving your memory", "number": 8} +{"topic": "creating a personal brand online", "number": 9} diff --git a/sglang/benchmark/tree_of_thought_deep/README.md b/sglang/benchmark/tree_of_thought_deep/README.md new file mode 100644 index 0000000..bf5ab16 --- /dev/null +++ b/sglang/benchmark/tree_of_thought_deep/README.md @@ -0,0 +1,51 @@ +## Download data +``` +wget https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl +``` + +## Run benchmark + +NOTE: This is an implementation for throughput/latency benchmark purposes. The prompts are not tuned to achieve good accuracy on the GSM-8K tasks. + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 32 +python3 bench_sglang.py --num-questions 16 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-questions 32 --backend vllm +``` + + +### Benchmark lightllm +``` +# A10G +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 +``` + +``` +python3 bench_other.py --num-questions 32 --backend lightllm +``` + + +### Benchmark guidance +``` +python3 bench_other.py --num-questions 8 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` + +### Benchmark lmql + +``` +python3 bench_other.py --num-questions 8 --backend lmql --parallel 1 +``` diff --git a/sglang/benchmark/tree_of_thought_deep/bench_other.py b/sglang/benchmark/tree_of_thought_deep/bench_other.py new file mode 100644 index 0000000..21c7df3 --- /dev/null +++ b/sglang/benchmark/tree_of_thought_deep/bench_other.py @@ -0,0 +1,222 @@ +import argparse +import ast +import json +import re +import time +from collections import Counter +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def most_frequent_number(numbers): + if not numbers: + return None + + frequency = Counter(numbers) + most_frequent = max(frequency, key=frequency.get) + return most_frequent + + +USER_PREFIX = "[INST] " +USER_SUFFIX = " [/INST]" +ASSISTANT_PREFIX = "" +ASSISTANT_SUFFIX = " " + +# Use a low temp to make the results more deterministic and the comparison more fair. +temp = 0.001 + + +def propose_plan(s, question, num_branches, call_generate): + s += ( + USER_PREFIX + + """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + + question + + USER_SUFFIX + ) + + s += ASSISTANT_PREFIX + comps = call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +def execute_plan(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +def reflect_solution(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """Okay. Now, evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +def get_final_answer(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """Based on your reflection, do you change your mind? Now, give me the final answer after careful consideration.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +def tree_search(question, num_branches, call_generate): + plan_forks = propose_plan("", question, num_branches, call_generate) + + sol_states = [] + for plan in plan_forks: + forks = execute_plan(plan, num_branches, call_generate) + sol_states.extend(forks) + + ref_states = [] + for sol in sol_states: + forks = reflect_solution(sol, num_branches, call_generate) + ref_states.extend(forks) + + solutions = [] + for sol in ref_states: + ans = get_final_answer(sol, num_branches, call_generate) + solutions.append(ans) + + return solutions + + +def main(args): + lines = read_jsonl(args.data_path) + + # Construct prompts + num_branches = 2 + questions = [] + labels = [] + for i in range(len(lines[: args.num_questions])): + questions.append(lines[i]["question"]) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q, "num_branches": num_branches} for q in questions] + + # Select backend + call_generate = get_call_generate(args) + + # Run requests + states = [None] * len(questions) + + tic = time.time() + if args.backend != "lmql": + + def get_one_answer(i): + states[i] = tree_search(**arguments[i], call_generate=call_generate) + + if args.parallel == 1: + for i in tqdm(range(len(questions))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(questions)))), + total=len(questions), + ) + ) + + else: + import asyncio + + from lmql_funcs import tree_search_async + + async def get_one_answer_async(i): + states[i] = await tree_search_async( + **arguments[i], call_generate=call_generate + ) + + batches = [ + [] for _ in range((len(questions) + args.parallel - 1) // args.parallel) + ] + for i in range(len(questions)): + batches[i // args.parallel].append(i) + + loop = asyncio.get_event_loop() + for bt in tqdm(batches): + tasks = [get_one_answer_async(k) for k in bt] + loop.run_until_complete(asyncio.gather(*tasks)) + + latency = time.time() - tic + + answers_text = [] + for s in states: + answers_text.append([x for xs in s for x in xs]) + + preds = [] + for i in range(len(states)): + answers = [get_answer_value(v) for v in answers_text[i]] + preds.append(most_frequent_number(answers)) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + print(f"Latency: {latency:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", answers_text) + + with open(args.result_file, "a") as fout: + value = { + "task": "tree_of_thought_gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/tree_of_thought_deep/bench_sglang.py b/sglang/benchmark/tree_of_thought_deep/bench_sglang.py new file mode 100644 index 0000000..bfb2a41 --- /dev/null +++ b/sglang/benchmark/tree_of_thought_deep/bench_sglang.py @@ -0,0 +1,171 @@ +import argparse +import ast +import json +import re +import time +from collections import Counter + +import numpy as np + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def most_frequent_number(numbers): + if not numbers: + return None + + frequency = Counter(numbers) + most_frequent = max(frequency, key=frequency.get) + return most_frequent + + +# Use a low temp to make the results more deterministic and the comparison more fair. +temp = 0.001 + + +def propose_plan(s, question, num_branches): + s += sgl.user( + """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + + question + ) + forks = s.fork(num_branches) + forks += sgl.assistant(sgl.gen("plan", max_tokens=256, temperature=temp)) + return forks + + +def execute_plan(s, num_branches): + s += sgl.user( + """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + ) + forks = s.fork(num_branches) + forks += sgl.assistant(sgl.gen("answer", max_tokens=256, temperature=temp)) + return forks + + +def reflect_solution(s, num_branches): + s += sgl.user( + """Okay. Now, evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + ) + forks = s.fork(num_branches) + forks += sgl.assistant(sgl.gen("score", max_tokens=256, temperature=temp)) + return forks + + +def get_final_answer(s, num_branches): + s += sgl.user( + """Based on your reflection, do you change your mind? Now, give me the final answer after careful consideration.""" + ) + forks = s.fork(num_branches) + forks += sgl.assistant(sgl.gen("final_answer", max_tokens=256, temperature=temp)) + return forks + + +@sgl.function +def tree_search(s, question, num_branches): + plan_forks = propose_plan(s, question, num_branches) + + sol_states = [] + for plan in plan_forks: + forks = execute_plan(plan, num_branches) + sol_states.extend(forks) + + ref_states = [] + for sol in sol_states: + forks = reflect_solution(sol, num_branches) + ref_states.extend(forks) + + solutions = [] + for sol in ref_states: + forks = get_final_answer(sol, num_branches) + solutions.append(forks) + solutions = [[s.text() for s in forks] for forks in solutions] + + return solutions + + +def main(args): + lines = read_jsonl(args.data_path) + lines = list(lines) + + # Construct prompts + num_branches = 2 + questions = [] + labels = [] + for i in range(len(lines[: args.num_questions])): + questions.append(lines[i]["question"]) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q, "num_branches": num_branches} for q in questions] + + # Select backend + backend = select_sglang_backend(args) + + # Run requests + tic = time.time() + states = tree_search.run_batch( + arguments, + temperature=0, + backend=backend, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.time() - tic + answers_text = [] + for s in states: + answers_text.append([x for xs in s.ret_value for x in xs]) + + preds = [] + for i in range(len(states)): + answers = [get_answer_value(v) for v in answers_text[i]] + preds.append(most_frequent_number(answers)) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + print(f"Latency: {latency:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", answers_text) + + with open(args.result_file, "a") as fout: + value = { + "task": "tree_of_thought_gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/tree_of_thought_deep/lmql_funcs.py b/sglang/benchmark/tree_of_thought_deep/lmql_funcs.py new file mode 100644 index 0000000..c783cdb --- /dev/null +++ b/sglang/benchmark/tree_of_thought_deep/lmql_funcs.py @@ -0,0 +1,82 @@ +from bench_other import ( + ASSISTANT_PREFIX, + ASSISTANT_SUFFIX, + USER_PREFIX, + USER_SUFFIX, + temp, +) + + +async def propose_plan_async(s, question, num_branches, call_generate): + s += ( + USER_PREFIX + + """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + + question + + USER_SUFFIX + ) + + s += ASSISTANT_PREFIX + comps = await call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +async def execute_plan_async(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = await call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +async def reflect_solution_async(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """Okay. Now, evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = await call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +async def get_final_answer_async(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """Based on your reflection, do you change your mind? Now, give me the final answer after careful consideration.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = await call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +async def tree_search_async(question, num_branches, call_generate): + plan_forks = await propose_plan_async("", question, num_branches, call_generate) + + sol_states = [] + for plan in plan_forks: + forks = await execute_plan_async(plan, num_branches, call_generate) + sol_states.extend(forks) + + ref_states = [] + for sol in sol_states: + forks = await reflect_solution_async(sol, num_branches, call_generate) + ref_states.extend(forks) + + solutions = [] + for sol in ref_states: + ans = await get_final_answer_async(sol, num_branches, call_generate) + solutions.append(ans) + + return solutions diff --git a/sglang/benchmark/tree_of_thought_v0/README.md b/sglang/benchmark/tree_of_thought_v0/README.md new file mode 100644 index 0000000..821bb20 --- /dev/null +++ b/sglang/benchmark/tree_of_thought_v0/README.md @@ -0,0 +1,43 @@ +## Download data +``` +wget https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl +``` + +## Run benchmark + +### Benchmark sglang +``` +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +``` + +``` +python3 bench_sglang.py --num-questions 32 --parallel 16 +python3 bench_sglang.py --num-questions 10 --parallel 1 +``` + + +### Benchmark vllm +``` +python3 -m vllm.entrypoints.api_server --tokenizer-mode auto --model meta-llama/Llama-2-7b-chat-hf --disable-log-requests --port 21000 +``` + +``` +python3 bench_other.py --num-questions 32 --backend vllm +``` + + +### Benchmark lightllm +``` +# A10G +python -m lightllm.server.api_server --tokenizer_mode auto --model_dir ~/model_weights/llama-2-7b-chat-hf --max_total_token_num 16000 --port 22000 +``` + +``` +python3 bench_other.py --num-questions 32 --backend lightllm +``` + + +### Benchmark guidance +``` +python3 bench_other.py --num-questions 32 --backend guidance --parallel 1 --n-ctx 4096 --model-path path/to/gguf +``` diff --git a/sglang/benchmark/tree_of_thought_v0/bench_other.py b/sglang/benchmark/tree_of_thought_v0/bench_other.py new file mode 100644 index 0000000..86e1335 --- /dev/null +++ b/sglang/benchmark/tree_of_thought_v0/bench_other.py @@ -0,0 +1,179 @@ +import argparse +import ast +import json +import re +import time +from collections import Counter +from concurrent.futures import ThreadPoolExecutor + +import numpy as np +from tqdm import tqdm + +from sglang.test.test_utils import add_common_other_args_and_parse, get_call_generate +from sglang.utils import dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def most_frequent_number(numbers): + if not numbers: + return None + + frequency = Counter(numbers) + most_frequent = max(frequency, key=frequency.get) + return most_frequent + + +USER_PREFIX = "[INST] " +USER_SUFFIX = " [/INST]" +ASSISTANT_PREFIX = "" +ASSISTANT_SUFFIX = " " + +# Use a low temp to make the results more deterministic and the comparison more fair. +temp = 0.3 + + +def propose_plan(s, question, num_branches, call_generate): + s += ( + USER_PREFIX + + """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + + question + + USER_SUFFIX + ) + + s += ASSISTANT_PREFIX + comps = call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +def execute_plan(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +def reflect_solution(s, num_branches, call_generate): + s += ( + USER_PREFIX + + """Okay. Now you evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + + USER_SUFFIX + ) + s += ASSISTANT_PREFIX + comps = call_generate( + s, max_tokens=256, temperature=temp, stop=None, n=num_branches + ) + return [s + comp + ASSISTANT_SUFFIX for comp in comps] + + +def tree_search(question, num_branches, call_generate): + s = "" + solutions = [] + + plan_forks = propose_plan(s, question, num_branches, call_generate) + for plan in plan_forks: + sol_forks = execute_plan(plan, num_branches, call_generate) + for sol in sol_forks: + score_forks = reflect_solution(sol, num_branches, call_generate) + solutions.append(sol_forks) + + return solutions + + +def main(args): + lines = read_jsonl(args.data_path) + + # Construct prompts + num_branches = 3 + questions = [] + labels = [] + for i in range(len(lines[: args.num_questions])): + questions.append(lines[i]["question"]) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q, "num_branches": num_branches} for q in questions] + + # Select backend + call_generate = get_call_generate(args) + + # Run requests + states = [None] * len(questions) + + def get_one_answer(i): + states[i] = tree_search(**arguments[i], call_generate=call_generate) + + tic = time.time() + if args.parallel == 1: + for i in tqdm(range(len(questions))): + get_one_answer(i) + else: + with ThreadPoolExecutor(args.parallel) as executor: + list( + tqdm( + executor.map(get_one_answer, list(range(len(questions)))), + total=len(questions), + ) + ) + + latency = time.time() - tic + + answers_text = [] + for s in states: + answers_text.append([x for xs in s for x in xs]) + + preds = [] + for i in range(len(states)): + answers = [get_answer_value(v) for v in answers_text[i]] + preds.append(most_frequent_number(answers)) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + print(f"Latency: {latency:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", answers_text) + + with open(args.result_file, "a") as fout: + value = { + "task": "tree_of_thought_gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_other_args_and_parse(parser) + main(args) diff --git a/sglang/benchmark/tree_of_thought_v0/bench_sglang.py b/sglang/benchmark/tree_of_thought_v0/bench_sglang.py new file mode 100644 index 0000000..f0d1307 --- /dev/null +++ b/sglang/benchmark/tree_of_thought_v0/bench_sglang.py @@ -0,0 +1,159 @@ +import argparse +import ast +import json +import re +import time +from collections import Counter + +import numpy as np + +import sglang as sgl +from sglang.test.test_utils import ( + add_common_sglang_args_and_parse, + select_sglang_backend, +) +from sglang.utils import dump_state_text, read_jsonl + +INVALID = -9999999 + + +def get_answer_value(answer_str): + answer_str = answer_str.replace(",", "") + numbers = re.findall(r"\d+", answer_str) + if len(numbers) < 1: + return INVALID + try: + return ast.literal_eval(numbers[-1]) + except SyntaxError: + return INVALID + + +def most_frequent_number(numbers): + if not numbers: + return None + + frequency = Counter(numbers) + most_frequent = max(frequency, key=frequency.get) + return most_frequent + + +# Use a low temp to make the results more deterministic and the comparison more fair. +temp = 0.3 + + +def propose_plan(s, question, num_branches): + s += sgl.user( + """Please generate a high-level plan for solving the following question. As the first step, just say what method and idea you will use to solve the question. You can reorganize the information in the question. Do not do the actual calculation. Keep your response concise and within 80 words. Question: """ + + question + ) + forks = s.fork(num_branches) + forks += sgl.assistant(sgl.gen("plan", max_tokens=256, temperature=temp)) + return forks + + +def execute_plan(s, num_branches): + s += sgl.user( + """The plan looks good! Now, use real numbers and do the calculation. Please solve the question step-by-step according to the high-level plan. Give me the final answer. Make your response short.""" + ) + forks = s.fork(num_branches) + forks += sgl.assistant(sgl.gen("answer", max_tokens=256, temperature=temp)) + return forks + + +def reflect_solution(s, num_branches): + s += sgl.user( + """Okay. Now you evaluate your own solution and give it a score on a scale of 1 to 5. Please do rigorous check of the correctness.""" + ) + forks = s.fork(num_branches) + forks += sgl.assistant(sgl.gen("score", max_tokens=256, temperature=temp)) + return forks + + +@sgl.function +def tree_search(s, question, num_branches): + forks_to_join = [] + + plan_forks = propose_plan(s, question, num_branches) + forks_to_join.append(plan_forks) + + sol_states = [] + for plan in plan_forks: + forks = execute_plan(plan, num_branches) + forks_to_join.append(forks) + sol_states.extend(forks) + + for sol in sol_states: + forks = reflect_solution(sol, num_branches) + forks_to_join.append(forks) + + for f in reversed(forks_to_join): + f.join() + + +def main(args): + lines = read_jsonl(args.data_path) + + # Construct prompts + num_branches = 3 + questions = [] + labels = [] + for i in range(len(lines[: args.num_questions])): + questions.append(lines[i]["question"]) + labels.append(get_answer_value(lines[i]["answer"])) + assert all(l != INVALID for l in labels) + arguments = [{"question": q, "num_branches": num_branches} for q in questions] + + # Select backend + backend = select_sglang_backend(args) + + # Run requests + tic = time.time() + states = tree_search.run_batch( + arguments, + temperature=0, + backend=backend, + num_threads=args.parallel, + progress_bar=True, + ) + latency = time.time() - tic + answers_text = [] + for s in states: + answers_text.append([x for xs in s["answer"] for x in xs]) + + preds = [] + for i in range(len(states)): + answers = [get_answer_value(v) for v in answers_text[i]] + preds.append(most_frequent_number(answers)) + + # Compute accuracy + acc = np.mean(np.array(preds) == np.array(labels)) + invalid = np.mean(np.array(preds) == INVALID) + print(f"Latency: {latency:.3f}") + print(f"Invalid: {invalid:.3f}") + print(f"Accuracy: {acc:.3f}") + + # Write results + dump_state_text(f"tmp_output_{args.backend}.txt", answers_text) + + with open(args.result_file, "a") as fout: + value = { + "task": "tree_of_thought_gsm8k", + "backend": args.backend, + "num_gpus": 1, + "latency": round(latency, 3), + "accuracy": round(acc, 3), + "num_requests": args.num_questions, + "other": { + "num_questions": args.num_questions, + "parallel": args.parallel, + }, + } + fout.write(json.dumps(value) + "\n") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--data-path", type=str, default="test.jsonl") + parser.add_argument("--num-questions", type=int, default=200) + args = add_common_sglang_args_and_parse(parser) + main(args) diff --git a/sglang/docker/Dockerfile b/sglang/docker/Dockerfile new file mode 100644 index 0000000..e70b580 --- /dev/null +++ b/sglang/docker/Dockerfile @@ -0,0 +1,78 @@ +ARG CUDA_VERSION=12.5.1 + +FROM nvcr.io/nvidia/tritonserver:24.04-py3-min + +ARG BUILD_TYPE=all +ENV DEBIAN_FRONTEND=noninteractive + +# 安装依赖（强制 IPv4） +RUN echo 'tzdata tzdata/Areas select Asia' | debconf-set-selections \ + && echo 'tzdata tzdata/Zones/Asia select Shanghai' | debconf-set-selections \ + && apt -o Acquire::ForceIPv4=true update -y \ + && apt -o Acquire::ForceIPv4=true install software-properties-common -y \ + && add-apt-repository ppa:deadsnakes/ppa -y \ + && apt -o Acquire::ForceIPv4=true update \ + && apt -o Acquire::ForceIPv4=true install python3.10 python3.10-dev -y \ + && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 \ + && update-alternatives --set python3 /usr/bin/python3.10 \ + && apt -o Acquire::ForceIPv4=true install python3.10-distutils -y \ + && apt -o Acquire::ForceIPv4=true install curl git sudo libibverbs-dev -y \ + && apt -o Acquire::ForceIPv4=true install -y rdma-core infiniband-diags openssh-server perftest ibverbs-providers libibumad3 libibverbs1 libnl-3-200 libnl-route-3-200 librdmacm1 \ + && curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py \ + && python3 get-pip.py \ + && python3 --version \ + && python3 -m pip --version \ + && rm -rf /var/lib/apt/lists/* \ + && apt clean + +# 安装 datamodel_code_generator（用于 MiniCPM 模型） +RUN pip3 install datamodel_code_generator + +WORKDIR /sgl-workspace + +# 拷贝 sglang 源代码并构建包 +COPY ./sglang /sgl-workspace/sglang + +# 拷贝模型文件（修正方式） +#COPY ./Alibaba/QwQ-32B /root/.cradle/Alibaba/QwQ-32B +COPY ./DeepSeek/DeepSeek-R1-Distill-Llama-70B /root/.cradle/DeepSeek/DeepSeek-R1-Distill-Llama-70B + +ARG CUDA_VERSION + +# 安装依赖、安装 sglang、安装 transformers，并清理源码 +RUN python3 -m pip install --upgrade pip setuptools wheel html5lib six \ + && if [ "$CUDA_VERSION" = "12.1.1" ]; then \ + CUINDEX=121; \ + elif [ "$CUDA_VERSION" = "12.4.1" ]; then \ + CUINDEX=124; \ + elif [ "$CUDA_VERSION" = "12.5.1" ]; then \ + CUINDEX=124; \ + elif [ "$CUDA_VERSION" = "11.8.0" ]; then \ + CUINDEX=118; \ + python3 -m pip install --no-cache-dir sgl-kernel -i https://docs.sglang.ai/whl/cu118; \ + else \ + echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1; \ + fi \ + && python3 -m pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu${CUINDEX} \ + && python3 -m pip install --no-cache-dir psutil pyzmq pynvml \ + && cd /sgl-workspace/sglang/python \ + && python3 -m pip install --no-cache-dir '.[srt,openai]' --find-links https://flashinfer.ai/whl/cu${CUINDEX}/torch2.5/flashinfer-python \ + && cd / && rm -rf /sgl-workspace/sglang \ + && python3 -m pip install --no-cache-dir transformers==4.48.3 \ + && python3 -c "import sglang; print('✅ sglang module installed successfully')" + +# 暴露端口 +EXPOSE 30000 + + +# 启动命令（QwQ 模型为例） +CMD ["python3", "-m", "sglang.launch_server", \ + "--host","0.0.0.0", \ + "--port","30000", \ + "--model-path", "/root/.cradle/DeepSeek/DeepSeek-R1-Distill-Llama-70B/", \ + "--cpu-offload-gb","60", \ + "--tp", "4", \ + "--api-key", "token-abc123"] + +# 设置回默认交互模式 +ENV DEBIAN_FRONTEND=interactive diff --git a/sglang/docker/Dockerfile.bak b/sglang/docker/Dockerfile.bak new file mode 100644 index 0000000..075b1e8 --- /dev/null +++ b/sglang/docker/Dockerfile.bak @@ -0,0 +1,48 @@ +ARG CUDA_VERSION=12.5.1 + +FROM nvcr.io/nvidia/tritonserver:24.04-py3-min + +ARG BUILD_TYPE=all +ENV DEBIAN_FRONTEND=noninteractive + +RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \ + && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \ + && apt update -y \ + && apt install software-properties-common -y \ + && add-apt-repository ppa:deadsnakes/ppa -y && apt update \ + && apt install python3.10 python3.10-dev -y \ + && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 \ + && update-alternatives --set python3 /usr/bin/python3.10 && apt install python3.10-distutils -y \ + && apt install curl git sudo libibverbs-dev -y \ + && apt install -y rdma-core infiniband-diags openssh-server perftest ibverbs-providers libibumad3 libibverbs1 libnl-3-200 libnl-route-3-200 librdmacm1 \ + && curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3 get-pip.py \ + && python3 --version \ + && python3 -m pip --version \ + && rm -rf /var/lib/apt/lists/* \ + && apt clean + +# For openbmb/MiniCPM models +RUN pip3 install datamodel_code_generator + +WORKDIR /sgl-workspace + +ARG CUDA_VERSION +RUN python3 -m pip install --upgrade pip setuptools wheel html5lib six \ + && git clone --depth=1 https://github.com/sgl-project/sglang.git \ + && if [ "$CUDA_VERSION" = "12.1.1" ]; then \ + export CUINDEX=121; \ + elif [ "$CUDA_VERSION" = "12.4.1" ]; then \ + export CUINDEX=124; \ + elif [ "$CUDA_VERSION" = "12.5.1" ]; then \ + export CUINDEX=124; \ + elif [ "$CUDA_VERSION" = "11.8.0" ]; then \ + export CUINDEX=118; \ + python3 -m pip install --no-cache-dir sgl-kernel -i https://docs.sglang.ai/whl/cu118; \ + else \ + echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1; \ + fi \ + && python3 -m pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu${CUINDEX} \ + && cd sglang \ + && python3 -m pip --no-cache-dir install -e "python[${BUILD_TYPE}]" --find-links https://flashinfer.ai/whl/cu${CUINDEX}/torch2.5/flashinfer-python + +ENV DEBIAN_FRONTEND=interactive diff --git a/sglang/docker/Dockerfile.deepep b/sglang/docker/Dockerfile.deepep new file mode 100644 index 0000000..bd32e37 --- /dev/null +++ b/sglang/docker/Dockerfile.deepep @@ -0,0 +1,77 @@ +FROM lmsysorg/sglang:latest + +# CMake +RUN apt-get update \ +&& apt-get install -y --no-install-recommends \ +build-essential \ +wget \ +libssl-dev \ +&& wget https://github.com/Kitware/CMake/releases/download/v3.27.4/cmake-3.27.4-linux-x86_64.sh \ +&& chmod +x cmake-3.27.4-linux-x86_64.sh \ +&& ./cmake-3.27.4-linux-x86_64.sh --skip-license --prefix=/usr/local \ +&& rm cmake-3.27.4-linux-x86_64.sh + +# Python +RUN apt-get update \ + && apt-get install -y --no-install-recommends \ + python3 \ + python3-pip \ + && ln -s /usr/bin/python3 /usr/bin/python + +# GDRCopy +WORKDIR /tmp +RUN git clone https://github.com/NVIDIA/gdrcopy.git +WORKDIR /tmp/gdrcopy +RUN git checkout v2.4.4 + +RUN apt update +RUN apt install -y nvidia-dkms-535 +RUN apt install -y build-essential devscripts debhelper fakeroot pkg-config dkms +RUN apt install -y check libsubunit0 libsubunit-dev + +WORKDIR /tmp/gdrcopy/packages +RUN CUDA=/usr/local/cuda ./build-deb-packages.sh +RUN dpkg -i gdrdrv-dkms_*.deb +RUN dpkg -i libgdrapi_*.deb +RUN dpkg -i gdrcopy-tests_*.deb +RUN dpkg -i gdrcopy_*.deb + +ENV GDRCOPY_HOME=/usr/src/gdrdrv-2.4.4/ + +# IBGDA dependency +RUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so +RUN apt-get install -y libfabric-dev + +# DeepEP +WORKDIR /sgl-workspace +RUN git clone https://github.com/deepseek-ai/DeepEP.git + +# NVSHMEM +WORKDIR /sgl-workspace +RUN wget https://developer.download.nvidia.com/compute/redist/nvshmem/3.2.5/source/nvshmem_src_3.2.5-1.txz +RUN tar -xf nvshmem_src_3.2.5-1.txz \ + && mv nvshmem_src nvshmem + +WORKDIR /sgl-workspace/nvshmem +RUN git apply /sgl-workspace/DeepEP/third-party/nvshmem.patch + +WORKDIR /sgl-workspace/nvshmem +ENV CUDA_HOME=/usr/local/cuda +RUN NVSHMEM_SHMEM_SUPPORT=0 \ + NVSHMEM_UCX_SUPPORT=0 \ + NVSHMEM_USE_NCCL=0 \ + NVSHMEM_MPI_SUPPORT=0 \ + NVSHMEM_IBGDA_SUPPORT=1 \ + NVSHMEM_PMIX_SUPPORT=0 \ + NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \ + NVSHMEM_USE_GDRCOPY=1 \ + cmake -S . -B build/ -DCMAKE_INSTALL_PREFIX=/sgl-workspace/nvshmem/install -DCMAKE_CUDA_ARCHITECTURES=90 \ + && cd build \ + && make install -j + +WORKDIR /sgl-workspace/DeepEP +ENV NVSHMEM_DIR=/sgl-workspace/nvshmem/install +RUN NVSHMEM_DIR=/sgl-workspace/nvshmem/install python setup.py install + +# Set workspace +WORKDIR /sgl-workspace diff --git a/sglang/docker/Dockerfile.dev b/sglang/docker/Dockerfile.dev new file mode 100644 index 0000000..6faa6ff --- /dev/null +++ b/sglang/docker/Dockerfile.dev @@ -0,0 +1,224 @@ +# Extend from the base sglang image +FROM lmsysorg/sglang:latest + +# Install development tools and utilities +RUN apt-get update && apt-get install -y \ + gdb \ + ninja-build \ + vim \ + tmux \ + htop \ + wget \ + curl \ + locales \ + git \ + git-lfs \ + zsh \ + tree \ + silversearcher-ag \ + cloc \ + unzip \ + pkg-config \ + libssl-dev \ + bear \ + && apt install -y rdma-core infiniband-diags openssh-server perftest ibverbs-providers libibumad3 libibverbs1 libnl-3-200 libnl-route-3-200 librdmacm1 \ + && rm -rf /var/lib/apt/lists/* \ + && apt-get clean + +RUN apt update -y \ + && apt install -y --no-install-recommends gnupg \ + && echo "deb http://developer.download.nvidia.com/devtools/repos/ubuntu2004/amd64 /" | tee /etc/apt/sources.list.d/nvidia-devtools.list \ + && apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub \ + && apt update -y \ + && apt install nsight-systems-cli -y + +# Set up locale +RUN locale-gen en_US.UTF-8 +ENV LANG en_US.UTF-8 +ENV LANGUAGE en_US:en +ENV LC_ALL en_US.UTF-8 + +# Install minimal Python packages +RUN python3 -m pip install --no-cache-dir \ + pytest \ + black \ + isort \ + icdiff \ + pre-commit + +# Install diff-so-fancy +RUN curl -LSso /usr/local/bin/diff-so-fancy https://github.com/so-fancy/diff-so-fancy/releases/download/v1.4.4/diff-so-fancy \ + && chmod +x /usr/local/bin/diff-so-fancy + +# Install clang-format +RUN curl -LSso /usr/local/bin/clang-format https://github.com/muttleyxd/clang-tools-static-binaries/releases/download/master-32d3ac78/clang-format-16_linux-amd64 \ + && chmod +x /usr/local/bin/clang-format + +# Install clangd +RUN curl -L https://github.com/clangd/clangd/releases/download/18.1.3/clangd-linux-18.1.3.zip -o clangd.zip \ + && unzip clangd.zip \ + && cp -r clangd_18.1.3/bin/* /usr/local/bin/ \ + && cp -r clangd_18.1.3/lib/* /usr/local/lib/ \ + && rm -rf clangd_18.1.3 clangd.zip + +# Install CMake +RUN wget https://github.com/Kitware/CMake/releases/download/v3.31.1/cmake-3.31.1-linux-x86_64.tar.gz \ + && tar -xzf cmake-3.31.1-linux-x86_64.tar.gz \ + && cp -r cmake-3.31.1-linux-x86_64/bin/* /usr/local/bin/ \ + && cp -r cmake-3.31.1-linux-x86_64/share/* /usr/local/share/ \ + && rm -rf cmake-3.31.1-linux-x86_64 cmake-3.31.1-linux-x86_64.tar.gz + +# Add yank script +COPY --chown=root:root <<-"EOF" /usr/local/bin/yank +#!/bin/bash +put() { + esc=$1 + test -n "$TMUX" -o -z "${TERM##screen*}" && esc="\033Ptmux;\033$esc\033\\" + printf "$esc" +} +put "\033]52;c;!\a" +buf=$( cat "$@" ) +len=$( printf %s "$buf" | wc -c ) max=74994 +test $len -gt $max && echo "$0: input is $(( len - max )) bytes too long" >&2 +put "\033]52;c;$( printf %s "$buf" | head -c $max | base64 | tr -d '\r\n' )\a" +test -n "$TMUX" && tmux set-buffer "$buf" ||: +EOF + +RUN chmod +x /usr/local/bin/yank + +# Install oh-my-zsh and plugins +RUN sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" "" --unattended \ + && git clone https://github.com/zsh-users/zsh-autosuggestions ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-autosuggestions \ + && git clone https://github.com/zsh-users/zsh-syntax-highlighting.git ${ZSH_CUSTOM:-~/.oh-my-zsh/custom}/plugins/zsh-syntax-highlighting + +# Configure Vim +COPY --chown=root:root <<-"EOF" /root/.vimrc +function! Yank(text) abort + let escape = system('yank', a:text) + if v:shell_error + echoerr escape + else + call writefile([escape], '/dev/tty', 'b') + endif +endfunction + +noremap y y:call Yank(@0) + +" automatically run yank(1) whenever yanking in Vim +function! CopyYank() abort + call Yank(join(v:event.regcontents, "\n")) +endfunction + +autocmd TextYankPost * call CopyYank() + +" Basic settings +set number +syntax on +set mouse=a +filetype indent on + +" Indentation +set autoindent nosmartindent +set smarttab +set expandtab +set shiftwidth=4 +set softtabstop=4 + +" Visual guides +set colorcolumn=120 +highlight ColorColumn ctermbg=5 + +" Status line +set laststatus=2 +set statusline=%<%f\ %h%m%r%=%{\"[\".(&fenc==\"\"?&enc:&fenc).((exists(\"+bomb\")\ &&\ &bomb)?\",B\":\"\").\"]\ \"}%k\ %-14.(%l,%c%V%)\ %P + +" Backspace behavior +set backspace=2 + +" Encoding +set encoding=utf-8 +set fileencoding=utf-8 +EOF + +# Configure tmux +COPY --chown=root:root <<-"EOF" /root/.tmux.conf +# Pane border styling +set -g pane-border-style fg='#742727',bg=black +set -g pane-active-border-style fg=red,bg=black + +# Status bar styling +set -g status-style bg='#0C8A92',fg=black + +# Change prefix key to backtick +set-option -g prefix ` +unbind C-b +bind-key ` send-prefix + +# Split panes using - and = with current path +unbind '"' +bind - splitw -v -c '#{pane_current_path}' +unbind '%' +bind = splitw -h -c '#{pane_current_path}' + +# Vi mode settings +bind-key -T copy-mode-vi Y send-keys -X copy-pipe 'yank > #{pane_tty}' +set-window-option -g mode-keys vi + +# Other settings +set-option -g escape-time 0 +set-option -g base-index 1 +set-window-option -g mouse on +EOF + +# Configure Git +RUN git config --global core.editor "vim" \ + && git config --global core.whitespace "fix,-indent-with-non-tab,trailing-space,cr-at-eol" \ + && git config --global core.pager "diff-so-fancy | less --tabs=4 -RFX" \ + && git config --global color.ui true \ + && git config --global color."diff-highlight".oldNormal "red bold" \ + && git config --global color."diff-highlight".oldHighlight "red bold 52" \ + && git config --global color."diff-highlight".newNormal "green bold" \ + && git config --global color."diff-highlight".newHighlight "green bold 22" \ + && git config --global color.diff.meta "11" \ + && git config --global color.diff.frag "magenta bold" \ + && git config --global color.diff.commit "yellow bold" \ + && git config --global color.diff.old "red bold" \ + && git config --global color.diff.new "green bold" \ + && git config --global color.diff.whitespace "red reverse" \ + && git config --global alias.lg "log --color --graph --pretty=format:'%Cred%h%Creset - %s %Cgreen(%cr) %C(bold blue)<%an>%Creset%C(auto)%d%Creset' --abbrev-commit --" \ + && git config --global http.sslVerify false \ + && git config --global pull.rebase true + +# Configure zsh +COPY --chown=root:root <<-"EOF" /root/.zshrc +export ZSH="/root/.oh-my-zsh" + +# Theme +ZSH_THEME="robbyrussell" + +# Plugins +plugins=( + git + z + zsh-autosuggestions + zsh-syntax-highlighting +) + +source $ZSH/oh-my-zsh.sh + +# Aliases +alias ll='ls -alF' +alias la='ls -A' +alias l='ls -CF' +alias vi='vim' + +# Enhanced history +HISTSIZE=10000 +SAVEHIST=10000 +setopt HIST_IGNORE_ALL_DUPS +setopt HIST_FIND_NO_DUPS +setopt INC_APPEND_HISTORY +EOF + +# Set workspace directory +WORKDIR /sgl-workspace/sglang diff --git a/sglang/docker/Dockerfile.ds-14b b/sglang/docker/Dockerfile.ds-14b new file mode 100644 index 0000000..6674347 --- /dev/null +++ b/sglang/docker/Dockerfile.ds-14b @@ -0,0 +1,77 @@ +ARG CUDA_VERSION=12.5.1 + +FROM nvcr.io/nvidia/tritonserver:24.04-py3-min + +ARG BUILD_TYPE=all +ENV DEBIAN_FRONTEND=noninteractive + +# 安装依赖（强制 IPv4） +RUN echo 'tzdata tzdata/Areas select Asia' | debconf-set-selections \ + && echo 'tzdata tzdata/Zones/Asia select Shanghai' | debconf-set-selections \ + && apt -o Acquire::ForceIPv4=true update -y \ + && apt -o Acquire::ForceIPv4=true install software-properties-common -y \ + && add-apt-repository ppa:deadsnakes/ppa -y \ + && apt -o Acquire::ForceIPv4=true update \ + && apt -o Acquire::ForceIPv4=true install python3.10 python3.10-dev -y \ + && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 \ + && update-alternatives --set python3 /usr/bin/python3.10 \ + && apt -o Acquire::ForceIPv4=true install python3.10-distutils -y \ + && apt -o Acquire::ForceIPv4=true install curl git sudo libibverbs-dev -y \ + && apt -o Acquire::ForceIPv4=true install -y rdma-core infiniband-diags openssh-server perftest ibverbs-providers libibumad3 libibverbs1 libnl-3-200 libnl-route-3-200 librdmacm1 \ + && curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py \ + && python3 get-pip.py \ + && python3 --version \ + && python3 -m pip --version \ + && rm -rf /var/lib/apt/lists/* \ + && apt clean + +# 安装 datamodel_code_generator（用于 MiniCPM 模型） +RUN pip3 install datamodel_code_generator + +WORKDIR /sgl-workspace + +# 拷贝 sglang 源代码并构建包 +COPY ./sglang /sgl-workspace/sglang + +# 拷贝模型文件（修正方式） +#COPY ./Alibaba/QwQ-32B /root/.cradle/Alibaba/QwQ-32B +COPY ./DeepSeek/DeepSeek-R1-Distill-Qwen-14B /root/.cradle/DeepSeek/DeepSeek-R1-Distill-Qwen-14B + +ARG CUDA_VERSION + +# 安装依赖、安装 sglang、安装 transformers，并清理源码 +RUN python3 -m pip install --upgrade pip setuptools wheel html5lib six \ + && if [ "$CUDA_VERSION" = "12.1.1" ]; then \ + CUINDEX=121; \ + elif [ "$CUDA_VERSION" = "12.4.1" ]; then \ + CUINDEX=124; \ + elif [ "$CUDA_VERSION" = "12.5.1" ]; then \ + CUINDEX=124; \ + elif [ "$CUDA_VERSION" = "11.8.0" ]; then \ + CUINDEX=118; \ + python3 -m pip install --no-cache-dir sgl-kernel -i https://docs.sglang.ai/whl/cu118; \ + else \ + echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1; \ + fi \ + && python3 -m pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu${CUINDEX} \ + && python3 -m pip install --no-cache-dir psutil pyzmq pynvml \ + && cd /sgl-workspace/sglang/python \ + && python3 -m pip install --no-cache-dir '.[srt,openai]' --find-links https://flashinfer.ai/whl/cu${CUINDEX}/torch2.5/flashinfer-python \ + && cd / && rm -rf /sgl-workspace/sglang \ + && python3 -m pip install --no-cache-dir transformers==4.48.3 \ + && python3 -c "import sglang; print('✅ sglang module installed successfully')" + +# 暴露端口 +EXPOSE 30000 + + +# 启动命令（QwQ 模型为例） +CMD ["python3", "-m", "sglang.launch_server", \ + "--host","0.0.0.0", \ + "--port","30000", \ + "--model-path", "/root/.cradle/DeepSeek/DeepSeek-R1-Distill-Qwen-14B/", \ + "--tp", "2", \ + "--api-key", "token-abc123"] + +# 设置回默认交互模式 +ENV DEBIAN_FRONTEND=interactive diff --git a/sglang/docker/Dockerfile.ds-8b b/sglang/docker/Dockerfile.ds-8b new file mode 100644 index 0000000..a6c93be --- /dev/null +++ b/sglang/docker/Dockerfile.ds-8b @@ -0,0 +1,77 @@ +ARG CUDA_VERSION=12.5.1 + +FROM nvcr.io/nvidia/tritonserver:24.04-py3-min + +ARG BUILD_TYPE=all +ENV DEBIAN_FRONTEND=noninteractive + +# 安装依赖（强制 IPv4） +RUN echo 'tzdata tzdata/Areas select Asia' | debconf-set-selections \ + && echo 'tzdata tzdata/Zones/Asia select Shanghai' | debconf-set-selections \ + && apt -o Acquire::ForceIPv4=true update -y \ + && apt -o Acquire::ForceIPv4=true install software-properties-common -y \ + && add-apt-repository ppa:deadsnakes/ppa -y \ + && apt -o Acquire::ForceIPv4=true update \ + && apt -o Acquire::ForceIPv4=true install python3.10 python3.10-dev -y \ + && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 \ + && update-alternatives --set python3 /usr/bin/python3.10 \ + && apt -o Acquire::ForceIPv4=true install python3.10-distutils -y \ + && apt -o Acquire::ForceIPv4=true install curl git sudo libibverbs-dev -y \ + && apt -o Acquire::ForceIPv4=true install -y rdma-core infiniband-diags openssh-server perftest ibverbs-providers libibumad3 libibverbs1 libnl-3-200 libnl-route-3-200 librdmacm1 \ + && curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py \ + && python3 get-pip.py \ + && python3 --version \ + && python3 -m pip --version \ + && rm -rf /var/lib/apt/lists/* \ + && apt clean + +# 安装 datamodel_code_generator（用于 MiniCPM 模型） +RUN pip3 install datamodel_code_generator + +WORKDIR /sgl-workspace + +# 拷贝 sglang 源代码并构建包 +COPY ./sglang /sgl-workspace/sglang + +# 拷贝模型文件（修正方式） +#COPY ./Alibaba/QwQ-32B /root/.cradle/Alibaba/QwQ-32B +COPY ./DeepSeek/DeepSeek-R1-Distill-Llama-8B /root/.cradle/DeepSeek/DeepSeek-R1-Distill-Llama-8B + +ARG CUDA_VERSION + +# 安装依赖、安装 sglang、安装 transformers，并清理源码 +RUN python3 -m pip install --upgrade pip setuptools wheel html5lib six \ + && if [ "$CUDA_VERSION" = "12.1.1" ]; then \ + CUINDEX=121; \ + elif [ "$CUDA_VERSION" = "12.4.1" ]; then \ + CUINDEX=124; \ + elif [ "$CUDA_VERSION" = "12.5.1" ]; then \ + CUINDEX=124; \ + elif [ "$CUDA_VERSION" = "11.8.0" ]; then \ + CUINDEX=118; \ + python3 -m pip install --no-cache-dir sgl-kernel -i https://docs.sglang.ai/whl/cu118; \ + else \ + echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1; \ + fi \ + && python3 -m pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu${CUINDEX} \ + && python3 -m pip install --no-cache-dir psutil pyzmq pynvml \ + && cd /sgl-workspace/sglang/python \ + && python3 -m pip install --no-cache-dir '.[srt,openai]' --find-links https://flashinfer.ai/whl/cu${CUINDEX}/torch2.5/flashinfer-python \ + && cd / && rm -rf /sgl-workspace/sglang \ + && python3 -m pip install --no-cache-dir transformers==4.48.3 \ + && python3 -c "import sglang; print('✅ sglang module installed successfully')" + +# 暴露端口 +EXPOSE 30000 + + +# 启动命令（QwQ 模型为例） +CMD ["python3", "-m", "sglang.launch_server", \ + "--host","0.0.0.0", \ + "--port","30000", \ + "--model-path", "/root/.cradle/DeepSeek/DeepSeek-R1-Distill-Llama-8B/", \ + "--tp", "1", \ + "--api-key", "token-abc123"] + +# 设置回默认交互模式 +ENV DEBIAN_FRONTEND=interactive diff --git a/sglang/docker/Dockerfile.qwq-32b b/sglang/docker/Dockerfile.qwq-32b new file mode 100644 index 0000000..1e10283 --- /dev/null +++ b/sglang/docker/Dockerfile.qwq-32b @@ -0,0 +1,77 @@ +ARG CUDA_VERSION=12.5.1 + +FROM nvcr.io/nvidia/tritonserver:24.04-py3-min + +ARG BUILD_TYPE=all +ENV DEBIAN_FRONTEND=noninteractive + +# 安装依赖（强制 IPv4） +RUN echo 'tzdata tzdata/Areas select Asia' | debconf-set-selections \ + && echo 'tzdata tzdata/Zones/Asia select Shanghai' | debconf-set-selections \ + && apt -o Acquire::ForceIPv4=true update -y \ + && apt -o Acquire::ForceIPv4=true install software-properties-common -y \ + && add-apt-repository ppa:deadsnakes/ppa -y \ + && apt -o Acquire::ForceIPv4=true update \ + && apt -o Acquire::ForceIPv4=true install python3.10 python3.10-dev -y \ + && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 \ + && update-alternatives --set python3 /usr/bin/python3.10 \ + && apt -o Acquire::ForceIPv4=true install python3.10-distutils -y \ + && apt -o Acquire::ForceIPv4=true install curl git sudo libibverbs-dev -y \ + && apt -o Acquire::ForceIPv4=true install -y rdma-core infiniband-diags openssh-server perftest ibverbs-providers libibumad3 libibverbs1 libnl-3-200 libnl-route-3-200 librdmacm1 \ + && curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py \ + && python3 get-pip.py \ + && python3 --version \ + && python3 -m pip --version \ + && rm -rf /var/lib/apt/lists/* \ + && apt clean + +# 安装 datamodel_code_generator（用于 MiniCPM 模型） +RUN pip3 install datamodel_code_generator + +WORKDIR /sgl-workspace + +# 拷贝 sglang 源代码并构建包 +COPY ./sglang /sgl-workspace/sglang + +# 拷贝模型文件（修正方式） +#COPY ./Alibaba/QwQ-32B /root/.cradle/Alibaba/QwQ-32B +COPY ./Alibaba/QwQ-32B /root/.cradle/Alibaba/QwQ-32B + +ARG CUDA_VERSION + +# 安装依赖、安装 sglang、安装 transformers，并清理源码 +RUN python3 -m pip install --upgrade pip setuptools wheel html5lib six \ + && if [ "$CUDA_VERSION" = "12.1.1" ]; then \ + CUINDEX=121; \ + elif [ "$CUDA_VERSION" = "12.4.1" ]; then \ + CUINDEX=124; \ + elif [ "$CUDA_VERSION" = "12.5.1" ]; then \ + CUINDEX=124; \ + elif [ "$CUDA_VERSION" = "11.8.0" ]; then \ + CUINDEX=118; \ + python3 -m pip install --no-cache-dir sgl-kernel -i https://docs.sglang.ai/whl/cu118; \ + else \ + echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1; \ + fi \ + && python3 -m pip install --no-cache-dir torch --index-url https://download.pytorch.org/whl/cu${CUINDEX} \ + && python3 -m pip install --no-cache-dir psutil pyzmq pynvml \ + && cd /sgl-workspace/sglang/python \ + && python3 -m pip install --no-cache-dir '.[srt,openai]' --find-links https://flashinfer.ai/whl/cu${CUINDEX}/torch2.5/flashinfer-python \ + && cd / && rm -rf /sgl-workspace/sglang \ + && python3 -m pip install --no-cache-dir transformers==4.48.3 \ + && python3 -c "import sglang; print('✅ sglang module installed successfully')" + +# 暴露端口 +EXPOSE 30000 + + +# 启动命令（QwQ 模型为例） +CMD ["python3", "-m", "sglang.launch_server", \ + "--host","0.0.0.0", \ + "--port","30000", \ + "--model-path", "/root/.cradle/Alibaba/QwQ-32B/", \ + "--tp", "4", \ + "--api-key", "token-abc123"] + +# 设置回默认交互模式 +ENV DEBIAN_FRONTEND=interactive diff --git a/sglang/docker/Dockerfile.rocm b/sglang/docker/Dockerfile.rocm new file mode 100644 index 0000000..0dde646 --- /dev/null +++ b/sglang/docker/Dockerfile.rocm @@ -0,0 +1,85 @@ +# Usage (to build SGLang ROCm docker image): +# docker build --build-arg SGL_BRANCH=v0.4.4.post3 -t v0.4.4.post3-rocm630 -f Dockerfile.rocm . + +# default base image +ARG BASE_IMAGE="rocm/sgl-dev:vllm20250114" + +FROM $BASE_IMAGE AS base +USER root + +WORKDIR /sgl-workspace +ARG BUILD_TYPE=all +ARG SGL_REPO="https://github.com/sgl-project/sglang" +ENV SGL_DEFAULT="main" +ARG SGL_BRANCH=${SGL_DEFAULT} + +ARG TRITON_REPO="https://github.com/ROCm/triton.git" +ARG TRITON_COMMIT="improve_fa_decode_3.0.0" + + +ARG AITER_REPO="https://github.com/ROCm/aiter.git" +ARG AITER_COMMIT="testx" + +RUN git clone ${SGL_REPO} \ + && cd sglang \ + && if [ "${SGL_BRANCH}" = ${SGL_DEFAULT} ]; then \ + echo "Using ${SGL_DEFAULT}, default branch."; \ + else \ + echo "Using ${SGL_BRANCH} branch."; \ + git checkout ${SGL_BRANCH}; \ + fi \ + && cd sgl-kernel \ + && rm -f pyproject.toml \ + && mv pyproject_rocm.toml pyproject.toml \ + && python setup_rocm.py install \ + && cd .. \ + && if [ "$BUILD_TYPE" = "srt" ]; then \ + python -m pip --no-cache-dir install -e "python[srt_hip]"; \ + else \ + python -m pip --no-cache-dir install -e "python[all_hip]"; \ + fi + +RUN cp -r /sgl-workspace/sglang /sglang +RUN python -m pip cache purge + +RUN pip install IPython \ + && pip install orjson \ + && pip install python-multipart \ + && pip install torchao \ + && pip install pybind11 + +RUN pip uninstall -y triton +RUN git clone ${TRITON_REPO} \ + && cd triton \ + && git checkout ${TRITON_COMMIT} \ + && cd python \ + && python3 setup.py install + +RUN git clone ${AITER_REPO} \ + && cd aiter \ + && git checkout ${AITER_COMMIT} \ + && git submodule update --init --recursive \ + && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py develop + +# Copy config files to support MI300X in virtualized environments (MI300X_VF). Symlinks will not be created in image build. +RUN find /sgl-workspace/sglang/python/sglang/srt/layers/quantization/configs/ \ + /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs/ \ + -type f -name '*MI300X*' | xargs -I {} sh -c 'vf_config=$(echo "$1" | sed "s/MI300X/MI300X_VF/"); cp "$1" "$vf_config"' -- {} + +# Performance environment variable. + +ENV HIP_FORCE_DEV_KERNARG=1 +ENV HSA_NO_SCRATCH_RECLAIM=1 +ENV SGLANG_SET_CPU_AFFINITY=1 +ENV SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1 +ENV NCCL_MIN_NCHANNELS=112 + +ENV MOE_PADDING=1 +ENV VLLM_FP8_PADDING=1 +ENV VLLM_FP8_ACT_PADDING=1 +ENV VLLM_FP8_WEIGHT_PADDING=1 +ENV VLLM_FP8_REDUCE_CONV=1 +ENV TORCHINDUCTOR_MAX_AUTOTUNE=1 +ENV TORCHINDUCTOR_MAX_AUTOTUNE_POINTWISE=1 + +CMD ["/bin/bash"] diff --git a/sglang/docker/Dockerfile.sagemaker b/sglang/docker/Dockerfile.sagemaker new file mode 100644 index 0000000..fde8d55 --- /dev/null +++ b/sglang/docker/Dockerfile.sagemaker @@ -0,0 +1,78 @@ +ARG CUDA_VERSION=12.5.1 + +FROM nvcr.io/nvidia/tritonserver:24.04-py3-min + +ARG BUILD_TYPE=all +ENV DEBIAN_FRONTEND=noninteractive + +RUN echo 'tzdata tzdata/Areas select America' | debconf-set-selections \ + && echo 'tzdata tzdata/Zones/America select Los_Angeles' | debconf-set-selections \ + && apt update -y \ + && apt install software-properties-common -y \ + && add-apt-repository ppa:deadsnakes/ppa -y && apt update \ + && apt install python3.10 python3.10-dev -y \ + && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1 \ + && update-alternatives --set python3 /usr/bin/python3.10 && apt install python3.10-distutils -y \ + && apt install curl git sudo libibverbs-dev -y \ + && apt install -y rdma-core infiniband-diags openssh-server perftest ibverbs-providers libibumad3 libibverbs1 libnl-3-200 libnl-route-3-200 librdmacm1 \ + && curl https://bootstrap.pypa.io/get-pip.py -o get-pip.py && python3 get-pip.py \ + && python3 --version \ + && python3 -m pip --version \ + && rm -rf /var/lib/apt/lists/* \ + && apt clean + +# For openbmb/MiniCPM models +RUN pip3 install datamodel_code_generator + +WORKDIR /sgl-workspace + +ARG CUDA_VERSION +RUN python3 -m pip install --upgrade pip setuptools wheel html5lib six \ + && git clone --depth=1 https://github.com/sgl-project/sglang.git \ + && if [ "$CUDA_VERSION" = "12.1.1" ]; then \ + python3 -m pip install torch --index-url https://download.pytorch.org/whl/cu121; \ + elif [ "$CUDA_VERSION" = "12.4.1" ]; then \ + python3 -m pip install torch --index-url https://download.pytorch.org/whl/cu124; \ + elif [ "$CUDA_VERSION" = "12.5.1" ]; then \ + python3 -m pip install torch --index-url https://download.pytorch.org/whl/cu124; \ + elif [ "$CUDA_VERSION" = "11.8.0" ]; then \ + python3 -m pip install torch --index-url https://download.pytorch.org/whl/cu118; \ + python3 -m pip install sgl-kernel -i https://docs.sglang.ai/whl/cu118; \ + else \ + echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1; \ + fi \ + && cd sglang \ + && if [ "$BUILD_TYPE" = "srt" ]; then \ + if [ "$CUDA_VERSION" = "12.1.1" ]; then \ + python3 -m pip --no-cache-dir install -e "python[srt]" --find-links https://flashinfer.ai/whl/cu121/torch2.5/flashinfer-python; \ + elif [ "$CUDA_VERSION" = "12.4.1" ]; then \ + python3 -m pip --no-cache-dir install -e "python[srt]" --find-links https://flashinfer.ai/whl/cu124/torch2.5/flashinfer-python; \ + elif [ "$CUDA_VERSION" = "12.5.1" ]; then \ + python3 -m pip --no-cache-dir install -e "python[srt]" --find-links https://flashinfer.ai/whl/cu124/torch2.5/flashinfer-python; \ + elif [ "$CUDA_VERSION" = "11.8.0" ]; then \ + python3 -m pip --no-cache-dir install -e "python[srt]" --find-links https://flashinfer.ai/whl/cu118/torch2.5/flashinfer-python; \ + python3 -m pip install sgl-kernel -i https://docs.sglang.ai/whl/cu118; \ + else \ + echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1; \ + fi; \ + else \ + if [ "$CUDA_VERSION" = "12.1.1" ]; then \ + python3 -m pip --no-cache-dir install -e "python[all]" --find-links https://flashinfer.ai/whl/cu121/torch2.5/flashinfer-python; \ + elif [ "$CUDA_VERSION" = "12.4.1" ]; then \ + python3 -m pip --no-cache-dir install -e "python[all]" --find-links https://flashinfer.ai/whl/cu124/torch2.5/flashinfer-python; \ + elif [ "$CUDA_VERSION" = "12.5.1" ]; then \ + python3 -m pip --no-cache-dir install -e "python[all]" --find-links https://flashinfer.ai/whl/cu124/torch2.5/flashinfer-python; \ + elif [ "$CUDA_VERSION" = "11.8.0" ]; then \ + python3 -m pip --no-cache-dir install -e "python[all]" --find-links https://flashinfer.ai/whl/cu118/torch2.5/flashinfer-python; \ + python3 -m pip install sgl-kernel -i https://docs.sglang.ai/whl/cu118; \ + else \ + echo "Unsupported CUDA version: $CUDA_VERSION" && exit 1; \ + fi; \ + fi + +ENV DEBIAN_FRONTEND=interactive + +COPY serve /usr/bin/serve +RUN chmod 777 /usr/bin/serve + +ENTRYPOINT [ "/usr/bin/serve" ] diff --git a/sglang/docker/compose.yaml b/sglang/docker/compose.yaml new file mode 100644 index 0000000..f7ff1fb --- /dev/null +++ b/sglang/docker/compose.yaml @@ -0,0 +1,35 @@ +services: + sglang: + image: lmsysorg/sglang:latest + container_name: sglang + volumes: + - ${HOME}/.cache/huggingface:/root/.cache/huggingface + # If you use modelscope, you need mount this directory + # - ${HOME}/.cache/modelscope:/root/.cache/modelscope + restart: always + network_mode: host # required by RDMA + privileged: true # required by RDMA + # Or you can only publish port 30000 + # ports: + # - 30000:30000 + environment: + HF_TOKEN: + # if you use modelscope to download model, you need set this environment + # - SGLANG_USE_MODELSCOPE: true + entrypoint: python3 -m sglang.launch_server + command: --model-path meta-llama/Llama-3.1-8B-Instruct + --host 0.0.0.0 + --port 30000 + ulimits: + memlock: -1 + stack: 67108864 + ipc: host + healthcheck: + test: ["CMD-SHELL", "curl -f http://localhost:30000/health || exit 1"] + deploy: + resources: + reservations: + devices: + - driver: nvidia + device_ids: ["0"] + capabilities: [gpu] diff --git a/sglang/docker/k8s-sglang-distributed-sts.yaml b/sglang/docker/k8s-sglang-distributed-sts.yaml new file mode 100644 index 0000000..6b81d9b --- /dev/null +++ b/sglang/docker/k8s-sglang-distributed-sts.yaml @@ -0,0 +1,104 @@ +# Two Nodes Sglang example + +apiVersion: apps/v1 +kind: StatefulSet +metadata: + name: distributed-sglang +spec: + replicas: 2 # number of nodes/pods to run distributed sglang + selector: + matchLabels: + app: distributed-sglang + serviceName: "" + template: + metadata: + labels: + app: distributed-sglang + spec: + containers: + - name: sglang-container + image: docker.io/lmsysorg/sglang:latest + imagePullPolicy: Always # image may be replaced by official CI versioned image + command: + - /bin/bash + - -c + # please modify the sglang serving arguments below, as necessary. + # NOTE: the --expert-parallel-size and --enable-ep-moe are for MoE model like DeepSeek-R1 + args: + - | + python3 -m sglang.launch_server \ + --model /llm-folder \ + --dist-init-addr sglang-master-pod:5000 \ + --tensor-parallel-size 16 \ + --nnodes 2 \ + --node-rank $POD_INDEX \ + --trust-remote-code \ + --host 0.0.0.0 \ + --port 8000 \ + --enable-metrics \ + --enable-ep-moe \ + --expert-parallel-size 16 + env: + - name: POD_INDEX # reflects the node-rank + valueFrom: + fieldRef: + apiVersion: v1 + fieldPath: metadata.labels['apps.kubernetes.io/pod-index'] + - name: NCCL_DEBUG + value: INFO + resources: + limits: + nvidia.com/gpu: "8" + requests: + volumeMounts: + - mountPath: /dev/shm + name: dshm + - mountPath: /llm-folder + name: llm + securityContext: + privileged: true # to leverage RDMA/InfiniBand device, co-work with HostNetwork=true + hostNetwork: true + volumes: + - emptyDir: + medium: Memory + sizeLimit: 10Gi + name: dshm + - hostPath: + path: /llm-folder # replace with PVC or hostPath with your model weights + type: DirectoryOrCreate + name: llm + #- persistentVolumeClaim: + # claimName: llm-pvc + # name: llm +--- +apiVersion: v1 +kind: Service +metadata: + name: sglang-master-pod +spec: + type: ClusterIP + selector: + app: distributed-sglang + apps.kubernetes.io/pod-index: "0" + ports: + - name: dist-port + port: 5000 + targetPort: 5000 +--- +# the serving service +apiVersion: v1 +kind: Service +metadata: + name: sglang-serving-on-master +spec: + type: NodePort + selector: + app: distributed-sglang + apps.kubernetes.io/pod-index: "0" + ports: + - name: serving + port: 8000 + targetPort: 8000 + - name: metrics + port: 8080 + targetPort: 8080 diff --git a/sglang/docker/k8s-sglang-service.yaml b/sglang/docker/k8s-sglang-service.yaml new file mode 100644 index 0000000..6f6f9af --- /dev/null +++ b/sglang/docker/k8s-sglang-service.yaml @@ -0,0 +1,81 @@ +apiVersion: node.k8s.io/v1 +kind: RuntimeClass +metadata: + name: nvidia +handler: nvidia +--- +apiVersion: apps/v1 +kind: Deployment +metadata: + name: meta-llama-31-8b-instruct-sglang +spec: + replicas: 1 + strategy: + type: Recreate + selector: + matchLabels: + app: meta-llama-31-8b-instruct-sglang + template: + metadata: + labels: + app: meta-llama-31-8b-instruct-sglang + model: meta-llama-31-8b-instruct + engine: sglang + spec: + restartPolicy: Always + runtimeClassName: nvidia + containers: + - name: meta-llama-31-8b-instruct-sglang + image: docker.io/lmsysorg/sglang:latest + imagePullPolicy: Always # IfNotPresent or Never + ports: + - containerPort: 30000 + command: ["python3", "-m", "sglang.launch_server"] + args: ["--model-path", "meta-llama/Llama-3.1-8B-Instruct", "--host", "0.0.0.0", "--port", "30000"] + env: + - name: HF_TOKEN + value: + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: shm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + readOnly: true + - name: localtime + mountPath: /etc/localtime + readOnly: true + livenessProbe: + httpGet: + path: /health + port: 30000 + initialDelaySeconds: 30 + periodSeconds: 10 + volumes: + - name: shm + emptyDir: + medium: Memory + sizeLimit: 10Gi + - name: hf-cache + hostPath: + path: /root/.cache/huggingface + type: Directory + - name: localtime + hostPath: + path: /etc/localtime + type: File +--- +apiVersion: v1 +kind: Service +metadata: + name: meta-llama-31-8b-instruct-sglang +spec: + selector: + app: meta-llama-31-8b-instruct-sglang + ports: + - protocol: TCP + port: 30000 # port on host + targetPort: 30000 # port in container + type: LoadBalancer diff --git a/sglang/docker/serve b/sglang/docker/serve new file mode 100644 index 0000000..493ecbd --- /dev/null +++ b/sglang/docker/serve @@ -0,0 +1,31 @@ +#!/bin/bash + +echo "Starting server" + +SERVER_ARGS="--host 0.0.0.0 --port 8080" + +if [ -n "$TENSOR_PARALLEL_DEGREE" ]; then + SERVER_ARGS="${SERVER_ARGS} --tp-size ${TENSOR_PARALLEL_DEGREE}" +fi + +if [ -n "$DATA_PARALLEL_DEGREE" ]; then + SERVER_ARGS="${SERVER_ARGS} --dp-size ${DATA_PARALLEL_DEGREE}" +fi + +if [ -n "$EXPERT_PARALLEL_DEGREE" ]; then + SERVER_ARGS="${SERVER_ARGS} --ep-size ${EXPERT_PARALLEL_DEGREE}" +fi + +if [ -n "$MEM_FRACTION_STATIC" ]; then + SERVER_ARGS="${SERVER_ARGS} --mem-fraction-static ${MEM_FRACTION_STATIC}" +fi + +if [ -n "$QUANTIZATION" ]; then + SERVER_ARGS="${SERVER_ARGS} --quantization ${QUANTIZATION}" +fi + +if [ -n "$CHUNKED_PREFILL_SIZE" ]; then + SERVER_ARGS="${SERVER_ARGS} --chunked-prefill-size ${CHUNKED_PREFILL_SIZE}" +fi + +python3 -m sglang.launch_server --model-path /opt/ml/model $SERVER_ARGS diff --git a/sglang/docker/start.sh b/sglang/docker/start.sh new file mode 100755 index 0000000..492b40a --- /dev/null +++ b/sglang/docker/start.sh @@ -0,0 +1,6 @@ +docker run -d \ + -p 30000:30000 \ + --gpus all \ + --shm-size=2g \ + --name alibaba_qwq-32b \ + alibaba_qwq-32b diff --git a/sglang/docs/Makefile b/sglang/docs/Makefile new file mode 100644 index 0000000..0d6cad1 --- /dev/null +++ b/sglang/docs/Makefile @@ -0,0 +1,57 @@ +# Minimal Makefile for Sphinx documentation +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SPHINXAUTOBUILD ?= sphinx-autobuild +SOURCEDIR = . +BUILDDIR = _build +PORT ?= 8003 + +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + @echo "" + @echo "Additional targets:" + @echo " serve to build and serve documentation with auto-build and live reload" + +# Compile Notebook files and record execution time +compile: + @set -e; \ + echo "Starting Notebook compilation..."; \ + mkdir -p logs; \ + echo "Notebook execution timings:" > logs/timing.log; \ + START_TOTAL=$$(date +%s); \ + find $(SOURCEDIR) -path "*/_build/*" -prune -o -name "*.ipynb" -print0 | \ + parallel -0 -j3 --halt soon,fail=1 ' \ + NB_NAME=$$(basename {}); \ + START_TIME=$$(date +%s); \ + retry --delay=0 --times=3 -- \ + jupyter nbconvert --to notebook --execute --inplace "{}" \ + --ExecutePreprocessor.timeout=600 \ + --ExecutePreprocessor.kernel_name=python3; \ + RET_CODE=$$?; \ + END_TIME=$$(date +%s); \ + ELAPSED_TIME=$$((END_TIME - START_TIME)); \ + echo "$${NB_NAME}: $${ELAPSED_TIME}s" >> logs/timing.log; \ + exit $$RET_CODE' || exit 1; \ + END_TOTAL=$$(date +%s); \ + TOTAL_ELAPSED=$$((END_TOTAL - START_TOTAL)); \ + echo "---------------------------------" >> logs/timing.log; \ + echo "Total execution time: $${TOTAL_ELAPSED}s" >> logs/timing.log; \ + echo "All Notebook execution timings:" && cat logs/timing.log + +# Serve documentation with auto-build and live reload +serve: + @echo "Starting auto-build server at http://localhost:$(PORT)" + @$(SPHINXAUTOBUILD) "$(SOURCEDIR)" "$(BUILDDIR)/html" \ + --port $(PORT) \ + --watch $(SOURCEDIR) \ + --re-ignore ".*\.(ipynb_checkpoints|pyc|pyo|pyd|git)" + +.PHONY: help Makefile compile clean serve + +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +clean: + find . -name "*.ipynb" -exec nbstripout {} \; + rm -rf $(BUILDDIR) + rm -rf logs diff --git a/sglang/docs/README.md b/sglang/docs/README.md new file mode 100644 index 0000000..0dfadb1 --- /dev/null +++ b/sglang/docs/README.md @@ -0,0 +1,110 @@ +# SGLang Documentation + +We recommend new contributors start from writing documentation, which helps you quickly understand SGLang codebase. Most documentation files are located under the `docs/` folder. We prefer **Jupyter Notebooks** over Markdown so that all examples can be executed and validated by our docs CI pipeline. + +## Docs Workflow + +### Install Dependency + +```bash +pip install -r requirements.txt +``` + +### Update Documentation + +Update your Jupyter notebooks in the appropriate subdirectories under `docs/`. If you add new files, remember to update `index.rst` (or relevant `.rst` files) accordingly. + +- **`pre-commit run --all-files`** manually runs all configured checks, applying fixes if possible. If it fails the first time, re-run it to ensure lint errors are fully resolved. Make sure your code passes all checks **before** creating a Pull Request. +- **Do not commit** directly to the `main` branch. Always create a new branch (e.g., `feature/my-new-feature`), push your changes, and open a PR from that branch. + +```bash +# 1) Compile all Jupyter notebooks +make compile + +# 2) Compile and Preview documentation locally with auto-build +# This will automatically rebuild docs when files change +# Open your browser at the displayed port to view the docs +bash serve.sh + +# 2a) Alternative ways to serve documentation +# Directly use make serve +make serve +# With custom port +PORT=8080 make serve + +# 3) Clean notebook outputs +# nbstripout removes notebook outputs so your PR stays clean +pip install nbstripout +find . -name '*.ipynb' -exec nbstripout {} \; + +# 4) Pre-commit checks and create a PR +# After these checks pass, push your changes and open a PR on your branch +pre-commit run --all-files +``` +--- + +### **Port Allocation and CI Efficiency** + +**To launch and kill the server:** + +```python +from sglang.test.test_utils import is_in_ci +from sglang.utils import wait_for_server, print_highlight, terminate_process + +if is_in_ci(): + from patch import launch_server_cmd +else: + from sglang.utils import launch_server_cmd + +server_process, port = launch_server_cmd( + """ +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ + --host 0.0.0.0 +""" +) + +wait_for_server(f"http://localhost:{port}") + +# Terminate Server +terminate_process(server_process) +``` + +**To launch and kill the engine:** + +```python +# Launch Engine +import sglang as sgl +import asyncio +from sglang.test.test_utils import is_in_ci + +if is_in_ci(): + import patch + +llm = sgl.Engine(model_path="meta-llama/Meta-Llama-3.1-8B-Instruct") + +# Terminalte Engine +llm.shutdown() +``` + +### **Why this approach?** + +- **Dynamic Port Allocation**: Avoids port conflicts by selecting an available port at runtime, enabling multiple server instances to run in parallel. +- **Optimized for CI**: The `patch` version of `launch_server_cmd` and `sgl.Engine()` in CI environments helps manage GPU memory dynamically, preventing conflicts and improving test parallelism. +- **Better Parallel Execution**: Ensures smooth concurrent tests by avoiding fixed port collisions and optimizing memory usage. + +### **Model Selection** + +For demonstrations in the docs, **prefer smaller models** to reduce memory consumption and speed up inference. Running larger models in CI can lead to instability due to memory constraints. + +### **Prompt Alignment Example** + +When designing prompts, ensure they align with SGLang's structured formatting. For example: + +```python +prompt = """You are an AI assistant. Answer concisely and accurately. + +User: What is the capital of France? +Assistant: The capital of France is Paris.""" +``` + +This keeps responses aligned with expected behavior and improves reliability across different files. diff --git a/sglang/docs/_static/css/custom_log.css b/sglang/docs/_static/css/custom_log.css new file mode 100644 index 0000000..61f65d0 --- /dev/null +++ b/sglang/docs/_static/css/custom_log.css @@ -0,0 +1,29 @@ +.output_area { + color: #615656; +} + +table.autosummary td { + width: 50% + } + + img.align-center { + display: block; + margin-left: auto; + margin-right: auto; +} + +.output_area.stderr { + color: #d3d3d3 !important; +} + +.output_area.stdout { + color: #d3d3d3 !important; +} + +div.output_area.stderr { + color: #d3d3d3 !important; +} + +div.output_area.stdout { + color: #d3d3d3 !important; +} diff --git a/sglang/docs/_static/css/readthedocs.css b/sglang/docs/_static/css/readthedocs.css new file mode 100644 index 0000000..aca6649 --- /dev/null +++ b/sglang/docs/_static/css/readthedocs.css @@ -0,0 +1,9 @@ +table.autosummary td { + width: 50% +} + +img.align-center { + display: block; + margin-left: auto; + margin-right: auto; +} diff --git a/sglang/docs/_static/image/logo.ico b/sglang/docs/_static/image/logo.ico new file mode 100644 index 0000000..7a0b5b2 Binary files /dev/null and b/sglang/docs/_static/image/logo.ico differ diff --git a/sglang/docs/_static/image/logo.png b/sglang/docs/_static/image/logo.png new file mode 100644 index 0000000..2a8bc25 Binary files /dev/null and b/sglang/docs/_static/image/logo.png differ diff --git a/sglang/docs/backend/__init__.py b/sglang/docs/backend/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sglang/docs/backend/custom_chat_template.md b/sglang/docs/backend/custom_chat_template.md new file mode 100644 index 0000000..557af5b --- /dev/null +++ b/sglang/docs/backend/custom_chat_template.md @@ -0,0 +1,42 @@ +# Custom Chat Template + +**NOTE**: There are two chat template systems in SGLang project. This document is about setting a custom chat template for the OpenAI-compatible API server (defined at [conversation.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/conversation.py)). It is NOT related to the chat template used in the SGLang language frontend (defined at [chat_template.py](https://github.com/sgl-project/sglang/blob/main/python/sglang/lang/chat_template.py)). + +By default, the server uses the chat template specified in the model tokenizer from Hugging Face. +It should just work for most official models such as Llama-2/Llama-3. + +If needed, you can also override the chat template when launching the server: + +```bash +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --chat-template llama-2 +``` + +If the chat template you are looking for is missing, you are welcome to contribute it or load it from a file. + +## JSON Format + +You can load the JSON format, which is defined by `conversation.py`. + +```json +{ + "name": "my_model", + "system": "<|im_start|>system", + "user": "<|im_start|>user", + "assistant": "<|im_start|>assistant", + "sep_style": "CHATML", + "sep": "<|im_end|>", + "stop_str": ["<|im_end|>", "<|im_start|>"] +} +``` + +```bash +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --chat-template ./my_model_template.json +``` + +## Jinja Format + +You can also use the [Jinja template format](https://huggingface.co/docs/transformers/main/en/chat_templating) as defined by Hugging Face Transformers. + +```bash +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --chat-template ./my_model_template.jinja +``` diff --git a/sglang/docs/backend/function_calling.ipynb b/sglang/docs/backend/function_calling.ipynb new file mode 100644 index 0000000..99d11f5 --- /dev/null +++ b/sglang/docs/backend/function_calling.ipynb @@ -0,0 +1,544 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Tool and Function Calling\n", + "\n", + "This guide demonstrates how to use SGLang’s [Funcion calling](https://platform.openai.com/docs/guides/function-calling) functionality." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## OpenAI Compatible API" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Launching the Server" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from openai import OpenAI\n", + "import json\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "from sglang.test.test_utils import is_in_ci\n", + "\n", + "if is_in_ci():\n", + " from patch import launch_server_cmd\n", + "else:\n", + " from sglang.utils import launch_server_cmd\n", + "\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"python3 -m sglang.launch_server --model-path Qwen/Qwen2.5-7B-Instruct --tool-call-parser qwen25 --host 0.0.0.0\" # qwen25\n", + ")\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Note that `--tool-call-parser` defines the parser used to interpret responses. Currently supported parsers include:\n", + "\n", + "- llama3: Llama 3.1 / 3.2 (e.g. meta-llama/Llama-3.1-8B-Instruct, meta-llama/Llama-3.2-1B-Instruct).\n", + "- mistral: Mistral (e.g. mistralai/Mistral-7B-Instruct-v0.3, mistralai/Mistral-Nemo-Instruct-2407, mistralai/\n", + "Mistral-Nemo-Instruct-2407, mistralai/Mistral-7B-v0.3).\n", + "- qwen25: Qwen 2.5 (e.g. Qwen/Qwen2.5-1.5B-Instruct, Qwen/Qwen2.5-7B-Instruct) and QwQ (i.e. Qwen/QwQ-32B). Especially, for QwQ, we can enable the reasoning parser together with tool call parser, details about reasoning parser can be found in [reasoning parser](https://docs.sglang.ai/backend/separate_reasoning.html)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Define Tools for Function Call\n", + "Below is a Python snippet that shows how to define a tool as a dictionary. The dictionary includes a tool name, a description, and property defined Parameters." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Define tools\n", + "tools = [\n", + " {\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": \"get_current_weather\",\n", + " \"description\": \"Get the current weather in a given location\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"city\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The city to find the weather for, e.g. 'San Francisco'\",\n", + " },\n", + " \"state\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"the two-letter abbreviation for the state that the city is\"\n", + " \" in, e.g. 'CA' which would mean 'California'\",\n", + " },\n", + " \"unit\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The unit to fetch the temperature in\",\n", + " \"enum\": [\"celsius\", \"fahrenheit\"],\n", + " },\n", + " },\n", + " \"required\": [\"city\", \"state\", \"unit\"],\n", + " },\n", + " },\n", + " }\n", + "]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Define Messages" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def get_messages():\n", + " return [\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"What's the weather like in Boston today? Output a reasoning before act, then use the tools to help you.\",\n", + " }\n", + " ]\n", + "\n", + "\n", + "messages = get_messages()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Initialize the Client" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Initialize OpenAI-like client\n", + "client = OpenAI(api_key=\"None\", base_url=f\"http://0.0.0.0:{port}/v1\")\n", + "model_name = client.models.list().data[0].id" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Non-Streaming Request" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Non-streaming mode test\n", + "response_non_stream = client.chat.completions.create(\n", + " model=model_name,\n", + " messages=messages,\n", + " temperature=0.1,\n", + " top_p=0.95,\n", + " max_tokens=1024,\n", + " stream=False, # Non-streaming\n", + " tools=tools,\n", + ")\n", + "print_highlight(\"Non-stream response:\")\n", + "print(response_non_stream)\n", + "print_highlight(\"==== content ====\")\n", + "print(response_non_stream.choices[0].message.content)\n", + "print_highlight(\"==== tool_calls ====\")\n", + "print(response_non_stream.choices[0].message.tool_calls)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Handle Tools\n", + "When the engine determines it should call a particular tool, it will return arguments or partial arguments through the response. You can parse these arguments and later invoke the tool accordingly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "name_non_stream = response_non_stream.choices[0].message.tool_calls[0].function.name\n", + "arguments_non_stream = (\n", + " response_non_stream.choices[0].message.tool_calls[0].function.arguments\n", + ")\n", + "\n", + "print_highlight(f\"Final streamed function call name: {name_non_stream}\")\n", + "print_highlight(f\"Final streamed function call arguments: {arguments_non_stream}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Streaming Request" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Streaming mode test\n", + "print_highlight(\"Streaming response:\")\n", + "response_stream = client.chat.completions.create(\n", + " model=model_name,\n", + " messages=messages,\n", + " temperature=0.1,\n", + " top_p=0.95,\n", + " max_tokens=1024,\n", + " stream=True, # Enable streaming\n", + " tools=tools,\n", + ")\n", + "\n", + "texts = \"\"\n", + "tool_calls = []\n", + "name = \"\"\n", + "arguments = \"\"\n", + "for chunk in response_stream:\n", + " if chunk.choices[0].delta.content:\n", + " texts += chunk.choices[0].delta.content\n", + " if chunk.choices[0].delta.tool_calls:\n", + " tool_calls.append(chunk.choices[0].delta.tool_calls[0])\n", + "print_highlight(\"==== Text ====\")\n", + "print(texts)\n", + "\n", + "print_highlight(\"==== Tool Call ====\")\n", + "for tool_call in tool_calls:\n", + " print(tool_call)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Handle Tools\n", + "When the engine determines it should call a particular tool, it will return arguments or partial arguments through the response. You can parse these arguments and later invoke the tool accordingly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Parse and combine function call arguments\n", + "arguments = []\n", + "for tool_call in tool_calls:\n", + " if tool_call.function.name:\n", + " print_highlight(f\"Streamed function call name: {tool_call.function.name}\")\n", + "\n", + " if tool_call.function.arguments:\n", + " arguments.append(tool_call.function.arguments)\n", + "\n", + "# Combine all fragments into a single JSON string\n", + "full_arguments = \"\".join(arguments)\n", + "print_highlight(f\"streamed function call arguments: {full_arguments}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Define a Tool Function" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# This is a demonstration, define real function according to your usage.\n", + "def get_current_weather(city: str, state: str, unit: \"str\"):\n", + " return (\n", + " f\"The weather in {city}, {state} is 85 degrees {unit}. It is \"\n", + " \"partly cloudly, with highs in the 90's.\"\n", + " )\n", + "\n", + "\n", + "available_tools = {\"get_current_weather\": get_current_weather}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "## Execute the Tool" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "call_data = json.loads(full_arguments)\n", + "\n", + "messages.append(\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"\",\n", + " \"tool_calls\": {\"name\": \"get_current_weather\", \"arguments\": full_arguments},\n", + " }\n", + ")\n", + "\n", + "# Call the corresponding tool function\n", + "tool_name = messages[-1][\"tool_calls\"][\"name\"]\n", + "tool_to_call = available_tools[tool_name]\n", + "result = tool_to_call(**call_data)\n", + "print_highlight(f\"Function call result: {result}\")\n", + "messages.append({\"role\": \"tool\", \"content\": result, \"name\": tool_name})\n", + "\n", + "print_highlight(f\"Updated message history: {messages}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Send Results Back to Model" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "final_response = client.chat.completions.create(\n", + " model=model_name,\n", + " messages=messages,\n", + " temperature=0.1,\n", + " top_p=0.95,\n", + " stream=False,\n", + " tools=tools,\n", + ")\n", + "print_highlight(\"Non-stream response:\")\n", + "print(final_response)\n", + "\n", + "print_highlight(\"==== Text ====\")\n", + "print(final_response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Native API and SGLang Runtime (SRT)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from transformers import AutoTokenizer\n", + "import requests\n", + "\n", + "# generate an answer\n", + "tokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen2.5-7B-Instruct\")\n", + "\n", + "messages = get_messages()\n", + "\n", + "input = tokenizer.apply_chat_template(\n", + " messages,\n", + " tokenize=False,\n", + " add_generation_prompt=True,\n", + " tools=tools,\n", + ")\n", + "\n", + "gen_url = f\"http://localhost:{port}/generate\"\n", + "gen_data = {\n", + " \"text\": input,\n", + " \"sampling_params\": {\n", + " \"skip_special_tokens\": False,\n", + " \"max_new_tokens\": 1024,\n", + " \"temperature\": 0.1,\n", + " \"top_p\": 0.95,\n", + " },\n", + "}\n", + "gen_response = requests.post(gen_url, json=gen_data).json()[\"text\"]\n", + "print_highlight(\"==== Reponse ====\")\n", + "print(gen_response)\n", + "\n", + "# parse the response\n", + "parse_url = f\"http://localhost:{port}/parse_function_call\"\n", + "\n", + "function_call_input = {\n", + " \"text\": gen_response,\n", + " \"tool_call_parser\": \"qwen25\",\n", + " \"tools\": tools,\n", + "}\n", + "\n", + "function_call_response = requests.post(parse_url, json=function_call_input)\n", + "function_call_response_json = function_call_response.json()\n", + "\n", + "print_highlight(\"==== Text ====\")\n", + "print(function_call_response_json[\"normal_text\"])\n", + "print_highlight(\"==== Calls ====\")\n", + "print(\"function name: \", function_call_response_json[\"calls\"][0][\"name\"])\n", + "print(\"function arguments: \", function_call_response_json[\"calls\"][0][\"parameters\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Offline Engine API" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import sglang as sgl\n", + "from sglang.srt.function_call_parser import FunctionCallParser\n", + "from sglang.srt.managers.io_struct import Tool, Function\n", + "\n", + "llm = sgl.Engine(model_path=\"Qwen/Qwen2.5-7B-Instruct\")\n", + "tokenizer = llm.tokenizer_manager.tokenizer\n", + "input_ids = tokenizer.apply_chat_template(\n", + " messages, tokenize=True, add_generation_prompt=True, tools=tools\n", + ")\n", + "\n", + "sampling_params = {\n", + " \"max_new_tokens\": 1024,\n", + " \"temperature\": 0.1,\n", + " \"top_p\": 0.95,\n", + " \"skip_special_tokens\": False,\n", + "}\n", + "\n", + "# 1) Offline generation\n", + "result = llm.generate(input_ids=input_ids, sampling_params=sampling_params)\n", + "generated_text = result[\"text\"] # Assume there is only one prompt\n", + "\n", + "print(\"=== Offline Engine Output Text ===\")\n", + "print(generated_text)\n", + "\n", + "\n", + "# 2) Parse using FunctionCallParser\n", + "def convert_dict_to_tool(tool_dict: dict) -> Tool:\n", + " function_dict = tool_dict.get(\"function\", {})\n", + " return Tool(\n", + " type=tool_dict.get(\"type\", \"function\"),\n", + " function=Function(\n", + " name=function_dict.get(\"name\"),\n", + " description=function_dict.get(\"description\"),\n", + " parameters=function_dict.get(\"parameters\"),\n", + " ),\n", + " )\n", + "\n", + "\n", + "tools = [convert_dict_to_tool(raw_tool) for raw_tool in tools]\n", + "\n", + "parser = FunctionCallParser(tools=tools, tool_call_parser=\"qwen25\")\n", + "normal_text, calls = parser.parse_non_stream(generated_text)\n", + "\n", + "print(\"=== Parsing Result ===\")\n", + "print(\"Normal text portion:\", normal_text)\n", + "print(\"Function call portion:\")\n", + "for call in calls:\n", + " # call: ToolCallItem\n", + " print(f\" - tool name: {call.name}\")\n", + " print(f\" parameters: {call.parameters}\")\n", + "\n", + "# 3) If needed, perform additional logic on the parsed functions, such as automatically calling the corresponding function to obtain a return value, etc." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "llm.shutdown()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## How to support a new model?\n", + "1. Update the TOOLS_TAG_LIST in sglang/srt/function_call_parser.py with the model’s tool tags. Currently supported tags include:\n", + "```\n", + "\tTOOLS_TAG_LIST = [\n", + "\t “<|plugin|>“,\n", + "\t ““,\n", + "\t “<|python_tag|>“,\n", + "\t “[TOOL_CALLS]”\n", + "\t]\n", + "```\n", + "2. Create a new detector class in sglang/srt/function_call_parser.py that inherits from BaseFormatDetector. The detector should handle the model’s specific function call format. For example:\n", + "```\n", + " class NewModelDetector(BaseFormatDetector):\n", + "```\n", + "3. Add the new detector to the MultiFormatParser class that manages all the format detectors." + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/sglang/docs/backend/hyperparameter_tuning.md b/sglang/docs/backend/hyperparameter_tuning.md new file mode 100644 index 0000000..ca2f52b --- /dev/null +++ b/sglang/docs/backend/hyperparameter_tuning.md @@ -0,0 +1,63 @@ +# Hyperparameter Tuning + +## Achieving Peak Throughput + +Achieving a large batch size is the most important thing for attaining high throughput. + +When the server is running at full load, look for the following in the log: + +```Decode batch. #running-req: 233, #token: 370959, token usage: 0.82, gen throughput (token/s): 4594.01, #queue-req: 317``` + +### Tune Your Request Submission Speed + +`#queue-req` indicates the number of requests in the queue. If you frequently see `#queue-req == 0`, it suggests you are bottlenecked by the request submission speed. + +A healthy range for `#queue-req` is `50 - 500`. + +On the other hand, do not make `#queue-req` too large because it will also increase the scheduling overhead on the server, especially when using the default longest-prefix-match schedule policy (`--schedule-policy lpm`). + +### Tune `--schedule-conservativeness` + +`token usage` indicates the KV cache memory utilization of the server. `token usage > 0.9` means good utilization. +If you frequently see `token usage < 0.9` and `#queue-req > 0`, it means the server is too conservative about taking in new requests. You can decrease `--schedule-conservativeness` to a value like 0.3. +The case of server being too conservative can happen when users send many requests with a large `max_new_tokens` but the requests stop very early due to EOS or stop strings. + +On the other hand, if you see `token usage` very high and you frequently see warnings like +`decode out of memory happened, #retracted_reqs: 1, #new_token_ratio: 0.9998 -> 1.0000`, you can increase `--schedule-conservativeness` to a value like 1.3. +If you see `decode out of memory happened` occasionally but not frequently, it is okay. + +### Tune `--dp-size` and `--tp-size` + +Data parallelism is better for throughput. When there is enough GPU memory, always favor data parallelism for throughput. Refer to [sglang router](../router/router.md) for a better data parallelism rather than using `dp_size` parameter. + +## Avoid out-of-memory by Tuning `--chunked-prefill-size`, `--mem-fraction-static`, `--max-running-requests` + +If you see out of memory (OOM) errors, you can try to tune the following parameters. + +- If OOM happens during prefill, try to decrease `--chunked-prefill-size` to `4096` or `2048`. +- If OOM happens during decoding, try to decrease `--max-running-requests`. +- You can also try to decrease `--mem-fraction-static`, which reduces the memory usage of the KV cache memory pool and helps both prefill and decoding. + +## Enabling cache for `torch.compile` + +To enable `torch.compile` acceleration, add `--enable-torch-compile`. It accelerates small models on small batch sizes. By default, `torch.compile` will automatically cache the FX graph and Triton in `/tmp/torchinductor_root`, which might be cleared according to the [system policy](https://serverfault.com/questions/377348/when-does-tmp-get-cleared). You can export the environment variable `TORCHINDUCTOR_CACHE_DIR` to save compilation cache in your desired directory to avoid unwanted deletion. You can also share the cache with other machines to reduce the compilation time. + +SGLang uses `max-autotune-no-cudagraphs` mode of `torch.compile`. The auto-tuning can be slow. +If you want to deploy a model on many different machines, you can ship the `torch.compile` cache to these machines and skip the compilation steps. This is based on [PyTorch official documentation](https://pytorch.org/tutorials/recipes/torch_compile_caching_tutorial.html). + +*Examples*： + +1. Generate the cache by setting `TORCHINDUCTOR_CACHE_DIR` and running the model once. + + ```bash + TORCHINDUCTOR_CACHE_DIR=/root/inductor_root_cache python3 -m sglang.launch_server --model meta-llama/Llama-3.1-8B-Instruct --enable-torch-compile + ``` + +2. Copy the cache folder to other machines and launch the server with `TORCHINDUCTOR_CACHE_DIR`. + +## Tune `--schedule-policy` + +If the workload has many shared prefixes, use the default `--schedule-policy lpm`. Where `lpm` stands for longest prefix match. + +When you have no shared prefixes at all or you always send the requests with the shared prefixes together, +you can try `--schedule-policy fcfs`. Where `fcfs` stands for first come first serve. This policy has a lower scheduling overhead. diff --git a/sglang/docs/backend/native_api.ipynb b/sglang/docs/backend/native_api.ipynb new file mode 100644 index 0000000..72b65c6 --- /dev/null +++ b/sglang/docs/backend/native_api.ipynb @@ -0,0 +1,515 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# SGLang Native APIs\n", + "\n", + "Apart from the OpenAI compatible APIs, the SGLang Runtime also provides its native server APIs. We introduce these following APIs:\n", + "\n", + "- `/generate` (text generation model)\n", + "- `/get_model_info`\n", + "- `/get_server_info`\n", + "- `/health`\n", + "- `/health_generate`\n", + "- `/flush_cache`\n", + "- `/update_weights`\n", + "- `/encode`(embedding model)\n", + "- `/classify`(reward model)\n", + "- `/start_expert_distribution_record`\n", + "- `/stop_expert_distribution_record`\n", + "- `/dump_expert_distribution_record`\n", + "\n", + "We mainly use `requests` to test these APIs in the following examples. You can also use `curl`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Launch A Server" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "from sglang.test.test_utils import is_in_ci\n", + "\n", + "if is_in_ci():\n", + " from patch import launch_server_cmd\n", + "else:\n", + " from sglang.utils import launch_server_cmd\n", + "\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"python -m sglang.launch_server --model-path meta-llama/Llama-3.2-1B-Instruct --host 0.0.0.0\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Generate (text generation model)\n", + "Generate completions. This is similar to the `/v1/completions` in OpenAI API. Detailed parameters can be found in the [sampling parameters](./sampling_params.md)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "url = f\"http://localhost:{port}/generate\"\n", + "data = {\"text\": \"What is the capital of France?\"}\n", + "\n", + "response = requests.post(url, json=data)\n", + "print_highlight(response.json())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Get Model Info\n", + "\n", + "Get the information of the model.\n", + "\n", + "- `model_path`: The path/name of the model.\n", + "- `is_generation`: Whether the model is used as generation model or embedding model.\n", + "- `tokenizer_path`: The path/name of the tokenizer." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "url = f\"http://localhost:{port}/get_model_info\"\n", + "\n", + "response = requests.get(url)\n", + "response_json = response.json()\n", + "print_highlight(response_json)\n", + "assert response_json[\"model_path\"] == \"meta-llama/Llama-3.2-1B-Instruct\"\n", + "assert response_json[\"is_generation\"] is True\n", + "assert response_json[\"tokenizer_path\"] == \"meta-llama/Llama-3.2-1B-Instruct\"\n", + "assert response_json.keys() == {\"model_path\", \"is_generation\", \"tokenizer_path\"}" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Get Server Info\n", + "Gets the server information including CLI arguments, token limits, and memory pool sizes.\n", + "- Note: `get_server_info` merges the following deprecated endpoints:\n", + " - `get_server_args`\n", + " - `get_memory_pool_size` \n", + " - `get_max_total_num_tokens`" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# get_server_info\n", + "\n", + "url = f\"http://localhost:{port}/get_server_info\"\n", + "\n", + "response = requests.get(url)\n", + "print_highlight(response.text)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Health Check\n", + "- `/health`: Check the health of the server.\n", + "- `/health_generate`: Check the health of the server by generating one token." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "url = f\"http://localhost:{port}/health_generate\"\n", + "\n", + "response = requests.get(url)\n", + "print_highlight(response.text)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "url = f\"http://localhost:{port}/health\"\n", + "\n", + "response = requests.get(url)\n", + "print_highlight(response.text)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Flush Cache\n", + "\n", + "Flush the radix cache. It will be automatically triggered when the model weights are updated by the `/update_weights` API." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# flush cache\n", + "\n", + "url = f\"http://localhost:{port}/flush_cache\"\n", + "\n", + "response = requests.post(url)\n", + "print_highlight(response.text)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Update Weights From Disk\n", + "\n", + "Update model weights from disk without restarting the server. Only applicable for models with the same architecture and parameter size.\n", + "\n", + "SGLang support `update_weights_from_disk` API for continuous evaluation during training (save checkpoint to disk and update weights from disk).\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# successful update with same architecture and size\n", + "\n", + "url = f\"http://localhost:{port}/update_weights_from_disk\"\n", + "data = {\"model_path\": \"meta-llama/Llama-3.2-1B\"}\n", + "\n", + "response = requests.post(url, json=data)\n", + "print_highlight(response.text)\n", + "assert response.json()[\"success\"] is True\n", + "assert response.json()[\"message\"] == \"Succeeded to update model weights.\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# failed update with different parameter size or wrong name\n", + "\n", + "url = f\"http://localhost:{port}/update_weights_from_disk\"\n", + "data = {\"model_path\": \"meta-llama/Llama-3.2-1B-wrong\"}\n", + "\n", + "response = requests.post(url, json=data)\n", + "response_json = response.json()\n", + "print_highlight(response_json)\n", + "assert response_json[\"success\"] is False\n", + "assert response_json[\"message\"] == (\n", + " \"Failed to get weights iterator: \"\n", + " \"meta-llama/Llama-3.2-1B-wrong\"\n", + " \" (repository not found).\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Encode (embedding model)\n", + "\n", + "Encode text into embeddings. Note that this API is only available for [embedding models](openai_api_embeddings.html#openai-apis-embedding) and will raise an error for generation models.\n", + "Therefore, we launch a new server to server an embedding model." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)\n", + "\n", + "embedding_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-7B-instruct \\\n", + " --host 0.0.0.0 --is-embedding\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# successful encode for embedding model\n", + "\n", + "url = f\"http://localhost:{port}/encode\"\n", + "data = {\"model\": \"Alibaba-NLP/gte-Qwen2-7B-instruct\", \"text\": \"Once upon a time\"}\n", + "\n", + "response = requests.post(url, json=data)\n", + "response_json = response.json()\n", + "print_highlight(f\"Text embedding (first 10): {response_json['embedding'][:10]}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(embedding_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Classify (reward model)\n", + "\n", + "SGLang Runtime also supports reward models. Here we use a reward model to classify the quality of pairwise generations." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(embedding_process)\n", + "\n", + "# Note that SGLang now treats embedding models and reward models as the same type of models.\n", + "# This will be updated in the future.\n", + "\n", + "reward_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python -m sglang.launch_server --model-path Skywork/Skywork-Reward-Llama-3.1-8B-v0.2 --host 0.0.0.0 --is-embedding\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from transformers import AutoTokenizer\n", + "\n", + "PROMPT = (\n", + " \"What is the range of the numeric output of a sigmoid node in a neural network?\"\n", + ")\n", + "\n", + "RESPONSE1 = \"The output of a sigmoid node is bounded between -1 and 1.\"\n", + "RESPONSE2 = \"The output of a sigmoid node is bounded between 0 and 1.\"\n", + "\n", + "CONVS = [\n", + " [{\"role\": \"user\", \"content\": PROMPT}, {\"role\": \"assistant\", \"content\": RESPONSE1}],\n", + " [{\"role\": \"user\", \"content\": PROMPT}, {\"role\": \"assistant\", \"content\": RESPONSE2}],\n", + "]\n", + "\n", + "tokenizer = AutoTokenizer.from_pretrained(\"Skywork/Skywork-Reward-Llama-3.1-8B-v0.2\")\n", + "prompts = tokenizer.apply_chat_template(CONVS, tokenize=False)\n", + "\n", + "url = f\"http://localhost:{port}/classify\"\n", + "data = {\"model\": \"Skywork/Skywork-Reward-Llama-3.1-8B-v0.2\", \"text\": prompts}\n", + "\n", + "responses = requests.post(url, json=data).json()\n", + "for response in responses:\n", + " print_highlight(f\"reward: {response['embedding'][0]}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(reward_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Capture expert selection distribution in MoE models\n", + "\n", + "SGLang Runtime supports recording the number of times an expert is selected in a MoE model run for each expert in the model. This is useful when analyzing the throughput of the model and plan for optimization." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "expert_record_server_process, port = launch_server_cmd(\n", + " \"python -m sglang.launch_server --model-path Qwen/Qwen1.5-MoE-A2.7B --host 0.0.0.0\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = requests.post(f\"http://localhost:{port}/start_expert_distribution_record\")\n", + "print_highlight(response)\n", + "\n", + "url = f\"http://localhost:{port}/generate\"\n", + "data = {\"text\": \"What is the capital of France?\"}\n", + "\n", + "response = requests.post(url, json=data)\n", + "print_highlight(response.json())\n", + "\n", + "response = requests.post(f\"http://localhost:{port}/stop_expert_distribution_record\")\n", + "print_highlight(response)\n", + "\n", + "response = requests.post(f\"http://localhost:{port}/dump_expert_distribution_record\")\n", + "print_highlight(response)\n", + "\n", + "import glob\n", + "\n", + "output_file = glob.glob(\"expert_distribution_*.csv\")[0]\n", + "with open(output_file, \"r\") as f:\n", + " print_highlight(\"Content of dumped record:\")\n", + " for line in f:\n", + " print_highlight(line.strip())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(expert_record_server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Skip Tokenizer and Detokenizer\n", + "\n", + "SGLang Runtime also supports skip tokenizer and detokenizer. This is useful in cases like integrating with RLHF workflow." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "tokenizer_free_server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model-path meta-llama/Llama-3.2-1B-Instruct --skip-tokenizer-init\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from transformers import AutoTokenizer\n", + "\n", + "tokenizer = AutoTokenizer.from_pretrained(\"meta-llama/Llama-3.2-1B-Instruct\")\n", + "\n", + "input_text = \"What is the capital of France?\"\n", + "\n", + "input_tokens = tokenizer.encode(input_text)\n", + "print_highlight(f\"Input Text: {input_text}\")\n", + "print_highlight(f\"Tokenized Input: {input_tokens}\")\n", + "\n", + "response = requests.post(\n", + " f\"http://localhost:{port}/generate\",\n", + " json={\n", + " \"input_ids\": input_tokens,\n", + " \"sampling_params\": {\n", + " \"temperature\": 0,\n", + " \"max_new_tokens\": 256,\n", + " \"stop_token_ids\": [tokenizer.eos_token_id],\n", + " },\n", + " \"stream\": False,\n", + " },\n", + ")\n", + "output = response.json()\n", + "output_tokens = output[\"output_ids\"]\n", + "\n", + "output_text = tokenizer.decode(output_tokens, skip_special_tokens=False)\n", + "print_highlight(f\"Tokenized Output: {output_tokens}\")\n", + "print_highlight(f\"Decoded Output: {output_text}\")\n", + "print_highlight(f\"Output Text: {output['meta_info']['finish_reason']}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(tokenizer_free_server_process)" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/sglang/docs/backend/offline_engine_api.ipynb b/sglang/docs/backend/offline_engine_api.ipynb new file mode 100644 index 0000000..22ec2e5 --- /dev/null +++ b/sglang/docs/backend/offline_engine_api.ipynb @@ -0,0 +1,231 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Offline Engine API\n", + "\n", + "SGLang provides a direct inference engine without the need for an HTTP server, especially for use cases where additional HTTP server adds unnecessary complexity or overhead. Here are two general use cases:\n", + "\n", + "- Offline Batch Inference\n", + "- Custom Server on Top of the Engine\n", + "\n", + "This document focuses on the offline batch inference, demonstrating four different inference modes:\n", + "\n", + "- Non-streaming synchronous generation\n", + "- Streaming synchronous generation\n", + "- Non-streaming asynchronous generation\n", + "- Streaming asynchronous generation\n", + "\n", + "Additionally, you can easily build a custom server on top of the SGLang offline engine. A detailed example working in a python script can be found in [custom_server](https://github.com/sgl-project/sglang/blob/main/examples/runtime/engine/custom_server.py).\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Advanced Usage\n", + "\n", + "The engine supports [vlm inference](https://github.com/sgl-project/sglang/blob/main/examples/runtime/engine/offline_batch_inference_vlm.py) as well as [extracting hidden states](https://github.com/sgl-project/sglang/blob/main/examples/runtime/hidden_states). \n", + "\n", + "Please see [the examples](https://github.com/sgl-project/sglang/tree/main/examples/runtime/engine) for further use cases." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Offline Batch Inference\n", + "\n", + "SGLang offline engine supports batch inference with efficient scheduling." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# launch the offline engine\n", + "import asyncio\n", + "import io\n", + "import os\n", + "\n", + "from PIL import Image\n", + "import requests\n", + "import sglang as sgl\n", + "\n", + "from sglang.srt.conversation import chat_templates\n", + "from sglang.test.test_utils import is_in_ci\n", + "from sglang.utils import async_stream_and_merge, stream_and_merge\n", + "\n", + "if is_in_ci():\n", + " import patch\n", + "\n", + "\n", + "llm = sgl.Engine(model_path=\"meta-llama/Meta-Llama-3.1-8B-Instruct\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Non-streaming Synchronous Generation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prompts = [\n", + " \"Hello, my name is\",\n", + " \"The president of the United States is\",\n", + " \"The capital of France is\",\n", + " \"The future of AI is\",\n", + "]\n", + "\n", + "sampling_params = {\"temperature\": 0.8, \"top_p\": 0.95}\n", + "\n", + "outputs = llm.generate(prompts, sampling_params)\n", + "for prompt, output in zip(prompts, outputs):\n", + " print(\"===============================\")\n", + " print(f\"Prompt: {prompt}\\nGenerated text: {output['text']}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Streaming Synchronous Generation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prompts = [\n", + " \"Write a short, neutral self-introduction for a fictional character. Hello, my name is\",\n", + " \"Provide a concise factual statement about France’s capital city. The capital of France is\",\n", + " \"Explain possible future trends in artificial intelligence. The future of AI is\",\n", + "]\n", + "\n", + "sampling_params = {\n", + " \"temperature\": 0.2,\n", + " \"top_p\": 0.9,\n", + "}\n", + "\n", + "print(\"\\n=== Testing synchronous streaming generation with overlap removal ===\\n\")\n", + "\n", + "for prompt in prompts:\n", + " print(f\"Prompt: {prompt}\")\n", + " merged_output = stream_and_merge(llm, prompt, sampling_params)\n", + " print(\"Generated text:\", merged_output)\n", + " print()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Non-streaming Asynchronous Generation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prompts = [\n", + " \"Write a short, neutral self-introduction for a fictional character. Hello, my name is\",\n", + " \"Provide a concise factual statement about France’s capital city. The capital of France is\",\n", + " \"Explain possible future trends in artificial intelligence. The future of AI is\",\n", + "]\n", + "\n", + "sampling_params = {\"temperature\": 0.8, \"top_p\": 0.95}\n", + "\n", + "print(\"\\n=== Testing asynchronous batch generation ===\")\n", + "\n", + "\n", + "async def main():\n", + " outputs = await llm.async_generate(prompts, sampling_params)\n", + "\n", + " for prompt, output in zip(prompts, outputs):\n", + " print(f\"\\nPrompt: {prompt}\")\n", + " print(f\"Generated text: {output['text']}\")\n", + "\n", + "\n", + "asyncio.run(main())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Streaming Asynchronous Generation" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prompts = [\n", + " \"Write a short, neutral self-introduction for a fictional character. Hello, my name is\",\n", + " \"Provide a concise factual statement about France’s capital city. The capital of France is\",\n", + " \"Explain possible future trends in artificial intelligence. The future of AI is\",\n", + "]\n", + "\n", + "sampling_params = {\"temperature\": 0.8, \"top_p\": 0.95}\n", + "\n", + "print(\"\\n=== Testing asynchronous streaming generation (no repeats) ===\")\n", + "\n", + "\n", + "async def main():\n", + " for prompt in prompts:\n", + " print(f\"\\nPrompt: {prompt}\")\n", + " print(\"Generated text: \", end=\"\", flush=True)\n", + "\n", + " # Replace direct calls to async_generate with our custom overlap-aware version\n", + " async for cleaned_chunk in async_stream_and_merge(llm, prompt, sampling_params):\n", + " print(cleaned_chunk, end=\"\", flush=True)\n", + "\n", + " print() # New line after each prompt\n", + "\n", + "\n", + "asyncio.run(main())" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "llm.shutdown()" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/sglang/docs/backend/openai_api_completions.ipynb b/sglang/docs/backend/openai_api_completions.ipynb new file mode 100644 index 0000000..b874283 --- /dev/null +++ b/sglang/docs/backend/openai_api_completions.ipynb @@ -0,0 +1,534 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# OpenAI APIs - Completions\n", + "\n", + "SGLang provides OpenAI-compatible APIs to enable a smooth transition from OpenAI services to self-hosted local models.\n", + "A complete reference for the API is available in the [OpenAI API Reference](https://platform.openai.com/docs/api-reference).\n", + "\n", + "This tutorial covers the following popular APIs:\n", + "\n", + "- `chat/completions`\n", + "- `completions`\n", + "- `batches`\n", + "\n", + "Check out other tutorials to learn about [vision APIs](https://docs.sglang.ai/backend/openai_api_vision.html) for vision-language models and [embedding APIs](https://docs.sglang.ai/backend/openai_api_embeddings.html) for embedding models." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Launch A Server\n", + "\n", + "Launch the server in your terminal and wait for it to initialize." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sglang.test.test_utils import is_in_ci\n", + "\n", + "if is_in_ci():\n", + " from patch import launch_server_cmd\n", + "else:\n", + " from sglang.utils import launch_server_cmd\n", + "\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --host 0.0.0.0\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")\n", + "print(f\"Server started on http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Chat Completions\n", + "\n", + "### Usage\n", + "\n", + "The server fully implements the OpenAI API.\n", + "It will automatically apply the chat template specified in the Hugging Face tokenizer, if one is available.\n", + "You can also specify a custom chat template with `--chat-template` when launching the server." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=64,\n", + ")\n", + "\n", + "print_highlight(f\"Response: {response}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Parameters\n", + "\n", + "The chat completions API accepts OpenAI Chat Completions API's parameters. Refer to [OpenAI Chat Completions API](https://platform.openai.com/docs/api-reference/chat/create) for more details.\n", + "\n", + "Here is an example of a detailed chat completion request:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\n", + " \"role\": \"system\",\n", + " \"content\": \"You are a knowledgeable historian who provides concise responses.\",\n", + " },\n", + " {\"role\": \"user\", \"content\": \"Tell me about ancient Rome\"},\n", + " {\n", + " \"role\": \"assistant\",\n", + " \"content\": \"Ancient Rome was a civilization centered in Italy.\",\n", + " },\n", + " {\"role\": \"user\", \"content\": \"What were their major achievements?\"},\n", + " ],\n", + " temperature=0.3, # Lower temperature for more focused responses\n", + " max_tokens=128, # Reasonable length for a concise response\n", + " top_p=0.95, # Slightly higher for better fluency\n", + " presence_penalty=0.2, # Mild penalty to avoid repetition\n", + " frequency_penalty=0.2, # Mild penalty for more natural language\n", + " n=1, # Single response is usually more stable\n", + " seed=42, # Keep for reproducibility\n", + ")\n", + "\n", + "print_highlight(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Streaming mode is also supported." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "stream = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[{\"role\": \"user\", \"content\": \"Say this is a test\"}],\n", + " stream=True,\n", + ")\n", + "for chunk in stream:\n", + " if chunk.choices[0].delta.content is not None:\n", + " print(chunk.choices[0].delta.content, end=\"\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Completions\n", + "\n", + "### Usage\n", + "Completions API is similar to Chat Completions API, but without the `messages` parameter or chat templates." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = client.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " prompt=\"List 3 countries and their capitals.\",\n", + " temperature=0,\n", + " max_tokens=64,\n", + " n=1,\n", + " stop=None,\n", + ")\n", + "\n", + "print_highlight(f\"Response: {response}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Parameters\n", + "\n", + "The completions API accepts OpenAI Completions API's parameters. Refer to [OpenAI Completions API](https://platform.openai.com/docs/api-reference/completions/create) for more details.\n", + "\n", + "Here is an example of a detailed completions request:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = client.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " prompt=\"Write a short story about a space explorer.\",\n", + " temperature=0.7, # Moderate temperature for creative writing\n", + " max_tokens=150, # Longer response for a story\n", + " top_p=0.9, # Balanced diversity in word choice\n", + " stop=[\"\\n\\n\", \"THE END\"], # Multiple stop sequences\n", + " presence_penalty=0.3, # Encourage novel elements\n", + " frequency_penalty=0.3, # Reduce repetitive phrases\n", + " n=1, # Generate one completion\n", + " seed=123, # For reproducible results\n", + ")\n", + "\n", + "print_highlight(f\"Response: {response}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Structured Outputs (JSON, Regex, EBNF)\n", + "\n", + "For OpenAI compatible structed outputs API, refer to [Structured Outputs](https://docs.sglang.ai/backend/structured_outputs.html#OpenAI-Compatible-API) for more details.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Batches\n", + "\n", + "Batches API for chat completions and completions are also supported. You can upload your requests in `jsonl` files, create a batch job, and retrieve the results when the batch job is completed (which takes longer but costs less).\n", + "\n", + "The batches APIs are:\n", + "\n", + "- `batches`\n", + "- `batches/{batch_id}/cancel`\n", + "- `batches/{batch_id}`\n", + "\n", + "Here is an example of a batch job for chat completions, completions are similar.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import time\n", + "from openai import OpenAI\n", + "\n", + "client = OpenAI(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "requests = [\n", + " {\n", + " \"custom_id\": \"request-1\",\n", + " \"method\": \"POST\",\n", + " \"url\": \"/chat/completions\",\n", + " \"body\": {\n", + " \"model\": \"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " \"messages\": [\n", + " {\"role\": \"user\", \"content\": \"Tell me a joke about programming\"}\n", + " ],\n", + " \"max_tokens\": 50,\n", + " },\n", + " },\n", + " {\n", + " \"custom_id\": \"request-2\",\n", + " \"method\": \"POST\",\n", + " \"url\": \"/chat/completions\",\n", + " \"body\": {\n", + " \"model\": \"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " \"messages\": [{\"role\": \"user\", \"content\": \"What is Python?\"}],\n", + " \"max_tokens\": 50,\n", + " },\n", + " },\n", + "]\n", + "\n", + "input_file_path = \"batch_requests.jsonl\"\n", + "\n", + "with open(input_file_path, \"w\") as f:\n", + " for req in requests:\n", + " f.write(json.dumps(req) + \"\\n\")\n", + "\n", + "with open(input_file_path, \"rb\") as f:\n", + " file_response = client.files.create(file=f, purpose=\"batch\")\n", + "\n", + "batch_response = client.batches.create(\n", + " input_file_id=file_response.id,\n", + " endpoint=\"/v1/chat/completions\",\n", + " completion_window=\"24h\",\n", + ")\n", + "\n", + "print_highlight(f\"Batch job created with ID: {batch_response.id}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "while batch_response.status not in [\"completed\", \"failed\", \"cancelled\"]:\n", + " time.sleep(3)\n", + " print(f\"Batch job status: {batch_response.status}...trying again in 3 seconds...\")\n", + " batch_response = client.batches.retrieve(batch_response.id)\n", + "\n", + "if batch_response.status == \"completed\":\n", + " print(\"Batch job completed successfully!\")\n", + " print(f\"Request counts: {batch_response.request_counts}\")\n", + "\n", + " result_file_id = batch_response.output_file_id\n", + " file_response = client.files.content(result_file_id)\n", + " result_content = file_response.read().decode(\"utf-8\")\n", + "\n", + " results = [\n", + " json.loads(line) for line in result_content.split(\"\\n\") if line.strip() != \"\"\n", + " ]\n", + "\n", + " for result in results:\n", + " print_highlight(f\"Request {result['custom_id']}:\")\n", + " print_highlight(f\"Response: {result['response']}\")\n", + "\n", + " print_highlight(\"Cleaning up files...\")\n", + " # Only delete the result file ID since file_response is just content\n", + " client.files.delete(result_file_id)\n", + "else:\n", + " print_highlight(f\"Batch job failed with status: {batch_response.status}\")\n", + " if hasattr(batch_response, \"errors\"):\n", + " print_highlight(f\"Errors: {batch_response.errors}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "It takes a while to complete the batch job. You can use these two APIs to retrieve the batch job status or cancel the batch job.\n", + "\n", + "1. `batches/{batch_id}`: Retrieve the batch job status.\n", + "2. `batches/{batch_id}/cancel`: Cancel the batch job.\n", + "\n", + "Here is an example to check the batch job status." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import time\n", + "from openai import OpenAI\n", + "\n", + "client = OpenAI(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "requests = []\n", + "for i in range(20):\n", + " requests.append(\n", + " {\n", + " \"custom_id\": f\"request-{i}\",\n", + " \"method\": \"POST\",\n", + " \"url\": \"/chat/completions\",\n", + " \"body\": {\n", + " \"model\": \"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " \"messages\": [\n", + " {\n", + " \"role\": \"system\",\n", + " \"content\": f\"{i}: You are a helpful AI assistant\",\n", + " },\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"Write a detailed story about topic. Make it very long.\",\n", + " },\n", + " ],\n", + " \"max_tokens\": 64,\n", + " },\n", + " }\n", + " )\n", + "\n", + "input_file_path = \"batch_requests.jsonl\"\n", + "with open(input_file_path, \"w\") as f:\n", + " for req in requests:\n", + " f.write(json.dumps(req) + \"\\n\")\n", + "\n", + "with open(input_file_path, \"rb\") as f:\n", + " uploaded_file = client.files.create(file=f, purpose=\"batch\")\n", + "\n", + "batch_job = client.batches.create(\n", + " input_file_id=uploaded_file.id,\n", + " endpoint=\"/v1/chat/completions\",\n", + " completion_window=\"24h\",\n", + ")\n", + "\n", + "print_highlight(f\"Created batch job with ID: {batch_job.id}\")\n", + "print_highlight(f\"Initial status: {batch_job.status}\")\n", + "\n", + "time.sleep(10)\n", + "\n", + "max_checks = 5\n", + "for i in range(max_checks):\n", + " batch_details = client.batches.retrieve(batch_id=batch_job.id)\n", + "\n", + " print_highlight(\n", + " f\"Batch job details (check {i+1} / {max_checks}) // ID: {batch_details.id} // Status: {batch_details.status} // Created at: {batch_details.created_at} // Input file ID: {batch_details.input_file_id} // Output file ID: {batch_details.output_file_id}\"\n", + " )\n", + " print_highlight(\n", + " f\"Request counts: Total: {batch_details.request_counts.total} // Completed: {batch_details.request_counts.completed} // Failed: {batch_details.request_counts.failed}\"\n", + " )\n", + "\n", + " time.sleep(3)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Here is an example to cancel a batch job." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import time\n", + "from openai import OpenAI\n", + "import os\n", + "\n", + "client = OpenAI(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "requests = []\n", + "for i in range(5000):\n", + " requests.append(\n", + " {\n", + " \"custom_id\": f\"request-{i}\",\n", + " \"method\": \"POST\",\n", + " \"url\": \"/chat/completions\",\n", + " \"body\": {\n", + " \"model\": \"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " \"messages\": [\n", + " {\n", + " \"role\": \"system\",\n", + " \"content\": f\"{i}: You are a helpful AI assistant\",\n", + " },\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"Write a detailed story about topic. Make it very long.\",\n", + " },\n", + " ],\n", + " \"max_tokens\": 128,\n", + " },\n", + " }\n", + " )\n", + "\n", + "input_file_path = \"batch_requests.jsonl\"\n", + "with open(input_file_path, \"w\") as f:\n", + " for req in requests:\n", + " f.write(json.dumps(req) + \"\\n\")\n", + "\n", + "with open(input_file_path, \"rb\") as f:\n", + " uploaded_file = client.files.create(file=f, purpose=\"batch\")\n", + "\n", + "batch_job = client.batches.create(\n", + " input_file_id=uploaded_file.id,\n", + " endpoint=\"/v1/chat/completions\",\n", + " completion_window=\"24h\",\n", + ")\n", + "\n", + "print_highlight(f\"Created batch job with ID: {batch_job.id}\")\n", + "print_highlight(f\"Initial status: {batch_job.status}\")\n", + "\n", + "time.sleep(10)\n", + "\n", + "try:\n", + " cancelled_job = client.batches.cancel(batch_id=batch_job.id)\n", + " print_highlight(f\"Cancellation initiated. Status: {cancelled_job.status}\")\n", + " assert cancelled_job.status == \"cancelling\"\n", + "\n", + " # Monitor the cancellation process\n", + " while cancelled_job.status not in [\"failed\", \"cancelled\"]:\n", + " time.sleep(3)\n", + " cancelled_job = client.batches.retrieve(batch_job.id)\n", + " print_highlight(f\"Current status: {cancelled_job.status}\")\n", + "\n", + " # Verify final status\n", + " assert cancelled_job.status == \"cancelled\"\n", + " print_highlight(\"Batch job successfully cancelled\")\n", + "\n", + "except Exception as e:\n", + " print_highlight(f\"Error during cancellation: {e}\")\n", + " raise e\n", + "\n", + "finally:\n", + " try:\n", + " del_response = client.files.delete(uploaded_file.id)\n", + " if del_response.deleted:\n", + " print_highlight(\"Successfully cleaned up input file\")\n", + " if os.path.exists(input_file_path):\n", + " os.remove(input_file_path)\n", + " print_highlight(\"Successfully deleted local batch_requests.jsonl file\")\n", + " except Exception as e:\n", + " print_highlight(f\"Error cleaning up: {e}\")\n", + " raise e" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/sglang/docs/backend/openai_api_embeddings.ipynb b/sglang/docs/backend/openai_api_embeddings.ipynb new file mode 100644 index 0000000..38543fa --- /dev/null +++ b/sglang/docs/backend/openai_api_embeddings.ipynb @@ -0,0 +1,191 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# OpenAI APIs - Embedding\n", + "\n", + "SGLang provides OpenAI-compatible APIs to enable a smooth transition from OpenAI services to self-hosted local models.\n", + "A complete reference for the API is available in the [OpenAI API Reference](https://platform.openai.com/docs/guides/embeddings).\n", + "\n", + "This tutorial covers the embedding APIs for embedding models, such as \n", + "- [intfloat/e5-mistral-7b-instruct](https://huggingface.co/intfloat/e5-mistral-7b-instruct) \n", + "- [Alibaba-NLP/gte-Qwen2-7B-instruct](https://huggingface.co/Alibaba-NLP/gte-Qwen2-7B-instruct) \n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Launch A Server\n", + "\n", + "Launch the server in your terminal and wait for it to initialize. Remember to add `--is-embedding` to the command." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sglang.test.test_utils import is_in_ci\n", + "\n", + "if is_in_ci():\n", + " from patch import launch_server_cmd\n", + "else:\n", + " from sglang.utils import launch_server_cmd\n", + "\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "embedding_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-7B-instruct \\\n", + " --host 0.0.0.0 --is-embedding\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using cURL" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import subprocess, json\n", + "\n", + "text = \"Once upon a time\"\n", + "\n", + "curl_text = f\"\"\"curl -s http://localhost:{port}/v1/embeddings \\\n", + " -d '{{\"model\": \"Alibaba-NLP/gte-Qwen2-7B-instruct\", \"input\": \"{text}\"}}'\"\"\"\n", + "\n", + "text_embedding = json.loads(subprocess.check_output(curl_text, shell=True))[\"data\"][0][\n", + " \"embedding\"\n", + "]\n", + "\n", + "print_highlight(f\"Text embedding (first 10): {text_embedding[:10]}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Python Requests" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "\n", + "text = \"Once upon a time\"\n", + "\n", + "response = requests.post(\n", + " f\"http://localhost:{port}/v1/embeddings\",\n", + " json={\"model\": \"Alibaba-NLP/gte-Qwen2-7B-instruct\", \"input\": text},\n", + ")\n", + "\n", + "text_embedding = response.json()[\"data\"][0][\"embedding\"]\n", + "\n", + "print_highlight(f\"Text embedding (first 10): {text_embedding[:10]}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using OpenAI Python Client" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "# Text embedding example\n", + "response = client.embeddings.create(\n", + " model=\"Alibaba-NLP/gte-Qwen2-7B-instruct\",\n", + " input=text,\n", + ")\n", + "\n", + "embedding = response.data[0].embedding[:10]\n", + "print_highlight(f\"Text embedding (first 10): {embedding}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Input IDs\n", + "\n", + "SGLang also supports `input_ids` as input to get the embedding." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import os\n", + "from transformers import AutoTokenizer\n", + "\n", + "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n", + "\n", + "tokenizer = AutoTokenizer.from_pretrained(\"Alibaba-NLP/gte-Qwen2-7B-instruct\")\n", + "input_ids = tokenizer.encode(text)\n", + "\n", + "curl_ids = f\"\"\"curl -s http://localhost:{port}/v1/embeddings \\\n", + " -d '{{\"model\": \"Alibaba-NLP/gte-Qwen2-7B-instruct\", \"input\": {json.dumps(input_ids)}}}'\"\"\"\n", + "\n", + "input_ids_embedding = json.loads(subprocess.check_output(curl_ids, shell=True))[\"data\"][\n", + " 0\n", + "][\"embedding\"]\n", + "\n", + "print_highlight(f\"Input IDs embedding (first 10): {input_ids_embedding[:10]}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(embedding_process)" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/sglang/docs/backend/openai_api_vision.ipynb b/sglang/docs/backend/openai_api_vision.ipynb new file mode 100644 index 0000000..e4fb768 --- /dev/null +++ b/sglang/docs/backend/openai_api_vision.ipynb @@ -0,0 +1,291 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# OpenAI APIs - Vision\n", + "\n", + "SGLang provides OpenAI-compatible APIs to enable a smooth transition from OpenAI services to self-hosted local models.\n", + "A complete reference for the API is available in the [OpenAI API Reference](https://platform.openai.com/docs/guides/vision).\n", + "This tutorial covers the vision APIs for vision language models.\n", + "\n", + "SGLang supports various vision language models such as Llama 3.2, LLaVA-OneVision, Qwen2.5-VL, Gemma3 and [more](https://docs.sglang.ai/references/supported_models): \n", + "- [meta-llama/Llama-3.2-11B-Vision-Instruct](https://huggingface.co/meta-llama/Llama-3.2-11B-Vision-Instruct) \n", + "- [lmms-lab/llava-onevision-qwen2-72b-ov-chat](https://huggingface.co/lmms-lab/llava-onevision-qwen2-72b-ov-chat) \n", + "- [Qwen/Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct)\n", + "- [google/gemma-3-4b-it](https://huggingface.co/google/gemma-3-4b-it)\n", + "- [openbmb/MiniCPM-V](https://huggingface.co/openbmb/MiniCPM-V)\n", + "- [deepseek-ai/deepseek-vl2](https://huggingface.co/deepseek-ai/deepseek-vl2)\n", + "\n", + "As an alternative to the OpenAI API, you can also use the [SGLang offline engine](https://github.com/sgl-project/sglang/blob/main/examples/runtime/engine/offline_batch_inference_vlm.py)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Launch A Server\n", + "\n", + "Launch the server in your terminal and wait for it to initialize.\n", + "\n", + "**Remember to add** `--chat-template llama_3_vision` **to specify the [vision chat template](https://docs.sglang.ai/backend/openai_api_vision.html#Chat-Template), otherwise, the server will only support text (images won’t be passed in), which can lead to degraded performance.**\n", + "\n", + "We need to specify `--chat-template` for vision language models because the chat template provided in Hugging Face tokenizer only supports text." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sglang.test.test_utils import is_in_ci\n", + "\n", + "if is_in_ci():\n", + " from patch import launch_server_cmd\n", + "else:\n", + " from sglang.utils import launch_server_cmd\n", + "\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "vision_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model-path meta-llama/Llama-3.2-11B-Vision-Instruct \\\n", + " --chat-template=llama_3_vision\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using cURL\n", + "\n", + "Once the server is up, you can send test requests using curl or requests." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import subprocess\n", + "\n", + "curl_command = f\"\"\"\n", + "curl -s http://localhost:{port}/v1/chat/completions \\\\\n", + " -d '{{\n", + " \"model\": \"meta-llama/Llama-3.2-11B-Vision-Instruct\",\n", + " \"messages\": [\n", + " {{\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {{\n", + " \"type\": \"text\",\n", + " \"text\": \"What’s in this image?\"\n", + " }},\n", + " {{\n", + " \"type\": \"image_url\",\n", + " \"image_url\": {{\n", + " \"url\": \"https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true\"\n", + " }}\n", + " }}\n", + " ]\n", + " }}\n", + " ],\n", + " \"max_tokens\": 300\n", + " }}'\n", + "\"\"\"\n", + "\n", + "response = subprocess.check_output(curl_command, shell=True).decode()\n", + "print_highlight(response)\n", + "\n", + "\n", + "response = subprocess.check_output(curl_command, shell=True).decode()\n", + "print_highlight(response)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Python Requests" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "\n", + "url = f\"http://localhost:{port}/v1/chat/completions\"\n", + "\n", + "data = {\n", + " \"model\": \"meta-llama/Llama-3.2-11B-Vision-Instruct\",\n", + " \"messages\": [\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\"type\": \"text\", \"text\": \"What’s in this image?\"},\n", + " {\n", + " \"type\": \"image_url\",\n", + " \"image_url\": {\n", + " \"url\": \"https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true\"\n", + " },\n", + " },\n", + " ],\n", + " }\n", + " ],\n", + " \"max_tokens\": 300,\n", + "}\n", + "\n", + "response = requests.post(url, json=data)\n", + "print_highlight(response.text)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using OpenAI Python Client" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from openai import OpenAI\n", + "\n", + "client = OpenAI(base_url=f\"http://localhost:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Llama-3.2-11B-Vision-Instruct\",\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\n", + " \"type\": \"text\",\n", + " \"text\": \"What is in this image?\",\n", + " },\n", + " {\n", + " \"type\": \"image_url\",\n", + " \"image_url\": {\n", + " \"url\": \"https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true\"\n", + " },\n", + " },\n", + " ],\n", + " }\n", + " ],\n", + " max_tokens=300,\n", + ")\n", + "\n", + "print_highlight(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Multiple-Image Inputs\n", + "\n", + "The server also supports multiple images and interleaved text and images if the model supports it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from openai import OpenAI\n", + "\n", + "client = OpenAI(base_url=f\"http://localhost:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Llama-3.2-11B-Vision-Instruct\",\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": [\n", + " {\n", + " \"type\": \"image_url\",\n", + " \"image_url\": {\n", + " \"url\": \"https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true\",\n", + " },\n", + " },\n", + " {\n", + " \"type\": \"image_url\",\n", + " \"image_url\": {\n", + " \"url\": \"https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png\",\n", + " },\n", + " },\n", + " {\n", + " \"type\": \"text\",\n", + " \"text\": \"I have two very different images. They are not related at all. \"\n", + " \"Please describe the first image in one sentence, and then describe the second image in another sentence.\",\n", + " },\n", + " ],\n", + " }\n", + " ],\n", + " temperature=0,\n", + ")\n", + "\n", + "print_highlight(response.choices[0].message.content)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(vision_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Chat Template\n", + "\n", + "As mentioned before, if you do not specify a vision model's `--chat-template`, the server uses Hugging Face's default template, which only supports text.\n", + "\n", + "We list popular vision models with their chat templates:\n", + "\n", + "- [meta-llama/Llama-3.2-Vision](https://huggingface.co/meta-llama/Llama-3.2-11B-Vision-Instruct) uses `llama_3_vision`.\n", + "- [Qwen/Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) uses `qwen2-vl`.\n", + "- [google/gemma-3-4b-it](https://huggingface.co/google/gemma-3-4b-it) uses `gemma-it`.\n", + "- [openbmb/MiniCPM-V](https://huggingface.co/openbmb/MiniCPM-V) uses `minicpmv`.\n", + "- [deepseek-ai/deepseek-vl2](https://huggingface.co/deepseek-ai/deepseek-vl2) uses `deepseek-vl2`.\n", + "- [LlaVA-OneVision](https://huggingface.co/lmms-lab/llava-onevision-qwen2-7b-ov) uses `chatml-llava`.\n", + "- [LLaVA-NeXT](https://huggingface.co/collections/lmms-lab/llava-next-6623288e2d61edba3ddbf5ff) uses `chatml-llava`.\n", + "- [Llama3-LLaVA-NeXT](https://huggingface.co/lmms-lab/llama3-llava-next-8b) uses `llava_llama_3`.\n", + "- [LLaVA-v1.5 / 1.6](https://huggingface.co/liuhaotian/llava-v1.6-34b) uses `vicuna_v1.1`." + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/sglang/docs/backend/patch.py b/sglang/docs/backend/patch.py new file mode 100644 index 0000000..ce7ba37 --- /dev/null +++ b/sglang/docs/backend/patch.py @@ -0,0 +1,53 @@ +import weakref + +import nest_asyncio + +nest_asyncio.apply() + +from sglang.utils import execute_shell_command, reserve_port + +DEFAULT_MAX_RUNNING_REQUESTS = 200 +DEFAULT_MAX_TOTAL_TOKENS = 20480 + +import sglang.srt.server_args as server_args_mod + +_original_post_init = server_args_mod.ServerArgs.__post_init__ + + +def patched_post_init(self): + _original_post_init(self) + if self.max_running_requests is None: + self.max_running_requests = DEFAULT_MAX_RUNNING_REQUESTS + if self.max_total_tokens is None: + self.max_total_tokens = DEFAULT_MAX_TOTAL_TOKENS + self.disable_cuda_graph = True + + +server_args_mod.ServerArgs.__post_init__ = patched_post_init + +process_socket_map = weakref.WeakKeyDictionary() + + +def launch_server_cmd(command: str, host: str = "0.0.0.0", port: int = None): + """ + Launch the server using the given command. + If no port is specified, a free port is reserved. + """ + if port is None: + port, lock_socket = reserve_port(host) + else: + lock_socket = None + + extra_flags = ( + f"--max-running-requests {DEFAULT_MAX_RUNNING_REQUESTS} " + f"--max-total-tokens {DEFAULT_MAX_TOTAL_TOKENS} " + f"--disable-cuda-graph" + ) + + full_command = f"{command} --port {port} {extra_flags}" + process = execute_shell_command(full_command) + + if lock_socket is not None: + process_socket_map[process] = lock_socket + + return process, port diff --git a/sglang/docs/backend/quantization.md b/sglang/docs/backend/quantization.md new file mode 100644 index 0000000..3a229f8 --- /dev/null +++ b/sglang/docs/backend/quantization.md @@ -0,0 +1,152 @@ +# Quantization + +SGLang supports various quantization methods, including offline quantization and online dynamic quantization. + +Offline quantization loads pre-quantized model weights directly during inference. This is required for quantization methods +such as GPTQ and AWQ, which collect and pre-compute various statistics from the original weights using the calibration dataset. + +Online quantization dynamically computes scaling parameters—such as the maximum/minimum values of model weights—during runtime. +Like NVIDIA FP8 training's [delayed scaling](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html#Mixed-precision-training-with-FP8) mechanism, online quantization calculates the appropriate scaling factors +on-the-fly to convert high-precision weights into a lower-precision format. + +**Note: For better performance, usability and convenience, offline quantization is recommended over online quantization.** + +If you use a pre-quantized model, do not add `--quantization` to enable online quantization at the same time. +For popular pre-quantized models, please visit [ModelCloud](https://huggingface.co/collections/ModelCloud/vortex-673743382af0a52b2a8b9fe2) +or [NeuralMagic](https://huggingface.co/collections/neuralmagic) collections on HF for some +popular quality validated quantized models. Quantized models must be validated via benchmarks post-quantization +to guard against abnormal quantization loss regressions. + +## Offline Quantization + +To load already quantized models, simply load the model weights and config. **Again, if the model has been quantized offline, +there's no need to add `--quantization` argument when starting the engine. The quantization method will be parsed from the +downloaded Hugging Face config. For example, DeepSeek V3/R1 models are already in FP8, so do not add redundant parameters.** + +```bash +python3 -m sglang.launch_server \ + --model-path hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4 \ + --port 30000 --host 0.0.0.0 +``` + +Take note, if your model is **per-channel quantized (INT8 or FP8) with per-token dynamic quantization activation**, you can opt to include `--quantization w8a8_int8` or `--quantization w8a8_fp8` to invoke the corresponding CUTLASS int8_kernel or fp8_kernel in sgl-kernel. This action will ignore the Hugging Face config's quantization settings. For instance, with `neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8-dynamic`, if you execute with `--quantization w8a8_fp8`, the system will use the `W8A8Fp8Config` from SGLang to invoke the sgl-kernel, rather than the `CompressedTensorsConfig` for vLLM kernels. + +```bash +python3 -m sglang.launch_server \ + --model-path neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8-dynamic \ + --quantization w8a8_fp8 \ + --port 30000 --host 0.0.0.0 +``` + +### Examples of Offline Model Quantization + +#### Using [GPTQModel](https://github.com/ModelCloud/GPTQModel) + +```bash +# install +pip install gptqmodel --no-build-isolation -v +``` + +```py +from datasets import load_dataset +from gptqmodel import GPTQModel, QuantizeConfig + +model_id = "meta-llama/Llama-3.2-1B-Instruct" +quant_path = "Llama-3.2-1B-Instruct-gptqmodel-4bit" + +calibration_dataset = load_dataset( + "allenai/c4", data_files="en/c4-train.00001-of-01024.json.gz", + split="train" + ).select(range(1024))["text"] + +quant_config = QuantizeConfig(bits=4, group_size=128) # quantization config +model = GPTQModel.load(model_id, quant_config) # load model + +model.quantize(calibration_dataset, batch_size=2) # quantize +model.save(quant_path) # save model +``` + +#### Using [LLM Compressor](https://github.com/vllm-project/llm-compressor/) + +```bash +# install +pip install llmcompressor +``` + +Here, we take quantize `meta-llama/Meta-Llama-3-8B-Instruct` to `FP8` as an example to elaborate on how to do offline quantization. + +```python +from transformers import AutoTokenizer +from llmcompressor.transformers import SparseAutoModelForCausalLM +from llmcompressor.transformers import oneshot +from llmcompressor.modifiers.quantization import QuantizationModifier + +# Step 1: Load the original model. +MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct" + +model = SparseAutoModelForCausalLM.from_pretrained( + MODEL_ID, device_map="auto", torch_dtype="auto") +tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) + +# Step 2: Perform offline quantization. +# Step 2.1: Configure the simple PTQ quantization. +recipe = QuantizationModifier( + targets="Linear", scheme="FP8_DYNAMIC", ignore=["lm_head"]) + +# Step 2.2: Apply the quantization algorithm. +oneshot(model=model, recipe=recipe) + +# Step 3: Save the model. +SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-Dynamic" +model.save_pretrained(SAVE_DIR) +tokenizer.save_pretrained(SAVE_DIR) +``` + +Then, you can directly use the quantized model with `SGLang`, by using the following command: + +```bash +python3 -m sglang.launch_server \ + --model-path $PWD/Meta-Llama-3-8B-Instruct-FP8-Dynamic \ + --port 30000 --host 0.0.0.0 +``` + +## Online Quantization + +To enable online quantization, you can simply specify `--quantization` in the command line. For example, you can launch the server with the following command to enable `FP8` quantization for model `meta-llama/Meta-Llama-3.1-8B-Instruct`: + +```bash +python3 -m sglang.launch_server \ + --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ + --quantization fp8 \ + --port 30000 --host 0.0.0.0 +``` + +Our team is working on supporting more online quantization methods. SGLang will soon support methods including but not limited to `["awq", "gptq", "marlin", "gptq_marlin", "awq_marlin", "bitsandbytes", "gguf"]`. + +SGLang also supports quantization methods based on [torchao](https://github.com/pytorch/ao). You can simply specify `--torchao-config` in the command line to support this feature. For example, if you want to enable `int4wo-128` for model `meta-llama/Meta-Llama-3.1-8B-Instruct`, you can launch the server with the following command: + +```bash +python3 -m sglang.launch_server \ + --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ + --torchao-config int4wo-128 \ + --port 30000 --host 0.0.0.0 +``` + +SGLang supports the following quantization methods based on torchao `["int8dq", "int8wo", "fp8wo", "fp8dq-per_tensor", "fp8dq-per_row", "int4wo-32", "int4wo-64", "int4wo-128", "int4wo-256"]`. + +Note: According to [this issue](https://github.com/sgl-project/sglang/issues/2219#issuecomment-2561890230), `"int8dq"` method currently has some bugs when using together with cuda graph capture. So we suggest to disable cuda graph capture when using `"int8dq"` method. Namely, please use the following command: + +```bash +python3 -m sglang.launch_server \ + --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ + --torchao-config int8dq \ + --disable-cuda-graph \ + --port 30000 --host 0.0.0.0 +``` + +## Reference + +- [GPTQModel](https://github.com/ModelCloud/GPTQModel) +- [LLM Compressor](https://github.com/vllm-project/llm-compressor/) +- [Torchao: PyTorch Architecture Optimization](https://github.com/pytorch/ao) +- [vLLM Quantization](https://docs.vllm.ai/en/latest/quantization/) diff --git a/sglang/docs/backend/sampling_params.md b/sglang/docs/backend/sampling_params.md new file mode 100644 index 0000000..4db0f86 --- /dev/null +++ b/sglang/docs/backend/sampling_params.md @@ -0,0 +1,282 @@ +# Sampling Parameters + +This doc describes the sampling parameters of the SGLang Runtime. It is the low-level endpoint of the runtime. + +If you want a high-level endpoint that can automatically handle chat templates, consider using the [OpenAI Compatible API](./openai_api_completions.ipynb). + +## `/generate` Endpoint + +The `/generate` endpoint accepts the following parameters in JSON format. For in detail usage see the [native api doc](./native_api.ipynb). + +* `text: Optional[Union[List[str], str]] = None` The input prompt. Can be a single prompt or a batch of prompts. +* `input_ids: Optional[Union[List[List[int]], List[int]]] = None` Alternative to `text`. Specify the input as token IDs instead of text. +* `sampling_params: Optional[Union[List[Dict], Dict]] = None` The sampling parameters as described in the sections below. +* `return_logprob: Optional[Union[List[bool], bool]] = None` Whether to return log probabilities for tokens. +* `logprob_start_len: Optional[Union[List[int], int]] = None` If returning log probabilities, specifies the start position in the prompt. Default is "-1" which returns logprobs only for output tokens. +* `top_logprobs_num: Optional[Union[List[int], int]] = None` If returning log probabilities, specifies the number of top logprobs to return at each position. +* `stream: bool = False` Whether to stream the output. +* `lora_path: Optional[Union[List[Optional[str]], Optional[str]]] = None` Path to LoRA weights. +* `custom_logit_processor: Optional[Union[List[Optional[str]], str]] = None` Custom logit processor for advanced sampling control. For usage see below. +* `return_hidden_states: bool = False` Whether to return hidden states of the model. Note that each time it changes, the cuda graph will be recaptured, which might lead to a performance hit. See the [examples](https://github.com/sgl-project/sglang/blob/main/examples/runtime/hidden_states) for more information. + +## Sampling params + +### Core Parameters + +* `max_new_tokens: int = 128` The maximum output length measured in tokens. +* `stop: Optional[Union[str, List[str]]] = None` One or multiple [stop words](https://platform.openai.com/docs/api-reference/chat/create#chat-create-stop). Generation will stop if one of these words is sampled. +* `stop_token_ids: Optional[List[int]] = None` Provide stop words in form of token ids. Generation will stop if one of these token ids is sampled. +* `temperature: float = 1.0` [Temperature](https://platform.openai.com/docs/api-reference/chat/create#chat-create-temperature) when sampling the next token. `temperature = 0` corresponds to greedy sampling, higher temperature leads to more diversity. +* `top_p: float = 1.0` [Top-p](https://platform.openai.com/docs/api-reference/chat/create#chat-create-top_p) selects tokens from the smallest sorted set whose cumulative probability exceeds `top_p`. When `top_p = 1`, this reduces to unrestricted sampling from all tokens. +* `top_k: int = -1` [Top-k](https://developer.nvidia.com/blog/how-to-get-better-outputs-from-your-large-language-model/#predictability_vs_creativity) randomly selects from the `k` highest-probability tokens. +* `min_p: float = 0.0` [Min-p](https://github.com/huggingface/transformers/issues/27670) samples from tokens with probability larger than `min_p * highest_token_probability`. + +### Penalizers + +* `frequency_penalty: float = 0.0`: Penalizes tokens based on their frequency in generation so far. Must be between `-2` and `2` where negative numbers encourage repeatment of tokens and positive number encourages sampling of new tokens. The scaling of penalization grows linearly with each appearance of a token. +* `presence_penalty: float = 0.0`: Penalizes tokens if they appeared in the generation so far. Must be between `-2` and `2` where negative numbers encourage repeatment of tokens and positive number encourages sampling of new tokens. The scaling of the penalization is constant if a token occured. +* `repetition_penalty: float = 0.0`: Penalizes tokens if they appeared in prompt or generation so far. Must be between `0` and `2` where numbers smaller than `1` encourage repeatment of tokens and numbers larger than `1` encourages sampling of new tokens. The penalization scales multiplicatively. +* `min_new_tokens: int = 0`: Forces the model to generate at least `min_new_tokens` until a stop word or EOS token is sampled. Note that this might lead to unintended behavior for example if the distribution is highly skewed towards these tokens. + +### Constrained decoding + +Please refer to our dedicated guide on [constrained decoding](./structured_outputs.ipynb) for the following parameters. + +* `json_schema: Optional[str] = None`: JSON schema for structured outputs. +* `regex: Optional[str] = None`: Regex for structured outputs. +* `ebnf: Optional[str] = None`: EBNF for structured outputs. + +### Other options + +* `n: int = 1`: Specifies the number of output sequences to generate per request. (Generating multiple outputs in one request (n > 1) is discouraged; repeat the same prompts for several times offer better control and efficiency.) +* `spaces_between_special_tokens: bool = True`: Whether or not to add spaces between special tokens during detokenization. +* `no_stop_trim: bool = False`: Don't trim stop words or EOS token from the generated text. +* `ignore_eos: bool = False`: Don't stop generation when EOS token is sampled. +* `skip_special_tokens: bool = True`: Remove special tokens during decoding. +* `custom_params: Optional[List[Optional[Dict[str, Any]]]] = None`: Used when employing `CustomLogitProcessor`. For usage see below. + +## Examples + +### Normal + +Launch a server: + +``` +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --port 30000 +``` + +Send a request: + +```python +import requests + +response = requests.post( + "http://localhost:30000/generate", + json={ + "text": "The capital of France is", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 32, + }, + }, +) +print(response.json()) +``` + +Detailed example in [send request](./send_request.ipynb). + +### Streaming + +Send a request and stream the output: + +```python +import requests, json + +response = requests.post( + "http://localhost:30000/generate", + json={ + "text": "The capital of France is", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 32, + }, + "stream": True, + }, + stream=True, +) + +prev = 0 +for chunk in response.iter_lines(decode_unicode=False): + chunk = chunk.decode("utf-8") + if chunk and chunk.startswith("data:"): + if chunk == "data: [DONE]": + break + data = json.loads(chunk[5:].strip("\n")) + output = data["text"].strip() + print(output[prev:], end="", flush=True) + prev = len(output) +print("") +``` + +Detailed example in [openai compatible api](https://docs.sglang.ai/backend/openai_api_completions.html#id2). + +### Multi modal + +Launch a server: + +``` +python3 -m sglang.launch_server --model-path lmms-lab/llava-onevision-qwen2-7b-ov --chat-template chatml-llava +``` + +Download an image: + +``` +curl -o example_image.png -L https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true +``` + +Send a request: + +```python +import requests + +response = requests.post( + "http://localhost:30000/generate", + json={ + "text": "<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n" + "<|im_start|>user\n\nDescribe this image in a very short sentence.<|im_end|>\n" + "<|im_start|>assistant\n", + "image_data": "example_image.png", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 32, + }, + }, +) +print(response.json()) +``` + +The `image_data` can be a file name, a URL, or a base64 encoded string. See also `python/sglang/srt/utils.py:load_image`. + +Streaming is supported in a similar manner as [above](#streaming). + +Detailed example in [openai api vision](./openai_api_vision.ipynb). + +### Structured Outputs (JSON, Regex, EBNF) + +You can specify a JSON schema, regular expression or [EBNF](https://en.wikipedia.org/wiki/Extended_Backus%E2%80%93Naur_form) to constrain the model output. The model output will be guaranteed to follow the given constraints. Only one constraint parameter (`json_schema`, `regex`, or `ebnf`) can be specified for a request. + +SGLang supports two grammar backends: + +- [Outlines](https://github.com/dottxt-ai/outlines) (default): Supports JSON schema and regular expression constraints. +- [XGrammar](https://github.com/mlc-ai/xgrammar): Supports JSON schema, regular expression, and EBNF constraints. + - XGrammar currently uses the [GGML BNF format](https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md) + +Initialize the XGrammar backend using `--grammar-backend xgrammar` flag +```bash +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ +--port 30000 --host 0.0.0.0 --grammar-backend [xgrammar|outlines] # xgrammar or outlines (default: outlines) +``` + +```python +import json +import requests + +json_schema = json.dumps({ + "type": "object", + "properties": { + "name": {"type": "string", "pattern": "^[\\w]+$"}, + "population": {"type": "integer"}, + }, + "required": ["name", "population"], +}) + +# JSON (works with both Outlines and XGrammar) +response = requests.post( + "http://localhost:30000/generate", + json={ + "text": "Here is the information of the capital of France in the JSON format.\n", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 64, + "json_schema": json_schema, + }, + }, +) +print(response.json()) + +# Regular expression (Outlines backend only) +response = requests.post( + "http://localhost:30000/generate", + json={ + "text": "Paris is the capital of", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 64, + "regex": "(France|England)", + }, + }, +) +print(response.json()) + +# EBNF (XGrammar backend only) +response = requests.post( + "http://localhost:30000/generate", + json={ + "text": "Write a greeting.", + "sampling_params": { + "temperature": 0, + "max_new_tokens": 64, + "ebnf": 'root ::= "Hello" | "Hi" | "Hey"', + }, + }, +) +print(response.json()) +``` + +Detailed example in [structured outputs](./structured_outputs.ipynb). +### Custom Logit Processor +Launch a server with `--enable-custom-logit-processor` flag on. +``` +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --port 30000 --enable-custom-logit-processor +``` + +Define a custom logit processor that will always sample a specific token id. +```python +from sglang.srt.sampling.custom_logit_processor import CustomLogitProcessor + +class DeterministicLogitProcessor(CustomLogitProcessor): + """A dummy logit processor that changes the logits to always + sample the given token id. + """ + + def __call__(self, logits, custom_param_list): + # Check that the number of logits matches the number of custom parameters + assert logits.shape[0] == len(custom_param_list) + key = "token_id" + + for i, param_dict in enumerate(custom_param_list): + # Mask all other tokens + logits[i, :] = -float("inf") + # Assign highest probability to the specified token + logits[i, param_dict[key]] = 0.0 + return logits +``` + +Send a request +```python +import requests + +response = requests.post( + "http://localhost:30000/generate", + json={ + "text": "The capital of France is", + "custom_logit_processor": DeterministicLogitProcessor().to_str(), + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": 32, + "custom_params": {"token_id": 5}, + }, + }, +) +print(response.json()) +``` diff --git a/sglang/docs/backend/send_request.ipynb b/sglang/docs/backend/send_request.ipynb new file mode 100644 index 0000000..6b51b42 --- /dev/null +++ b/sglang/docs/backend/send_request.ipynb @@ -0,0 +1,259 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Sending Requests\n", + "This notebook provides a quick-start guide to use SGLang in chat completions after installation.\n", + "\n", + "- For Vision Language Models, see [OpenAI APIs - Vision](../backend/openai_api_vision.ipynb).\n", + "- For Embedding Models, see [OpenAI APIs - Embedding](../backend/openai_api_embeddings.ipynb) and [Encode (embedding model)](../backend/native_api.html#Encode-(embedding-model)).\n", + "- For Reward Models, see [Classify (reward model)](../backend/native_api.html#Classify-(reward-model))." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Launch A Server" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sglang.test.test_utils import is_in_ci\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "if is_in_ci():\n", + " from patch import launch_server_cmd\n", + "else:\n", + " from sglang.utils import launch_server_cmd\n", + "\n", + "# This is equivalent to running the following command in your terminal\n", + "\n", + "# python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --host 0.0.0.0\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \\\n", + " --host 0.0.0.0\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using cURL\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import subprocess, json\n", + "\n", + "curl_command = f\"\"\"\n", + "curl -s http://localhost:{port}/v1/chat/completions \\\n", + " -H \"Content-Type: application/json\" \\\n", + " -d '{{\"model\": \"meta-llama/Meta-Llama-3.1-8B-Instruct\", \"messages\": [{{\"role\": \"user\", \"content\": \"What is the capital of France?\"}}]}}'\n", + "\"\"\"\n", + "\n", + "response = json.loads(subprocess.check_output(curl_command, shell=True))\n", + "print_highlight(response)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Python Requests" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "\n", + "url = f\"http://localhost:{port}/v1/chat/completions\"\n", + "\n", + "data = {\n", + " \"model\": \"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " \"messages\": [{\"role\": \"user\", \"content\": \"What is the capital of France?\"}],\n", + "}\n", + "\n", + "response = requests.post(url, json=data)\n", + "print_highlight(response.json())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using OpenAI Python Client" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=64,\n", + ")\n", + "print_highlight(response)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Streaming" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "# Use stream=True for streaming responses\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=64,\n", + " stream=True,\n", + ")\n", + "\n", + "# Handle the streaming output\n", + "for chunk in response:\n", + " if chunk.choices[0].delta.content:\n", + " print(chunk.choices[0].delta.content, end=\"\", flush=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Using Native Generation APIs\n", + "\n", + "You can also use the native `/generate` endpoint with requests, which provides more flexiblity. An API reference is available at [Sampling Parameters](../references/sampling_params.md)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "\n", + "response = requests.post(\n", + " f\"http://localhost:{port}/generate\",\n", + " json={\n", + " \"text\": \"The capital of France is\",\n", + " \"sampling_params\": {\n", + " \"temperature\": 0,\n", + " \"max_new_tokens\": 32,\n", + " },\n", + " },\n", + ")\n", + "\n", + "print_highlight(response.json())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Streaming" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests, json\n", + "\n", + "response = requests.post(\n", + " f\"http://localhost:{port}/generate\",\n", + " json={\n", + " \"text\": \"The capital of France is\",\n", + " \"sampling_params\": {\n", + " \"temperature\": 0,\n", + " \"max_new_tokens\": 32,\n", + " },\n", + " \"stream\": True,\n", + " },\n", + " stream=True,\n", + ")\n", + "\n", + "prev = 0\n", + "for chunk in response.iter_lines(decode_unicode=False):\n", + " chunk = chunk.decode(\"utf-8\")\n", + " if chunk and chunk.startswith(\"data:\"):\n", + " if chunk == \"data: [DONE]\":\n", + " break\n", + " data = json.loads(chunk[5:].strip(\"\\n\"))\n", + " output = data[\"text\"]\n", + " print(output[prev:], end=\"\", flush=True)\n", + " prev = len(output)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/sglang/docs/backend/separate_reasoning.ipynb b/sglang/docs/backend/separate_reasoning.ipynb new file mode 100644 index 0000000..756ecba --- /dev/null +++ b/sglang/docs/backend/separate_reasoning.ipynb @@ -0,0 +1,427 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Reasoning Parser\n", + "\n", + "SGLang supports parsing reasoning content our from \"normal\" content for reasoning models such as [DeepSeek R1](https://huggingface.co/deepseek-ai/DeepSeek-R1).\n", + "\n", + "## Supported Models\n", + "\n", + "Currently, SGLang supports the following reasoning models:\n", + "- [DeepSeek R1 series](https://huggingface.co/collections/deepseek-ai/deepseek-r1-678e1e131c0169c0bc89728d): The reasoning content is wrapped with `` and `` tags.\n", + "- [QwQ](https://huggingface.co/Qwen/QwQ-32B): The reasoning content is wrapped with `` and `` tags." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Usage\n", + "\n", + "### Launching the Server" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Specify the `--reasoning-parser` option." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "from openai import OpenAI\n", + "from sglang.test.test_utils import is_in_ci\n", + "\n", + "if is_in_ci():\n", + " from patch import launch_server_cmd\n", + "else:\n", + " from sglang.utils import launch_server_cmd\n", + "\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-R1-Distill-Qwen-7B --host 0.0.0.0 --reasoning-parser deepseek-r1\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Note that `--reasoning-parser` defines the parser used to interpret responses. Currently supported parsers include:\n", + "\n", + "- deepseek-r1: DeepSeek R1 series and QwQ (e.g. deepseek-ai/DeepSeek-R1, Qwen/QwQ-32B)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### OpenAI Compatible API\n", + "\n", + "Using the OpenAI compatible API, the contract follows the [DeepSeek API design](https://api-docs.deepseek.com/guides/reasoning_model) established with the release of DeepSeek-R1:\n", + "\n", + "- `reasoning_content`: The content of the CoT.\n", + "- `content`: The content of the final answer." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Initialize OpenAI-like client\n", + "client = OpenAI(api_key=\"None\", base_url=f\"http://0.0.0.0:{port}/v1\")\n", + "model_name = client.models.list().data[0].id\n", + "\n", + "messages = [\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"What is 1+3?\",\n", + " }\n", + "]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Non-Streaming Request" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response_non_stream = client.chat.completions.create(\n", + " model=model_name,\n", + " messages=messages,\n", + " temperature=0.6,\n", + " top_p=0.95,\n", + " stream=False, # Non-streaming\n", + " extra_body={\"separate_reasoning\": True},\n", + ")\n", + "print_highlight(\"==== Reasoning ====\")\n", + "print_highlight(response_non_stream.choices[0].message.reasoning_content)\n", + "\n", + "print_highlight(\"==== Text ====\")\n", + "print_highlight(response_non_stream.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "#### Streaming Request" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response_stream = client.chat.completions.create(\n", + " model=model_name,\n", + " messages=messages,\n", + " temperature=0.6,\n", + " top_p=0.95,\n", + " stream=True, # Non-streaming\n", + " extra_body={\"separate_reasoning\": True},\n", + ")\n", + "\n", + "reasoning_content = \"\"\n", + "content = \"\"\n", + "for chunk in response_stream:\n", + " if chunk.choices[0].delta.content:\n", + " content += chunk.choices[0].delta.content\n", + " if chunk.choices[0].delta.reasoning_content:\n", + " reasoning_content += chunk.choices[0].delta.reasoning_content\n", + "\n", + "print_highlight(\"==== Reasoning ====\")\n", + "print_highlight(reasoning_content)\n", + "\n", + "print_highlight(\"==== Text ====\")\n", + "print_highlight(content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Optionally, you can buffer the reasoning content to the last reasoning chunk (or the first chunk after the reasoning content)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response_stream = client.chat.completions.create(\n", + " model=model_name,\n", + " messages=messages,\n", + " temperature=0.6,\n", + " top_p=0.95,\n", + " stream=True, # Non-streaming\n", + " extra_body={\"separate_reasoning\": True, \"stream_reasoning\": False},\n", + ")\n", + "\n", + "reasoning_content = \"\"\n", + "content = \"\"\n", + "for chunk in response_stream:\n", + " if chunk.choices[0].delta.content:\n", + " content += chunk.choices[0].delta.content\n", + " if chunk.choices[0].delta.reasoning_content:\n", + " reasoning_content = chunk.choices[0].delta.reasoning_content\n", + "\n", + "print_highlight(\"==== Reasoning ====\")\n", + "print_highlight(reasoning_content)\n", + "\n", + "print_highlight(\"==== Text ====\")\n", + "print_highlight(content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The reasoning separation is enable by default when specify . \n", + "**To disable it, set the `separate_reasoning` option to `False` in request.**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response_non_stream = client.chat.completions.create(\n", + " model=model_name,\n", + " messages=messages,\n", + " temperature=0.6,\n", + " top_p=0.95,\n", + " stream=False, # Non-streaming\n", + " extra_body={\"separate_reasoning\": False},\n", + ")\n", + "\n", + "print_highlight(\"==== Original Output ====\")\n", + "print_highlight(response_non_stream.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### SGLang Native API " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from transformers import AutoTokenizer\n", + "\n", + "tokenizer = AutoTokenizer.from_pretrained(\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\")\n", + "input = tokenizer.apply_chat_template(\n", + " messages,\n", + " tokenize=False,\n", + " add_generation_prompt=True,\n", + ")\n", + "\n", + "gen_url = f\"http://localhost:{port}/generate\"\n", + "gen_data = {\n", + " \"text\": input,\n", + " \"sampling_params\": {\n", + " \"skip_special_tokens\": False,\n", + " \"max_new_tokens\": 1024,\n", + " \"temperature\": 0.6,\n", + " \"top_p\": 0.95,\n", + " },\n", + "}\n", + "gen_response = requests.post(gen_url, json=gen_data).json()[\"text\"]\n", + "\n", + "print_highlight(\"==== Original Output ====\")\n", + "print_highlight(gen_response)\n", + "\n", + "parse_url = f\"http://localhost:{port}/separate_reasoning\"\n", + "separate_reasoning_data = {\n", + " \"text\": gen_response,\n", + " \"reasoning_parser\": \"deepseek-r1\",\n", + "}\n", + "separate_reasoning_response_json = requests.post(\n", + " parse_url, json=separate_reasoning_data\n", + ").json()\n", + "print_highlight(\"==== Reasoning ====\")\n", + "print_highlight(separate_reasoning_response_json[\"reasoning_text\"])\n", + "print_highlight(\"==== Text ====\")\n", + "print_highlight(separate_reasoning_response_json[\"text\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Offline Engine API" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import sglang as sgl\n", + "from sglang.srt.reasoning_parser import ReasoningParser\n", + "from sglang.utils import print_highlight\n", + "\n", + "llm = sgl.Engine(model_path=\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\")\n", + "tokenizer = AutoTokenizer.from_pretrained(\"deepseek-ai/DeepSeek-R1-Distill-Qwen-7B\")\n", + "input = tokenizer.apply_chat_template(\n", + " messages,\n", + " tokenize=False,\n", + " add_generation_prompt=True,\n", + ")\n", + "sampling_params = {\n", + " \"max_new_tokens\": 1024,\n", + " \"skip_special_tokens\": False,\n", + " \"temperature\": 0.6,\n", + " \"top_p\": 0.95,\n", + "}\n", + "result = llm.generate(prompt=input, sampling_params=sampling_params)\n", + "\n", + "generated_text = result[\"text\"] # Assume there is only one prompt\n", + "\n", + "print_highlight(\"==== Original Output ====\")\n", + "print_highlight(generated_text)\n", + "\n", + "parser = ReasoningParser(\"deepseek-r1\")\n", + "reasoning_text, text = parser.parse_non_stream(generated_text)\n", + "print_highlight(\"==== Reasoning ====\")\n", + "print_highlight(reasoning_text)\n", + "print_highlight(\"==== Text ====\")\n", + "print_highlight(text)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "llm.shutdown()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Supporting New Reasoning Model Schemas\n", + "\n", + "For future reasoning models, you can implement the reasoning parser as a subclass of `BaseReasoningFormatDetector` in `python/sglang/srt/reasoning_parser.py` and specify the reasoning parser for new reasoning model schemas accordingly." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "```python\n", + "class DeepSeekR1Detector(BaseReasoningFormatDetector):\n", + " \"\"\"\n", + " Detector for DeepSeek-R1 model.\n", + " Assumes reasoning format:\n", + " ()*(.*)\n", + " Returns all the text before the tag as `reasoning_text`\n", + " and the rest of the text as `normal_text`.\n", + "\n", + " Args:\n", + " stream_reasoning (bool): If False, accumulates reasoning content until the end tag.\n", + " If True, streams reasoning content as it arrives.\n", + " \"\"\"\n", + "\n", + " def __init__(self, stream_reasoning: bool = False):\n", + " # DeepSeek-R1 is assumed to be reasoning until `` token\n", + " super().__init__(\"\", \"\", True, stream_reasoning=stream_reasoning)\n", + " # https://github.com/sgl-project/sglang/pull/3202#discussion_r1950153599\n", + "\n", + "\n", + "class ReasoningParser:\n", + " \"\"\"\n", + " Parser that handles both streaming and non-streaming scenarios for extracting\n", + " reasoning content from model outputs.\n", + "\n", + " Args:\n", + " model_type (str): Type of model to parse reasoning from\n", + " stream_reasoning (bool): If Flase, accumulates reasoning content until complete.\n", + " If True, streams reasoning content as it arrives.\n", + " \"\"\"\n", + "\n", + " DetectorMap: Dict[str, BaseReasoningFormatDetector] = {\n", + " \"deepseek-r1\": DeepSeekR1Detector\n", + " }\n", + "\n", + " def __init__(self, model_type: str = None, stream_reasoning: bool = True):\n", + " if not model_type:\n", + " raise ValueError(\"Model type must be specified\")\n", + "\n", + " detector_class = self.DetectorMap.get(model_type.lower())\n", + " if not detector_class:\n", + " raise ValueError(f\"Unsupported model type: {model_type}\")\n", + "\n", + " self.detector = detector_class(stream_reasoning=stream_reasoning)\n", + "\n", + " def parse_non_stream(self, full_text: str) -> StreamingParseResult:\n", + " \"\"\"Non-streaming call: one-time parsing\"\"\"\n", + " ret = self.detector.detect_and_parse(full_text)\n", + " return ret.reasoning_text, ret.normal_text\n", + "\n", + " def parse_stream_chunk(self, chunk_text: str) -> StreamingParseResult:\n", + " \"\"\"Streaming call: incremental parsing\"\"\"\n", + " ret = self.detector.parse_streaming_increment(chunk_text)\n", + " return ret.reasoning_text, ret.normal_text\n", + "```" + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/sglang/docs/backend/server_arguments.md b/sglang/docs/backend/server_arguments.md new file mode 100644 index 0000000..3d2aae8 --- /dev/null +++ b/sglang/docs/backend/server_arguments.md @@ -0,0 +1,195 @@ +# Server Arguments + +## Common launch commands + +- To enable multi-GPU tensor parallelism, add `--tp 2`. If it reports the error "peer access is not supported between these two devices", add `--enable-p2p-check` to the server launch command. + +```bash +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --tp 2 +``` + +- To enable multi-GPU data parallelism, add `--dp 2`. Data parallelism is better for throughput if there is enough memory. It can also be used together with tensor parallelism. The following command uses 4 GPUs in total. We recommend [SGLang Router](../router/router.md) for data parallelism. + +```bash +python -m sglang_router.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --dp 2 --tp 2 +``` + +- If you see out-of-memory errors during serving, try to reduce the memory usage of the KV cache pool by setting a smaller value of `--mem-fraction-static`. The default value is `0.9`. + +```bash +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --mem-fraction-static 0.7 +``` + +- See [hyperparameter tuning](hyperparameter_tuning.md) on tuning hyperparameters for better performance. +- For docker and Kubernetes runs, you need to set up shared memory which is used for communication between processes. See `--shm-size` for docker and `/dev/shm` size update for Kubernetes manifests. +- If you see out-of-memory errors during prefill for long prompts, try to set a smaller chunked prefill size. + +```bash +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --chunked-prefill-size 4096 +``` +- To enable `torch.compile` acceleration, add `--enable-torch-compile`. It accelerates small models on small batch sizes. By default, the cache path is located at `/tmp/torchinductor_root`, you can customize it using environment variable `TORCHINDUCTOR_CACHE_DIR`. For more details, please refer to [PyTorch official documentation](https://pytorch.org/tutorials/recipes/torch_compile_caching_tutorial.html) and [Enabling cache for torch.compile](https://docs.sglang.ai/backend/hyperparameter_tuning.html#enabling-cache-for-torch-compile). +- To enable torchao quantization, add `--torchao-config int4wo-128`. It supports other [quantization strategies (INT8/FP8)](https://github.com/sgl-project/sglang/blob/v0.3.6/python/sglang/srt/server_args.py#L671) as well. +- To enable fp8 weight quantization, add `--quantization fp8` on a fp16 checkpoint or directly load a fp8 checkpoint without specifying any arguments. +- To enable fp8 kv cache quantization, add `--kv-cache-dtype fp8_e5m2`. +- If the model does not have a chat template in the Hugging Face tokenizer, you can specify a [custom chat template](custom_chat_template.md). +- To run tensor parallelism on multiple nodes, add `--nnodes 2`. If you have two nodes with two GPUs on each node and want to run TP=4, let `sgl-dev-0` be the hostname of the first node and `50000` be an available port, you can use the following commands. If you meet deadlock, please try to add `--disable-cuda-graph` + +```bash +# Node 0 +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --tp 4 --dist-init-addr sgl-dev-0:50000 --nnodes 2 --node-rank 0 + +# Node 1 +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct --tp 4 --dist-init-addr sgl-dev-0:50000 --nnodes 2 --node-rank 1 +``` + +Please consult the documentation below to learn more about the parameters you may provide when launching a server. + + +## Model and tokenizer + +* `model_path`: Path to the model that will be served. +* `tokenizer_path`: Defaults to the `model_path`. +* `tokenizer_mode`: By default `auto`, see [here](https://huggingface.co/docs/transformers/en/main_classes/tokenizer) for different mode. +* `load_format`: The format the weights are loaded in. Defaults to `*.safetensors`/`*.bin`. +* `trust_remote_code`: If `True`, will use locally cached config files, otherwise use remote configs in HuggingFace. +* `dtype`: Dtype used for the model, defaults to `bfloat16`. +* `kv_cache_dtype`: Dtype of the kv cache, defaults to the `dtype`. +* `context_length`: The number of tokens our model can process *including the input*. Note that extending the default might lead to strange behavior. +* `device`: The device we put the model, defaults to `cuda`. +* `chat_template`: The chat template to use. Deviating from the default might lead to unexpected responses. For multi-modal chat templates, refer to [here](https://docs.sglang.ai/backend/openai_api_vision.ipynb#Chat-Template). **Make sure the correct** `chat_template` **is passed, or performance degradation may occur!!!!** +* `is_embedding`: Set to true to perform [embedding](./openai_api_embeddings.ipynb) / [encode](https://docs.sglang.ai/backend/native_api#Encode-(embedding-model)) and [reward](https://docs.sglang.ai/backend/native_api#Classify-(reward-model)) tasks. +* `revision`: Adjust if a specific version of the model should be used. +* `skip_tokenizer_init`: Set to true to provide the tokens to the engine and get the output tokens directly, typically used in RLHF. Please see this [example for reference](https://github.com/sgl-project/sglang/blob/main/examples/runtime/token_in_token_out/). +* `json_model_override_args`: Override model config with the provided JSON. +* `delete_ckpt_after_loading`: Delete the model checkpoint after loading the model. + + +## Serving: HTTP & API + +### HTTP Server configuration + +* `port` and `host`: Setup the host for HTTP server. By default `host: str = "127.0.0.1"` and `port: int = 30000` + +### API configuration + +* `api_key`: Sets an API key for the server and the OpenAI-compatible API. +* `file_storage_path`: Directory for storing uploaded or generated files from API calls. +* `enable_cache_report`: If set, includes detailed usage of cached tokens in the response usage. + +## Parallelism + +### Tensor parallelism + +* `tp_size`: The number of GPUs the model weights get sharded over. Mainly for saving memory rather than for high throughput, see [this blogpost](https://pytorch.org/tutorials/intermediate/TP_tutorial.html#how-tensor-parallel-works). + +### Data parallelism + +* `dp_size`: Will be deprecated. The number of data-parallel copies of the model. [SGLang router](../router/router.md) is recommended instead of the current naive data parallel. +* `load_balance_method`: Will be deprecated. Load balancing strategy for data parallel requests. + +### Expert parallelism +* `enable_ep_moe`: Enables expert parallelism that distributes the experts onto multiple GPUs for MoE models. +* `ep_size`: The size of EP. Please shard the model weights with `tp_size=ep_size`, for detailed benchmarking refer to [this PR](https://github.com/sgl-project/sglang/pull/2203). If not set, `ep_size` will be automatically set to `tp_size`. +* `enable_deepep_moe`: Enables expert parallelism that distributes the experts onto multiple GPUs for DeepSeek-V3 model based on deepseek-ai/DeepEP. + +## Memory and scheduling + +* `mem_fraction_static`: Fraction of the free GPU memory used for static memory like model weights and KV cache. If building KV cache fails, it should be increased. If CUDA runs out of memory, it should be decreased. +* `max_running_requests`: The maximum number of requests to run concurrently. +* `max_total_tokens`: The maximum number of tokens that can be stored into the KV cache. Use mainly for debugging. +* `chunked_prefill_size`: Perform the prefill in chunks of these size. Larger chunk size speeds up the prefill phase but increases the VRAM consumption. If CUDA runs out of memory, it should be decreased. +* `max_prefill_tokens`: Token budget of how many tokens to accept in one prefill batch. The actual number is the max of this parameter and the `context_length`. +* `schedule_policy`: The scheduling policy to control the processing order of waiting prefill requests in a single engine. +* `schedule_conservativeness`: Can be used to decrease/increase the conservativeness of the server when taking new requests. Highly conservative behavior leads to starvation, but low conservativeness leads to slowed-down performance. +* `cpu_offload_gb`: Reserve this amount of RAM in GB for offloading of model parameters to the CPU. + +## Other runtime options + +* `stream_interval`: Interval (in tokens) for streaming responses. Smaller values lead to smoother streaming, and larger values lead to better throughput. +* `random_seed`: Can be used to enforce more deterministic behavior. +* `watchdog_timeout`: Adjusts the watchdog thread's timeout before killing the server if batch generation takes too long. +* `download_dir`: Use to override the default Hugging Face cache directory for model weights. +* `base_gpu_id`: Use to adjust first GPU used to distribute the model across available GPUs. +* `allow_auto_truncate`: Automatically truncate requests that exceed the maximum input length. + +## Logging + +* `log_level`: Global log verbosity. +* `log_level_http`: Separate verbosity level for the HTTP server logs (if unset, defaults to `log_level`). +* `log_requests`: Logs the inputs and outputs of all requests for debugging. +* `log_requests_level`: Ranges from 0 to 2: level 0 only shows some basic metadata in requests, level 1 and 2 show request details (e.g., text, images), and level 1 limits output to 2048 characters (if unset, defaults to `0`). +* `show_time_cost`: Prints or logs detailed timing info for internal operations (helpful for performance tuning). +* `enable_metrics`: Exports Prometheus-like metrics for request usage and performance. +* `decode_log_interval`: How often (in tokens) to log decode progress. + +## Multi-node distributed serving + +* `dist_init_addr`: The TCP address used for initializing PyTorch's distributed backend (e.g. `192.168.0.2:25000`). +* `nnodes`: Total number of nodes in the cluster. Refer to how to run the [Llama 405B model](https://docs.sglang.ai/references/llama_405B.html#run-405b-fp16-on-two-nodes). +* `node_rank`: Rank (ID) of this node among the `nnodes` in the distributed setup. + + +## LoRA + +* `lora_paths`: You may provide a list of adapters to your model as a list. Each batch element will get model response with the corresponding lora adapter applied. Currently `cuda_graph` and `radix_attention` are not supported with this option so you need to disable them manually. We are still working on through these [issues](https://github.com/sgl-project/sglang/issues/2929). +* `max_loras_per_batch`: Maximum number of LoRAs in a running batch including base model. +* `lora_backend`: The backend of running GEMM kernels for Lora modules, can be one of `triton` or `flashinfer`. Defaults to be `triton`. + +## Kernel backend + +* `attention_backend`: The backend for attention computation and KV cache management. +* `sampling_backend`: The backend for sampling. + +## Constrained Decoding + +* `grammar_backend`: The grammar backend for constraint decoding. Detailed usage can be found in this [document](https://docs.sglang.ai/backend/structured_outputs.html). +* `constrained_json_whitespace_pattern`: Use with `Outlines` grammar backend to allow JSON with syntatic newlines, tabs or multiple spaces. Details can be found [here](https://dottxt-ai.github.io/outlines/latest/reference/generation/json/#using-pydantic). + +## Speculative decoding + +* `speculative_draft_model_path`: The draft model path for speculative decoding. +* `speculative_algorithm`: The algorithm for speculative decoding. Currently only [Eagle](https://arxiv.org/html/2406.16858v1) is supported. Note that the radix cache, chunked prefill, and overlap scheduler are disabled when using eagle speculative decoding. +* `speculative_num_steps`: How many draft passes we run before verifying. +* `speculative_num_draft_tokens`: The number of tokens proposed in a draft. +* `speculative_eagle_topk`: The number of top candidates we keep for verification at each step for [Eagle](https://arxiv.org/html/2406.16858v1). +* `speculative_token_map`: Optional, the path to the high frequency token list of [FR-Spec](https://arxiv.org/html/2502.14856v1), used for accelerating [Eagle](https://arxiv.org/html/2406.16858v1). + + +## Double Sparsity + +* `enable_double_sparsity`: Enables [double sparsity](https://arxiv.org/html/2408.07092v2) which increases throughput. +* `ds_channel_config_path`: The double sparsity config. For a guide on how to generate the config for your model see [this repo](https://github.com/andy-yang-1/DoubleSparse/tree/main/config). +* `ds_heavy_channel_num`: Number of channel indices to keep for each layer. +* `ds_heavy_token_num`: Number of tokens used for attention during decode. Skip sparse decoding if `min_seq_len` in batch < this number. +* `ds_heavy_channel_type`: The type of heavy channels. Either `q`, `k` or `qk`. +* `ds_sparse_decode_threshold`: Don't apply sparse decoding if `max_seq_len` in batch < this threshold. + +## Debug options + +*Note: We recommend to stay with the defaults and only use these options for debugging for best possible performance.* + +* `disable_radix_cache`: Disable [Radix](https://lmsys.org/blog/2024-01-17-sglang/) backend for prefix caching. +* `disable_cuda_graph`: Disable [cuda graph](https://pytorch.org/blog/accelerating-pytorch-with-cuda-graphs/) for model forward. Use if encountering uncorrectable CUDA ECC errors. +* `disable_cuda_graph_padding`: Disable cuda graph when padding is needed. In other case still use cuda graph. +* `disable_outlines_disk_cache`: Disable disk cache for outlines grammar backend. +* `disable_custom_all_reduce`: Disable usage of custom all reduce kernel. +* `disable_mla`: Disable [Multi-Head Latent Attention](https://arxiv.org/html/2405.04434v5) for Deepseek model. +* `disable_overlap_schedule`: Disable the [Overhead-Scheduler](https://lmsys.org/blog/2024-12-04-sglang-v0-4/#zero-overhead-batch-scheduler). +* `enable_nan_detection`: Turning this on makes the sampler print a warning if the logits contain `NaN`. +* `enable_p2p_check`: Turns off the default of allowing always p2p check when accessing GPU. +* `triton_attention_reduce_in_fp32`: In triton kernels this will cast the intermediate attention result to `float32`. + +## Optimization + +*Note: Some of these options are still in experimental stage.* + +* `enable_mixed_chunk`: Enables mixing prefill and decode, see [this discussion](https://github.com/sgl-project/sglang/discussions/1163). +* `enable_dp_attention`: Enable [Data Parallelism Attention](https://lmsys.org/blog/2024-12-04-sglang-v0-4/#data-parallelism-attention-for-deepseek-models) for Deepseek models. +* `enable_torch_compile`: Torch compile the model. Note that compiling a model takes a long time but have a great performance boost. The compiled model can also be [cached for future use](https://docs.sglang.ai/backend/hyperparameter_tuning.html#enabling-cache-for-torch-compile). +* `torch_compile_max_bs`: The maximum batch size when using `torch_compile`. +* `cuda_graph_max_bs`: Adjust the maximum batchsize when using cuda graph. By default this is chosen for you based on GPU specifics. +* `cuda_graph_bs`: The batch sizes to capture by `CudaGraphRunner`. By default this is done for you. +* `torchao_config`: Experimental feature that optimizes the model with [torchao](https://github.com/pytorch/ao). Possible choices are: int8dq, int8wo, int4wo-, fp8wo, fp8dq-per_tensor, fp8dq-per_row. +* `triton_attention_num_kv_splits`: Use to adjust the number of KV splits in triton kernels. Default is 8. +* `enable_flashinfer_mla`: Use the attention backend with flashinfer MLA wrapper for deepseek models. When providing this argument, `attention_backend` argument is overridden. +* `flashinfer_mla_disable_ragged`: Disable usage of ragged prefill wrapper for flashinfer mla attention backend. Should be used when `enable_flashinfer_mla` is turned on. diff --git a/sglang/docs/backend/speculative_decoding.ipynb b/sglang/docs/backend/speculative_decoding.ipynb new file mode 100644 index 0000000..061f415 --- /dev/null +++ b/sglang/docs/backend/speculative_decoding.ipynb @@ -0,0 +1,305 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Speculative Decoding\n", + "\n", + "SGLang now provides an EAGLE2-based speculative decoding option. Our implementation aims to maximize speed and efficiency and is considered to be among the fastest in open-source LLM engines.\n", + "\n", + "**Note:** Currently, Speculative Decoding in SGLang is compatible with radix cache and chunked prefill.\n", + "\n", + "### Performance Highlights\n", + "\n", + "- Official EAGLE code ([SafeAILab/EAGLE](https://github.com/SafeAILab/EAGLE)): ~200 tokens/s\n", + "- Standard SGLang Decoding: ~156 tokens/s\n", + "- EAGLE Decoding in SGLang: ~297 tokens/s\n", + "- EAGLE Decoding in SGLang (w/ `torch.compile`): ~316 tokens/s\n", + "\n", + "All benchmarks below were run on a single H100." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## EAGLE Decoding\n", + "\n", + "To enable EAGLE speculative decoding the following parameters are relevant:\n", + "* `speculative_draft_model_path`: Specifies draft model. This parameter is required.\n", + "* `speculative_num_steps`: Depth of autoregressive drafting. Increases speculation range but risks rejection cascades. Default is 5.\n", + "\n", + "* `speculative_eagle_topk`: Branching factor per step. Improves candidate diversity, will lead to higher acceptance rate, but more lead to higher memory/compute consumption. Default is 4.\n", + "\n", + "* `speculative_num_draft_tokens`: Maximum parallel verification capacity. Allows deeper tree evaluation but will lead to higher GPU memory usage. Default is 8." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sglang.test.test_utils import is_in_ci\n", + "\n", + "if is_in_ci():\n", + " from patch import launch_server_cmd\n", + "else:\n", + " from sglang.utils import launch_server_cmd\n", + "\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model meta-llama/Llama-2-7b-chat-hf --speculative-algorithm EAGLE \\\n", + " --speculative-draft-model-path lmsys/sglang-EAGLE-llama2-chat-7B --speculative-num-steps 3 \\\n", + " --speculative-eagle-topk 4 --speculative-num-draft-tokens 16 --cuda-graph-max-bs 8\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=64,\n", + ")\n", + "\n", + "print_highlight(f\"Response: {response}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EAGLE Decoding with `torch.compile`\n", + "\n", + "You can also enable `torch.compile` for further optimizations and optionally set `--cuda-graph-max-bs`:\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model meta-llama/Llama-2-7b-chat-hf --speculative-algorithm EAGLE \\\n", + " --speculative-draft-model-path lmsys/sglang-EAGLE-llama2-chat-7B --speculative-num-steps 5 \\\n", + " --speculative-eagle-topk 8 --speculative-num-draft-tokens 64 --mem-fraction 0.6 \\\n", + " --enable-torch-compile --cuda-graph-max-bs 2\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=64,\n", + ")\n", + "\n", + "print_highlight(f\"Response: {response}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EAGLE Decoding via Frequency-Ranked Speculative Sampling\n", + "\n", + "By employing a truncated high-frequency token vocabulary in the draft model, Eagle speculative decoding reduces `lm_head` computational overhead while accelerating the pipeline without quality degradation. For more details, checkout [the paper](https://arxiv.org/pdf/arXiv:2502.14856).\n", + "\n", + "In our implementation, set `--speculative-token-map` to enable the optimization. You can get the high-frequency token in FR-Spec from [this model](https://huggingface.co/thunlp/LLaMA3-Instruct-8B-FR-Spec). Or you can obtain high-frequency token by directly downloading these token from [this repo](https://github.com/thunlp/FR-Spec/tree/main?tab=readme-ov-file#prepare-fr-spec-vocabulary-subset).\n", + "\n", + "Thanks for the contribution from [Weilin Zhao](https://github.com/https://github.com/Achazwl) and [Zhousx](https://github.com/Zhou-sx). " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model meta-llama/Meta-Llama-3-8B-Instruct --speculative-algorithm EAGLE \\\n", + " --speculative-draft-model-path lmsys/sglang-EAGLE-LLaMA3-Instruct-8B --speculative-num-steps 5 \\\n", + " --speculative-eagle-topk 8 --speculative-num-draft-tokens 64 --speculative-token-map thunlp/LLaMA3-Instruct-8B-FR-Spec/freq_32768.pt \\\n", + " --mem-fraction 0.7 --cuda-graph-max-bs 2 --dtype float16 \n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3-8B-Instruct\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=64,\n", + ")\n", + "\n", + "print_highlight(f\"Response: {response}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EAGLE-3 Decoding\n", + "\n", + "You can enable EAGLE-3 decoding by setting `--speculative_draft_model_path: EAGLE3`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"\"\"\n", + "python3 -m sglang.launch_server --model meta-llama/Llama-3.1-8B-Instruct --speculative-algorithm EAGLE3 \\\n", + " --speculative-draft-model-path jamesliu1/sglang-EAGLE3-Llama-3.1-Instruct-8B --speculative-num-steps 5 \\\n", + " --speculative-eagle-topk 8 --speculative-num-draft-tokens 32 --mem-fraction 0.6 \\\n", + " --cuda-graph-max-bs 2 --dtype float16\n", + "\"\"\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"List 3 countries and their capitals.\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=64,\n", + ")\n", + "\n", + "print_highlight(f\"Response: {response}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## References\n", + "\n", + "EAGLE process is as follows:\n", + "\n", + "- Within EAGLE the draft model predicts the next feature vector, i.e. the last hidden state of the original LLM, using the feature sequence $(f_1, ..., f_k)$ and the token sequence $(t_2, ..., t_{k+1})$. \n", + "- The next token is then sampled from $p_{k+2}=\\text{LMHead}(f_{k+1})$. Afterwards, the two sequences are extended in a tree style—branching out multiple potential continuations, with the branching factor per step controlled by the `speculative_eagle_topk` parameter—to ensure a more coherent connection of context, and are given as input again.\n", + "- EAGLE-2 additionally uses the draft model to evaluate how probable certain branches in the draft tree are, dynamically stopping the expansion of unlikely branches. After the expansion phase, reranking is employed to select only the top `speculative_num_draft_tokens` final nodes as draft tokens.\n", + "- EAGLE-3 removes the feature prediction objective, incorporates low and mid-layer features, and is trained in an on-policy manner.\n", + "\n", + "This enhances drafting accuracy by operating on the features instead of tokens for more regular inputs and passing the tokens from the next timestep additionaly to minimize randomness effects from sampling. Furthermore the dynamic adjustment of the draft tree and selection of reranked final nodes increases acceptance rate of draft tokens further. For more details see [the paper](https://arxiv.org/abs/2406.16858)." + ] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/sglang/docs/backend/structured_outputs.ipynb b/sglang/docs/backend/structured_outputs.ipynb new file mode 100644 index 0000000..acb2620 --- /dev/null +++ b/sglang/docs/backend/structured_outputs.ipynb @@ -0,0 +1,827 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Structured Outputs" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "You can specify a JSON schema, [regular expression](https://en.wikipedia.org/wiki/Regular_expression) or [EBNF](https://en.wikipedia.org/wiki/Extended_Backus%E2%80%93Naur_form) to constrain the model output. The model output will be guaranteed to follow the given constraints. Only one constraint parameter (`json_schema`, `regex`, or `ebnf`) can be specified for a request.\n", + "\n", + "SGLang supports three grammar backends:\n", + "\n", + "- [Outlines](https://github.com/dottxt-ai/outlines): Supports JSON schema and regular expression constraints.\n", + "- [XGrammar](https://github.com/mlc-ai/xgrammar)(default): Supports JSON schema, regular expression, and EBNF constraints.\n", + "- [Llguidance](https://github.com/guidance-ai/llguidance): Supports JSON schema, regular expression, and EBNF constraints.\n", + "\n", + "We suggest using XGrammar for its better performance and utility. XGrammar currently uses the [GGML BNF format](https://github.com/ggerganov/llama.cpp/blob/master/grammars/README.md). For more details, see [XGrammar technical overview](https://blog.mlc.ai/2024/11/22/achieving-efficient-flexible-portable-structured-generation-with-xgrammar).\n", + "\n", + "To use Outlines, simply add `--grammar-backend outlines` when launching the server.\n", + "To use llguidance, add `--grammar-backend llguidance` when launching the server.\n", + "If no backend is specified, XGrammar will be used as the default.\n", + "\n", + "For better output quality, **It's advisable to explicitly include instructions in the prompt to guide the model to generate the desired format.** For example, you can specify, 'Please generate the output in the following JSON format: ...'.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## OpenAI Compatible API" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import openai\n", + "import os\n", + "from sglang.test.test_utils import is_in_ci\n", + "\n", + "if is_in_ci():\n", + " from patch import launch_server_cmd\n", + "else:\n", + " from sglang.utils import launch_server_cmd\n", + "\n", + "from sglang.utils import wait_for_server, print_highlight, terminate_process\n", + "\n", + "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n", + "\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --host 0.0.0.0\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")\n", + "client = openai.Client(base_url=f\"http://127.0.0.1:{port}/v1\", api_key=\"None\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### JSON\n", + "\n", + "you can directly define a JSON schema or use [Pydantic](https://docs.pydantic.dev/latest/) to define and validate the response." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Using Pydantic**" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from pydantic import BaseModel, Field\n", + "\n", + "\n", + "# Define the schema using Pydantic\n", + "class CapitalInfo(BaseModel):\n", + " name: str = Field(..., pattern=r\"^\\w+$\", description=\"Name of the capital city\")\n", + " population: int = Field(..., description=\"Population of the capital city\")\n", + "\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"Please generate the information of the capital of France in the JSON format.\",\n", + " },\n", + " ],\n", + " temperature=0,\n", + " max_tokens=128,\n", + " response_format={\n", + " \"type\": \"json_schema\",\n", + " \"json_schema\": {\n", + " \"name\": \"foo\",\n", + " # convert the pydantic model to json schema\n", + " \"schema\": CapitalInfo.model_json_schema(),\n", + " },\n", + " },\n", + ")\n", + "\n", + "response_content = response.choices[0].message.content\n", + "# validate the JSON response by the pydantic model\n", + "capital_info = CapitalInfo.model_validate_json(response_content)\n", + "print_highlight(f\"Validated response: {capital_info.model_dump_json()}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**JSON Schema Directly**\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "\n", + "json_schema = json.dumps(\n", + " {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"name\": {\"type\": \"string\", \"pattern\": \"^[\\\\w]+$\"},\n", + " \"population\": {\"type\": \"integer\"},\n", + " },\n", + " \"required\": [\"name\", \"population\"],\n", + " }\n", + ")\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"Give me the information of the capital of France in the JSON format.\",\n", + " },\n", + " ],\n", + " temperature=0,\n", + " max_tokens=128,\n", + " response_format={\n", + " \"type\": \"json_schema\",\n", + " \"json_schema\": {\"name\": \"foo\", \"schema\": json.loads(json_schema)},\n", + " },\n", + ")\n", + "\n", + "print_highlight(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### EBNF" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "ebnf_grammar = \"\"\"\n", + "root ::= city | description\n", + "city ::= \"London\" | \"Paris\" | \"Berlin\" | \"Rome\"\n", + "description ::= city \" is \" status\n", + "status ::= \"the capital of \" country\n", + "country ::= \"England\" | \"France\" | \"Germany\" | \"Italy\"\n", + "\"\"\"\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\"role\": \"system\", \"content\": \"You are a helpful geography bot.\"},\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"Give me the information of the capital of France.\",\n", + " },\n", + " ],\n", + " temperature=0,\n", + " max_tokens=32,\n", + " extra_body={\"ebnf\": ebnf_grammar},\n", + ")\n", + "\n", + "print_highlight(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Regular expression" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=[\n", + " {\"role\": \"user\", \"content\": \"What is the capital of France?\"},\n", + " ],\n", + " temperature=0,\n", + " max_tokens=128,\n", + " extra_body={\"regex\": \"(Paris|London)\"},\n", + ")\n", + "\n", + "print_highlight(response.choices[0].message.content)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Structural Tag" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "tool_get_current_weather = {\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": \"get_current_weather\",\n", + " \"description\": \"Get the current weather in a given location\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"city\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The city to find the weather for, e.g. 'San Francisco'\",\n", + " },\n", + " \"state\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"the two-letter abbreviation for the state that the city is\"\n", + " \" in, e.g. 'CA' which would mean 'California'\",\n", + " },\n", + " \"unit\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The unit to fetch the temperature in\",\n", + " \"enum\": [\"celsius\", \"fahrenheit\"],\n", + " },\n", + " },\n", + " \"required\": [\"city\", \"state\", \"unit\"],\n", + " },\n", + " },\n", + "}\n", + "\n", + "tool_get_current_date = {\n", + " \"type\": \"function\",\n", + " \"function\": {\n", + " \"name\": \"get_current_date\",\n", + " \"description\": \"Get the current date and time for a given timezone\",\n", + " \"parameters\": {\n", + " \"type\": \"object\",\n", + " \"properties\": {\n", + " \"timezone\": {\n", + " \"type\": \"string\",\n", + " \"description\": \"The timezone to fetch the current date and time for, e.g. 'America/New_York'\",\n", + " }\n", + " },\n", + " \"required\": [\"timezone\"],\n", + " },\n", + " },\n", + "}\n", + "\n", + "schema_get_current_weather = tool_get_current_weather[\"function\"][\"parameters\"]\n", + "schema_get_current_date = tool_get_current_date[\"function\"][\"parameters\"]\n", + "\n", + "\n", + "def get_messages():\n", + " return [\n", + " {\n", + " \"role\": \"system\",\n", + " \"content\": f\"\"\"\n", + "# Tool Instructions\n", + "- Always execute python code in messages that you share.\n", + "- When looking for real time information use relevant functions if available else fallback to brave_search\n", + "You have access to the following functions:\n", + "Use the function 'get_current_weather' to: Get the current weather in a given location\n", + "{tool_get_current_weather[\"function\"]}\n", + "Use the function 'get_current_date' to: Get the current date and time for a given timezone\n", + "{tool_get_current_date[\"function\"]}\n", + "If a you choose to call a function ONLY reply in the following format:\n", + "<{{start_tag}}={{function_name}}>{{parameters}}{{end_tag}}\n", + "where\n", + "start_tag => ` a JSON dict with the function argument name as key and function argument value as value.\n", + "end_tag => ``\n", + "Here is an example,\n", + "{{\"example_name\": \"example_value\"}}\n", + "Reminder:\n", + "- Function calls MUST follow the specified format\n", + "- Required parameters MUST be specified\n", + "- Only call one function at a time\n", + "- Put the entire function call reply on one line\n", + "- Always add your sources when using search results to answer the user query\n", + "You are a helpful assistant.\"\"\",\n", + " },\n", + " {\n", + " \"role\": \"user\",\n", + " \"content\": \"You are in New York. Please get the current date and time, and the weather.\",\n", + " },\n", + " ]\n", + "\n", + "\n", + "messages = get_messages()\n", + "\n", + "response = client.chat.completions.create(\n", + " model=\"meta-llama/Meta-Llama-3.1-8B-Instruct\",\n", + " messages=messages,\n", + " response_format={\n", + " \"type\": \"structural_tag\",\n", + " \"structures\": [\n", + " {\n", + " \"begin\": \"\",\n", + " \"schema\": schema_get_current_weather,\n", + " \"end\": \"\",\n", + " },\n", + " {\n", + " \"begin\": \"\",\n", + " \"schema\": schema_get_current_date,\n", + " \"end\": \"\",\n", + " },\n", + " ],\n", + " \"triggers\": [\"\",\n", + " \"schema\": schema_get_current_weather,\n", + " \"end\": \"\",\n", + " },\n", + " {\n", + " \"begin\": \"\",\n", + " \"schema\": schema_get_current_date,\n", + " \"end\": \"\",\n", + " },\n", + " ],\n", + " \"triggers\": [\"\",\n", + " \"schema\": schema_get_current_weather,\n", + " \"end\": \"\",\n", + " },\n", + " {\n", + " \"begin\": \"\",\n", + " \"schema\": schema_get_current_date,\n", + " \"end\": \"\",\n", + " },\n", + " ],\n", + " \"triggers\": [\">> |\.\.\. " +copybutton_prompt_is_regexp = True + +autodoc_preserve_defaults = True +navigation_with_keys = False + +autodoc_mock_imports = [ + "torch", + "transformers", + "triton", +] + +intersphinx_mapping = { + "python": ("https://docs.python.org/3.12", None), + "typing_extensions": ("https://typing-extensions.readthedocs.io/en/latest", None), + "pillow": ("https://pillow.readthedocs.io/en/stable", None), + "numpy": ("https://numpy.org/doc/stable", None), + "torch": ("https://pytorch.org/docs/stable", None), +} + +html_theme = "sphinx_book_theme" + + +nbsphinx_prolog = """ +.. raw:: html + + +""" diff --git a/sglang/docs/deploy.py b/sglang/docs/deploy.py new file mode 100644 index 0000000..75b7ea7 --- /dev/null +++ b/sglang/docs/deploy.py @@ -0,0 +1,22 @@ +# Deploy the documents + +import os +from datetime import datetime + + +def run_cmd(cmd): + print(cmd) + os.system(cmd) + + +run_cmd("cd $DOC_SITE_PATH; git pull") + +# (Optional) Remove old files +# run_cmd("rm -rf $ALPA_SITE_PATH/*") + +run_cmd("cp -r _build/html/* $DOC_SITE_PATH") + +cmd_message = f"Update {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}" +run_cmd( + f"cd $DOC_SITE_PATH; git add .; git commit -m '{cmd_message}'; git push origin main" +) diff --git a/sglang/docs/developer/development_guide_using_docker.md b/sglang/docs/developer/development_guide_using_docker.md new file mode 100644 index 0000000..963d04f --- /dev/null +++ b/sglang/docs/developer/development_guide_using_docker.md @@ -0,0 +1,52 @@ +# Development Guide Using Docker + +## Setup VSCode + +Download `code` from `Https://code.visualstudio.com/docs/?dv=linux64cli` + +```bash +wget https://vscode.download.prss.microsoft.com/dbazure/download/stable/fabdb6a30b49f79a7aba0f2ad9df9b399473380f/vscode_cli_alpine_x64_cli.tar.gz +tar xf vscode_cli_alpine_x64_cli.tar.gz + +# https://code.visualstudio.com/docs/remote/tunnels +./code tunnel +``` + +## Setup Docker Container + +The following startup command is an example for internal development by the SGLang team. You can **modify or add directory mappings as needed**, especially for model weight downloads, to prevent repeated downloads by different Docker containers. + +❗️ **Note on RDMA** + + 1. `--network host` and `--privileged` are required by RDMA. If you don't need RDMA, you can remove them but keeping them there does not harm. Thus, we enable these two flags by default in the commands below. + 2. You may need to set `NCCL_IB_GID_INDEX` if you are using RoCE, for example: `export NCCL_IB_GID_INDEX=3`. + +### H100 + +```bash +# Change the name to yours +docker run -itd --shm-size 32g --gpus all -v /opt/dlami/nvme/.cache:/root/.cache --ipc=host --network=host --privileged --name sglang_zhyncs lmsysorg/sglang:dev /bin/zsh +docker exec -it sglang_zhyncs /bin/zsh +``` + +### H200 + +```bash +docker run -itd --shm-size 32g --gpus all -v /mnt/co-research/shared-models:/root/.cache/huggingface --ipc=host --network=host --privileged --name sglang_zhyncs lmsysorg/sglang:dev /bin/zsh +docker exec -it sglang_zhyncs /bin/zsh +``` + +## Profile + +```bash +# Change batch size, input, output and add `disable-cuda-graph` (for easier analysis) +# e.g. DeepSeek V3 +nsys profile -o deepseek_v3 python3 -m sglang.bench_one_batch --batch-size 1 --input 128 --output 256 --model deepseek-ai/DeepSeek-V3 --trust-remote-code --tp 8 --disable-cuda-graph +``` + +## Evaluation + +```bash +# e.g. gsm8k 8 shot +python3 benchmark/gsm8k/bench_sglang.py --num-questions 2000 --parallel 2000 --num-shots 8 +``` diff --git a/sglang/docs/developer/release_process.md b/sglang/docs/developer/release_process.md new file mode 100644 index 0000000..e817a25 --- /dev/null +++ b/sglang/docs/developer/release_process.md @@ -0,0 +1,18 @@ +# PyPI Package Release Process + +## Update the version in code +Update the package version in `python/pyproject.toml` and `python/sglang/__init__.py`. + +## Upload the PyPI package + +``` +pip install build twine +``` + +``` +cd python +bash upload_pypi.sh +``` + +## Make a release in GitHub +Make a new release https://github.com/sgl-project/sglang/releases/new. diff --git a/sglang/docs/developer/setup_github_runner.md b/sglang/docs/developer/setup_github_runner.md new file mode 100644 index 0000000..55240e7 --- /dev/null +++ b/sglang/docs/developer/setup_github_runner.md @@ -0,0 +1,49 @@ +# Set Up Self-Hosted Runners for GitHub Action + +## Add a Runner + +### Step 1: Start a docker container. + +You can mount a folder for the shared huggingface model weights cache. The command below uses `/tmp/huggingface` as an example. + +``` +docker pull nvidia/cuda:12.1.1-devel-ubuntu22.04 +# Nvidia +docker run --shm-size 128g -it -v /tmp/huggingface:/hf_home --gpus all nvidia/cuda:12.1.1-devel-ubuntu22.04 /bin/bash +# AMD +docker run --rm --device=/dev/kfd --device=/dev/dri --group-add video --shm-size 128g -it -v /tmp/huggingface:/hf_home lmsysorg/sglang:v0.4.4.post3-rocm630 /bin/bash +# AMD just the last 2 GPUs +docker run --rm --device=/dev/kfd --device=/dev/dri/renderD176 --device=/dev/dri/renderD184 --group-add video --shm-size 128g -it -v /tmp/huggingface:/hf_home lmsysorg/sglang:v0.4.4.post3-rocm630 /bin/bash +``` + +### Step 2: Configure the runner by `config.sh` + +Run these commands inside the container. + +``` +apt update && apt install -y curl python3-pip git +export RUNNER_ALLOW_RUNASROOT=1 +``` + +Then follow https://github.com/sgl-project/sglang/settings/actions/runners/new?arch=x64&os=linux to run `config.sh` + +**Notes** +- Do not need to specify the runner group +- Give it a name (e.g., `test-sgl-gpu-0`) and some labels (e.g., `1-gpu-runner`). The labels can be editted later in Github Settings. +- Do not need to change the work folder. + +### Step 3: Run the runner by `run.sh` + +- Set up environment variables +``` +export HF_HOME=/hf_home +export SGLANG_IS_IN_CI=true +export HF_TOKEN=hf_xxx +export OPENAI_API_KEY=sk-xxx +export CUDA_VISIBLE_DEVICES=0 +``` + +- Run it forever +``` +while true; do ./run.sh; echo "Restarting..."; sleep 2; done +``` diff --git a/sglang/docs/frontend/__init__.py b/sglang/docs/frontend/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sglang/docs/frontend/choices_methods.md b/sglang/docs/frontend/choices_methods.md new file mode 100644 index 0000000..30a0a18 --- /dev/null +++ b/sglang/docs/frontend/choices_methods.md @@ -0,0 +1,77 @@ +# Choices Methods in SGLang +This doc describes the choices methods supported by SGLang. + +The optional `choices_method` arg determines how options supplied to SGLang's `choices` primitive are selected. Only the `RuntimeEndpoint` backend supports the `choices_method` arg. Other backends, such as `OpenAI`, have bespoke selection implementations due to API limitations. + +## Methods + +### Token Length Normalized + +Token length normalized is the default SGLang choices method. It selects the option with the highest average logprob across all of its tokens. + +Usage example (alternatively, simply omit the `choices_method` arg): +```python +@sgl.function +def example(s): + s += sgl.user("What is the capital of France?") + s += sgl.assistant( + sgl.gen( + "answer", + choices=["London", "Paris", "Berlin"], + choices_method=sgl.token_length_normalized, + ) + ) +``` + + +This can perform poorly if an option contains many tokens, where its later tokens are predicted with high confidence based on its earlier tokens. For instance, even strong models will fail the above example if the specified options are `["Paris", "Antidisestablishmentarianism"]`. + +### Greedy Token Selection + +Greedy token selection simply selects the option with the highest logprob for its initial token. For overlapping options where one option is a subset of a longer option, the logprobs of the shorter option are extended using its average logprob for comparison against the longer option. + +Usage example: +```python +@sgl.function +def example(s): + s += sgl.user("What is the capital of France?") + s += sgl.assistant( + sgl.gen( + "answer", + choices=["London", "Paris", "Berlin"], + choices_method=sgl.greedy_token_selection, + ) + ) +``` + +This can perform poorly if an option misleads the model down a bad path based on an attractive initial token. For instance, greedy selection will result in an incorrect response for this example: +```python +@sgl.function +def us_president_example(s): + s += sgl.user("Name a US president.") + s += sgl.assistant( + sgl.gen( + "answer", + choices=["Donald Duck", "Millard Fillmore"], + choices_method=sgl.greedy_token_selection, + ) + ) +``` + +### Unconditional Likelihood Normalized + +Unconditional likelihood normalized selects the option with the highest average token logprob once normalized by the unconditional token logprobs, as described in [this EleutherAI blogpost](https://blog.eleuther.ai/multiple-choice-normalization/). This method incurs an additional LLM call to obtain the unconditional likelihoods. + +Usage example: +```python +@sgl.function +def example(s): + s += sgl.user("What is the capital of France?") + s += sgl.assistant( + sgl.gen( + "answer", + choices=["London", "Paris", "Berlin"], + choices_method=sgl.unconditional_likelihood_normalized, + ) + ) +``` diff --git a/sglang/docs/frontend/frontend.ipynb b/sglang/docs/frontend/frontend.ipynb new file mode 100644 index 0000000..7bc07f6 --- /dev/null +++ b/sglang/docs/frontend/frontend.ipynb @@ -0,0 +1,464 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# SGLang Frontend Language" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "SGLang frontend language can be used to define simple and easy prompts in a convenient, structured way." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Launch A Server\n", + "\n", + "Launch the server in your terminal and wait for it to initialize." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import requests\n", + "import os\n", + "\n", + "from sglang import assistant_begin, assistant_end\n", + "from sglang import assistant, function, gen, system, user\n", + "from sglang import image\n", + "from sglang import RuntimeEndpoint, set_default_backend\n", + "from sglang.srt.utils import load_image\n", + "from sglang.test.test_utils import is_in_ci\n", + "from sglang.utils import print_highlight, terminate_process, wait_for_server\n", + "\n", + "if is_in_ci():\n", + " from patch import launch_server_cmd\n", + "else:\n", + " from sglang.utils import launch_server_cmd\n", + "\n", + "\n", + "server_process, port = launch_server_cmd(\n", + " \"python -m sglang.launch_server --model-path Qwen/Qwen2.5-7B-Instruct --host 0.0.0.0\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")\n", + "print(f\"Server started on http://localhost:{port}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Set the default backend. Note: Besides the local server, you may use also `OpenAI` or other API endpoints." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "set_default_backend(RuntimeEndpoint(f\"http://localhost:{port}\"))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Basic Usage\n", + "\n", + "The most simple way of using SGLang frontend language is a simple question answer dialog between a user and an assistant." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def basic_qa(s, question):\n", + " s += system(f\"You are a helpful assistant than can answer questions.\")\n", + " s += user(question)\n", + " s += assistant(gen(\"answer\", max_tokens=512))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "state = basic_qa(\"List 3 countries and their capitals.\")\n", + "print_highlight(state[\"answer\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Multi-turn Dialog\n", + "\n", + "SGLang frontend language can also be used to define multi-turn dialogs." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def multi_turn_qa(s):\n", + " s += system(f\"You are a helpful assistant than can answer questions.\")\n", + " s += user(\"Please give me a list of 3 countries and their capitals.\")\n", + " s += assistant(gen(\"first_answer\", max_tokens=512))\n", + " s += user(\"Please give me another list of 3 countries and their capitals.\")\n", + " s += assistant(gen(\"second_answer\", max_tokens=512))\n", + " return s\n", + "\n", + "\n", + "state = multi_turn_qa()\n", + "print_highlight(state[\"first_answer\"])\n", + "print_highlight(state[\"second_answer\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Control flow\n", + "\n", + "You may use any Python code within the function to define more complex control flows." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def tool_use(s, question):\n", + " s += assistant(\n", + " \"To answer this question: \"\n", + " + question\n", + " + \". I need to use a \"\n", + " + gen(\"tool\", choices=[\"calculator\", \"search engine\"])\n", + " + \". \"\n", + " )\n", + "\n", + " if s[\"tool\"] == \"calculator\":\n", + " s += assistant(\"The math expression is: \" + gen(\"expression\"))\n", + " elif s[\"tool\"] == \"search engine\":\n", + " s += assistant(\"The key word to search is: \" + gen(\"word\"))\n", + "\n", + "\n", + "state = tool_use(\"What is 2 * 2?\")\n", + "print_highlight(state[\"tool\"])\n", + "print_highlight(state[\"expression\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Parallelism\n", + "\n", + "Use `fork` to launch parallel prompts. Because `sgl.gen` is non-blocking, the for loop below issues two generation calls in parallel." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def tip_suggestion(s):\n", + " s += assistant(\n", + " \"Here are two tips for staying healthy: \"\n", + " \"1. Balanced Diet. 2. Regular Exercise.\\n\\n\"\n", + " )\n", + "\n", + " forks = s.fork(2)\n", + " for i, f in enumerate(forks):\n", + " f += assistant(\n", + " f\"Now, expand tip {i+1} into a paragraph:\\n\"\n", + " + gen(\"detailed_tip\", max_tokens=256, stop=\"\\n\\n\")\n", + " )\n", + "\n", + " s += assistant(\"Tip 1:\" + forks[0][\"detailed_tip\"] + \"\\n\")\n", + " s += assistant(\"Tip 2:\" + forks[1][\"detailed_tip\"] + \"\\n\")\n", + " s += assistant(\n", + " \"To summarize the above two tips, I can say:\\n\" + gen(\"summary\", max_tokens=512)\n", + " )\n", + "\n", + "\n", + "state = tip_suggestion()\n", + "print_highlight(state[\"summary\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Constrained Decoding\n", + "\n", + "Use `regex` to specify a regular expression as a decoding constraint. This is only supported for local models." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def regular_expression_gen(s):\n", + " s += user(\"What is the IP address of the Google DNS servers?\")\n", + " s += assistant(\n", + " gen(\n", + " \"answer\",\n", + " temperature=0,\n", + " regex=r\"((25[0-5]|2[0-4]\\d|[01]?\\d\\d?).){3}(25[0-5]|2[0-4]\\d|[01]?\\d\\d?)\",\n", + " )\n", + " )\n", + "\n", + "\n", + "state = regular_expression_gen()\n", + "print_highlight(state[\"answer\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use `regex` to define a `JSON` decoding schema." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "character_regex = (\n", + " r\"\"\"\\{\\n\"\"\"\n", + " + r\"\"\" \"name\": \"[\\w\\d\\s]{1,16}\",\\n\"\"\"\n", + " + r\"\"\" \"house\": \"(Gryffindor|Slytherin|Ravenclaw|Hufflepuff)\",\\n\"\"\"\n", + " + r\"\"\" \"blood status\": \"(Pure-blood|Half-blood|Muggle-born)\",\\n\"\"\"\n", + " + r\"\"\" \"occupation\": \"(student|teacher|auror|ministry of magic|death eater|order of the phoenix)\",\\n\"\"\"\n", + " + r\"\"\" \"wand\": \\{\\n\"\"\"\n", + " + r\"\"\" \"wood\": \"[\\w\\d\\s]{1,16}\",\\n\"\"\"\n", + " + r\"\"\" \"core\": \"[\\w\\d\\s]{1,16}\",\\n\"\"\"\n", + " + r\"\"\" \"length\": [0-9]{1,2}\\.[0-9]{0,2}\\n\"\"\"\n", + " + r\"\"\" \\},\\n\"\"\"\n", + " + r\"\"\" \"alive\": \"(Alive|Deceased)\",\\n\"\"\"\n", + " + r\"\"\" \"patronus\": \"[\\w\\d\\s]{1,16}\",\\n\"\"\"\n", + " + r\"\"\" \"bogart\": \"[\\w\\d\\s]{1,16}\"\\n\"\"\"\n", + " + r\"\"\"\\}\"\"\"\n", + ")\n", + "\n", + "\n", + "@function\n", + "def character_gen(s, name):\n", + " s += user(\n", + " f\"{name} is a character in Harry Potter. Please fill in the following information about this character.\"\n", + " )\n", + " s += assistant(gen(\"json_output\", max_tokens=256, regex=character_regex))\n", + "\n", + "\n", + "state = character_gen(\"Harry Potter\")\n", + "print_highlight(state[\"json_output\"])" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Batching \n", + "\n", + "Use `run_batch` to run a batch of prompts." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def text_qa(s, question):\n", + " s += user(question)\n", + " s += assistant(gen(\"answer\", stop=\"\\n\"))\n", + "\n", + "\n", + "states = text_qa.run_batch(\n", + " [\n", + " {\"question\": \"What is the capital of the United Kingdom?\"},\n", + " {\"question\": \"What is the capital of France?\"},\n", + " {\"question\": \"What is the capital of Japan?\"},\n", + " ],\n", + " progress_bar=True,\n", + ")\n", + "\n", + "for i, state in enumerate(states):\n", + " print_highlight(f\"Answer {i+1}: {states[i]['answer']}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Streaming \n", + "\n", + "Use `stream` to stream the output to the user." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def text_qa(s, question):\n", + " s += user(question)\n", + " s += assistant(gen(\"answer\", stop=\"\\n\"))\n", + "\n", + "\n", + "state = text_qa.run(\n", + " question=\"What is the capital of France?\", temperature=0.1, stream=True\n", + ")\n", + "\n", + "for out in state.text_iter():\n", + " print(out, end=\"\", flush=True)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Complex Prompts\n", + "\n", + "You may use `{system|user|assistant}_{begin|end}` to define complex prompts." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def chat_example(s):\n", + " s += system(\"You are a helpful assistant.\")\n", + " # Same as: s += s.system(\"You are a helpful assistant.\")\n", + "\n", + " with s.user():\n", + " s += \"Question: What is the capital of France?\"\n", + "\n", + " s += assistant_begin()\n", + " s += \"Answer: \" + gen(\"answer\", max_tokens=100, stop=\"\\n\")\n", + " s += assistant_end()\n", + "\n", + "\n", + "state = chat_example()\n", + "print_highlight(state[\"answer\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Multi-modal Generation\n", + "\n", + "You may use SGLang frontend language to define multi-modal prompts.\n", + "See [here](https://docs.sglang.ai/references/supported_models.html) for supported models." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "server_process, port = launch_server_cmd(\n", + " \"python -m sglang.launch_server --model-path Qwen/Qwen2.5-VL-7B-Instruct --host 0.0.0.0\"\n", + ")\n", + "\n", + "wait_for_server(f\"http://localhost:{port}\")\n", + "print(f\"Server started on http://localhost:{port}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "set_default_backend(RuntimeEndpoint(f\"http://localhost:{port}\"))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Ask a question about an image." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def image_qa(s, image_file, question):\n", + " s += user(image(image_file) + question)\n", + " s += assistant(gen(\"answer\", max_tokens=256))\n", + "\n", + "\n", + "image_url = \"https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true\"\n", + "image_bytes, _ = load_image(image_url)\n", + "state = image_qa(image_bytes, \"What is in the image?\")\n", + "print_highlight(state[\"answer\"])" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "terminate_process(server_process)" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} diff --git a/sglang/docs/frontend/patch.py b/sglang/docs/frontend/patch.py new file mode 100644 index 0000000..b39951d --- /dev/null +++ b/sglang/docs/frontend/patch.py @@ -0,0 +1,49 @@ +import weakref + +from sglang.utils import execute_shell_command, reserve_port + +DEFAULT_MAX_RUNNING_REQUESTS = 200 +DEFAULT_MAX_TOTAL_TOKENS = 20480 + +import sglang.srt.server_args as server_args_mod + +_original_post_init = server_args_mod.ServerArgs.__post_init__ + + +def patched_post_init(self): + _original_post_init(self) + if self.max_running_requests is None: + self.max_running_requests = DEFAULT_MAX_RUNNING_REQUESTS + if self.max_total_tokens is None: + self.max_total_tokens = DEFAULT_MAX_TOTAL_TOKENS + self.disable_cuda_graph = True + + +server_args_mod.ServerArgs.__post_init__ = patched_post_init + +process_socket_map = weakref.WeakKeyDictionary() + + +def launch_server_cmd(command: str, host: str = "0.0.0.0", port: int = None): + """ + Launch the server using the given command. + If no port is specified, a free port is reserved. + """ + if port is None: + port, lock_socket = reserve_port(host) + else: + lock_socket = None + + extra_flags = ( + f"--max-running-requests {DEFAULT_MAX_RUNNING_REQUESTS} " + f"--max-total-tokens {DEFAULT_MAX_TOTAL_TOKENS} " + f"--disable-cuda-graph" + ) + + full_command = f"{command} --port {port} {extra_flags}" + process = execute_shell_command(full_command) + + if lock_socket is not None: + process_socket_map[process] = lock_socket + + return process, port diff --git a/sglang/docs/index.rst b/sglang/docs/index.rst new file mode 100644 index 0000000..085e722 --- /dev/null +++ b/sglang/docs/index.rst @@ -0,0 +1,65 @@ +SGLang Documentation +==================== + +SGLang is a fast serving framework for large language models and vision language models. +It makes your interaction with models faster and more controllable by co-designing the backend runtime and frontend language. +The core features include: + +- **Fast Backend Runtime**: Provides efficient serving with RadixAttention for prefix caching, zero-overhead CPU scheduler, continuous batching, token attention (paged attention), speculative decoding, tensor parallelism, chunked prefill, structured outputs, and quantization (FP8/INT4/AWQ/GPTQ). +- **Flexible Frontend Language**: Offers an intuitive interface for programming LLM applications, including chained generation calls, advanced prompting, control flow, multi-modal inputs, parallelism, and external interactions. +- **Extensive Model Support**: Supports a wide range of generative models (Llama, Gemma, Mistral, QWen, DeepSeek, LLaVA, etc.), embedding models (e5-mistral, gte, mcdse) and reward models (Skywork), with easy extensibility for integrating new models. +- **Active Community**: SGLang is open-source and backed by an active community with industry adoption. + +.. toctree:: + :maxdepth: 1 + :caption: Installation + + start/install.md + +.. toctree:: + :maxdepth: 1 + :caption: Backend Tutorial + + references/deepseek + backend/send_request.ipynb + backend/openai_api_completions.ipynb + backend/openai_api_vision.ipynb + backend/openai_api_embeddings.ipynb + backend/native_api.ipynb + backend/offline_engine_api.ipynb + backend/server_arguments.md + backend/sampling_params.md + backend/hyperparameter_tuning.md + +.. toctree:: + :maxdepth: 1 + :caption: Advanced Features + + backend/speculative_decoding.ipynb + backend/structured_outputs.ipynb + backend/function_calling.ipynb + backend/separate_reasoning.ipynb + backend/custom_chat_template.md + backend/quantization.md + +.. toctree:: + :maxdepth: 1 + :caption: Frontend Tutorial + + frontend/frontend.ipynb + frontend/choices_methods.md + +.. toctree:: + :maxdepth: 1 + :caption: SGLang Router + + router/router.md + +.. toctree:: + :maxdepth: 1 + :caption: References + + references/general + references/hardware + references/advanced_deploy + references/performance_tuning diff --git a/sglang/docs/references/accuracy_evaluation.md b/sglang/docs/references/accuracy_evaluation.md new file mode 100644 index 0000000..84c1edd --- /dev/null +++ b/sglang/docs/references/accuracy_evaluation.md @@ -0,0 +1,60 @@ +# Measuring Model Accuracy in SGLang + +This guide shows how to evaluate model accuracy using SGLang's [built-in benchmarks](https://github.com/sgl-project/sglang/tree/b045841baeff37a5601fcde23fa98bd09d942c36/benchmark). Please include accuracy on crucial benchmarks in your PR if you make modifications on the model side, like the kernel and model architecture. + +## Benchmarking Model Accuracy + +This is a reference workflow for the [MMLU benchmark](https://github.com/sgl-project/sglang/tree/main/benchmark/mmlu). For more details or other benchmarks, please refer to the README in each specific benchmark folder under [sglang/benchmark](https://github.com/sgl-project/sglang/tree/b045841baeff37a5601fcde23fa98bd09d942c36/benchmark). + +```bash +# Step 1: Download the dataset +bash download_data.sh + +# Step 2: Launch the server +python3 -m sglang.launch_server \ + --model-path Qwen/Qwen2.5-Math-1.5B-Instruct \ # Model selection + --port 30000 \ # Network configuration + --mem-fraction-static 0.8 # Memory optimization + +# Step 3: Run the benchmark script +python3 bench_sglang.py --nsub 10 # Test 10 subjects + +# Step 4: Extract the accuracy +cat result.jsonl | grep -oP '"accuracy": \K\d+\.\d+' +``` + +## Customizing Benchmark Scripts + +Some benchmark implementations may differ from ours, causing accuracy discrepancies. To match [[Qwen2.5-Math]](https://github.com/QwenLM/Qwen2.5-Math)'s reported 76.8% GSM8K accuracy, customization is required. + +```python +# The GSM8K benchmark script includes few shot examples for evaluation by default. +# Here we exclude them. +for i in range(len(lines[num_shots:num_questions])): + questions.append(get_one_example(lines, i, False)) + labels.append(get_answer_value(lines[i]["answer"])) +``` + +```python +@sgl.function +def few_shot_gsm8k(s, question): + # System prompt given in https://github.com/QwenLM/Qwen2.5-Math + s += sgl.system("Please reason step by step, and put your final answer within \\boxed{}.") # Include system prompt + s += few_shot_examples + question + # Stopwords given in evaluation/math_eval.py of the Qwen2.5-Math repo + s += sgl.gen( + "answer", max_tokens=2048, stop=["Question", "Assistant:", "", "<|im_end|>", "<|endoftext|>"] + ) +``` + +These adjustments should return the desired accuracy. + +## Extending Evaluation Capabilities + +1. **Contribute New Benchmarks** + * Follow our [contribution guidelines](../references/contribution_guide.md) to add new test scripts +2. **Request Implementations** + * Feel free to open an issue describing your evaluation needs +3. **Use Alternative Tools** + * [OpenCompass](https://opencompass.org.cn) + * [LM Evaluation Harness](https://github.com/EleutherAI/lm-evaluation-harness) diff --git a/sglang/docs/references/advanced_deploy.rst b/sglang/docs/references/advanced_deploy.rst new file mode 100644 index 0000000..06a1002 --- /dev/null +++ b/sglang/docs/references/advanced_deploy.rst @@ -0,0 +1,7 @@ +Multi-Node Deployment +========================== +.. toctree:: + :maxdepth: 1 + + multi_node.md + deploy_on_k8s.md diff --git a/sglang/docs/references/amd.md b/sglang/docs/references/amd.md new file mode 100644 index 0000000..c4e53eb --- /dev/null +++ b/sglang/docs/references/amd.md @@ -0,0 +1,146 @@ +# SGLang on AMD + +This document describes how to set up an AMD-based environment for [SGLang](https://github.com/sgl-project/sglang). If you encounter issues or have questions, please [open an issue](https://github.com/sgl-project/sglang/issues) on the SGLang repository. + +## System Configuration + +When using AMD GPUs (such as MI300X), certain system-level optimizations help ensure stable performance. Here we take MI300X as an example. AMD provides official documentation for MI300X optimization and system tuning: + +- [AMD MI300X Tuning Guides](https://rocm.docs.amd.com/en/latest/how-to/tuning-guides/mi300x/index.html) + - [LLM inference performance validation on AMD Instinct MI300X](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/inference/vllm-benchmark.html) + - [AMD Instinct MI300X System Optimization](https://rocm.docs.amd.com/en/latest/how-to/system-optimization/mi300x.html) + - [AMD Instinct MI300X Workload Optimization](https://rocm.docs.amd.com/en/latest/how-to/rocm-for-ai/inference-optimization/workload.html) + +**NOTE:** We strongly recommend reading these docs and guides entirely to fully utilize your system. + +Below are a few key settings to confirm or enable for SGLang: + +### Update GRUB Settings + +In `/etc/default/grub`, append the following to `GRUB_CMDLINE_LINUX`: + +```text +pci=realloc=off iommu=pt +``` + +Afterward, run `sudo update-grub` (or your distro’s equivalent) and reboot. + +### Disable NUMA Auto-Balancing + +```bash +sudo sh -c 'echo 0 > /proc/sys/kernel/numa_balancing' +``` + +You can automate or verify this change using [this helpful script](https://github.com/ROCm/triton/blob/rocm_env/scripts/amd/env_check.sh). + +Again, please go through the entire documentation to confirm your system is using the recommended configuration. + +## Installing SGLang + +For general installation instructions, see the official [SGLang Installation Docs](../start/install.md). Below are the AMD-specific steps summarized for convenience. + +### Install from Source + +```bash +git clone https://github.com/sgl-project/sglang.git +cd sglang + +pip install --upgrade pip +pip install sgl-kernel --force-reinstall --no-deps +pip install -e "python[all_hip]" +``` + +### Install Using Docker (Recommended) + +1. Build the docker image. + + ```bash + docker build -t sglang_image -f Dockerfile.rocm . + ``` + +2. Create a convenient alias. + + ```bash + alias drun='docker run -it --rm --network=host --privileged --device=/dev/kfd --device=/dev/dri \ + --ipc=host --shm-size 16G --group-add video --cap-add=SYS_PTRACE \ + --security-opt seccomp=unconfined \ + -v $HOME/dockerx:/dockerx \ + -v /data:/data' + ``` + +If you are using RDMA, please note that: + +1. `--network host` and `--privileged` are required by RDMA. If you don't need RDMA, you can remove them. +2. You may need to set `NCCL_IB_GID_INDEX` if you are using RoCE, for example: `export NCCL_IB_GID_INDEX=3`. +3. Launch the server. + + **NOTE:** Replace `` below with your [huggingface hub token](https://huggingface.co/docs/hub/en/security-tokens). + + ```bash + drun -p 30000:30000 \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --env "HF_TOKEN=" \ + sglang_image \ + python3 -m sglang.launch_server \ + --model-path NousResearch/Meta-Llama-3.1-8B \ + --host 0.0.0.0 \ + --port 30000 + ``` + +4. To verify the utility, you can run a benchmark in another terminal or refer to [other docs](https://docs.sglang.ai/backend/openai_api_completions.html) to send requests to the engine. + + ```bash + drun sglang_image \ + python3 -m sglang.bench_serving \ + --backend sglang \ + --dataset-name random \ + --num-prompts 4000 \ + --random-input 128 \ + --random-output 128 + ``` + +With your AMD system properly configured and SGLang installed, you can now fully leverage AMD hardware to power SGLang’s machine learning capabilities. + +## Examples + +### Running DeepSeek-V3 + +The only difference when running DeepSeek-V3 is in how you start the server. Here's an example command: + +```bash +drun -p 30000:30000 \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --ipc=host \ + --env "HF_TOKEN=" \ + sglang_image \ + python3 -m sglang.launch_server \ + --model-path deepseek-ai/DeepSeek-V3 \ # <- here + --tp 8 \ + --trust-remote-code \ + --host 0.0.0.0 \ + --port 30000 +``` + +[Running DeepSeek-R1 on a single NDv5 MI300X VM](https://techcommunity.microsoft.com/blog/azurehighperformancecomputingblog/running-deepseek-r1-on-a-single-ndv5-mi300x-vm/4372726) could also be a good reference. + +### Running Llama3.1 + +Running Llama3.1 is nearly identical to running DeepSeek-V3. The only difference is in the model specified when starting the server, shown by the following example command: + +```bash +drun -p 30000:30000 \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --ipc=host \ + --env "HF_TOKEN=" \ + sglang_image \ + python3 -m sglang.launch_server \ + --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ # <- here + --tp 8 \ + --trust-remote-code \ + --host 0.0.0.0 \ + --port 30000 +``` + +### Warmup Step + +When the server displays `The server is fired up and ready to roll!`, it means the startup is successful. diff --git a/sglang/docs/references/benchmark_and_profiling.md b/sglang/docs/references/benchmark_and_profiling.md new file mode 100644 index 0000000..37c4545 --- /dev/null +++ b/sglang/docs/references/benchmark_and_profiling.md @@ -0,0 +1,142 @@ +# Benchmark and Profiling + +## Benchmark + +- Benchmark the latency of running a single static batch without a server. The arguments are the same as for `launch_server.py`. + Note that this is a simplified test script without a dynamic batching server, so it may run out of memory for a batch size that a real server can handle. A real server truncates the prefill into several batches, while this simplified script does not. + + ```bash + python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --batch 32 --input-len 256 --output-len 32 + ``` + +- Benchmark offline processing. This script will start an offline engine and run the benchmark. + + ```bash + python3 -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --num-prompts 10 + ``` + +- Benchmark online serving. Please use `sglang.launch_server` to launch a server first and run the following command. + + ```bash + python3 -m sglang.bench_serving --backend sglang --num-prompt 10 + ``` + +## Profile with PyTorch Profiler + +[Pytorch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) is a convenient basic tool to inspect kernel execution time, call stack, and kernel overlap and occupancy. + +- To profile a server + + ```bash + # set trace path + export SGLANG_TORCH_PROFILER_DIR=/root/sglang/profile_log + + # start server + python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct + + # send profiling request from client + python -m sglang.bench_serving --backend sglang --model meta-llama/Llama-3.1-8B-Instruct --num-prompts 10 --sharegpt-output-len 100 --profile + ``` + + Please make sure that the `SGLANG_TORCH_PROFILER_DIR` should be set at both server and client side, otherwise the trace file cannot be generated correctly . A secure way will be setting `SGLANG_TORCH_PROFILER_DIR` in the `.*rc` file of shell (e.g. `~/.bashrc` for bash shells). + +- To profile offline + + ```bash + export SGLANG_TORCH_PROFILER_DIR=/root/sglang/profile_log + python -m sglang.bench_offline_throughput --model-path meta-llama/Llama-3.1-8B-Instruct --dataset-name random --num-prompts 10 --profile --mem-frac=0.8 + ``` + +- View Traces + + Trace files can be loaded and visualized from: + + 1. https://ui.perfetto.dev/ (any browser) + 2. chrome://tracing (Chrome browser only) + + If browser cannot open trace file due to its large size, + client can generate a small trace file (<100MB) by controlling number of prompts and lengths of prompt outputs. + For example, when profiling a server, + + ```bash + python -m sglang.bench_serving --backend sglang --model meta-llama/Llama-3.1-8B-Instruct --num-prompts 2 --sharegpt-output-len 100 --profile + ``` + + This command sets the number of prompts to 2 with `--num-prompts` argument and limits the length of output sequences to 100 with `--sharegpt-output-len` argument, which can generate a small trace file for browser to open smoothly. + +## Profile with Nsight + +[Nsight systems](https://docs.nvidia.com/nsight-systems/) is an advanced tool that exposes more profiling details, such as register and shared memory usage, annotated code regions and low-level CUDA APIs and events. + +1. Prerequisite: + + Install using apt, or run inside a [NVIDIA Docker container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch/tags) or [SGLang Docker container](https://github.com/sgl-project/sglang/tree/main/docker). + + ```bash + # install nsys + # https://docs.nvidia.com/nsight-systems/InstallationGuide/index.html + apt update + apt install -y --no-install-recommends gnupg + echo "deb http://developer.download.nvidia.com/devtools/repos/ubuntu$(source /etc/lsb-release; echo "$DISTRIB_RELEASE" | tr -d .)/$(dpkg --print-architecture) /" | tee /etc/apt/sources.list.d/nvidia-devtools.list + apt-key adv --fetch-keys http://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub + apt update + apt install nsight-systems-cli + ``` + +2. To profile a single batch, use + + ```bash + nsys profile --trace-fork-before-exec=true --cuda-graph-trace=node python3 -m sglang.bench_one_batch --model meta-llama/Meta-Llama-3-8B --batch-size 64 --input-len 512 + ``` + +3. To profile a server, e.g. + + ```bash + # launch the server, set the delay and duration times according to needs + # after the duration time has been used up, server will be killed by nsys + + nsys profile --trace-fork-before-exec=true --cuda-graph-trace=node -o sglang.out --delay 60 --duration 70 python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disable-radix-cache + + # client + python3 -m sglang.bench_serving --backend sglang --num-prompts 1000 --dataset-name random --random-input 1024 --random-output 512 + ``` + + In practice, we recommend users to set `--duration` argument to a large value. Whenever user wants the server to stop profiling. Firstly run: + + ```bash + nsys sessions list + ``` + + to get the session id in the form of `profile-XXXXX`, then run: + + ```bash + nsys stop --session=profile-XXXXX + ``` + + to manually kill the profiler and generate `nsys-rep` files instantly. + +4. Use NVTX to annotate code regions, e.g. to see their execution time. + + ```bash + # install nvtx + pip install nvtx + ``` + + ```python + # code snippets + import nvtx + with nvtx.annotate("description", color="color"): + # some critical code + ``` + +## Other tips + +1. You can benchmark a model using dummy weights by only providing the config.json file. This allows for quick testing of model variants without training. To do so, add `--load-format dummy` to the above commands and then you only need a correct `config.json` under the checkpoint folder. +2. You can benchmark a model with modified configs (e.g., less layers) by using `--json-model-override-args`. For example, you can benchmark a model with only 2 layers and 2 kv heads using: + + ```bash + python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --batch 32 --input-len 256 --output-len 32 --load-format dummy --json-model-override-args '{"num_hidden_layers": 1, "num_key_value_heads": 1}' + ``` + +3. You can use `--python-backtrace=cuda` to see python call stack for all CUDA kernels, as in PyTorch Profiler. (Caveat: this can cause inaccurately long kernel runtimes for CUDA event based timing) +4. For more arguments see [Nsight Systems User Guide](https://docs.nvidia.com/nsight-systems/UserGuide/index.html). diff --git a/sglang/docs/references/contribution_guide.md b/sglang/docs/references/contribution_guide.md new file mode 100644 index 0000000..3d5e331 --- /dev/null +++ b/sglang/docs/references/contribution_guide.md @@ -0,0 +1,46 @@ +# Contribution Guide + +Welcome to **SGLang**! We appreciate your interest in contributing. This guide provides a concise overview of how to set up your environment, run tests, build documentation, and open a Pull Request (PR). Whether you’re fixing a small bug or developing a major feature, we encourage following these steps for a smooth contribution process. + +## Setting Up & Building from Source + +### Fork and Clone the Repository + +**Note**: New contributors do **not** have the write permission to push to the official SGLang repo. Please fork the repository under your GitHub account, then clone your fork locally. + +```bash +git clone https://github.com//sglang.git +``` + +### Install Dependencies & Build + +Refer to [Install SGLang from Source](https://docs.sglang.ai/start/install.html#method-2-from-source) documentation for more details on setting up the necessary dependencies. + +## Code Formatting with Pre-Commit + +We use [pre-commit](https://pre-commit.com/) to maintain consistent code style checks. Before pushing your changes, please run: + +```bash +pip3 install pre-commit +pre-commit install +pre-commit run --all-files +``` + +- **`pre-commit run --all-files`** manually runs all configured checks, applying fixes if possible. If it fails the first time, re-run it to ensure lint errors are fully resolved. Make sure your code passes all checks **before** creating a Pull Request. +- **Do not commit** directly to the `main` branch. Always create a new branch (e.g., `feature/my-new-feature`), push your changes, and open a PR from that branch. + +## Running Unit Tests & Adding to CI + +SGLang uses Python's built-in [unittest](https://docs.python.org/3/library/unittest.html) framework. For detailed instructions on running tests and adding them to CI, please refer to [test/README.md](https://github.com/sgl-project/sglang/tree/main/test/README.md). + +## Writing Documentation & Running Docs CI + +We recommend new contributors start from writing documentation, which helps you quickly understand SGLang codebase. For more details, please refer to [docs/README.md](https://github.com/sgl-project/sglang/tree/main/docs/README.md). + +## Tips for Newcomers + +If you want to contribute but don’t have a specific idea in mind, pick issues labeled [“good first issue” or “help wanted”](https://github.com/sgl-project/sglang/issues?q=is%3Aissue+label%3A%22good+first+issue%22%2C%22help+wanted%22). These tasks typically have lower complexity and provide an excellent introduction to the codebase. Also check out this [code walk-through](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/tree/main/sglang/code-walk-through) for a deeper look into SGLang’s workflow. + +If you have any questions or want to start a discussion, please feel free to ask in our [Slack channel](https://join.slack.com/t/sgl-fru7574/shared_invite/zt-2um0ad92q-LkU19KQTxCGzlCgRiOiQEw). + +Thank you for your interest in SGLang. Happy coding! diff --git a/sglang/docs/references/deepseek.md b/sglang/docs/references/deepseek.md new file mode 100644 index 0000000..9a079b9 --- /dev/null +++ b/sglang/docs/references/deepseek.md @@ -0,0 +1,166 @@ +# DeepSeek Usage + +SGLang provides several optimizations specifically designed for the DeepSeek model to boost its inference speed. This document outlines current optimizations for DeepSeek. Additionally, the SGLang team is actively developing enhancements for [DeepSeek V3](https://github.com/sgl-project/sglang/issues/2591). + +## Launch DeepSeek V3 with SGLang + +SGLang is recognized as one of the top engines for [DeepSeek model inference](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3). To run DeepSeek V3/R1 models, the requirements are as follows: + +| Weight Type | Configuration | +|------------|-------------------| +| **Full precision FP8**
*(recommended)* | 8 x H200 | +| | 8 x MI300X | +| | 2 x 8 x H100/800/20 | +| **Full precision BF16** | 2 x 8 x H200 | +| | 2 x 8 x MI300X | +| | 4 x 8 x H100/800/20 | +| | 4 x 8 x A100/A800 | +| **Quantized weights (AWQ)** | 8 x H100/800/20 | +| | 8 x A100/A800 | +| **Quantized weights (int8)** | 16 x A100/800 | +| | 32 x L40S | + + + +Detailed commands for reference: + +- [8 x H200](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#using-docker-recommended) +- [8 x MI300X](https://docs.sglang.ai/references/amd.html#running-deepseek-v3) +- [2 x 8 x H200](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-two-h208-nodes) +- [4 x 8 x A100](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-four-a1008-nodes) +- [8 x A100 (AWQ)](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-8-a100a800-with-awq-quantization) +- [16 x A100 (int8)](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-16-a100a800-with-int8-quantization) +- [32 x L40S (int8)](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-32-l40s-with-int8-quantization) + +### Download Weights + +If you encounter errors when starting the server, ensure the weights have finished downloading. It's recommended to download them beforehand or restart multiple times until all weights are downloaded. Please refer to [DeepSeek V3](https://huggingface.co/deepseek-ai/DeepSeek-V3-Base#61-inference-with-deepseek-infer-demo-example-only) official guide to download the weights. + +### Caching `torch.compile` + +The DeepSeek series have huge model weights, it takes some time to compile the model with `torch.compile` for the first time if you have added the flag `--enable-torch-compile`. You can refer [here](https://docs.sglang.ai/backend/hyperparameter_tuning.html#try-advanced-options) to optimize the caching of compilation results, so that the cache can be used to speed up the next startup. +### Launch with One node of 8 H200 + +Please refer to [the example](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#using-docker-recommended). **Note that Deepseek V3 is already in FP8. So we should not run it with any quantization arguments like `--quantization fp8 --kv-cache-dtype fp8_e5m2`.** Also, `--enable-dp-attention` can be useful to improve for Deepseek V3/R1's throughput. Please refer to [Data Parallelism Attention](https://docs.sglang.ai/references/deepseek.html#multi-head-latent-attention-mla-throughput-optimizations) for detail. + +### Running examples on Multi-node + +- [Serving with two H20*8 nodes](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-two-h208-nodes). + +- [Serving with two H200*8 nodes and docker](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-two-h2008-nodes-and-docker). + +- [Serving with four A100*8 nodes](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-four-a1008-nodes). + +## Optimizations + +### Multi-head Latent Attention (MLA) Throughput Optimizations + +**Description**: [MLA](https://arxiv.org/pdf/2405.04434) is an innovative attention mechanism introduced by the DeepSeek team, aimed at improving inference efficiency. SGLang has implemented specific optimizations for this, including: + +- **Weight Absorption**: By applying the associative law of matrix multiplication to reorder computation steps, this method balances computation and memory access and improves efficiency in the decoding phase. + +- **Flashinfer MLA Wrapper**: By providing `--enable-flashinfer-mla` argument, the server will use MLA kernels customized by Flashinfer. More details can be referred to [this document](https://docs.flashinfer.ai/api/mla.html). Under long input scenarios, flashinfer mla can improve performance significantly. Optimized triton kernels will be used when flashinfer mla is turned off. + +- **FP8 Quantization**: W8A8 FP8 and KV Cache FP8 quantization enables efficient FP8 inference. Additionally, we have implemented Batched Matrix Multiplication (BMM) operator to facilitate FP8 inference in MLA with weight absorption. + +- **CUDA Graph & Torch.compile**: Both MLA and Mixture of Experts (MoE) are compatible with CUDA Graph and Torch.compile, which reduces latency and accelerates decoding speed for small batch sizes. + +Overall, with these optimizations, we have achieved up to **7x** acceleration in output throughput compared to the previous version. + +

+ Multi-head Latent Attention for DeepSeek Series Models +

+ +**Usage**: MLA optimization is enabled by default, to disable, use `--disable-mla`. + +**Reference**: Check [Blog](https://lmsys.org/blog/2024-09-04-sglang-v0-3/#deepseek-multi-head-latent-attention-mla-throughput-optimizations) and [Slides](https://github.com/sgl-project/sgl-learning-materials/blob/main/slides/lmsys_1st_meetup_deepseek_mla.pdf) for more details. + +### Data Parallelism Attention + +**Description**: This optimization involves data parallelism (DP) for the MLA attention mechanism of DeepSeek Series Models, which allows for a significant reduction in the KV cache size, enabling larger batch sizes. Each DP worker independently handles different types of batches (prefill, decode, idle), which are then synchronized before and after processing through the Mixture-of-Experts (MoE) layer. + +

+ Data Parallelism Attention for DeepSeek Series Models +

+ +With data parallelism attention enabled, we have achieved up to **1.9x** decoding throughput improvement compared to the previous version. + +

+ Data Parallelism Attention Performance Comparison +

+ +**Usage**: +- This optimization is aimed at improving throughput and should be used for scenarios with high QPS (Queries Per Second). It can be enabled by `--enable-dp-attention` for DeepSeek models. +- Since v0.4.4, DP and TP attention can be flexibly combined. For example, to deploy DeepSeek-V3/R1 on 2 node with 8*H100, you can specify `--tp 16` and `--dp 2`, which means for attention part there are 2 DP groups, and in each DP group there are 8 TP groups. + +**Reference**: Check [Blog](https://lmsys.org/blog/2024-12-04-sglang-v0-4/#data-parallelism-attention-for-deepseek-models). + +### Multi Node Tensor Parallelism + +**Description**: For users with limited memory on a single node, SGLang supports serving DeepSeek Series Models, including DeepSeek V3, across multiple nodes using tensor parallelism. This approach partitions the model parameters across multiple GPUs or nodes to handle models that are too large for one node's memory. + +**Usage**: Check [here](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#example-serving-with-2-h208) for usage examples. + +### Block-wise FP8 + +**Description**: SGLang implements block-wise FP8 quantization with two key optimizations: + +- **Activation**: E4M3 format using per-token-per-128-channel sub-vector scales with online casting. + +- **Weight**: Per-128x128-block quantization for better numerical stability. + +**Usage**: Turn on by default for DeepSeek V3 models. + +### Multi-token Prediction +**Description**: SGLang implements DeepSeek V3 Multi-Token Prediction (MTP) based on [EAGLE speculative decoding](https://docs.sglang.ai/backend/speculative_decoding.html#EAGLE-Decoding). With this optimization, the decoding speed can be improved by **1.8x** for batch size 1 and **1.5x** for batch size 32 respectively on H200 TP8 setting. + +**Usage**: +Add arguments `--speculative-algorithm`, `--speculative-draft-model-path`, +`--speculative-num-steps`, `--speculative-eagle-topk` and `--speculative-num-draft-tokens` to enable this feature. For example: +``` +python3 -m sglang.launch_server --model-path deepseek-ai/DeepSeek-V3-0324 --speculative-algorithm EAGLE --speculative-draft-model-path lmsys/DeepSeek-V3-0324-NextN --speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2 --trust-remote-code --tp 8 +``` +- The draft model are available at huggingface: [lmsys/DeepSeek-V3-0324-NextN](https://huggingface.co/lmsys/DeepSeek-V3-0324-NextN), [lmsys/DeepSeek-R1-NextN](https://huggingface.co/lmsys/DeepSeek-R1-NextN). It can also be exported from original DeepSeek-V3/R1 model with [export_deepseek_nextn.py](https://github.com/sgl-project/sglang/blob/main/scripts/export_deepseek_nextn.py) script. +- The best configuratin for `--speculative-num-steps`, `--speculative-eagle-topk` and `--speculative-num-draft-tokens` can be searched with [bench_speculative.py](https://github.com/sgl-project/sglang/blob/main/scripts/playground/bench_speculative.py) script for given batch size. The minimum configuration is `--speculative-num-steps 1 --speculative-eagle-topk 1 --speculative-num-draft-tokens 2`, which can achieve speedup for larger batch sizes. +- Currently when using flashinfer mla wrapper (`--enable-flashinfer-mla`) and speculative decoding together, the `--speculative-eagle-topk` parameter should be set to `1`. +- To enable DeepSeek MTP for large batch sizes (>32), there are some parameters should be changed (Reference [this discussion](https://github.com/sgl-project/sglang/issues/4543#issuecomment-2737413756)): + - Adjust `--max-running-requests` to a larger number. The default value is `32` for MTP. For larger batch sizes, you should increase this value beyond the default value. + - Set `--cuda-graph-bs`. It's a list of batch sizes for cuda graph capture. The default captured batch sizes for speculative decoding is set [here](https://github.com/sgl-project/sglang/blob/49420741746c8f3e80e0eb17e7d012bfaf25793a/python/sglang/srt/model_executor/cuda_graph_runner.py#L126). You can include more batch sizes into it. + + +### Reasoning Content for DeepSeek R1 + +See [Separate Reasoning](https://docs.sglang.ai/backend/separate_reasoning.html). + +## FAQ + +1. **Question**: What should I do if model loading takes too long and NCCL timeout occurs? + + **Answer**: You can try to add `--dist-timeout 3600` when launching the model, this allows for 1-hour timeout. diff --git a/sglang/docs/references/deepseek.rst b/sglang/docs/references/deepseek.rst new file mode 100644 index 0000000..b45383a --- /dev/null +++ b/sglang/docs/references/deepseek.rst @@ -0,0 +1,6 @@ +Multi-Node Deployment +========================== +.. toctree:: + :maxdepth: 1 + + deepseek.md diff --git a/sglang/docs/references/deploy_on_k8s.md b/sglang/docs/references/deploy_on_k8s.md new file mode 100644 index 0000000..6fb1afb --- /dev/null +++ b/sglang/docs/references/deploy_on_k8s.md @@ -0,0 +1,345 @@ +# Deploy On Kubernetes + +This document is for deploying a RoCE network-based SGLang two-node inference service on a Kubernetes (K8S) cluster. + +[LeaderWorkerSet (LWS)](https://github.com/kubernetes-sigs/lws) is a Kubernetes API that aims to address common deployment patterns of AI/ML inference workloads. A major use case is for multi-host/multi-node distributed inference. + +SGLang can also be deployed with LWS on Kubernetes for distributed model serving. + +Please see this guide for more details on deploying SGLang on Kubernetes using LWS. + +Here we take the deployment of DeepSeek-R1 as an example. + +## Prerequisites + +1. At least two Kubernetes nodes, each with two H20 systems and eight GPUs, are required. + +2. Make sure your K8S cluster has LWS correctly installed. If it hasn't been set up yet, please follow the [installation instructions](https://github.com/kubernetes-sigs/lws/blob/main/site/content/en/docs/installation/_index.md). + +## Basic example + +For the basic example documentation, refer to [Deploy Distributed Inference Service with SGLang and LWS on GPUs](https://github.com/kubernetes-sigs/lws/tree/main/docs/examples/sglang). + +However, that document only covers the basic NCCL socket mode. + +In this section, we’ll make some simple modifications to adapt the setup to the RDMA scenario. + +## RDMA RoCE case + +* Check your env: + +```bash +[root@node1 ~]# ibstatus +Infiniband device 'mlx5_bond_0' port 1 status: + default gid: fe80:0000:0000:0000:0225:9dff:fe64:c79a + base lid: 0x0 + sm lid: 0x0 + state: 4: ACTIVE + phys state: 5: LinkUp + rate: 200 Gb/sec (2X NDR) + link_layer: Ethernet + +Infiniband device 'mlx5_bond_1' port 1 status: + default gid: fe80:0000:0000:0000:0225:9dff:fe6e:c3ec + base lid: 0x0 + sm lid: 0x0 + state: 4: ACTIVE + phys state: 5: LinkUp + rate: 200 Gb/sec (2X NDR) + link_layer: Ethernet + +Infiniband device 'mlx5_bond_2' port 1 status: + default gid: fe80:0000:0000:0000:0225:9dff:fe73:0dd7 + base lid: 0x0 + sm lid: 0x0 + state: 4: ACTIVE + phys state: 5: LinkUp + rate: 200 Gb/sec (2X NDR) + link_layer: Ethernet + +Infiniband device 'mlx5_bond_3' port 1 status: + default gid: fe80:0000:0000:0000:0225:9dff:fe36:f7ff + base lid: 0x0 + sm lid: 0x0 + state: 4: ACTIVE + phys state: 5: LinkUp + rate: 200 Gb/sec (2X NDR) + link_layer: Ethernet +``` + +* Prepare the `lws.yaml` file for deploying on k8s. + +```yaml +apiVersion: leaderworkerset.x-k8s.io/v1 +kind: LeaderWorkerSet +metadata: + name: sglang +spec: + replicas: 1 + leaderWorkerTemplate: + size: 2 + restartPolicy: RecreateGroupOnPodRestart + leaderTemplate: + metadata: + labels: + role: leader + spec: + dnsPolicy: ClusterFirstWithHostNet + hostNetwork: true + hostIPC: true + containers: + - name: sglang-leader + image: sglang:latest + securityContext: + privileged: true + env: + - name: NCCL_IB_GID_INDEX + value: "3" + - name: LWS_WORKER_INDEX + valueFrom: + fieldRef: + fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] + command: + - python3 + - -m + - sglang.launch_server + - --model-path + - /work/models + - --mem-fraction-static + - "0.93" + - --torch-compile-max-bs + - "8" + - --max-running-requests + - "20" + - --tp + - "16" # Size of Tensor Parallelism + - --dist-init-addr + - $(LWS_LEADER_ADDRESS):20000 + - --nnodes + - $(LWS_GROUP_SIZE) + - --node-rank + - $(LWS_WORKER_INDEX) + - --trust-remote-code + - --host + - "0.0.0.0" + - --port + - "40000" + resources: + limits: + nvidia.com/gpu: "8" + ports: + - containerPort: 40000 + readinessProbe: + tcpSocket: + port: 40000 + initialDelaySeconds: 15 + periodSeconds: 10 + volumeMounts: + - mountPath: /dev/shm + name: dshm + - name: model + mountPath: /work/models + - name: ib + mountPath: /dev/infiniband + volumes: + - name: dshm + emptyDir: + medium: Memory + - name: model + hostPath: + path: '< your models dir >' # modify it according your models dir + - name: ib + hostPath: + path: /dev/infiniband + workerTemplate: + spec: + dnsPolicy: ClusterFirstWithHostNet + hostNetwork: true + hostIPC: true + containers: + - name: sglang-worker + image: sglang:latest + securityContext: + privileged: true + env: + - name: NCCL_IB_GID_INDEX + value: "3" + - name: LWS_WORKER_INDEX + valueFrom: + fieldRef: + fieldPath: metadata.labels['leaderworkerset.sigs.k8s.io/worker-index'] + command: + - python3 + - -m + - sglang.launch_server + - --model-path + - /work/models + - --mem-fraction-static + - "0.93" + - --torch-compile-max-bs + - "8" + - --max-running-requests + - "20" + - --tp + - "16" # Size of Tensor Parallelism + - --dist-init-addr + - $(LWS_LEADER_ADDRESS):20000 + - --nnodes + - $(LWS_GROUP_SIZE) + - --node-rank + - $(LWS_WORKER_INDEX) + - --trust-remote-code + resources: + limits: + nvidia.com/gpu: "8" + volumeMounts: + - mountPath: /dev/shm + name: dshm + - name: model + mountPath: /work/models + - name: ib + mountPath: /dev/infiniband + volumes: + - name: dshm + emptyDir: + medium: Memory + - name: ib + hostPath: + path: /dev/infiniband + - name: model + hostPath: + path: /data1/models/deepseek_v3_moe +--- +apiVersion: v1 +kind: Service +metadata: + name: sglang-leader +spec: + selector: + leaderworkerset.sigs.k8s.io/name: sglang + role: leader + ports: + - protocol: TCP + port: 40000 + targetPort: 40000 + +``` + +* Then use `kubectl apply -f lws.yaml` you will get this output. + +```text +NAME READY STATUS RESTARTS AGE +sglang-0 0/1 Running 0 9s +sglang-0-1 1/1 Running 0 9s +``` + +Wait for the sglang leader (`sglang-0`) status to change to 1/1, which indicates it is `Ready`. + +You can use the command `kubectl logs -f sglang-0` to view the logs of the leader node. + +Once successful, you should see output like this: + +```text +[2025-02-17 05:27:24 TP1] Capture cuda graph end. Time elapsed: 84.89 s +[2025-02-17 05:27:24 TP6] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24 TP0] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24 TP7] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24 TP3] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24 TP2] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24 TP4] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24 TP1] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24 TP5] max_total_num_tokens=712400, chunked_prefill_size=8192, max_prefill_tokens=16384, max_running_requests=50, context_len=163840 +[2025-02-17 05:27:24] INFO: Started server process [1] +[2025-02-17 05:27:24] INFO: Waiting for application startup. +[2025-02-17 05:27:24] INFO: Application startup complete. +[2025-02-17 05:27:24] INFO: Uvicorn running on http://0.0.0.0:40000 (Press CTRL+C to quit) +[2025-02-17 05:27:25] INFO: 127.0.0.1:48908 - "GET /get_model_info HTTP/1.1" 200 OK +[2025-02-17 05:27:25 TP0] Prefill batch. #new-seq: 1, #new-token: 7, #cached-token: 0, cache hit rate: 0.00%, token usage: 0.00, #running-req: 0, #queue-req: 0 +[2025-02-17 05:27:32] INFO: 127.0.0.1:48924 - "POST /generate HTTP/1.1" 200 OK +[2025-02-17 05:27:32] The server is fired up and ready to roll! +``` + +If it doesn’t start up successfully, please follow these steps to check for any remaining issues. Thanks! + +### Debug + +* Set `NCCL_DEBUG=TRACE` to check if it is a NCCL communication problem. + +This should resolve most NCCL-related issues. + +***Notice: If you find that NCCL_DEBUG=TRACE is not effective in the container environment, but the process is stuck or you encounter hard-to-diagnose issues, try switching to a different container image. Some images may not handle standard error output properly.*** + +#### RoCE scenario + +* Please make sure that RDMA devices are available in the cluster environment. +* Please make sure that the nodes in the cluster have Mellanox NICs with RoCE. In this example, we use Mellanox ConnectX 5 model NICs, and the proper OFED driver has been installed. If not, please refer to the document [Install OFED Driver](https://docs.nvidia.com/networking/display/mlnxofedv461000/installing+mellanox+ofed) to install the driver. +* Check your env: + + ```shell + $ lspci -nn | grep Eth | grep Mellanox + 0000:7f:00.0 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + 0000:7f:00.1 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + 0000:c7:00.0 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + 0000:c7:00.1 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + 0001:08:00.0 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + 0001:08:00.1 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + 0001:a2:00.0 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + 0001:a2:00.1 Ethernet controller [0200]: Mellanox Technologies MT43244 BlueField-3 integrated ConnectX-7 network controller [15b3:a2dc] (rev 01) + ``` + +* Check the OFED driver: + + ```shell + ofed_info -s + OFED-internal-23.07-0.5.0: + ``` + +* Show RDMA link status and check IB devices: + + ```shell + $ rdma link show + 8/1: mlx5_bond_0/1: state ACTIVE physical_state LINK_UP netdev reth0 + 9/1: mlx5_bond_1/1: state ACTIVE physical_state LINK_UP netdev reth2 + 10/1: mlx5_bond_2/1: state ACTIVE physical_state LINK_UP netdev reth4 + 11/1: mlx5_bond_3/1: state ACTIVE physical_state LINK_UP netdev reth6 + + $ ibdev2netdev + 8/1: mlx5_bond_0/1: state ACTIVE physical_state LINK_UP netdev reth0 + 9/1: mlx5_bond_1/1: state ACTIVE physical_state LINK_UP netdev reth2 + 10/1: mlx5_bond_2/1: state ACTIVE physical_state LINK_UP netdev reth4 + 11/1: mlx5_bond_3/1: state ACTIVE physical_state LINK_UP netdev reth6 + ``` + +* Test RoCE network speed on the host: + + ```shell + yum install qperf + # for server： + execute qperf + # for client + qperf -t 60 -cm1 rc_rdma_write_bw + ``` + +* Check RDMA accessible in your container: + + ```shell + # ibv_devices + # ibv_devinfo + ``` + +## Keys to success + +* In the YAML configuration above, pay attention to the NCCL environment variable. For older versions of NCCL, you should check the NCCL_IB_GID_INDEX environment setting. +* NCCL_SOCKET_IFNAME is also crucial, but in a containerized environment, this typically isn’t an issue. +* In some cases, it’s necessary to configure GLOO_SOCKET_IFNAME correctly. +* NCCL_DEBUG is essential for troubleshooting, but I've found that sometimes it doesn't show error logs within containers. This could be related to the Docker image you're using. You may want to try switching images if needed. +* Avoid using Docker images based on Ubuntu 18.04, as they tend to have compatibility issues. + +## Remaining issues + +* In Kubernetes, Docker, or Containerd environments, we use hostNetwork to prevent performance degradation. +* We utilize privileged mode, which isn’t secure. Additionally, in containerized environments, full GPU isolation cannot be achieved. + +## TODO + +* Integrated with [k8s-rdma-shared-dev-plugin](https://github.com/Mellanox/k8s-rdma-shared-dev-plugin). diff --git a/sglang/docs/references/faq.md b/sglang/docs/references/faq.md new file mode 100644 index 0000000..eb43ee6 --- /dev/null +++ b/sglang/docs/references/faq.md @@ -0,0 +1,15 @@ +# Frequently Asked Questions + +## The results are not deterministic, even with a temperature of 0 + +You may notice that when you send the same request twice, the results from the engine will be slightly different, even when the temperature is set to 0. + +From our initial investigation, this indeterminism arises from two factors: dynamic batching and prefix caching. Roughly speaking, dynamic batching accounts for about 95% of the indeterminism, while prefix caching accounts for the remaining portion. The server runs dynamic batching under the hood. Different batch sizes can cause PyTorch/CuBLAS to dispatch to different CUDA kernels, which can lead to slight numerical differences. This difference accumulates across many layers, resulting in nondeterministic output when the batch size changes. Similarly, when prefix caching is enabled, it can also dispatch to different kernels. Even when the computations are mathematically equivalent, small numerical differences from different kernel implementations lead to the final nondeterministic outputs. + +To achieve more deterministic outputs in the current code, you can add `--disable-radix-cache` and send only one request at a time. The results will be mostly deterministic under this setting. + +We are still investigating the root causes and potential solutions. In the short term, we may introduce a "deterministic mode" that uses more padding to address the variance caused by dynamic batching. This mode will be more deterministic but slower. + +We have two issues to track our progress: +- The deterministic mode is tracked at [https://github.com/sgl-project/sglang/issues/1729](https://github.com/sgl-project/sglang/issues/1729). +- The per-request random seed is tracked at [https://github.com/sgl-project/sglang/issues/1335](https://github.com/sgl-project/sglang/issues/1335). diff --git a/sglang/docs/references/general.rst b/sglang/docs/references/general.rst new file mode 100644 index 0000000..8ea335d --- /dev/null +++ b/sglang/docs/references/general.rst @@ -0,0 +1,14 @@ + +General Guidance +========== + +.. toctree:: + :maxdepth: 1 + + supported_models.md + contribution_guide.md + troubleshooting.md + faq.md + learn_more.md + modelscope.md + production_metrics.md diff --git a/sglang/docs/references/hardware.rst b/sglang/docs/references/hardware.rst new file mode 100644 index 0000000..0500e54 --- /dev/null +++ b/sglang/docs/references/hardware.rst @@ -0,0 +1,7 @@ +Hardware Supports +========== +.. toctree:: + :maxdepth: 1 + + amd.md + nvidia_jetson.md diff --git a/sglang/docs/references/learn_more.md b/sglang/docs/references/learn_more.md new file mode 100644 index 0000000..d5c430e --- /dev/null +++ b/sglang/docs/references/learn_more.md @@ -0,0 +1,3 @@ +# Learn more + +You can find more blogs, slides, and videos about SGLang at [https://github.com/sgl-project/sgl-learning-materials](https://github.com/sgl-project/sgl-learning-materials). diff --git a/sglang/docs/references/modelscope.md b/sglang/docs/references/modelscope.md new file mode 100644 index 0000000..4740c27 --- /dev/null +++ b/sglang/docs/references/modelscope.md @@ -0,0 +1,28 @@ +# Use Models From ModelScope + +To use a model from [ModelScope](https://www.modelscope.cn), set the environment variable `SGLANG_USE_MODELSCOPE`. + +```bash +export SGLANG_USE_MODELSCOPE=true +``` + +We take [Qwen2-7B-Instruct](https://www.modelscope.cn/models/qwen/qwen2-7b-instruct) as an example. + +Launch the Server: +```bash +python -m sglang.launch_server --model-path qwen/Qwen2-7B-Instruct --port 30000 +``` + +Or start it by docker: + +```bash +docker run --gpus all \ + -p 30000:30000 \ + -v ~/.cache/modelscope:/root/.cache/modelscope \ + --env "SGLANG_USE_MODELSCOPE=true" \ + --ipc=host \ + lmsysorg/sglang:latest \ + python3 -m sglang.launch_server --model-path Qwen/Qwen2.5-7B-Instruct --host 0.0.0.0 --port 30000 +``` + +Note that modelscope uses a different cache directory than huggingface. You may need to set it manually to avoid running out of disk space. diff --git a/sglang/docs/references/multi_node.md b/sglang/docs/references/multi_node.md new file mode 100644 index 0000000..0d9abc1 --- /dev/null +++ b/sglang/docs/references/multi_node.md @@ -0,0 +1,92 @@ +# Multi-Node Deployment + +## Llama 3.1 405B + +**Run 405B (fp16) on Two Nodes** + +```bash +# replace 172.16.4.52:20000 with your own node ip address and port of the first node + +python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-405B-Instruct --tp 16 --dist-init-addr 172.16.4.52:20000 --nnodes 2 --node-rank 0 + +# replace 172.18.45.52:20000 with your own node ip address and port of the second node + +python3 -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-405B-Instruct --tp 16 --dist-init-addr 172.18.45.52:20000 --nnodes 2 --node-rank 1 +``` + +Note that LLama 405B (fp8) can also be launched on a single node. + +```bash +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-405B-Instruct-FP8 --tp 8 +``` + +## DeepSeek V3/R1 + +Please refer to [DeepSeek documents for reference](https://docs.sglang.ai/references/deepseek.html#running-examples-on-multi-node). + +## Multi-Node Inference on SLURM + +This example showcases how to serve SGLang server across multiple nodes by SLURM. Submit the following job to the SLURM cluster. + +``` +#!/bin/bash -l + +#SBATCH -o SLURM_Logs/%x_%j_master.out +#SBATCH -e SLURM_Logs/%x_%j_master.err +#SBATCH -D ./ +#SBATCH -J Llama-405B-Online-Inference-TP16-SGL + +#SBATCH --nodes=2 +#SBATCH --ntasks=2 +#SBATCH --ntasks-per-node=1 # Ensure 1 task per node +#SBATCH --cpus-per-task=18 +#SBATCH --mem=224GB +#SBATCH --partition="lmsys.org" +#SBATCH --gres=gpu:8 +#SBATCH --time=12:00:00 + +echo "[INFO] Activating environment on node $SLURM_PROCID" +if ! source ENV_FOLDER/bin/activate; then + echo "[ERROR] Failed to activate environment" >&2 + exit 1 +fi + +# Define parameters +model=MODEL_PATH +tp_size=16 + +echo "[INFO] Running inference" +echo "[INFO] Model: $model" +echo "[INFO] TP Size: $tp_size" + +# Set NCCL initialization address using the hostname of the head node +HEAD_NODE=$(scontrol show hostname "$SLURM_NODELIST" | head -n 1) +NCCL_INIT_ADDR="${HEAD_NODE}:8000" +echo "[INFO] NCCL_INIT_ADDR: $NCCL_INIT_ADDR" + +# Launch the model server on each node using SLURM +srun --ntasks=2 --nodes=2 --output="SLURM_Logs/%x_%j_node$SLURM_NODEID.out" \ + --error="SLURM_Logs/%x_%j_node$SLURM_NODEID.err" \ + python3 -m sglang.launch_server \ + --model-path "$model" \ + --grammar-backend "xgrammar" \ + --tp "$tp_size" \ + --dist-init-addr "$NCCL_INIT_ADDR" \ + --nnodes 2 \ + --node-rank "$SLURM_NODEID" & + +# Wait for the NCCL server to be ready on port 30000 +while ! nc -z "$HEAD_NODE" 30000; do + sleep 1 + echo "[INFO] Waiting for $HEAD_NODE:30000 to accept connections" +done + +echo "[INFO] $HEAD_NODE:30000 is ready to accept connections" + +# Keep the script running until the SLURM job times out +wait +``` + +Then, you can test the server by sending requests following other [documents](https://docs.sglang.ai/backend/openai_api_completions.html). + +Thanks for [aflah02](https://github.com/aflah02) for providing the example, based on his [blog post](https://aflah02.substack.com/p/multi-node-llm-inference-with-sglang). diff --git a/sglang/docs/references/nvidia_jetson.md b/sglang/docs/references/nvidia_jetson.md new file mode 100644 index 0000000..1dc4850 --- /dev/null +++ b/sglang/docs/references/nvidia_jetson.md @@ -0,0 +1,76 @@ +# Apply SGLang on NVIDIA Jetson Orin + +## Prerequisites + +Before starting, ensure the following: + +- [**NVIDIA Jetson AGX Orin Devkit**](https://www.nvidia.com/en-us/autonomous-machines/embedded-systems/jetson-orin/) is set up with **JetPack 6.1** or later. +- **CUDA Toolkit** and **cuDNN** are installed. +- Verify that the Jetson AGX Orin is in **high-performance mode**: +```bash +sudo nvpmodel -m 0 +``` +* * * * * +## Installing and running SGLang with Jetson Containers +Clone the jetson-containers github repository: +``` +git clone https://github.com/dusty-nv/jetson-containers.git +``` +Run the installation script: +``` +bash jetson-containers/install.sh +``` +Build the container: +``` +CUDA_VERSION=12.6 jetson-containers build sglang +``` +Run the container: +``` +docker run --runtime nvidia -it --rm --network=host IMAGE_NAME +``` +* * * * * + +Running Inference +----------------------------------------- + +Launch the server: +```bash +python -m sglang.launch_server \ + --model-path deepseek-ai/DeepSeek-R1-Distill-Llama-8B \ + --device cuda \ + --dtype half \ + --attention-backend flashinfer \ + --mem-fraction-static 0.8 \ + --context-length 8192 +``` +The quantization and limited context length (`--dtype half --context-length 8192`) are due to the limited computational resources in [Nvidia jetson kit](https://www.nvidia.com/en-us/autonomous-machines/embedded-systems/jetson-orin/). A detailed explanation can be found in [Server Arguments](../backend/server_arguments.md). + +After launching the engine, refer to [Chat completions](https://docs.sglang.ai/backend/openai_api_completions.html#Usage) to test the usability. +* * * * * +Running quantization with TorchAO +------------------------------------- +TorchAO is suggested to NVIDIA Jetson Orin. +```bash +python -m sglang.launch_server \ + --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ + --device cuda \ + --dtype bfloat16 \ + --attention-backend flashinfer \ + --mem-fraction-static 0.8 \ + --context-length 8192 \ + --torchao-config int4wo-128 +``` +This enables TorchAO's int4 weight-only quantization with a 128-group size. The usage of `--torchao-config int4wo-128` is also for memory efficiency. + + +* * * * * +Structured output with XGrammar +------------------------------- +Please refer to [SGLang doc structured output](../backend/structured_outputs.ipynb). +* * * * * + +Thanks to the support from [shahizat](https://github.com/shahizat). + +References +---------- +- [NVIDIA Jetson AGX Orin Documentation](https://developer.nvidia.com/embedded/jetson-agx-orin) diff --git a/sglang/docs/references/performance_tuning.rst b/sglang/docs/references/performance_tuning.rst new file mode 100644 index 0000000..6cc20e0 --- /dev/null +++ b/sglang/docs/references/performance_tuning.rst @@ -0,0 +1,7 @@ +Performance Tuning +==================== +.. toctree:: + :maxdepth: 1 + + benchmark_and_profiling.md + accuracy_evaluation.md diff --git a/sglang/docs/references/production_metrics.md b/sglang/docs/references/production_metrics.md new file mode 100644 index 0000000..7183158 --- /dev/null +++ b/sglang/docs/references/production_metrics.md @@ -0,0 +1,173 @@ +# Production Metrics + +SGLang exposes the following metrics via Prometheus. The metrics are namespaced by `$name` (the model name). + +An example of the monitoring dashboard is available in [examples/monitoring/grafana.json](../examples/monitoring/grafana.json). + +Here is an example of the metrics: + +``` +$ curl http://localhost:30000/metrics +# HELP sglang:prompt_tokens_total Number of prefill tokens processed. +# TYPE sglang:prompt_tokens_total counter +sglang:prompt_tokens_total{model_name="meta-llama/Llama-3.1-8B-Instruct"} 8.128902e+06 +# HELP sglang:generation_tokens_total Number of generation tokens processed. +# TYPE sglang:generation_tokens_total counter +sglang:generation_tokens_total{model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.557572e+06 +# HELP sglang:token_usage The token usage +# TYPE sglang:token_usage gauge +sglang:token_usage{model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.28 +# HELP sglang:cache_hit_rate The cache hit rate +# TYPE sglang:cache_hit_rate gauge +sglang:cache_hit_rate{model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.007507552643049313 +# HELP sglang:time_to_first_token_seconds Histogram of time to first token in seconds. +# TYPE sglang:time_to_first_token_seconds histogram +sglang:time_to_first_token_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct"} 2.3518979474117756e+06 +sglang:time_to_first_token_seconds_bucket{le="0.001",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 +sglang:time_to_first_token_seconds_bucket{le="0.005",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 +sglang:time_to_first_token_seconds_bucket{le="0.01",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 +sglang:time_to_first_token_seconds_bucket{le="0.02",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 +sglang:time_to_first_token_seconds_bucket{le="0.04",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1.0 +sglang:time_to_first_token_seconds_bucket{le="0.06",model_name="meta-llama/Llama-3.1-8B-Instruct"} 3.0 +sglang:time_to_first_token_seconds_bucket{le="0.08",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:time_to_first_token_seconds_bucket{le="0.1",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:time_to_first_token_seconds_bucket{le="0.25",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:time_to_first_token_seconds_bucket{le="0.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:time_to_first_token_seconds_bucket{le="0.75",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:time_to_first_token_seconds_bucket{le="1.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 27.0 +sglang:time_to_first_token_seconds_bucket{le="2.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 140.0 +sglang:time_to_first_token_seconds_bucket{le="5.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 314.0 +sglang:time_to_first_token_seconds_bucket{le="7.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 941.0 +sglang:time_to_first_token_seconds_bucket{le="10.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1330.0 +sglang:time_to_first_token_seconds_bucket{le="15.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1970.0 +sglang:time_to_first_token_seconds_bucket{le="20.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 2326.0 +sglang:time_to_first_token_seconds_bucket{le="25.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 2417.0 +sglang:time_to_first_token_seconds_bucket{le="30.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 2513.0 +sglang:time_to_first_token_seconds_bucket{le="+Inf",model_name="meta-llama/Llama-3.1-8B-Instruct"} 11008.0 +sglang:time_to_first_token_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct"} 11008.0 +# HELP sglang:e2e_request_latency_seconds Histogram of End-to-end request latency in seconds +# TYPE sglang:e2e_request_latency_seconds histogram +sglang:e2e_request_latency_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct"} 3.116093850019932e+06 +sglang:e2e_request_latency_seconds_bucket{le="0.3",model_name="meta-llama/Llama-3.1-8B-Instruct"} 0.0 +sglang:e2e_request_latency_seconds_bucket{le="0.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:e2e_request_latency_seconds_bucket{le="0.8",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:e2e_request_latency_seconds_bucket{le="1.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:e2e_request_latency_seconds_bucket{le="1.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:e2e_request_latency_seconds_bucket{le="2.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:e2e_request_latency_seconds_bucket{le="2.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 6.0 +sglang:e2e_request_latency_seconds_bucket{le="5.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.0 +sglang:e2e_request_latency_seconds_bucket{le="10.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 10.0 +sglang:e2e_request_latency_seconds_bucket{le="15.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 11.0 +sglang:e2e_request_latency_seconds_bucket{le="20.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 14.0 +sglang:e2e_request_latency_seconds_bucket{le="30.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 247.0 +sglang:e2e_request_latency_seconds_bucket{le="40.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 486.0 +sglang:e2e_request_latency_seconds_bucket{le="50.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 845.0 +sglang:e2e_request_latency_seconds_bucket{le="60.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1513.0 +sglang:e2e_request_latency_seconds_bucket{le="+Inf",model_name="meta-llama/Llama-3.1-8B-Instruct"} 11228.0 +sglang:e2e_request_latency_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct"} 11228.0 +# HELP sglang:time_per_output_token_seconds Histogram of time per output token in seconds. +# TYPE sglang:time_per_output_token_seconds histogram +sglang:time_per_output_token_seconds_sum{model_name="meta-llama/Llama-3.1-8B-Instruct"} 866964.5791549598 +sglang:time_per_output_token_seconds_bucket{le="0.005",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1.0 +sglang:time_per_output_token_seconds_bucket{le="0.01",model_name="meta-llama/Llama-3.1-8B-Instruct"} 73.0 +sglang:time_per_output_token_seconds_bucket{le="0.015",model_name="meta-llama/Llama-3.1-8B-Instruct"} 382.0 +sglang:time_per_output_token_seconds_bucket{le="0.02",model_name="meta-llama/Llama-3.1-8B-Instruct"} 593.0 +sglang:time_per_output_token_seconds_bucket{le="0.025",model_name="meta-llama/Llama-3.1-8B-Instruct"} 855.0 +sglang:time_per_output_token_seconds_bucket{le="0.03",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1035.0 +sglang:time_per_output_token_seconds_bucket{le="0.04",model_name="meta-llama/Llama-3.1-8B-Instruct"} 1815.0 +sglang:time_per_output_token_seconds_bucket{le="0.05",model_name="meta-llama/Llama-3.1-8B-Instruct"} 11685.0 +sglang:time_per_output_token_seconds_bucket{le="0.075",model_name="meta-llama/Llama-3.1-8B-Instruct"} 433413.0 +sglang:time_per_output_token_seconds_bucket{le="0.1",model_name="meta-llama/Llama-3.1-8B-Instruct"} 4.950195e+06 +sglang:time_per_output_token_seconds_bucket{le="0.15",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.039435e+06 +sglang:time_per_output_token_seconds_bucket{le="0.2",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.171662e+06 +sglang:time_per_output_token_seconds_bucket{le="0.3",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.266055e+06 +sglang:time_per_output_token_seconds_bucket{le="0.4",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.296752e+06 +sglang:time_per_output_token_seconds_bucket{le="0.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.312226e+06 +sglang:time_per_output_token_seconds_bucket{le="0.75",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.339675e+06 +sglang:time_per_output_token_seconds_bucket{le="1.0",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.357747e+06 +sglang:time_per_output_token_seconds_bucket{le="2.5",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.389414e+06 +sglang:time_per_output_token_seconds_bucket{le="+Inf",model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.400757e+06 +sglang:time_per_output_token_seconds_count{model_name="meta-llama/Llama-3.1-8B-Instruct"} 7.400757e+06 +# HELP sglang:func_latency_seconds Function latency in seconds +# TYPE sglang:func_latency_seconds histogram +sglang:func_latency_seconds_sum{name="generate_request"} 4.514771912145079 +sglang:func_latency_seconds_bucket{le="0.05",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="0.07500000000000001",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="0.1125",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="0.16875",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="0.253125",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="0.3796875",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="0.56953125",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="0.8542968750000001",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="1.2814453125",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="1.9221679687500002",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="2.8832519531250003",name="generate_request"} 14006.0 +sglang:func_latency_seconds_bucket{le="4.3248779296875",name="generate_request"} 14007.0 +sglang:func_latency_seconds_bucket{le="6.487316894531251",name="generate_request"} 14007.0 +sglang:func_latency_seconds_bucket{le="9.730975341796876",name="generate_request"} 14007.0 +sglang:func_latency_seconds_bucket{le="14.596463012695313",name="generate_request"} 14007.0 +sglang:func_latency_seconds_bucket{le="21.89469451904297",name="generate_request"} 14007.0 +sglang:func_latency_seconds_bucket{le="32.84204177856446",name="generate_request"} 14007.0 +sglang:func_latency_seconds_bucket{le="49.26306266784668",name="generate_request"} 14007.0 +sglang:func_latency_seconds_bucket{le="+Inf",name="generate_request"} 14007.0 +sglang:func_latency_seconds_count{name="generate_request"} 14007.0 +# HELP sglang:num_running_reqs The number of running requests +# TYPE sglang:num_running_reqs gauge +sglang:num_running_reqs{model_name="meta-llama/Llama-3.1-8B-Instruct"} 162.0 +# HELP sglang:num_used_tokens The number of used tokens +# TYPE sglang:num_used_tokens gauge +sglang:num_used_tokens{model_name="meta-llama/Llama-3.1-8B-Instruct"} 123859.0 +# HELP sglang:gen_throughput The generate throughput (token/s) +# TYPE sglang:gen_throughput gauge +sglang:gen_throughput{model_name="meta-llama/Llama-3.1-8B-Instruct"} 86.50814177726902 +# HELP sglang:num_queue_reqs The number of requests in the waiting queue +# TYPE sglang:num_queue_reqs gauge +sglang:num_queue_reqs{model_name="meta-llama/Llama-3.1-8B-Instruct"} 2826.0 +``` + +## Setup Guide + +To setup a monitoring dashboard, you can use the following docker compose file: [examples/monitoring/docker-compose.yaml](../examples/monitoring/docker-compose.yaml). + +Assume you have sglang server running at `localhost:30000`, to start the server, ensure you have `--enable-metrics` flag enabled: + +```bash +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct \ +--port 30000 --host 0.0.0.0 --enable-metrics +``` + +To start the monitoring dashboard (prometheus + grafana), cd to `examples/monitoring` and run: + +```bash +docker compose -f compose.yaml -p monitoring up +``` + +Then you can access the Grafana dashboard at http://localhost:3000. + +### Grafana Dashboard + +In a new Grafana setup, ensure that you have the `Prometheus` data source enabled. To check that, go to `http://localhost:3000/connections/datasources` and ensure that `Prometheus` is enabled. + +If not, click `Add data source` -> `Prometheus`, set Prometheus URL to `http://localhost:9090`, and click `Save & Test`. + +To import the Grafana dashboard, click `+` -> `Import` -> `Upload JSON file` -> `Upload` and select [grafana.json](../examples/monitoring/grafana.json). + +### Troubleshooting + +#### Check if the variables are created + +The example dashboard assume you have the following variables avaliable: +- `model_name` (name: `model_name`, label: `model name`, Data source: `Prometheus`, Type: `Label values`) +- `instance` (name: `instance`, label: `instance`, Data source: `Prometheus`, Type: `Label values`) + +If you don't have these variables, you can create them manually. + +To create a variable, go to dashboard settings, `Variables` -> `New variable`. + +You should be able to see the preview the values (e.g. `meta-llama/Llama-3.1-8B-Instruct` for `model_name`). + +#### Check if the metrics are being collected + +Run `python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 3000 --random-input 1024 --random-output 1024 --random-range-ratio 0.5` to generate some requests. + +Then you should be able to see the metrics in the Grafana dashboard. diff --git a/sglang/docs/references/supported_models.md b/sglang/docs/references/supported_models.md new file mode 100644 index 0000000..790e9a8 --- /dev/null +++ b/sglang/docs/references/supported_models.md @@ -0,0 +1,145 @@ +# Supported Models + +## Generative Models +- Llama / Llama 2 / Llama 3 / Llama 3.1 / Llama 3.2 / Llama 3.3 +- Mistral / Mixtral / Mistral NeMo / Mistral Small 3 +- Gemma / Gemma 2 / Gemma3 +- Qwen / Qwen 2 / Qwen 2 MoE / Qwen 2 VL / Qwen 2.5 VL / Olympic Coder +- DeepSeek / DeepSeek 2 / [DeepSeek 3](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3) +- OLMoE +- [LLaVA-OneVision](https://llava-vl.github.io/blog/2024-08-05-llava-onevision/) + - `python3 -m sglang.launch_server --model-path lmms-lab/llava-onevision-qwen2-7b-ov --port=30000 --chat-template=chatml-llava` + - `python3 -m sglang.launch_server --model-path lmms-lab/llava-onevision-qwen2-72b-ov --port=30000 --tp-size=8 --chat-template=chatml-llava` + - Query the server with the [OpenAI Vision API](https://platform.openai.com/docs/guides/vision). See examples at [test/srt/test_vision_openai_server.py](https://github.com/sgl-project/sglang/blob/main/test/srt/test_vision_openai_server.py) +- LLaVA 1.5 / 1.6 / NeXT + - `python -m sglang.launch_server --model-path lmms-lab/llama3-llava-next-8b --port=30000 --tp-size=1 --chat-template=llava_llama_3` + - `python -m sglang.launch_server --model-path lmms-lab/llava-next-72b --port=30000 --tp-size=8 --chat-template=chatml-llava` + - Query the server with the [OpenAI Vision API](https://platform.openai.com/docs/guides/vision). See examples at [test/srt/test_vision_openai_server.py](https://github.com/sgl-project/sglang/blob/main/test/srt/test_vision_openai_server.py) +- Yi-VL +- StableLM +- Command-R +- DBRX +- Grok +- ChatGLM +- InternLM 2 +- Exaone 3 +- BaiChuan2 +- MiniCPM / MiniCPM 3 / MiniCPM-v / MiniCPM-o +- XVERSE / XVERSE MoE +- SmolLM +- GLM-4 +- Phi-3 / Phi-4 +- Phi-3-Small +- IBM Granite 3 +- Janus-Pro-1B / Janus-Pro-7B +- Deepseek-VL2 / Deepseek-VL2-small +- Gemma 3 (it) + +## Embedding Models + +- LlamaEmbeddingModel +- Mistral embedding models +- Qwen embedding models + - `python -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-7B-instruct --is-embedding` +- Multi-modal embedding models + - `python -m sglang.launch_server --model-path Alibaba-NLP/gme-Qwen2-VL-2B-Instruct --is-embedding --chat-template gme-qwen2-vl` +- CLIP + - `python -m sglang.launch_server --model-path openai/clip-vit-large-patch14-336 --is-embedding` + +## Reward Models + +- LlamaForSequenceClassification + - `python -m sglang.launch_server --model-path Skywork/Skywork-Reward-Llama-3.1-8B-v0.2 --is-embedding` +- Gemma2ForSequenceClassification + - `python -m sglang.launch_server --model-path Skywork/Skywork-Reward-Gemma-2-27B-v0.2 --is-embedding` +- InternLM2ForRewardModel + - `python -m sglang.launch_server --model-path internlm/internlm2-7b-reward --is-embedding --trust-remote-code` +- Qwen2ForRewardModel + - `python -m sglang.launch_server --model-path Qwen/Qwen2.5-Math-RM-72B --is-embedding --trust-remote-code --tp-size=4` +- Qwen2ForSequenceClassification + - `python -m sglang.launch_server --model-path jason9693/Qwen2.5-1.5B-apeach --is-embedding --trust-remote-code` +## How to Support a New Language Model + +To support a new model in SGLang, you only need to add a single file under [SGLang Models Directory](https://github.com/sgl-project/sglang/tree/main/python/sglang/srt/models). +You can learn from existing model implementations and create new files for the new models. +For most models, you should be able to find a similar model to start with (e.g., starting from Llama). + +## How to Support a New vLM + +To support a new vision-language model (vLM) in SGLang, there are several key components in addition to the standard +LLM. + +1. **Register your new model as multimodal**: Extend `is_multimodal_model` in [ + `model_config.py`](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/configs/model_config.py) to + return True for your model. +2. **Process Images**: Define a new `Processor` class that inherits from `BaseProcessor` and register this + processor as your model's dedicated processor. See [ + `multimodal_processor.py`](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/managers/multimodal_processor.py) + for more details. +3. **Handle Image Tokens**: Implement a `pad_input_ids` function for your new model, in which image tokens in the prompt + should be expanded and replaced with image-hashes, so that SGLang can recognize different images for + `RadixAttention`. +4. Replace Multi-headed `Attention` of ViT with SGLang's `VisionAttention`. + +You can refer [Qwen2VL](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/models/qwen2_vl.py) or other +vLMs. These models demonstrate how to properly handle both multimodal and textual inputs. + +You should test the new vLM locally against hf models. See [`mmmu`](https://github.com/sgl-project/sglang/tree/main/benchmark/mmmu) for an example. + +### Test the correctness + +#### Interactive debugging +For interactive debugging, you can compare the outputs of huggingface/transformers and SGLang. +The following two commands should give the same text output and very similar prefill logits. + +- Get the reference output by `python3 scripts/playground/reference_hf.py --model-path [new model] --model-type {text,vlm}` +- Get the SGLang output by `python3 -m sglang.bench_one_batch --correct --model [new model]` + +#### Add the model to the test suite +To make sure the new model is well maintained in the future, it is better to add it to the test suite. +You can add it to the `ALL_OTHER_MODELS` list in the [test_generation_models.py](https://github.com/sgl-project/sglang/blob/main/test/srt/models/test_generation_models.py) and run the following command to test it. + +For example, if the model is Qwen/Qwen2-1.5B +``` +ONLY_RUN=Qwen/Qwen2-1.5B python3 -m unittest test_generation_models.TestGenerationModels.test_others +``` + +### Port a model from vLLM to SGLang +Another valuable resource is the [vLLM Models Directory](https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/models). vLLM has extensive coverage of models, and SGLang reuses vLLM's interface and some layers to implement the models. This similarity makes it easy to port many models from vLLM to SGLang. + +To port a model from vLLM to SGLang, you can compare these two files [SGLang Llama Implementation](https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/models/llama.py) and [vLLM Llama Implementation](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/models/llama.py). This comparison will help you understand how to convert a model implementation from vLLM to SGLang. The major difference is the replacement of Attention with RadixAttention. The other parts are almost identical. Specifically, + - Replace vllm's `Attention` with `RadixAttention`. Note that you need to pass `layer_id` all the way to `RadixAttention`. + - Replace vllm's `LogitsProcessor` with SGLang's `LogitsProcessor`. + - Replace Multi-headed `Attention` of ViT with SGLang's `VisionAttention`. + - Replace other vLLM layers with SGLang layers (e.g., `RMSNorm`, `SiluAndMul`). + - Remove `Sample`. + - Change `forward()` functions, and add `forward_batch`. + - Add `EntryClass` at the end. + - Please ensure the new implementation uses **only SGLang components and does not rely on any vLLM components**. + +### Registering an external model implementation + +In addition to the methods described above, you can also register your new model with the `ModelRegistry` before launching the server. This approach is useful if you want to integrate your model without needing to modify the source code. + +Here is how you can do it: + +```python +from sglang.srt.models.registry import ModelRegistry +from sglang.srt.entrypoints.http_server import launch_server + +# for a single model, you can add it to the registry +ModelRegistry.models[model_name] = model_class + +# for multiple models, you can imitate the import_model_classes() function in sglang/srt/models/registry.py +from functools import lru_cache + +@lru_cache() +def import_new_model_classes(): + model_arch_name_to_cls = {} + ... + return model_arch_name_to_cls + +ModelRegistry.models.update(import_new_model_classes()) + +launch_server(server_args) +``` diff --git a/sglang/docs/references/troubleshooting.md b/sglang/docs/references/troubleshooting.md new file mode 100644 index 0000000..5a7646a --- /dev/null +++ b/sglang/docs/references/troubleshooting.md @@ -0,0 +1,14 @@ +# Troubleshooting + +This page lists some common errors and tips for fixing them. + +## CUDA out of memory +If you see out of memory (OOM) errors, you can try to tune the following parameters. +- If OOM happens during prefill, try to decrease `--chunked-prefill-size` to `4096` or `2048`. +- If OOM happens during decoding, try to decrease `--max-running-requests`. +- You can also try to decrease `--mem-fraction-static`, which reduces the memory usage of the KV cache memory pool and helps both prefill and decoding. + +## CUDA error: an illegal memory access was encountered +This error may be due to kernel errors or out-of-memory issues. +- If it is a kernel error, it is not easy to fix. Please file an issue on the GitHub. +- If it is out-of-memory, sometimes it will report this error instead of "Out-of-memory." Please refer to the above section to avoid the OOM. diff --git a/sglang/docs/requirements.txt b/sglang/docs/requirements.txt new file mode 100644 index 0000000..1a7e5d4 --- /dev/null +++ b/sglang/docs/requirements.txt @@ -0,0 +1,20 @@ +ipykernel +ipywidgets +jupyter_client +markdown>=3.4.0 +matplotlib +myst-parser +nbconvert +nbsphinx +pandoc +pillow +pydantic +sphinx +sphinx-book-theme +sphinx-copybutton +sphinx-tabs +nbstripout +sphinxcontrib-mermaid +urllib3<2.0.0 +gguf>=0.10.0 +sphinx-autobuild diff --git a/sglang/docs/router/router.md b/sglang/docs/router/router.md new file mode 100644 index 0000000..26fa942 --- /dev/null +++ b/sglang/docs/router/router.md @@ -0,0 +1,165 @@ +# Router for Data Parallelism + +Given multiple GPUs running multiple SGLang Runtimes, SGLang Router distributes the requests to different Runtimes with its unique cache-aware load-balancing algorithm. + +The router is an independent Python package, and it can be used as a drop-in replacement for the OpenAI API. + +## Installation + +```bash +pip install sglang-router +``` + +Detailed usage of the router can be found in [launch_router](https://github.com/sgl-project/sglang/blob/main/sgl-router/py_src/sglang_router/launch_router.py) and [launch_server](https://github.com/sgl-project/sglang/blob/main/sgl-router/py_src/sglang_router/launch_server.py). Also, you can directly run the following command to see the usage of the router. + +```bash +python -m sglang_router.launch_server --help +python -m sglang_router.launch_router --help +``` + +The router supports two working modes: + +1. Co-launch Router and Runtimes +2. Launch Runtimes and Router separately + +## Co-launch Router and Runtimes + +This will be a drop-in replacement for the existing `--dp-size` argument of SGLang Runtime. Under the hood, it uses multi-processes to launch multiple workers, wait for them to be ready, then connect the router to all workers. + +```bash +python -m sglang_router.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --dp-size 4 --host 0.0.0.0 +``` + +After the server is ready, you can directly send requests to the router as the same way as sending requests to each single worker. + +Please adjust the batchsize accordingly to archieve maximum throughput. + +```python +import requests + +url = "http://localhost:30000/generate" +data = {"text": "What is the capital of France?"} + +response = requests.post(url, json=data) +print(response.json()) +``` + +## Launch Runtimes and Router Separately + +This is useful for multi-node DP. First, launch workers on multiple nodes, then launch a router on the main node, and connect the router to all workers. + +```bash +python -m sglang_router.launch_router --worker-urls http://worker_url_1 http://worker_url_2 +``` + +## Dynamic Scaling APIs + +We offer `/add_worker` and `/remove_worker` APIs to dynamically add or remove workers from the router. + +- `/add_worker` + +Usage: + +```bash +curl -X POST http://localhost:30000/add_worker?url=http://worker_url_1 +``` + +Example: + +```bash +python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --port 30001 + +curl -X POST http://localhost:30000/add_worker?url=http://127.0.0.1:30001 + +# Successfully added worker: http://127.0.0.1:30001 +``` + +- `/remove_worker` + +Usage: + +```bash +curl -X POST http://localhost:30000/remove_worker?url=http://worker_url_1 +``` + +Example: + +```bash +curl -X POST http://localhost:30000/remove_worker?url=http://127.0.0.1:30001 + +# Successfully removed worker: http://127.0.0.1:30001 +``` + +Note: + +- For cache-aware router, the worker will be removed from the tree and the queues. + +## Fault Tolerance + +We provide retries based for failure tolerance. + +1. If the request to a worker fails for `max_worker_retries` times, the router will remove the worker from the router and move on to the next worker. +2. If the total number of retries exceeds `max_total_retries`, the router will return an error. + +Note: + +- `max_worker_retries` is 3 and `max_total_retries` is 6 by default. + +## Routing Strategies + +### Cache-Aware Load-Balancing Router + +The native router combines two strategies to optimize both cache utilization and request distribution: + +1. Cache-Aware Routing (Approximate Tree) +2. Load-Balancing Routing (Shortest Queue with Balance Thresholds) + +The router dynamically switches between these strategies based on load conditions: + +- Uses load balancing when the system is imbalanced +- Uses cache-aware routing when the system is balanced + +A system is considered imbalanced if both conditions are met: + +1. (max_load - min_load) > balance_abs_threshold +2. max_load > balance_rel_threshold * min_load + +***Cache-Aware Routing (Approximate Tree)*** + +When the workers are considered to be balanced, the router maintains an approximate radix tree for each worker based on request history, eliminating the need for direct cache state queries on each worker. The tree stores raw text characters instead of token IDs to avoid tokenization overhead. + +Process: + +1. For each request, find the worker with the highest prefix match. + + - If match rate > cache_threshold, route the request to the worker with highest match (likely has relevant data cached) + - If match rate ≤ cache_threshold, route the request to the worker with smallest tree size (most available cache capacity) + +2. Background maintenance: Periodically evict least recently used leaf nodes on the approximate tree to prevent memory overflow. + +***Load-Balancing (Shortest Queue)*** + +For unbalanced systems, this strategy tracks pending request counts per worker and routes new requests to the least busy worker. This helps maintain optimal load distribution across workers. + +## Configuration Parameters + +1. `cache_threshold`: (float, 0.0 to 1.0, default: 0.5) + - Minimum prefix match ratio to use highest-match routing. + - Below this threshold, the request will be routed to the worker with most available cache space. + +2. `balance_abs_threshold`: (integer, default: 32) + - Absolute difference threshold for load imbalance detection. + - The system is potentially imbalanced if (max_load - min_load) > abs_threshold. + +3. `balance_rel_threshold`: (float, default: 1.0001) + - Relative ratio threshold for load imbalance detection. + - The system is potentially imbalanced if max_load > min_load * rel_threshold. + - Used in conjunction with `balance_abs_threshold` to determine the final imbalance state. + +4. `eviction_interval`: (integer, default: 60) + - Interval in seconds between LRU eviction cycles for the approximate trees. + - Background thread periodically evicts least recently used nodes to maintain tree size. + +5. `max_tree_size`: (integer, default: 16777216) + - Maximum nodes on the approximate tree. + - When exceeded, LRU leaf nodes are evicted during the next eviction cycle. diff --git a/sglang/docs/serve.sh b/sglang/docs/serve.sh new file mode 100644 index 0000000..049f767 --- /dev/null +++ b/sglang/docs/serve.sh @@ -0,0 +1,3 @@ +# Clean and serve documentation with auto-build +make clean +make serve diff --git a/sglang/docs/start/install.md b/sglang/docs/start/install.md new file mode 100644 index 0000000..fff91d0 --- /dev/null +++ b/sglang/docs/start/install.md @@ -0,0 +1,160 @@ +# Install SGLang + +You can install SGLang using any of the methods below. + +For running DeepSeek V3/R1, refer to [DeepSeek V3 Support](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3). It is recommended to use the latest version and deploy it with [Docker](https://github.com/sgl-project/sglang/tree/main/benchmark/deepseek_v3#using-docker-recommended) to avoid environment-related issues. + +It is recommended to use uv to install the dependencies for faster installation: + +## Method 1: With pip or uv + +```bash +pip install --upgrade pip +pip install uv +uv pip install "sglang[all]>=0.4.4.post3" --find-links https://flashinfer.ai/whl/cu124/torch2.5/flashinfer-python +``` + +**Quick Fixes to Common Problems** + +- SGLang currently uses torch 2.5, so you need to install flashinfer for torch 2.5. If you want to install flashinfer separately, please refer to [FlashInfer installation doc](https://docs.flashinfer.ai/installation.html). Please note that the FlashInfer pypi package is called `flashinfer-python` instead of `flashinfer`. + +- If you encounter `OSError: CUDA_HOME environment variable is not set`. Please set it to your CUDA install root with either of the following solutions: + 1. Use `export CUDA_HOME=/usr/local/cuda-` to set the `CUDA_HOME` environment variable. + 2. Install FlashInfer first following [FlashInfer installation doc](https://docs.flashinfer.ai/installation.html), then install SGLang as described above. + +- If you encounter `ImportError; cannot import name 'is_valid_list_of_images' from 'transformers.models.llama.image_processing_llama'`, try to use the specified version of `transformers` in [pyproject.toml](https://github.com/sgl-project/sglang/blob/main/python/pyproject.toml). Currently, just running `pip install transformers==4.48.3`. + +## Method 2: From source +``` +# Use the last release branch +git clone -b v0.4.4.post3 https://github.com/sgl-project/sglang.git +cd sglang + +pip install --upgrade pip +pip install -e "python[all]" --find-links https://flashinfer.ai/whl/cu124/torch2.5/flashinfer-python +``` + +Note: SGLang currently uses torch 2.5, so you need to install flashinfer for torch 2.5. If you want to install flashinfer separately, please refer to [FlashInfer installation doc](https://docs.flashinfer.ai/installation.html). + +If you want to develop SGLang, it is recommended to use docker. Please refer to [setup docker container](https://github.com/sgl-project/sglang/blob/main/docs/developer/development_guide_using_docker.md#setup-docker-container) for guidance. The docker image is `lmsysorg/sglang:dev`. + +Note: For AMD ROCm system with Instinct/MI GPUs, do following instead: + +``` +# Use the last release branch +git clone -b v0.4.4.post3 https://github.com/sgl-project/sglang.git +cd sglang + +pip install --upgrade pip +cd sgl-kernel +python setup_rocm.py install +cd .. +pip install -e "python[all_hip]" +``` + +## Method 3: Using docker +The docker images are available on Docker Hub as [lmsysorg/sglang](https://hub.docker.com/r/lmsysorg/sglang/tags), built from [Dockerfile](https://github.com/sgl-project/sglang/tree/main/docker). +Replace `` below with your huggingface hub [token](https://huggingface.co/docs/hub/en/security-tokens). + +```bash +docker run --gpus all \ + --shm-size 32g \ + -p 30000:30000 \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --env "HF_TOKEN=" \ + --ipc=host \ + lmsysorg/sglang:latest \ + python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --host 0.0.0.0 --port 30000 +``` + +Note: For AMD ROCm system with Instinct/MI GPUs, it is recommended to use `docker/Dockerfile.rocm` to build images, example and usage as below: + +```bash +docker build --build-arg SGL_BRANCH=v0.4.4.post3 -t v0.4.4.post3-rocm630 -f Dockerfile.rocm . + +alias drun='docker run -it --rm --network=host --device=/dev/kfd --device=/dev/dri --ipc=host \ + --shm-size 16G --group-add video --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ + -v $HOME/dockerx:/dockerx -v /data:/data' + +drun -p 30000:30000 \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + --env "HF_TOKEN=" \ + v0.4.4.post3-rocm630 \ + python3 -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --host 0.0.0.0 --port 30000 + +# Till flashinfer backend available, --attention-backend triton --sampling-backend pytorch are set by default +drun v0.4.4.post3-rocm630 python3 -m sglang.bench_one_batch --batch-size 32 --input 1024 --output 128 --model amd/Meta-Llama-3.1-8B-Instruct-FP8-KV --tp 8 --quantization fp8 +``` + +## Method 4: Using docker compose + +
+More + +> This method is recommended if you plan to serve it as a service. +> A better approach is to use the [k8s-sglang-service.yaml](https://github.com/sgl-project/sglang/blob/main/docker/k8s-sglang-service.yaml). + +1. Copy the [compose.yml](https://github.com/sgl-project/sglang/blob/main/docker/compose.yaml) to your local machine +2. Execute the command `docker compose up -d` in your terminal. +
+ +## Method 5: Using Kubernetes + +
+More + +1. Option 1: For single node serving (typically when the model size fits into GPUs on one node) + Execute command `kubectl apply -f docker/k8s-sglang-service.yaml`, to create k8s deployment and service, with llama-31-8b as example. + +2. Option 2: For multi-node serving (usually when a large model requires more than one GPU node, such as `DeepSeek-R1`) + Modify the LLM model path and arguments as necessary, then execute command `kubectl apply -f docker/k8s-sglang-distributed-sts.yaml`, to create two nodes k8s statefulset and serving service. +
+ + + +## Method 6: Run on Kubernetes or Clouds with SkyPilot + +
+More + +To deploy on Kubernetes or 12+ clouds, you can use [SkyPilot](https://github.com/skypilot-org/skypilot). + +1. Install SkyPilot and set up Kubernetes cluster or cloud access: see [SkyPilot's documentation](https://skypilot.readthedocs.io/en/latest/getting-started/installation.html). +2. Deploy on your own infra with a single command and get the HTTP API endpoint: +
+SkyPilot YAML: sglang.yaml + +```yaml +# sglang.yaml +envs: + HF_TOKEN: null + +resources: + image_id: docker:lmsysorg/sglang:latest + accelerators: A100 + ports: 30000 + +run: | + conda deactivate + python3 -m sglang.launch_server \ + --model-path meta-llama/Llama-3.1-8B-Instruct \ + --host 0.0.0.0 \ + --port 30000 +``` +
+ +```bash +# Deploy on any cloud or Kubernetes cluster. Use --cloud to select a specific cloud provider. +HF_TOKEN= sky launch -c sglang --env HF_TOKEN sglang.yaml + +# Get the HTTP API endpoint +sky status --endpoint 30000 sglang +``` +3. To further scale up your deployment with autoscaling and failure recovery, check out the [SkyServe + SGLang guide](https://github.com/skypilot-org/skypilot/tree/master/llm/sglang#serving-llama-2-with-sglang-for-more-traffic-using-skyserve). +
+ +## Common Notes +- [FlashInfer](https://github.com/flashinfer-ai/flashinfer) is the default attention kernel backend. It only supports sm75 and above. If you encounter any FlashInfer-related issues on sm75+ devices (e.g., T4, A10, A100, L4, L40S, H100), please switch to other kernels by adding `--attention-backend triton --sampling-backend pytorch` and open an issue on GitHub. +- If you only need to use OpenAI models with the frontend language, you can avoid installing other dependencies by using `pip install "sglang[openai]"`. +- The language frontend operates independently of the backend runtime. You can install the frontend locally without needing a GPU, while the backend can be set up on a GPU-enabled machine. To install the frontend, run `pip install sglang`, and for the backend, use `pip install sglang[srt]`. `srt` is the abbreviation of SGLang runtime. +- To reinstall flashinfer locally, use the following command: `pip install "flashinfer-python>=0.2.3" -i https://flashinfer.ai/whl/cu124/torch2.5 --force-reinstall --no-deps` and then delete the cache with `rm -rf ~/.cache/flashinfer`. diff --git a/sglang/examples/frontend_language/quick_start/anthropic_example_chat.py b/sglang/examples/frontend_language/quick_start/anthropic_example_chat.py new file mode 100644 index 0000000..03d699b --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/anthropic_example_chat.py @@ -0,0 +1,73 @@ +""" +Usage: +export ANTHROPIC_API_KEY=sk-****** +python3 anthropic_example_chat.py +""" + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print() + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.Anthropic("claude-3-haiku-20240307")) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/sglang/examples/frontend_language/quick_start/anthropic_example_complete.py b/sglang/examples/frontend_language/quick_start/anthropic_example_complete.py new file mode 100644 index 0000000..bce2a61 --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/anthropic_example_complete.py @@ -0,0 +1,68 @@ +""" +Usage: +export ANTHROPIC_API_KEY=sk-****** +python3 anthropic_example_complete.py +""" + +import sglang as sgl + + +@sgl.function +def few_shot_qa(s, question): + s += """ +\n\nHuman: What is the capital of France? +\n\nAssistant: Paris +\n\nHuman: What is the capital of Germany? +\n\nAssistant: Berlin +\n\nHuman: What is the capital of Italy? +\n\nAssistant: Rome +""" + s += "\n\nHuman: " + question + "\n" + s += "\n\nAssistant:" + sgl.gen("answer", temperature=0) + + +def single(): + state = few_shot_qa.run(question="What is the capital of the United States?") + answer = state["answer"].strip().lower() + + assert "washington" in answer, f"answer: {state['answer']}" + + print(state.text()) + + +def stream(): + state = few_shot_qa.run( + question="What is the capital of the United States?", stream=True + ) + + for out in state.text_iter("answer"): + print(out, end="", flush=True) + print() + + +def batch(): + states = few_shot_qa.run_batch( + [ + {"question": "What is the capital of the United States?"}, + {"question": "What is the capital of China?"}, + ] + ) + + for s in states: + print(s["answer"]) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.Anthropic("claude-3-haiku-20240307")) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/sglang/examples/frontend_language/quick_start/azure_openai_example_chat.py b/sglang/examples/frontend_language/quick_start/azure_openai_example_chat.py new file mode 100644 index 0000000..d53f935 --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/azure_openai_example_chat.py @@ -0,0 +1,83 @@ +""" +Usage: +export AZURE_OPENAI_API_KEY=sk-****** +python3 openai_example_chat.py +""" + +import os + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print() + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + backend = sgl.OpenAI( + model_name="azure-gpt-4", + api_version="2023-07-01-preview", + azure_endpoint="https://oai-arena-sweden.openai.azure.com/", + api_key=os.environ["AZURE_OPENAI_API_KEY"], + is_azure=True, + ) + sgl.set_default_backend(backend) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/sglang/examples/frontend_language/quick_start/gemini_example_chat.py b/sglang/examples/frontend_language/quick_start/gemini_example_chat.py new file mode 100644 index 0000000..0ae6231 --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/gemini_example_chat.py @@ -0,0 +1,73 @@ +""" +Usage: +export GCP_PROJECT_ID=****** +python3 gemini_example_chat.py +""" + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print() + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.VertexAI("gemini-pro")) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/sglang/examples/frontend_language/quick_start/gemini_example_complete.py b/sglang/examples/frontend_language/quick_start/gemini_example_complete.py new file mode 100644 index 0000000..5188bf4 --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/gemini_example_complete.py @@ -0,0 +1,68 @@ +""" +Usage: +export GCP_PROJECT_ID=****** +python3 gemini_example_complete.py +""" + +import sglang as sgl + + +@sgl.function +def few_shot_qa(s, question): + s += """The following are questions with answers. +Q: What is the capital of France? +A: Paris +Q: What is the capital of Germany? +A: Berlin +Q: What is the capital of Italy? +A: Rome +""" + s += "Q: " + question + "\n" + s += "A:" + sgl.gen("answer", stop="\n", temperature=0) + + +def single(): + state = few_shot_qa.run(question="What is the capital of the United States?") + answer = state["answer"].strip().lower() + + assert "washington" in answer, f"answer: {state['answer']}" + + print(state.text()) + + +def stream(): + state = few_shot_qa.run( + question="What is the capital of the United States?", stream=True + ) + + for out in state.text_iter("answer"): + print(out, end="", flush=True) + print() + + +def batch(): + states = few_shot_qa.run_batch( + [ + {"question": "What is the capital of the United States?"}, + {"question": "What is the capital of China?"}, + ] + ) + + for s in states: + print(s["answer"]) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.VertexAI("gemini-pro")) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/sglang/examples/frontend_language/quick_start/gemini_example_multimodal_chat.py b/sglang/examples/frontend_language/quick_start/gemini_example_multimodal_chat.py new file mode 100644 index 0000000..afe0c72 --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/gemini_example_multimodal_chat.py @@ -0,0 +1,30 @@ +""" +Usage: +export GCP_PROJECT_ID=****** +python3 gemini_example_multimodal_chat.py +""" + +import sglang as sgl + + +@sgl.function +def image_qa(s, image_file1, image_file2, question): + s += sgl.user(sgl.image(image_file1) + sgl.image(image_file2) + question) + s += sgl.assistant(sgl.gen("answer", max_tokens=256)) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.VertexAI("gemini-pro-vision")) + + state = image_qa.run( + image_file1="./images/cat.jpeg", + image_file2="./images/dog.jpeg", + question="Describe difference of the two images in one sentence.", + stream=True, + ) + + for out in state.text_iter("answer"): + print(out, end="", flush=True) + print() + + print(state["answer"]) diff --git a/sglang/examples/frontend_language/quick_start/images/cat.jpeg b/sglang/examples/frontend_language/quick_start/images/cat.jpeg new file mode 100644 index 0000000..a6a8e48 Binary files /dev/null and b/sglang/examples/frontend_language/quick_start/images/cat.jpeg differ diff --git a/sglang/examples/frontend_language/quick_start/images/dog.jpeg b/sglang/examples/frontend_language/quick_start/images/dog.jpeg new file mode 100644 index 0000000..dc4b40e Binary files /dev/null and b/sglang/examples/frontend_language/quick_start/images/dog.jpeg differ diff --git a/sglang/examples/frontend_language/quick_start/local_example_chat.py b/sglang/examples/frontend_language/quick_start/local_example_chat.py new file mode 100644 index 0000000..e1e4b62 --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/local_example_chat.py @@ -0,0 +1,75 @@ +""" +Usage: +python3 local_example_chat.py +""" + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print() + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + runtime = sgl.Runtime(model_path="meta-llama/Llama-2-7b-chat-hf") + sgl.set_default_backend(runtime) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() + + runtime.shutdown() diff --git a/sglang/examples/frontend_language/quick_start/local_example_complete.py b/sglang/examples/frontend_language/quick_start/local_example_complete.py new file mode 100644 index 0000000..00a451c --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/local_example_complete.py @@ -0,0 +1,70 @@ +""" +Usage: +python3 local_example_complete.py +""" + +import sglang as sgl + + +@sgl.function +def few_shot_qa(s, question): + s += """The following are questions with answers. +Q: What is the capital of France? +A: Paris +Q: What is the capital of Germany? +A: Berlin +Q: What is the capital of Italy? +A: Rome +""" + s += "Q: " + question + "\n" + s += "A:" + sgl.gen("answer", stop="\n", temperature=0) + + +def single(): + state = few_shot_qa.run(question="What is the capital of the United States?") + answer = state["answer"].strip().lower() + + assert "washington" in answer, f"answer: {state['answer']}" + + print(state.text()) + + +def stream(): + state = few_shot_qa.run( + question="What is the capital of the United States?", stream=True + ) + + for out in state.text_iter("answer"): + print(out, end="", flush=True) + print() + + +def batch(): + states = few_shot_qa.run_batch( + [ + {"question": "What is the capital of the United States?"}, + {"question": "What is the capital of China?"}, + ] + ) + + for s in states: + print(s["answer"]) + + +if __name__ == "__main__": + runtime = sgl.Runtime(model_path="meta-llama/Llama-2-7b-chat-hf") + sgl.set_default_backend(runtime) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() + + runtime.shutdown() diff --git a/sglang/examples/frontend_language/quick_start/local_example_llava_next.py b/sglang/examples/frontend_language/quick_start/local_example_llava_next.py new file mode 100644 index 0000000..c941a54 --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/local_example_llava_next.py @@ -0,0 +1,78 @@ +""" +Usage: python3 local_example_llava_next.py +""" + +import sglang as sgl +from sglang.lang.chat_template import get_chat_template + + +@sgl.function +def image_qa(s, image_path, question): + s += sgl.user(sgl.image(image_path) + question) + s += sgl.assistant(sgl.gen("answer")) + + +def single(): + state = image_qa.run( + image_path="images/cat.jpeg", question="What is this?", max_new_tokens=128 + ) + print(state["answer"], "\n") + + +def stream(): + state = image_qa.run( + image_path="images/cat.jpeg", + question="What is this?", + max_new_tokens=64, + stream=True, + ) + + for out in state.text_iter("answer"): + print(out, end="", flush=True) + print() + + +def batch(): + states = image_qa.run_batch( + [ + {"image_path": "images/cat.jpeg", "question": "What is this?"}, + {"image_path": "images/dog.jpeg", "question": "What is this?"}, + ], + max_new_tokens=128, + ) + for s in states: + print(s["answer"], "\n") + + +if __name__ == "__main__": + import multiprocessing as mp + + mp.set_start_method("spawn", force=True) + + runtime = sgl.Runtime(model_path="lmms-lab/llama3-llava-next-8b") + runtime.endpoint.chat_template = get_chat_template("llama-3-instruct-llava") + + # Or you can use the 72B model + # runtime = sgl.Runtime(model_path="lmms-lab/llava-next-72b", tp_size=8) + # runtime.endpoint.chat_template = get_chat_template("chatml-llava") + + sgl.set_default_backend(runtime) + print(f"chat template: {runtime.endpoint.chat_template.name}") + + # Or you can use API models + # sgl.set_default_backend(sgl.OpenAI("gpt-4-vision-preview")) + # sgl.set_default_backend(sgl.VertexAI("gemini-pro-vision")) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() + + runtime.shutdown() diff --git a/sglang/examples/frontend_language/quick_start/openai_example_chat.py b/sglang/examples/frontend_language/quick_start/openai_example_chat.py new file mode 100644 index 0000000..9511e21 --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/openai_example_chat.py @@ -0,0 +1,74 @@ +""" +Usage: +export OPENAI_API_KEY=sk-****** +python3 openai_example_chat.py +""" + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print() + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo")) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/sglang/examples/frontend_language/quick_start/openai_example_complete.py b/sglang/examples/frontend_language/quick_start/openai_example_complete.py new file mode 100644 index 0000000..d64bcaf --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/openai_example_complete.py @@ -0,0 +1,68 @@ +""" +Usage: +export OPENAI_API_KEY=sk-****** +python3 openai_example_complete.py +""" + +import sglang as sgl + + +@sgl.function +def few_shot_qa(s, question): + s += """The following are questions with answers. +Q: What is the capital of France? +A: Paris +Q: What is the capital of Germany? +A: Berlin +Q: What is the capital of Italy? +A: Rome +""" + s += "Q: " + question + "\n" + s += "A:" + sgl.gen("answer", stop="\n", temperature=0) + + +def single(): + state = few_shot_qa.run(question="What is the capital of the United States?") + answer = state["answer"].strip().lower() + + assert "washington" in answer, f"answer: {state['answer']}" + + print(state.text()) + + +def stream(): + state = few_shot_qa.run( + question="What is the capital of the United States?", stream=True + ) + + for out in state.text_iter("answer"): + print(out, end="", flush=True) + print() + + +def batch(): + states = few_shot_qa.run_batch( + [ + {"question": "What is the capital of the United States?"}, + {"question": "What is the capital of China?"}, + ] + ) + + for s in states: + print(s["answer"]) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo-instruct")) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/sglang/examples/frontend_language/quick_start/openai_example_n.py b/sglang/examples/frontend_language/quick_start/openai_example_n.py new file mode 100644 index 0000000..25372b9 --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/openai_example_n.py @@ -0,0 +1,71 @@ +""" +Usage: +export OPENAI_API_KEY=sk-****** +python3 openai_example_chat.py +""" + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=1024, n=2)) + s += sgl.user(question_2) + s += sgl.assistant( + sgl.gen( + "answer_2", + max_tokens=1024, + ) + ) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + print("\n-- answer_2 --\n", state["answer_2"]) + assert isinstance(state["answer_1"], list) + assert len(state["answer_1"]) == 2 + assert isinstance(state["answer_2"], str) + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + print("\n-- answer_1 --\n", s["answer_1"]) + print("\n-- answer_2 --\n", s["answer_2"]) + assert isinstance(s["answer_1"], list) + assert len(s["answer_1"]) == 2 + assert isinstance(s["answer_2"], str) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.OpenAI("o1")) + + # Run a single request + print("\n========== single ==========\n") + single() + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/sglang/examples/frontend_language/quick_start/openai_example_o1.py b/sglang/examples/frontend_language/quick_start/openai_example_o1.py new file mode 100644 index 0000000..2e5c140 --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/openai_example_o1.py @@ -0,0 +1,57 @@ +""" +Usage: +export OPENAI_API_KEY=sk-****** +python3 openai_example_chat.py +""" + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=100)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2")) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.OpenAI("o1")) + + # Run a single request + print("\n========== single ==========\n") + single() + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/sglang/examples/frontend_language/quick_start/openrouter_example_chat.py b/sglang/examples/frontend_language/quick_start/openrouter_example_chat.py new file mode 100644 index 0000000..a0b6f15 --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/openrouter_example_chat.py @@ -0,0 +1,81 @@ +""" +Usage: +export OPENROUTER_API_KEY=sk-****** +python3 together_example_chat.py +""" + +import os + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print() + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + backend = sgl.OpenAI( + model_name="google/gemma-7b-it:free", + base_url="https://openrouter.ai/api/v1", + api_key=os.environ.get("OPENROUTER_API_KEY"), + ) + sgl.set_default_backend(backend) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/sglang/examples/frontend_language/quick_start/together_example_chat.py b/sglang/examples/frontend_language/quick_start/together_example_chat.py new file mode 100644 index 0000000..2d20590 --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/together_example_chat.py @@ -0,0 +1,81 @@ +""" +Usage: +export TOGETHER_API_KEY=sk-****** +python3 together_example_chat.py +""" + +import os + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +def single(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + + +def stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print() + + +def batch(): + states = multi_turn_question.run_batch( + [ + { + "question_1": "What is the capital of the United States?", + "question_2": "List two local attractions.", + }, + { + "question_1": "What is the capital of France?", + "question_2": "What is the population of this city?", + }, + ] + ) + + for s in states: + print(s.messages()) + + +if __name__ == "__main__": + backend = sgl.OpenAI( + model_name="mistralai/Mixtral-8x7B-Instruct-v0.1", + base_url="https://api.together.xyz/v1", + api_key=os.environ.get("TOGETHER_API_KEY"), + ) + sgl.set_default_backend(backend) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/sglang/examples/frontend_language/quick_start/together_example_complete.py b/sglang/examples/frontend_language/quick_start/together_example_complete.py new file mode 100644 index 0000000..d9119ed --- /dev/null +++ b/sglang/examples/frontend_language/quick_start/together_example_complete.py @@ -0,0 +1,76 @@ +""" +Usage: +export TOGETHER_API_KEY=sk-****** +python3 together_example_complete.py +""" + +import os + +import sglang as sgl + + +@sgl.function +def few_shot_qa(s, question): + s += """The following are questions with answers. +Q: What is the capital of France? +A: Paris +Q: What is the capital of Germany? +A: Berlin +Q: What is the capital of Italy? +A: Rome +""" + s += "Q: " + question + "\n" + s += "A:" + sgl.gen("answer", stop="\n", temperature=0) + + +def single(): + state = few_shot_qa.run(question="What is the capital of the United States?") + answer = state["answer"].strip().lower() + + assert "washington" in answer, f"answer: {state['answer']}" + + print(state.text()) + + +def stream(): + state = few_shot_qa.run( + question="What is the capital of the United States?", stream=True + ) + + for out in state.text_iter("answer"): + print(out, end="", flush=True) + print() + + +def batch(): + states = few_shot_qa.run_batch( + [ + {"question": "What is the capital of the United States?"}, + {"question": "What is the capital of China?"}, + ] + ) + + for s in states: + print(s["answer"]) + + +if __name__ == "__main__": + backend = sgl.OpenAI( + model_name="mistralai/Mixtral-8x7B-Instruct-v0.1", + is_chat_model=False, + base_url="https://api.together.xyz/v1", + api_key=os.environ.get("TOGETHER_API_KEY"), + ) + sgl.set_default_backend(backend) + + # Run a single request + print("\n========== single ==========\n") + single() + + # Stream output + print("\n========== stream ==========\n") + stream() + + # Run a batch of requests + print("\n========== batch ==========\n") + batch() diff --git a/sglang/examples/frontend_language/usage/chinese_regex.py b/sglang/examples/frontend_language/usage/chinese_regex.py new file mode 100644 index 0000000..78e9c7e --- /dev/null +++ b/sglang/examples/frontend_language/usage/chinese_regex.py @@ -0,0 +1,53 @@ +import sglang as sgl + +character_regex = ( + r"""\{\n""" + + r""" "姓名": "[^"]{1,32}",\n""" + + r""" "学院": "(格兰芬多|赫奇帕奇|拉文克劳|斯莱特林)",\n""" + + r""" "血型": "(纯血|混血|麻瓜)",\n""" + + r""" "职业": "(学生|教师|傲罗|魔法部|食死徒|凤凰社成员)",\n""" + + r""" "魔杖": \{\n""" + + r""" "材质": "[^"]{1,32}",\n""" + + r""" "杖芯": "[^"]{1,32}",\n""" + + r""" "长度": [0-9]{1,2}\.[0-9]{0,2}\n""" + + r""" \},\n""" + + r""" "存活": "(存活|死亡)",\n""" + + r""" "守护神": "[^"]{1,32}",\n""" + + r""" "博格特": "[^"]{1,32}"\n""" + + r"""\}""" +) + + +@sgl.function +def character_gen(s, name): + s += name + " 是一名哈利波特系列小说中的角色。请填写以下关于这个角色的信息。" + s += """\ +这是一个例子 +{ + "姓名": "哈利波特", + "学院": "格兰芬多", + "血型": "混血", + "职业": "学生", + "魔杖": { + "材质": "冬青木", + "杖芯": "凤凰尾羽", + "长度": 11.0 + }, + "存活": "存活", + "守护神": "麋鹿", + "博格特": "摄魂怪" +} +""" + s += f"现在请你填写{name}的信息：\n" + s += sgl.gen("json_output", max_tokens=256, regex=character_regex) + + +def main(): + backend = sgl.RuntimeEndpoint("http://localhost:30000") + sgl.set_default_backend(backend) + ret = character_gen.run(name="赫敏格兰杰", temperature=0) + print(ret.text()) + + +if __name__ == "__main__": + main() diff --git a/sglang/examples/frontend_language/usage/choices_logprob.py b/sglang/examples/frontend_language/usage/choices_logprob.py new file mode 100644 index 0000000..6cd733f --- /dev/null +++ b/sglang/examples/frontend_language/usage/choices_logprob.py @@ -0,0 +1,44 @@ +""" +Usage: +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +python choices_logprob.py +""" + +import sglang as sgl + + +@sgl.function +def tool_use(s, question): + s += "To answer this question: " + question + ", " + s += "I need to use a " + sgl.gen("tool", choices=["calculator", "search engine"]) + + +def main(): + # Run one case + question = "What is 5 + 5?" + state = tool_use.run(question) + print("questions:", question) + print("choice:", state["tool"]) + meta_info = state.get_meta_info("tool") + print("logprobs of choice 1", meta_info["input_token_logprobs"][0]) + print("logprobs of choice 2", meta_info["input_token_logprobs"][1]) + print("-" * 50) + + # Run a batch + questions = [ + "What is 5 + 6?", + "Who is Michael Jordan?", + ] + states = tool_use.run_batch([{"question": q} for q in questions]) + for question, state in zip(questions, states): + print("questions:", question) + print("choice:", state["tool"]) + meta_info = state.get_meta_info("tool") + print("logprobs of choice 1", meta_info["input_token_logprobs"][0]) + print("logprobs of choice 2", meta_info["input_token_logprobs"][1]) + print("-" * 50) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) + main() diff --git a/sglang/examples/frontend_language/usage/cot_decoding.py b/sglang/examples/frontend_language/usage/cot_decoding.py new file mode 100644 index 0000000..7a7a04b --- /dev/null +++ b/sglang/examples/frontend_language/usage/cot_decoding.py @@ -0,0 +1,115 @@ +from math import exp +from pprint import pformat + +import sglang as sgl + +YELLOW = "\033[1;33m" +GREEN = "\033[1;32m" +BLUE = "\033[1;34m" +CLEAR = "\033[1;0m" + + +@sgl.function +def cot_decoding(s, question, get_top_k, is_chat_model, verbose): + """CoT Decoding: http://arxiv.org/abs/2402.10200""" + + if is_chat_model: + s += sgl.user("Question: " + question + "\nAnswer:") + s += sgl.assistant_begin() + else: + s += "Question: " + question + "\nAnswer:" + + step_0 = s.fork(1)[0] + forks = s.fork(get_top_k) + answer_forks = s.fork(get_top_k) + + # decoding step 0 + step_0 += sgl.gen( + "get_top_k", + max_tokens=0, + return_logprob=True, + top_logprobs_num=get_top_k, + return_text_in_logprobs=True, + ) + logprobs = step_0.get_meta_info("get_top_k")["output_top_logprobs"][0] + + print("Decoding step 0:", ", ".join(pformat(token[2]) for token in logprobs)) + for idx, (f, token) in enumerate(zip(forks, logprobs)): + logprob, token_id, text = token + f += text + + if text == "<|end_of_text|>": + print( + f"{YELLOW}Path #{idx} {pformat(text)}[{exp(logprob):.3f}] (score=nan, answer=nan){CLEAR}" + ) + continue + + # continue greedy decoding + f += sgl.gen( + "answer", + temperature=0, + max_tokens=1024, + return_logprob=True, + top_logprobs_num=2, + return_text_in_logprobs=True, + ) + + # calculate probability disparity between the top and secondary tokens + x1s = [exp(xt[0][0]) for xt in f.get_meta_info("answer")["output_top_logprobs"]] + x2s = [exp(xt[1][0]) for xt in f.get_meta_info("answer")["output_top_logprobs"]] + tokens = [xt[0][2] for xt in f.get_meta_info("answer")["output_top_logprobs"]] + delta = (sum(x1s) - sum(x2s)) / len(x1s) + + # extract the answer span (without the '<|end_of_text|>' token) + answer_forks[idx] += text + f["answer"] + "\nSo the answer is" + answer_forks[idx] += sgl.gen( + "answer_span", + temperature=0, + max_tokens=64, + return_logprob=True, + top_logprobs_num=2, + return_text_in_logprobs=True, + ) + answer = answer_forks[idx]["answer_span"].replace("\n", " ").strip(":") + print( + f"{YELLOW}Path #{idx} {pformat(text)}[{exp(logprob):.3f}] (score={delta}, answer={answer}){CLEAR}" + ) + generated_text = str(answer_forks[idx])[len("ProgramState(") : -1] + print(f"{BLUE}{pformat(generated_text)}{CLEAR}") + + if verbose: + answer_tokens = [ + xt[0][2] + for xt in answer_forks[idx].get_meta_info("answer_span")[ + "output_top_logprobs" + ] + ] + answer_x1s = [ + exp(xt[0][0]) + for xt in answer_forks[idx].get_meta_info("answer_span")[ + "output_top_logprobs" + ] + ] + answer_x2s = [ + exp(xt[1][0]) + for xt in answer_forks[idx].get_meta_info("answer_span")[ + "output_top_logprobs" + ] + ] + + for token, x1, x2 in zip(tokens, x1s, x2s): + print(f" {GREEN}{pformat(token)}{CLEAR}({x1:.3f}-{x2:.3f})", end="") + print("\n===========") + for token, x1, x2 in zip(answer_tokens, answer_x1s, answer_x2s): + print(f" {GREEN}{pformat(token)}{CLEAR}({x1:.3f}-{x2:.3f})", end="") + print() + + +sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) + +state = cot_decoding.run( + question=r"Claire makes a 3 egg omelet every morning for breakfast. How many dozens of eggs will she eat in 4 weeks?", + get_top_k=10, + is_chat_model=True, + verbose=False, +) diff --git a/sglang/examples/frontend_language/usage/json_decode.py b/sglang/examples/frontend_language/usage/json_decode.py new file mode 100644 index 0000000..5dc3522 --- /dev/null +++ b/sglang/examples/frontend_language/usage/json_decode.py @@ -0,0 +1,83 @@ +""" +Usage: +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +python json_decode.py +""" + +from enum import Enum + +from pydantic import BaseModel + +import sglang as sgl +from sglang.srt.constrained.outlines_backend import build_regex_from_object + +character_regex = ( + r"""\{\n""" + + r""" "name": "[\w\d\s]{1,16}",\n""" + + r""" "house": "(Gryffindor|Slytherin|Ravenclaw|Hufflepuff)",\n""" + + r""" "blood status": "(Pure-blood|Half-blood|Muggle-born)",\n""" + + r""" "occupation": "(student|teacher|auror|ministry of magic|death eater|order of the phoenix)",\n""" + + r""" "wand": \{\n""" + + r""" "wood": "[\w\d\s]{1,16}",\n""" + + r""" "core": "[\w\d\s]{1,16}",\n""" + + r""" "length": [0-9]{1,2}\.[0-9]{0,2}\n""" + + r""" \},\n""" + + r""" "alive": "(Alive|Deceased)",\n""" + + r""" "patronus": "[\w\d\s]{1,16}",\n""" + + r""" "bogart": "[\w\d\s]{1,16}"\n""" + + r"""\}""" +) + + +@sgl.function +def character_gen(s, name): + s += ( + name + + " is a character in Harry Potter. Please fill in the following information about this character.\n" + ) + s += "The constrained regex is:\n" + s += character_regex + "\n" + s += "The JSON output is:\n" + s += sgl.gen("json_output", max_tokens=256, regex=character_regex) + + +def driver_character_gen(): + state = character_gen.run(name="Hermione Granger") + print(state.text()) + + +class Weapon(str, Enum): + sword = "sword" + axe = "axe" + mace = "mace" + spear = "spear" + bow = "bow" + crossbow = "crossbow" + + +class Wizard(BaseModel): + name: str + age: int + weapon: Weapon + + +@sgl.function +def pydantic_wizard_gen(s): + s += "Give me a description about a wizard in the JSON format.\n" + s += sgl.gen( + "character", + max_tokens=128, + temperature=0, + regex=build_regex_from_object(Wizard), # Requires pydantic >= 2.0 + ) + + +def driver_pydantic_wizard_gen(): + state = pydantic_wizard_gen.run() + print(state.text()) + + +if __name__ == "__main__": + sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) + driver_character_gen() + # driver_pydantic_wizard_gen() diff --git a/sglang/examples/frontend_language/usage/json_logprobs.py b/sglang/examples/frontend_language/usage/json_logprobs.py new file mode 100644 index 0000000..15206a6 --- /dev/null +++ b/sglang/examples/frontend_language/usage/json_logprobs.py @@ -0,0 +1,103 @@ +# NOTE: Currently this can only be run through HTTP requests. +from concurrent.futures import ThreadPoolExecutor + +from json_decode import character_regex + +from sglang.utils import http_request + +character_names = ["Hermione Granger", "Ron Weasley", "Harry Potter"] + +base_url = "http://localhost:30000" + +prompt = "is a character in Harry Potter. Please fill in the following information about this character.\n" + + +def openai_api_request(name): + data = { + "model": "", + "prompt": name + prompt, + "temperature": 0, + "max_tokens": 128, + "regex": character_regex, + "logprobs": 3, + } + res = http_request(base_url + "/v1/completions", json=data).json() + + # with open(f"json_logprobs_{name.replace(' ', '_')}_tmp.json", "w") as fout: + # fout.write(json.dumps(res, indent=4)) + + logprobs = res["choices"][0]["logprobs"] + usage = res["usage"] + assert len(logprobs["token_logprobs"]) == len(logprobs["tokens"]) + assert len(logprobs["token_logprobs"]) == len(logprobs["top_logprobs"]) + assert len(logprobs["token_logprobs"]) == usage["completion_tokens"] - 1 + + return res + + +def srt_api_request(name): + data = { + "text": name + prompt, + "sampling_params": { + "temperature": 0, + "max_new_tokens": 128, + "regex": character_regex, + }, + "return_logprob": True, + "logprob_start_len": 0, + "top_logprobs_num": 3, + "return_text_in_logprobs": True, + } + + res = http_request(base_url + "/generate", json=data).json() + + # with open(f"json_logprobs_{name.replace(' ', '_')}_tmp.json", "w") as fout: + # fout.write(json.dumps(res, indent=4)) + + meta_info = res["meta_info"] + assert len(meta_info["input_token_logprobs"]) == len( + meta_info["input_top_logprobs"] + ) + assert len(meta_info["output_token_logprobs"]) == len( + meta_info["output_top_logprobs"] + ) + assert len(meta_info["input_token_logprobs"]) == meta_info["prompt_tokens"] + assert len(meta_info["output_token_logprobs"]) == meta_info["completion_tokens"] - 1 + + return res + + +def pretty_print(res): + meta_info = res["meta_info"] + + print("\n\n", "=" * 30, "Prefill", "=" * 30) + for i in range(len(meta_info["input_token_logprobs"])): + print(f"{str(meta_info['input_token_logprobs'][i][2].encode()): <20}", end="") + top_ks = ( + [str(t[2].encode()) for t in meta_info["input_top_logprobs"][i]] + if meta_info["input_top_logprobs"][i] + else [] + ) + for top_k in top_ks: + print(f"{top_k: <15}", end="") + print() + + print("\n\n", "=" * 30, "Decode", "=" * 30) + for i in range(len(meta_info["output_token_logprobs"])): + print(f"{str(meta_info['output_token_logprobs'][i][2].encode()): <20}", end="") + top_ks = [str(t[2].encode()) for t in meta_info["output_top_logprobs"][i]] + for top_k in top_ks: + print(f"{top_k: <15}", end="") + print() + + print(res["text"]) + + +if __name__ == "__main__": + with ThreadPoolExecutor() as executor: + ress = executor.map(srt_api_request, character_names) + + for res in ress: + pretty_print(res) + + openai_api_request("Hermione Granger") diff --git a/sglang/examples/frontend_language/usage/llava_video/srt_example_llava_v.py b/sglang/examples/frontend_language/usage/llava_video/srt_example_llava_v.py new file mode 100644 index 0000000..bc56d42 --- /dev/null +++ b/sglang/examples/frontend_language/usage/llava_video/srt_example_llava_v.py @@ -0,0 +1,260 @@ +""" +Usage: +pip install opencv-python-headless + +python3 srt_example_llava_v.py +""" + +import argparse +import csv +import json +import os +import time + +import requests + +import sglang as sgl + + +@sgl.function +def video_qa(s, num_frames, video_path, question): + s += sgl.user(sgl.video(video_path, num_frames) + question) + s += sgl.assistant(sgl.gen("answer")) + + +def single(path, num_frames=16): + state = video_qa.run( + num_frames=num_frames, + video_path=path, + question="Please provide a detailed description of the video, focusing on the main subjects, their actions, the background scenes", + temperature=0.0, + max_new_tokens=1024, + ) + print(state["answer"], "\n") + + +def split_into_chunks(lst, num_chunks): + """Split a list into a specified number of chunks.""" + # Calculate the chunk size using integer division. Note that this may drop some items if not evenly divisible. + chunk_size = len(lst) // num_chunks + + if chunk_size == 0: + chunk_size = len(lst) + # Use list comprehension to generate chunks. The last chunk will take any remainder if the list size isn't evenly divisible. + chunks = [lst[i : i + chunk_size] for i in range(0, len(lst), chunk_size)] + # Ensure we have exactly num_chunks chunks, even if some are empty + chunks.extend([[] for _ in range(num_chunks - len(chunks))]) + return chunks + + +def save_batch_results(batch_video_files, states, cur_chunk, batch_idx, save_dir): + csv_filename = f"{save_dir}/chunk_{cur_chunk}_batch_{batch_idx}.csv" + with open(csv_filename, "w", newline="") as csvfile: + writer = csv.writer(csvfile) + writer.writerow(["video_name", "answer"]) + for video_path, state in zip(batch_video_files, states): + video_name = os.path.basename(video_path) + writer.writerow([video_name, state["answer"]]) + + +def compile_and_cleanup_final_results(cur_chunk, num_batches, save_dir): + final_csv_filename = f"{save_dir}/final_results_chunk_{cur_chunk}.csv" + with open(final_csv_filename, "w", newline="") as final_csvfile: + writer = csv.writer(final_csvfile) + writer.writerow(["video_name", "answer"]) + for batch_idx in range(num_batches): + batch_csv_filename = f"{save_dir}/chunk_{cur_chunk}_batch_{batch_idx}.csv" + with open(batch_csv_filename, "r") as batch_csvfile: + reader = csv.reader(batch_csvfile) + next(reader) # Skip header row + for row in reader: + writer.writerow(row) + os.remove(batch_csv_filename) + + +def find_video_files(video_dir): + # Check if the video_dir is actually a file + if os.path.isfile(video_dir): + # If it's a file, return it as a single-element list + return [video_dir] + + # Original logic to find video files in a directory + video_files = [] + for root, dirs, files in os.walk(video_dir): + for file in files: + if file.endswith((".mp4", ".avi", ".mov")): + video_files.append(os.path.join(root, file)) + return video_files + + +def batch(video_dir, save_dir, cur_chunk, num_chunks, num_frames=16, batch_size=64): + video_files = find_video_files(video_dir) + chunked_video_files = split_into_chunks(video_files, num_chunks)[cur_chunk] + num_batches = 0 + + for i in range(0, len(chunked_video_files), batch_size): + batch_video_files = chunked_video_files[i : i + batch_size] + print(f"Processing batch of {len(batch_video_files)} video(s)...") + + if not batch_video_files: + print("No video files found in the specified directory.") + return + + batch_input = [ + { + "num_frames": num_frames, + "video_path": video_path, + "question": "Please provide a detailed description of the video, focusing on the main subjects, their actions, the background scenes.", + } + for video_path in batch_video_files + ] + + start_time = time.time() + states = video_qa.run_batch(batch_input, max_new_tokens=512, temperature=0.2) + total_time = time.time() - start_time + average_time = total_time / len(batch_video_files) + print( + f"Number of videos in batch: {len(batch_video_files)}. Average processing time per video: {average_time:.2f} seconds. Total time for this batch: {total_time:.2f} seconds" + ) + + save_batch_results(batch_video_files, states, cur_chunk, num_batches, save_dir) + num_batches += 1 + + compile_and_cleanup_final_results(cur_chunk, num_batches, save_dir) + + +if __name__ == "__main__": + + url = "https://raw.githubusercontent.com/EvolvingLMMs-Lab/sglang/dev/onevision_local/assets/jobs.mp4" + + cache_dir = os.path.expanduser("~/.cache") + file_path = os.path.join(cache_dir, "jobs.mp4") + + os.makedirs(cache_dir, exist_ok=True) + + response = requests.get(url) + response.raise_for_status() # Raise an exception for bad responses + + with open(file_path, "wb") as f: + f.write(response.content) + + print(f"File downloaded and saved to: {file_path}") + # Create the parser + parser = argparse.ArgumentParser( + description="Run video processing with specified port." + ) + + # Add an argument for the port + parser.add_argument( + "--port", + type=int, + default=30000, + help="The master port for distributed serving.", + ) + parser.add_argument( + "--chunk-idx", type=int, default=0, help="The index of the chunk to process." + ) + parser.add_argument( + "--num-chunks", type=int, default=8, help="The number of chunks to process." + ) + parser.add_argument( + "--save-dir", + type=str, + default="./work_dirs/llava_video", + help="The directory to save the processed video files.", + ) + parser.add_argument( + "--video-dir", + type=str, + default=os.path.expanduser("~/.cache/jobs.mp4"), + help="The directory or path for the processed video files.", + ) + parser.add_argument( + "--model-path", + type=str, + default="lmms-lab/LLaVA-NeXT-Video-7B", + help="The model path for the video processing.", + ) + parser.add_argument( + "--num-frames", + type=int, + default=16, + help="The number of frames to process in each video.", + ) + parser.add_argument("--mm_spatial_pool_stride", type=int, default=2) + + # Parse the arguments + args = parser.parse_args() + cur_port = args.port + cur_chunk = args.chunk_idx + num_chunks = args.num_chunks + num_frames = args.num_frames + + if "34b" in args.model_path.lower(): + tokenizer_path = "liuhaotian/llava-v1.6-34b-tokenizer" + elif "7b" in args.model_path.lower(): + tokenizer_path = "llava-hf/llava-1.5-7b-hf" + else: + print("Invalid model path. Please specify a valid model path.") + exit() + + model_override_args = {} + model_override_args["mm_spatial_pool_stride"] = args.mm_spatial_pool_stride + model_override_args["architectures"] = ["LlavaVidForCausalLM"] + model_override_args["num_frames"] = args.num_frames + model_override_args["model_type"] = "llava" + + if "34b" in args.model_path.lower(): + model_override_args["image_token_index"] = 64002 + + if args.num_frames == 32: + model_override_args["rope_scaling"] = {"factor": 2.0, "rope_type": "linear"} + model_override_args["max_sequence_length"] = 4096 * 2 + model_override_args["tokenizer_model_max_length"] = 4096 * 2 + elif args.num_frames < 32: + pass + else: + print( + "The maximum number of frames to process is 32. Please specify a valid number of frames." + ) + exit() + + runtime = sgl.Runtime( + model_path=args.model_path, # "liuhaotian/llava-v1.6-vicuna-7b", + tokenizer_path=tokenizer_path, + port=cur_port, + json_model_override_args=json.dumps(model_override_args), + tp_size=1, + ) + sgl.set_default_backend(runtime) + print(f"chat template: {runtime.endpoint.chat_template.name}") + + # Run a single request + print("\n========== single ==========\n") + root = args.video_dir + if os.path.isfile(root): + video_files = [root] + else: + video_files = [ + os.path.join(root, f) + for f in os.listdir(root) + if f.endswith((".mp4", ".avi", ".mov")) + ] # Add more extensions if needed + start_time = time.time() # Start time for processing a single video + for cur_video in video_files[:1]: + print(cur_video) + single(cur_video, num_frames) + end_time = time.time() # End time for processing a single video + total_time = end_time - start_time + average_time = total_time / len( + video_files + ) # Calculate the average processing time + print(f"Average processing time per video: {average_time:.2f} seconds") + runtime.shutdown() + + # # Run a batch of requests + # print("\n========== batch ==========\n") + # if not os.path.exists(args.save_dir): + # os.makedirs(args.save_dir) + # batch(args.video_dir, args.save_dir, cur_chunk, num_chunks, num_frames, num_chunks) + # runtime.shutdown() diff --git a/sglang/examples/frontend_language/usage/llava_video/srt_example_llava_v.sh b/sglang/examples/frontend_language/usage/llava_video/srt_example_llava_v.sh new file mode 100644 index 0000000..ffb1af9 --- /dev/null +++ b/sglang/examples/frontend_language/usage/llava_video/srt_example_llava_v.sh @@ -0,0 +1,131 @@ +#!/bin/bash + +##### USAGE ##### +# - First node: +# ```sh +# bash examples/usage/llava_video/srt_example_llava_v.sh K 0 YOUR_VIDEO_PATH YOUR_MODEL_PATH FRAMES_PER_VIDEO +# ``` +# - Second node: +# ```sh +# bash examples/usage/llava_video/srt_example_llava_v.sh K 1 YOUR_VIDEO_PATH YOUR_MODEL_PATH FRAMES_PER_VIDEO +# ``` +# - The K node: +# ```sh +# bash examples/usage/llava_video/srt_example_llava_v.sh K K-1 YOUR_VIDEO_PATH YOUR_MODEL_PATH FRAMES_PER_VIDEO +# ``` + + +# Replace `K`, `YOUR_VIDEO_PATH`, `YOUR_MODEL_PATH`, and `FRAMES_PER_VIDEO` with your specific details. +# CURRENT_ROOT="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )" +CURRENT_ROOT=$(dirname "$0") + +echo ${CURRENT_ROOT} + +cd ${CURRENT_ROOT} + +export PYTHONWARNINGS=ignore + +START_TIME=$(date +%s) # Capture start time + +NUM_NODES=$1 + +CUR_NODES_IDX=$2 + +VIDEO_DIR=$3 + +MODEL_PATH=$4 + +NUM_FRAMES=$5 + + +# FRAME_FORMAT=$6 + +# FRAME_FORMAT=$(echo $FRAME_FORMAT | tr '[:lower:]' '[:upper:]') + +# # Check if FRAME_FORMAT is either JPEG or PNG +# if [[ "$FRAME_FORMAT" != "JPEG" && "$FRAME_FORMAT" != "PNG" ]]; then +# echo "Error: FRAME_FORMAT must be either JPEG or PNG." +# exit 1 +# fi + +# export TARGET_FRAMES=$TARGET_FRAMES + +echo "Each video you will sample $NUM_FRAMES frames" + +# export FRAME_FORMAT=$FRAME_FORMAT + +# echo "The frame format is $FRAME_FORMAT" + +# Assuming GPULIST is a bash array containing your GPUs +GPULIST=(0 1 2 3 4 5 6 7) +LOCAL_CHUNKS=${#GPULIST[@]} + +echo "Number of GPUs in GPULIST: $LOCAL_CHUNKS" + +ALL_CHUNKS=$((NUM_NODES * LOCAL_CHUNKS)) + +# Calculate GPUs per chunk +GPUS_PER_CHUNK=1 + +echo $GPUS_PER_CHUNK + +for IDX in $(seq 1 $LOCAL_CHUNKS); do + ( + START=$(((IDX-1) * GPUS_PER_CHUNK)) + LENGTH=$GPUS_PER_CHUNK # Length for slicing, not the end index + + CHUNK_GPUS=(${GPULIST[@]:$START:$LENGTH}) + + # Convert the chunk GPUs array to a comma-separated string + CHUNK_GPUS_STR=$(IFS=,; echo "${CHUNK_GPUS[*]}") + + LOCAL_IDX=$((CUR_NODES_IDX * LOCAL_CHUNKS + IDX)) + + echo "Chunk $(($LOCAL_IDX - 1)) will run on GPUs $CHUNK_GPUS_STR" + + # Calculate the port for this chunk. Ensure it's incremented by 5 for each chunk. + PORT=$((10000 + RANDOM % 55536)) + + MAX_RETRIES=10 + RETRY_COUNT=0 + COMMAND_STATUS=1 # Initialize as failed + + while [ $RETRY_COUNT -lt $MAX_RETRIES ] && [ $COMMAND_STATUS -ne 0 ]; do + echo "Running chunk $(($LOCAL_IDX - 1)) on GPUs $CHUNK_GPUS_STR with port $PORT. Attempt $(($RETRY_COUNT + 1))" + +#!/bin/bash + CUDA_VISIBLE_DEVICES=$CHUNK_GPUS_STR python3 srt_example_llava_v.py \ + --port $PORT \ + --num-chunks $ALL_CHUNKS \ + --chunk-idx $(($LOCAL_IDX - 1)) \ + --save-dir work_dirs/llava_next_video_inference_results \ + --video-dir $VIDEO_DIR \ + --model-path $MODEL_PATH \ + --num-frames $NUM_FRAMES #& + + wait $! # Wait for the process to finish and capture its exit status + COMMAND_STATUS=$? + + if [ $COMMAND_STATUS -ne 0 ]; then + echo "Execution failed for chunk $(($LOCAL_IDX - 1)), attempt $(($RETRY_COUNT + 1)). Retrying..." + RETRY_COUNT=$(($RETRY_COUNT + 1)) + sleep 180 # Wait a bit before retrying + else + echo "Execution succeeded for chunk $(($LOCAL_IDX - 1))." + fi + done + + if [ $COMMAND_STATUS -ne 0 ]; then + echo "Execution failed for chunk $(($LOCAL_IDX - 1)) after $MAX_RETRIES attempts." + fi + ) #& + sleep 2 # Slight delay to stagger the start times +done + +wait + +cat work_dirs/llava_next_video_inference_results/final_results_chunk_*.csv > work_dirs/llava_next_video_inference_results/final_results_node_${CUR_NODES_IDX}.csv + +END_TIME=$(date +%s) # Capture end time +ELAPSED_TIME=$(($END_TIME - $START_TIME)) +echo "Total execution time: $ELAPSED_TIME seconds." diff --git a/sglang/examples/frontend_language/usage/openai_chat_speculative.py b/sglang/examples/frontend_language/usage/openai_chat_speculative.py new file mode 100644 index 0000000..f3fd74e --- /dev/null +++ b/sglang/examples/frontend_language/usage/openai_chat_speculative.py @@ -0,0 +1,155 @@ +""" +Usage: +***Note: for speculative execution to work, user must put all "gen" in "assistant". +Show in "assistant" the desired answer format. Each "gen" term should have a stop token. +The stream mode is not supported in speculative execution. + +E.g. +correct: + sgl.assistant("\nName:" + sgl.gen("name", stop="\n") + "\nBirthday:" + sgl.gen("birthday", stop="\n") + "\nJob:" + sgl.gen("job", stop="\n")) +incorrect: + s += sgl.assistant("\nName:" + sgl.gen("name", stop="\n")) + s += sgl.assistant("\nBirthday:" + sgl.gen("birthday", stop="\n")) + s += sgl.assistant("\nJob:" + sgl.gen("job", stop="\n")) + +export OPENAI_API_KEY=sk-****** +python3 openai_chat_speculative.py +""" + +import sglang as sgl +from sglang import OpenAI, function, set_default_backend + + +@function(num_api_spec_tokens=256) +def gen_character_spec(s): + s += sgl.system("You are a helpful assistant.") + s += sgl.user("Construct a character within the following format:") + s += sgl.assistant( + "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\n" + ) + s += sgl.user("Please generate new Name, Birthday and Job.\n") + s += sgl.assistant( + "Name:" + + sgl.gen("name", stop="\n") + + "\nBirthday:" + + sgl.gen("birthday", stop="\n") + + "\nJob:" + + sgl.gen("job", stop="\n") + ) + + +@function(num_api_spec_tokens=256) +def gen_character_spec_no_few_shot(s): + s += sgl.user("Construct a character. For each field stop with a newline\n") + s += sgl.assistant( + "Name:" + + sgl.gen("name", stop="\n") + + "\nAge:" + + sgl.gen("age", stop="\n") + + "\nJob:" + + sgl.gen("job", stop="\n") + ) + + +@function +def gen_character_normal(s): + s += sgl.system("You are a helpful assistant.") + s += sgl.user("What's the answer of 23 + 8?") + s += sgl.assistant(sgl.gen("answer", max_tokens=64)) + + +@function(num_api_spec_tokens=1024) +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user("Answer questions in the following format:") + s += sgl.user( + "Question 1: What is the capital of France?\nQuestion 2: What is the population of this city?\n" + ) + s += sgl.assistant( + "Answer 1: The capital of France is Paris.\nAnswer 2: The population of Paris in 2024 is estimated to be around 2.1 million for the city proper.\n" + ) + s += sgl.user("Question 1: " + question_1 + "\nQuestion 2: " + question_2) + s += sgl.assistant( + "Answer 1: " + + sgl.gen("answer_1", stop="\n") + + "\nAnswer 2: " + + sgl.gen("answer_2", stop="\n") + ) + + +def test_spec_single_turn(): + backend.token_usage.reset() + + state = gen_character_spec.run() + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- name:", state["name"]) + print("-- birthday:", state["birthday"]) + print("-- job:", state["job"]) + print(backend.token_usage) + + +def test_inaccurate_spec_single_turn(): + state = gen_character_spec_no_few_shot.run() + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- name:", state["name"]) + print("\n-- age:", state["age"]) + print("\n-- job:", state["job"]) + + +def test_normal_single_turn(): + state = gen_character_normal.run() + for m in state.messages(): + print(m["role"], ":", m["content"]) + + +def test_spec_multi_turn(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions in the capital of the United States.", + ) + + for m in state.messages(): + print(m["role"], ":", m["content"]) + + print("\n-- answer_1 --\n", state["answer_1"]) + print("\n-- answer_2 --\n", state["answer_2"]) + + +def test_spec_multi_turn_stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + + +if __name__ == "__main__": + backend = OpenAI("gpt-4-turbo") + set_default_backend(backend) + + print("\n========== test spec single turn ==========\n") + # expect reasonable answer for each field + test_spec_single_turn() + + print("\n========== test inaccurate spec single turn ==========\n") + # expect incomplete or unreasonable answers + test_inaccurate_spec_single_turn() + + print("\n========== test normal single turn ==========\n") + # expect reasonable answer + test_normal_single_turn() + + print("\n========== test spec multi turn ==========\n") + # expect answer with same format as in the few shot + test_spec_multi_turn() + + print("\n========== test spec multi turn stream ==========\n") + # expect error in stream_executor: stream is not supported... + test_spec_multi_turn_stream() diff --git a/sglang/examples/frontend_language/usage/openai_speculative.py b/sglang/examples/frontend_language/usage/openai_speculative.py new file mode 100644 index 0000000..4389cb0 --- /dev/null +++ b/sglang/examples/frontend_language/usage/openai_speculative.py @@ -0,0 +1,54 @@ +""" +Usage: +python3 openai_speculative.py +""" + +from sglang import OpenAI, function, gen, set_default_backend + + +@function(num_api_spec_tokens=64) +def gen_character_spec(s): + s += "Construct a character within the following format:\n" + s += "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\n" + s += "\nPlease generate new Name, Birthday and Job.\n" + s += "Name:" + gen("name", stop="\n") + "\nBirthday:" + gen("birthday", stop="\n") + s += "\nJob:" + gen("job", stop="\n") + "\n" + + +@function +def gen_character_no_spec(s): + s += "Construct a character within the following format:\n" + s += "Name: Steve Jobs.\nBirthday: February 24, 1955.\nJob: Apple CEO.\n" + s += "\nPlease generate new Name, Birthday and Job.\n" + s += "Name:" + gen("name", stop="\n") + "\nBirthday:" + gen("birthday", stop="\n") + s += "\nJob:" + gen("job", stop="\n") + "\n" + + +@function(num_api_spec_tokens=64) +def gen_character_spec_no_few_shot(s): + # s += "Construct a character with name, birthday, and job:\n" + s += "Construct a character:\n" + s += "Name:" + gen("name", stop="\n") + "\nBirthday:" + gen("birthday", stop="\n") + s += "\nJob:" + gen("job", stop="\n") + "\n" + + +if __name__ == "__main__": + backend = OpenAI("gpt-3.5-turbo-instruct") + set_default_backend(backend) + + for function in [ + gen_character_spec, + gen_character_no_spec, + gen_character_spec_no_few_shot, + ]: + backend.token_usage.reset() + + print(f"function: {function.func.__name__}") + + state = function.run() + + print("...name:", state["name"]) + print("...birthday:", state["birthday"]) + print("...job:", state["job"]) + print(backend.token_usage) + print() diff --git a/sglang/examples/frontend_language/usage/parallel_sample.py b/sglang/examples/frontend_language/usage/parallel_sample.py new file mode 100644 index 0000000..0f3cf17 --- /dev/null +++ b/sglang/examples/frontend_language/usage/parallel_sample.py @@ -0,0 +1,40 @@ +""" +Usage: +python3 parallel_sample.py +""" + +import sglang as sgl + + +@sgl.function +def parallel_sample(s, question, n): + s += ( + "Question: Compute 1 + 2 + 3\n" + "Reasoning: I need to use a calculator.\n" + "Tool: calculator\n" + "Answer: 6\n" + "Question: Compute 3 + 2 + 2\n" + "Reasoning: I will try a calculator.\n" + "Tool: calculator\n" + "Answer: 7\n" + ) + s += "Question: " + question + "\n" + forks = s.fork(n) + forks += "Reasoning:" + sgl.gen("reasoning", stop="\n") + "\n" + forks += "Tool:" + sgl.gen("tool", choices=["calculator", "browser"]) + "\n" + forks += "Answer:" + sgl.gen("answer", stop="\n") + "\n" + forks.join() + + +sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo-instruct")) +# sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) + +state = parallel_sample.run(question="Compute 5 + 2 + 4.", n=5, temperature=1.0) + +for i in range(5): + obj = { + "reasoning": state["reasoning"][i], + "tool": state["tool"][i], + "answer": state["answer"][i], + } + print(f"[{i}], {obj}") diff --git a/sglang/examples/frontend_language/usage/rag_using_parea/trace_and_evaluate_rag_using_parea.ipynb b/sglang/examples/frontend_language/usage/rag_using_parea/trace_and_evaluate_rag_using_parea.ipynb new file mode 100644 index 0000000..3c1b2a6 --- /dev/null +++ b/sglang/examples/frontend_language/usage/rag_using_parea/trace_and_evaluate_rag_using_parea.ipynb @@ -0,0 +1,408 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# RAG Powered by SGLang & Chroma Evaluated using Parea\n", + "\n", + "In this notebook, we will build a simple RAG pipeline using SGLang to execute our LLM calls, Chroma as vector database for retrieval and [Parea](https://www.parea.ai) for tracing and evaluation. We will then evaluate the performance of our RAG pipeline. The dataset we will use was created by [Virat](https://twitter.com/virattt) and contains 100 questions, contexts and answers from the Airbnb 2023 10k filing.\n", + "\n", + "The RAG pipeline consists of two steps:\n", + "1. Retrieval: Given a question, we retrieve the relevant context from all provided contexts.\n", + "2. Generation: Given the question and the retrieved context, we generate an answer.\n", + "\n", + "ℹ️ This notebook requires an OpenAI API key.\n", + "\n", + "ℹ️ This notebook requires a Parea API key, which can be created [here](https://docs.parea.ai/api-reference/authentication#parea-api-key)." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setting up the environment\n", + "\n", + "We will first install the necessary packages: `sglang`, `parea-ai` and `chromadb`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# note, if you use a Mac M1 chip, you might need to install grpcio 1.59.0 first such that installing chromadb works\n", + "# !pip install grpcio==1.59.0\n", + "\n", + "!pip install sglang[openai] parea-ai chromadb" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Create a Parea API key as outlined [here](https://docs.parea.ai/api-reference/authentication#parea-api-key) and save it in a `.env` file as `PAREA_API_KEY=your-api-key`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Indexing the data\n", + "\n", + "Now it's time to download the data & index it! For that, we create a collection called `contexts` in Chroma and add the contexts as documents." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import json\n", + "import os\n", + "from typing import List\n", + "\n", + "import chromadb\n", + "\n", + "path_qca = \"airbnb-2023-10k-qca.json\"\n", + "\n", + "if not os.path.exists(path_qca):\n", + " !wget https://virattt.github.io/datasets/abnb-2023-10k.json -O airbnb-2023-10k-qca.json\n", + "\n", + "with open(path_qca, \"r\") as f:\n", + " question_context_answers = json.load(f)\n", + "\n", + "chroma_client = chromadb.PersistentClient()\n", + "collection = chroma_client.get_or_create_collection(name=\"contexts\")\n", + "if collection.count() == 0:\n", + " collection.add(\n", + " documents=[qca[\"context\"] for qca in question_context_answers],\n", + " ids=[str(i) for i in range(len(question_context_answers))],\n", + " )" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Defining the RAG pipeline\n", + "\n", + "We will start with importing the necessary packages, setting up tracing of OpenAI calls via Parea and setting OpenAI as the default backend for SGLang." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "import time\n", + "\n", + "from dotenv import load_dotenv\n", + "\n", + "from sglang import function, user, assistant, gen, set_default_backend, OpenAI\n", + "from sglang.lang.interpreter import ProgramState\n", + "from parea import Parea, trace\n", + "\n", + "\n", + "load_dotenv()\n", + "\n", + "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n", + "\n", + "p = Parea(api_key=os.getenv(\"PAREA_API_KEY\"), project_name=\"rag_sglang\")\n", + "p.integrate_with_sglang()\n", + "\n", + "set_default_backend(OpenAI(\"gpt-3.5-turbo\"))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we can define our retrieval step shown below. Notice, the `trace` decorator which will automatically trace inputs, output, latency, etc. of that call." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@trace\n", + "def retrieval(question: str) -> List[str]:\n", + " return collection.query(query_texts=[question], n_results=1)[\"documents\"][0]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Next we will define the generation step which uses SGLang to execute the LLM call." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@function\n", + "def generation_sglang(s, question: str, *context: str):\n", + " context = \"\\n\".join(context)\n", + " s += user(\n", + " f\"Given this question:\\n{question}\\n\\nAnd this context:\\n{context}\\n\\nAnswer the question.\"\n", + " )\n", + " s += assistant(gen(\"answer\"))\n", + "\n", + "\n", + "@trace\n", + "def generation(question: str, *context):\n", + " state: ProgramState = generation_sglang.run(question, *context)\n", + " while not state.stream_executor.is_finished:\n", + " time.sleep(1)\n", + " return state.stream_executor.variables[\"answer\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finally, we can tie it together and execute a sample query." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@trace\n", + "def rag_pipeline(question: str) -> str:\n", + " contexts = retrieval(question)\n", + " return generation(question, *contexts)\n", + "\n", + "\n", + "rag_pipeline(\n", + " \"When did the World Health Organization formally declare an end to the COVID-19 global health emergency?\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Debug Trace\n", + "\n", + "The output is unfortunately wrong! Using the traced pipeline, we can see that\n", + "\n", + "- the context is relevant to the question and contains the correct information\n", + "- but, the generation step is cut off as max tokens is set to 16\n", + "\n", + "When opening the generation step in the playground and rerunning the prompt with max. tokens set to 1000, the correct answer is produced.\n", + "\n", + "![RAG Trace](https://drive.google.com/uc?id=1QI243ogGjzbO01tUrR72g9rFoGzUJqVH)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Evaluating RAG Pipelines\n", + "\n", + "Before we apply above's fix, let's dive into evaluating RAG pipelines.\n", + "\n", + "RAG pipelines consist of a retrieval step to fetch relevant information and a generation step to generate a response to a users question. A RAG pipeline can fail at either step. E.g. the retrieval step can fail to find relevant information which makes generating the correct impossible. Another failure mode is that the generation step doesn't leverage the retrieved information correctly. We will apply the following evaluation metrics to understand different failure modes:\n", + "\n", + "- `context_relevancy`: measures how relevant the context is given the question\n", + "- `percent_target_supported_by_context`: measures how much of the target answer is supported by the context; this will give an upper ceiling of how well the generation step can perform\n", + "- `answer_context_faithfulness`: measures how much the generated answer utilizes the context\n", + "- `answer_matches_target`: measures how well the generated answer matches the target answer judged by a LLM and gives a sense of accuracy of our entire pipeline\n", + "\n", + "To use these evaluation metrics, we can import them from `parea.evals.rag` and `parea.evals.general` and apply them to a function by specifying in the `trace` decorator which evaluation metrics to use. The `@trace` decorator will automatically log the results of the evaluation metrics to the Parea dashboard.\n", + "\n", + "Applying them to the retrieval step:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from parea.evals.rag import (\n", + " context_query_relevancy_factory,\n", + " percent_target_supported_by_context_factory,\n", + ")\n", + "\n", + "\n", + "context_relevancy_eval = context_query_relevancy_factory()\n", + "percent_target_supported_by_context = percent_target_supported_by_context_factory()\n", + "\n", + "\n", + "@trace(eval_funcs=[context_relevancy_eval, percent_target_supported_by_context])\n", + "def retrieval(question: str) -> List[str]:\n", + " return collection.query(query_texts=[question], n_results=1)[\"documents\"][0]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Now we can apply `answer_context_faithfulness` and `answer_matches_target` to the generation step." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from parea.evals.general import answer_matches_target_llm_grader_factory\n", + "from parea.evals.rag import answer_context_faithfulness_statement_level_factory\n", + "\n", + "\n", + "answer_context_faithfulness = answer_context_faithfulness_statement_level_factory()\n", + "answer_matches_target_llm_grader = answer_matches_target_llm_grader_factory()\n", + "\n", + "\n", + "@function\n", + "def generation_sglang(s, question: str, *context: str):\n", + " context = \"\\n\".join(context)\n", + " s += user(\n", + " f\"Given this question:\\n{question}\\n\\nAnd this context:\\n{context}\\n\\nAnswer the question.\"\n", + " )\n", + " s += assistant(gen(\"answer\", max_tokens=1_000))\n", + "\n", + "\n", + "@trace(eval_funcs=[answer_context_faithfulness, answer_matches_target_llm_grader])\n", + "def generation(question: str, *context):\n", + " state: ProgramState = generation_sglang.run(question, *context)\n", + " while not state.stream_executor.is_finished:\n", + " time.sleep(1)\n", + " return state.stream_executor.variables[\"answer\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Finally, we tie them together & execute the original sample query." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "@trace\n", + "def rag_pipeline(question: str) -> str:\n", + " contexts = retrieval(question)\n", + " return generation(question, *contexts)\n", + "\n", + "\n", + "rag_pipeline(\n", + " \"When did the World Health Organization formally declare an end to the COVID-19 global health emergency?\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Great, the answer is correct! Can you spot the line where we fixed the output truncation issue?\n", + "\n", + "The evaluation scores appear in the bottom right of the logs (screenshot below). Note, that there is no score for `answer_matches_target_llm_grader` and `percent_target_supported_by_context` as these evals are automatically skipped if the target answer is not provided.\n", + "\n", + "![Fixed Max. Tokens](max-tokens-fixed-rag-trace.png)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Running an experiment\n", + "\n", + "Now we are (almost) ready to evaluate the performance of our RAG pipeline on the entire dataset. First, we will need to apply the `nest_asyncio` package to avoid issues with the Jupyter notebook event loop." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "!pip install nest-asyncio\n", + "import nest_asyncio\n", + "\n", + "nest_asyncio.apply()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Running the actual experiment is straight-forward. For that we use `p.experiment` to initialize the experiment with a name, the data (list of key-value pairs fed into our entry function) and the entry function. We then call `run` on the experiment to execute it. Note, that `target` is a reserved key in the data dictionary and will be used as the target answer for evaluation." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "e = p.experiment(\n", + " \"RAG\",\n", + " data=[\n", + " {\n", + " \"question\": qca[\"question\"],\n", + " \"target\": qca[\"answer\"],\n", + " }\n", + " for qca in question_context_answers\n", + " ],\n", + " func=rag_pipeline,\n", + ").run()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Analyzing the results\n", + "\n", + "When opening above experiment, we will see an overview of the experiment as shown below. The upper half shows a summary of the statistics on the left and charts to investigate the distribution and relationships of scores on the right. The lower half is a table with the individual traces which we can use to debug individual samples.\n", + "\n", + "When looking at the statistics, we can see that the accuracy of our RAG pipeline is 22% as measured by `answer_matches_target_llm_grader`. Though when checking the quality of our retrieval step (`context_query_relevancy`), we can see that our retrival step is fetching relevant information in only 27% of all samples. As shown in the GIF, we investigate the relationship between the two and see the two scores have 95% agreement. This confirms that the retrieval step is a major bottleneck for our RAG pipeline. So, now it's your turn to improve the retrieval step!\n", + "\n", + "Note, above link isn't publicly accessible but the experiment can be accessed through [here](https://app.parea.ai/public-experiments/parea/rag_sglang/30f0244a-d56c-44ff-bdfb-8f47626304b6).\n", + "\n", + "![Experiment Results](https://drive.google.com/uc?id=1KMtJBU47nPB02Pvv3SPPTK7RnHRh5YdA)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 2 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython2" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/sglang/examples/frontend_language/usage/readme_examples.py b/sglang/examples/frontend_language/usage/readme_examples.py new file mode 100644 index 0000000..7269ef1 --- /dev/null +++ b/sglang/examples/frontend_language/usage/readme_examples.py @@ -0,0 +1,109 @@ +""" +Usage: +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +python readme_examples.py +""" + +import sglang as sgl + + +@sgl.function +def tool_use(s, question): + s += "To answer this question: " + question + ". " + s += ( + "I need to use a " + + sgl.gen("tool", choices=["calculator", "search engine"]) + + ". " + ) + + if s["tool"] == "calculator": + s += "The math expression is" + sgl.gen("expression") + elif s["tool"] == "search engine": + s += "The key word to search is" + sgl.gen("word") + + +@sgl.function +def tip_suggestion(s): + s += ( + "Here are two tips for staying healthy: " + "1. Balanced Diet. 2. Regular Exercise.\n\n" + ) + + forks = s.fork(2) + for i, f in enumerate(forks): + f += f"Now, expand tip {i+1} into a paragraph:\n" + f += sgl.gen(f"detailed_tip", max_tokens=256, stop="\n\n") + + s += "Tip 1:" + forks[0]["detailed_tip"] + "\n" + s += "Tip 2:" + forks[1]["detailed_tip"] + "\n" + s += "In summary" + sgl.gen("summary") + + +@sgl.function +def regular_expression_gen(s): + s += "Q: What is the IP address of the Google DNS servers?\n" + s += "A: " + sgl.gen( + "answer", + temperature=0, + regex=r"((25[0-5]|2[0-4]\d|[01]?\d\d?).){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)", + ) + + +@sgl.function +def text_qa(s, question): + s += "Q: " + question + "\n" + s += "A:" + sgl.gen("answer", stop="\n") + + +def driver_tool_use(): + state = tool_use.run(question="What is the capital of the United States?") + print(state.text()) + print("\n") + + +def driver_tip_suggestion(): + state = tip_suggestion.run() + print(state.text()) + print("\n") + + +def driver_regex(): + state = regular_expression_gen.run() + print(state.text()) + print("\n") + + +def driver_batching(): + states = text_qa.run_batch( + [ + {"question": "What is the capital of the United Kingdom?"}, + {"question": "What is the capital of France?"}, + {"question": "What is the capital of Japan?"}, + ], + progress_bar=True, + ) + + for s in states: + print(s.text()) + print("\n") + + +def driver_stream(): + state = text_qa.run( + question="What is the capital of France?", temperature=0.1, stream=True + ) + + for out in state.text_iter(): + print(out, end="", flush=True) + print("\n") + + +if __name__ == "__main__": + # sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo-instruct")) + sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000")) + + driver_tool_use() + driver_tip_suggestion() + driver_regex() + driver_batching() + driver_stream() diff --git a/sglang/examples/frontend_language/usage/sgl_gen_min_tokens.py b/sglang/examples/frontend_language/usage/sgl_gen_min_tokens.py new file mode 100644 index 0000000..a508819 --- /dev/null +++ b/sglang/examples/frontend_language/usage/sgl_gen_min_tokens.py @@ -0,0 +1,35 @@ +""" +This example demonstrates how to use `min_tokens` to enforce sgl.gen to generate a longer sequence + +Usage: +python3 sgl_gen_min_tokens.py +""" + +import sglang as sgl + + +@sgl.function +def long_answer(s): + s += sgl.user("What is the capital of the United States?") + s += sgl.assistant(sgl.gen("answer", min_tokens=64, max_tokens=128)) + + +@sgl.function +def short_answer(s): + s += sgl.user("What is the capital of the United States?") + s += sgl.assistant(sgl.gen("answer")) + + +if __name__ == "__main__": + runtime = sgl.Runtime(model_path="meta-llama/Meta-Llama-3.1-8B-Instruct") + sgl.set_default_backend(runtime) + + state = long_answer.run() + print("=" * 20) + print("Longer Answer", state["answer"]) + + state = short_answer.run() + print("=" * 20) + print("Short Answer", state["answer"]) + + runtime.shutdown() diff --git a/sglang/examples/frontend_language/usage/streaming.py b/sglang/examples/frontend_language/usage/streaming.py new file mode 100644 index 0000000..506ee35 --- /dev/null +++ b/sglang/examples/frontend_language/usage/streaming.py @@ -0,0 +1,49 @@ +""" +Usage: +python3 streaming.py +""" + +import asyncio + +import sglang as sgl + + +@sgl.function +def multi_turn_question(s, question_1, question_2): + s += sgl.system("You are a helpful assistant.") + s += sgl.user(question_1) + s += sgl.assistant(sgl.gen("answer_1", max_tokens=256)) + s += sgl.user(question_2) + s += sgl.assistant(sgl.gen("answer_2", max_tokens=256)) + + +sgl.set_default_backend(sgl.OpenAI("gpt-3.5-turbo")) + + +def stream_a_variable(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + for out in state.text_iter(var_name="answer_2"): + print(out, end="", flush=True) + print("\n") + + +async def async_stream(): + state = multi_turn_question.run( + question_1="What is the capital of the United States?", + question_2="List two local attractions.", + stream=True, + ) + + async for out in state.text_async_iter(var_name="answer_2"): + print(out, end="", flush=True) + print("\n") + + +if __name__ == "__main__": + stream_a_variable() + asyncio.run(async_stream()) diff --git a/sglang/examples/frontend_language/usage/triton/Dockerfile b/sglang/examples/frontend_language/usage/triton/Dockerfile new file mode 100644 index 0000000..e4741a1 --- /dev/null +++ b/sglang/examples/frontend_language/usage/triton/Dockerfile @@ -0,0 +1,10 @@ +FROM nvcr.io/nvidia/tritonserver:24.01-py3 + +WORKDIR /opt + +RUN git clone https://github.com/sgl-project/sglang.git + +WORKDIR /opt/sglang +RUN pip install --upgrade pip && \ + pip install -e "python[all]" && \ + pip install datasets diff --git a/sglang/examples/frontend_language/usage/triton/README.md b/sglang/examples/frontend_language/usage/triton/README.md new file mode 100644 index 0000000..b2e5596 --- /dev/null +++ b/sglang/examples/frontend_language/usage/triton/README.md @@ -0,0 +1,35 @@ +# sglang_triton + +Build the docker image: +``` +docker build -t sglang-triton . +``` + +Then do: +``` +docker run -ti --gpus=all --network=host --name sglang-triton -v ./models:/mnt/models sglang-triton +``` + +inside the docker container: +``` +cd sglang +python3 -m sglang.launch_server --model-path mistralai/Mistral-7B-Instruct-v0.2 --port 30000 --mem-fraction-static 0.9 +``` + +with another shell, inside the docker container: +``` +docker exec -ti sglang-triton /bin/bash +cd /mnt +tritonserver --model-repository=/mnt/models +``` + + +Send request to the server: +``` +curl -X POST http://localhost:8000/v2/models/character_generation/generate \ +-H "Content-Type: application/json" \ +-d '{ + "INPUT_TEXT": ["harry"] +}' + +``` diff --git a/sglang/examples/frontend_language/usage/triton/models/character_generation/config.pbtxt b/sglang/examples/frontend_language/usage/triton/models/character_generation/config.pbtxt new file mode 100644 index 0000000..7546f99 --- /dev/null +++ b/sglang/examples/frontend_language/usage/triton/models/character_generation/config.pbtxt @@ -0,0 +1,23 @@ +name: "character_generation" +backend: "python" +input [ + { + name: "INPUT_TEXT" + data_type: TYPE_STRING + dims: [ -1 ] + } +] +output [ + { + name: "OUTPUT_TEXT" + data_type: TYPE_STRING + dims: [ -1 ] + } +] +instance_group [ + { + count: 1 + kind: KIND_GPU + gpus: [ 0 ] + } +] diff --git a/sglang/examples/monitoring/docker-compose.yaml b/sglang/examples/monitoring/docker-compose.yaml new file mode 100644 index 0000000..6c18b40 --- /dev/null +++ b/sglang/examples/monitoring/docker-compose.yaml @@ -0,0 +1,16 @@ +services: + prometheus: + image: prom/prometheus:latest + network_mode: host + ports: + - "9090:9090" + volumes: + - ${PWD}/prometheus.yaml:/etc/prometheus/prometheus.yml + + grafana: + image: grafana/grafana:latest + network_mode: host + depends_on: + - prometheus + ports: + - "3000:3000" diff --git a/sglang/examples/monitoring/grafana.json b/sglang/examples/monitoring/grafana.json new file mode 100644 index 0000000..5fb7fbe --- /dev/null +++ b/sglang/examples/monitoring/grafana.json @@ -0,0 +1,962 @@ +{ + "annotations": { + "list": [ + { + "builtIn": 1, + "datasource": { + "type": "grafana", + "uid": "-- Grafana --" + }, + "enable": true, + "hide": true, + "iconColor": "rgba(0, 211, 255, 1)", + "name": "Annotations & Alerts", + "type": "dashboard" + } + ] + }, + "editable": true, + "fiscalYearStartMonth": 0, + "graphTooltip": 0, + "id": 1, + "links": [], + "panels": [ + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green", + "value": null + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 0 + }, + "id": 14, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "single", + "sort": "none" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "histogram_quantile(0.99, sum by (le) (rate(sglang:e2e_request_latency_seconds_bucket[$__rate_interval])))\r\n", + "fullMetaSearch": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "P99", + "range": true, + "refId": "A", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "histogram_quantile(0.9, sum by (le) (rate(sglang:e2e_request_latency_seconds_bucket[$__rate_interval])))\r\n", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "P90", + "range": true, + "refId": "B", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "histogram_quantile(0.5, sum by (le) (rate(sglang:e2e_request_latency_seconds_bucket[$__rate_interval])))\r\n", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "P95", + "range": true, + "refId": "C", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "avg(rate(sglang:e2e_request_latency_seconds_sum[$__rate_interval]) / rate(sglang:e2e_request_latency_seconds_count[$__rate_interval]))\r\n", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "P50", + "range": true, + "refId": "D", + "useBackend": false + } + ], + "title": "End-to-End Request Latency", + "type": "timeseries" + }, + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "custom": { + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "scaleDistribution": { + "type": "linear" + } + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 0 + }, + "id": 17, + "options": { + "calculate": true, + "cellGap": 1, + "color": { + "exponent": 0.5, + "fill": "dark-orange", + "mode": "scheme", + "reverse": false, + "scale": "exponential", + "scheme": "Spectral", + "steps": 64 + }, + "exemplars": { + "color": "rgba(255,0,255,0.7)" + }, + "filterValues": { + "le": 1e-09 + }, + "legend": { + "show": true + }, + "rowsFrame": { + "layout": "auto" + }, + "tooltip": { + "mode": "single", + "showColorScale": true, + "yHistogram": false + }, + "yAxis": { + "axisPlacement": "left", + "reverse": false + } + }, + "pluginVersion": "11.2.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "rate(sglang:e2e_request_latency_seconds_bucket{model_name=~\"$model_name\"}[$__rate_interval])\r\n", + "fullMetaSearch": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "{{le}}", + "range": true, + "refId": "A", + "useBackend": false + } + ], + "title": "End-to-End Request Latency Heatmap", + "type": "heatmap" + }, + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green", + "value": null + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 8 + }, + "id": 20, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "single", + "sort": "none" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "histogram_quantile(0.99, sum by (le) (rate(sglang:time_to_first_token_seconds_bucket[$__rate_interval])))\r\n", + "fullMetaSearch": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "P99", + "range": true, + "refId": "A", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "histogram_quantile(0.9, sum by (le) (rate(sglang:time_to_first_token_seconds_bucket[$__rate_interval])))\r\n", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "P90", + "range": true, + "refId": "B", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "histogram_quantile(0.95, sum by (le) (rate(sglang:time_to_first_token_seconds_bucket[$__rate_interval])))\r\n", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "P95", + "range": true, + "refId": "C", + "useBackend": false + }, + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "avg(rate(sglang:time_to_first_token_seconds_bucket[$__rate_interval]) / rate(sglang:time_to_first_token_seconds_bucket[$__rate_interval]))\r\n", + "fullMetaSearch": false, + "hide": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "P50", + "range": true, + "refId": "D", + "useBackend": false + } + ], + "title": "End-to-End Request Latency", + "type": "timeseries" + }, + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "custom": { + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "scaleDistribution": { + "type": "linear" + } + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 8 + }, + "id": 19, + "options": { + "calculate": true, + "cellGap": 1, + "color": { + "exponent": 0.5, + "fill": "dark-orange", + "mode": "scheme", + "reverse": false, + "scale": "exponential", + "scheme": "Spectral", + "steps": 64 + }, + "exemplars": { + "color": "rgba(255,0,255,0.7)" + }, + "filterValues": { + "le": 1e-09 + }, + "legend": { + "show": true + }, + "rowsFrame": { + "layout": "auto" + }, + "tooltip": { + "mode": "single", + "showColorScale": true, + "yHistogram": false + }, + "yAxis": { + "axisPlacement": "left", + "reverse": false, + "unit": "s" + } + }, + "pluginVersion": "11.2.0", + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "rate(sglang:time_to_first_token_seconds_bucket{model_name=~\"$model_name\"}[$__rate_interval])\r\n", + "fullMetaSearch": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "{{le}}", + "range": true, + "refId": "A", + "useBackend": false + } + ], + "title": "Time-To-First-Token Seconds Heatmap", + "type": "heatmap" + }, + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green", + "value": null + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 16 + }, + "id": 7, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "single", + "sort": "none" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "sglang:num_running_reqs", + "fullMetaSearch": false, + "includeNullMetadata": true, + "instant": false, + "interval": "", + "legendFormat": "{{instance}}", + "range": true, + "refId": "A", + "useBackend": false + } + ], + "title": "Num Running Requests", + "type": "timeseries" + }, + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green", + "value": null + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 16 + }, + "id": 18, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "single", + "sort": "none" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "editorMode": "code", + "expr": "sglang:gen_throughput", + "instant": false, + "legendFormat": "{{instance}}", + "range": true, + "refId": "A" + } + ], + "title": "Token Generation Throughput (Tokens / S)", + "type": "timeseries" + }, + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green", + "value": null + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 0, + "y": 24 + }, + "id": 11, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "single", + "sort": "none" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "sglang:cache_hit_rate", + "fullMetaSearch": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "{{instance}}", + "range": true, + "refId": "A", + "useBackend": false + } + ], + "title": "Cache Hit Rate", + "type": "timeseries" + }, + { + "datasource": { + "default": true, + "type": "prometheus" + }, + "fieldConfig": { + "defaults": { + "color": { + "mode": "palette-classic" + }, + "custom": { + "axisBorderShow": false, + "axisCenteredZero": false, + "axisColorMode": "text", + "axisLabel": "", + "axisPlacement": "auto", + "barAlignment": 0, + "barWidthFactor": 0.6, + "drawStyle": "line", + "fillOpacity": 0, + "gradientMode": "none", + "hideFrom": { + "legend": false, + "tooltip": false, + "viz": false + }, + "insertNulls": false, + "lineInterpolation": "linear", + "lineWidth": 1, + "pointSize": 5, + "scaleDistribution": { + "type": "linear" + }, + "showPoints": "auto", + "spanNulls": false, + "stacking": { + "group": "A", + "mode": "none" + }, + "thresholdsStyle": { + "mode": "off" + } + }, + "mappings": [], + "thresholds": { + "mode": "absolute", + "steps": [ + { + "color": "green", + "value": null + }, + { + "color": "red", + "value": 80 + } + ] + } + }, + "overrides": [] + }, + "gridPos": { + "h": 8, + "w": 12, + "x": 12, + "y": 24 + }, + "id": 8, + "options": { + "legend": { + "calcs": [], + "displayMode": "list", + "placement": "bottom", + "showLegend": true + }, + "tooltip": { + "mode": "single", + "sort": "none" + } + }, + "targets": [ + { + "datasource": { + "type": "prometheus", + "uid": "ddyfngn31dg5cf" + }, + "disableTextWrap": false, + "editorMode": "code", + "expr": "sglang:num_queue_reqs", + "fullMetaSearch": false, + "includeNullMetadata": true, + "instant": false, + "legendFormat": "{{instance}}", + "range": true, + "refId": "A", + "useBackend": false + } + ], + "title": "Number Queued Requests", + "type": "timeseries" + } + ], + "refresh": "5s", + "schemaVersion": 39, + "tags": [], + "templating": { + "list": [ + { + "current": { + "selected": false, + "text": "127.0.0.1:30000", + "value": "127.0.0.1:30000" + }, + "datasource": { + "type": "prometheus" + }, + "definition": "label_values(instance)", + "hide": 0, + "includeAll": false, + "label": "instance", + "multi": false, + "name": "instance", + "options": [], + "query": { + "qryType": 1, + "query": "label_values(instance)", + "refId": "PrometheusVariableQueryEditor-VariableQuery" + }, + "refresh": 1, + "regex": "", + "skipUrlSync": false, + "sort": 0, + "type": "query" + }, + { + "current": { + "selected": false, + "text": "meta-llama/Llama-3.1-8B-Instruct", + "value": "meta-llama/Llama-3.1-8B-Instruct" + }, + "datasource": { + "type": "prometheus" + }, + "definition": "label_values(model_name)", + "hide": 0, + "includeAll": false, + "label": "model name", + "multi": false, + "name": "model_name", + "options": [], + "query": { + "qryType": 1, + "query": "label_values(model_name)", + "refId": "PrometheusVariableQueryEditor-VariableQuery" + }, + "refresh": 1, + "regex": "", + "skipUrlSync": false, + "sort": 0, + "type": "query" + } + ] + }, + "time": { + "from": "now-30m", + "to": "now" + }, + "timepicker": {}, + "timezone": "browser", + "title": "SGLang Dashboard", + "uid": "sglang-dashboard", + "version": 5, + "weekStart": "" +} diff --git a/sglang/examples/monitoring/prometheus.yaml b/sglang/examples/monitoring/prometheus.yaml new file mode 100644 index 0000000..ba16ac3 --- /dev/null +++ b/sglang/examples/monitoring/prometheus.yaml @@ -0,0 +1,10 @@ +# prometheus.yaml +global: + scrape_interval: 5s + evaluation_interval: 30s + +scrape_configs: + - job_name: sglang + static_configs: + - targets: + - '127.0.0.1:30000' diff --git a/sglang/examples/runtime/README.md b/sglang/examples/runtime/README.md new file mode 100644 index 0000000..ecf116d --- /dev/null +++ b/sglang/examples/runtime/README.md @@ -0,0 +1,42 @@ +# Runtime examples + +The below examples will mostly need you to start a server in a separate terminal before you can execute them. Please see in the code for detailed instruction. + +## Native API + +* `lora.py`: An example how to use LoRA adapters. +* `multimodal_embedding.py`: An example how perform [multi modal embedding](Alibaba-NLP/gme-Qwen2-VL-2B-Instruct). +* `openai_batch_chat.py`: An example how to process batch requests for chat completions. +* `openai_batch_complete.py`: An example how to process batch requests for text completions. +* `openai_chat_with_response_prefill.py`: An example how to [prefill](https://eugeneyan.com/writing/prompting/#prefill-claudes-responses) a response using OpenAI API. +* `reward_model.py`: An example how to extract scores from a reward model. +* `vertex_predict.py`: An example how to deploy a model to [Vertex AI](https://cloud.google.com/vertex-ai?hl=en). + +## Engine + +The `engine` folder contains that examples that show how to use [Offline Engine API](https://docs.sglang.ai/backend/offline_engine_api.html#Offline-Engine-API) for common workflows. + +* `custom_server.py`: An example how to deploy a custom server. +* `embedding.py`: An example how to extract embeddings. +* `launch_engine.py`: An example how to launch the Engine. +* `offline_batch_inference_eagle.py`: An example how to perform speculative decoding using [EAGLE](https://docs.sglang.ai/backend/speculative_decoding.html). +* `offline_batch_inference_torchrun.py`: An example how to perform inference using [torchrun](https://pytorch.org/docs/stable/elastic/run.html). +* `offline_batch_inference_vlm.py`: An example how to use VLMs with the engine. +* `offline_batch_inference.py`: An example how to use the engine to perform inference on a batch of examples. + +## Hidden States + +The `hidden_states` folder contains examples on how to extract hidden states using SGLang. Please note that this might degrade throughput due to cuda graph rebuilding. + +* `hidden_states_engine.py`: An example how to extract hidden states using the Engine API. +* `hidden_states_server.py`: An example how to extract hidden states using the Server API. + +## LLaVA-NeXT + +SGLang support LLaVA-OneVision with single-image, multi-image and video are supported. The folder `llava_onevision` shows how to do this. + +## Token In, Token Out + +The folder `token_in_token_out` shows how to perform inference, where we provide tokens and get tokens as response. + +* `token_in_token_out_{llm|vlm}_{engine|server}.py`: Shows how to perform token in, token out workflow for llm/vlm using either the engine or native API. diff --git a/sglang/examples/runtime/engine/custom_server.py b/sglang/examples/runtime/engine/custom_server.py new file mode 100644 index 0000000..b190a46 --- /dev/null +++ b/sglang/examples/runtime/engine/custom_server.py @@ -0,0 +1,53 @@ +from sanic import Sanic, text +from sanic.response import json + +import sglang as sgl + +engine = None + +# Create an instance of the Sanic app +app = Sanic("sanic-server") + + +# Define an asynchronous route handler +@app.route("/generate", methods=["POST"]) +async def generate(request): + prompt = request.json.get("prompt") + if not prompt: + return json({"error": "Prompt is required"}, status=400) + + # async_generate returns a dict + result = await engine.async_generate(prompt) + + return text(result["text"]) + + +@app.route("/generate_stream", methods=["POST"]) +async def generate_stream(request): + prompt = request.json.get("prompt") + + if not prompt: + return json({"error": "Prompt is required"}, status=400) + + # async_generate returns a dict + result = await engine.async_generate(prompt, stream=True) + + # https://sanic.dev/en/guide/advanced/streaming.md#streaming + # init the response + response = await request.respond() + + # result is an async generator + async for chunk in result: + await response.send(chunk["text"]) + + await response.eof() + + +def run_server(): + global engine + engine = sgl.Engine(model_path="meta-llama/Meta-Llama-3.1-8B-Instruct") + app.run(host="0.0.0.0", port=8000, single_process=True) + + +if __name__ == "__main__": + run_server() diff --git a/sglang/examples/runtime/engine/embedding.py b/sglang/examples/runtime/engine/embedding.py new file mode 100644 index 0000000..b927a18 --- /dev/null +++ b/sglang/examples/runtime/engine/embedding.py @@ -0,0 +1,27 @@ +import sglang as sgl + + +def main(): + # Sample prompts. + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + # Create an LLM. + llm = sgl.Engine( + model_path="Alibaba-NLP/gte-Qwen2-1.5B-instruct", is_embedding=True + ) + + outputs = llm.encode(prompts) + # Print the outputs (embedding vectors) + for prompt, output in zip(prompts, outputs): + print("===============================") + print(f"Prompt: {prompt}\nEmbedding vector: {output['embedding']}") + + +# The __main__ condition is necessary here because we use "spawn" to create subprocesses +# Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine +if __name__ == "__main__": + main() diff --git a/sglang/examples/runtime/engine/launch_engine.py b/sglang/examples/runtime/engine/launch_engine.py new file mode 100644 index 0000000..b97e0bc --- /dev/null +++ b/sglang/examples/runtime/engine/launch_engine.py @@ -0,0 +1,17 @@ +""" +This example demonstrates how to launch the offline engine. +""" + +import sglang as sgl + + +def main(): + llm = sgl.Engine(model_path="meta-llama/Meta-Llama-3.1-8B-Instruct") + llm.generate("What is the capital of France?") + llm.shutdown() + + +# The __main__ condition is necessary here because we use "spawn" to create subprocesses +# Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine +if __name__ == "__main__": + main() diff --git a/sglang/examples/runtime/engine/offline_batch_inference.py b/sglang/examples/runtime/engine/offline_batch_inference.py new file mode 100644 index 0000000..92e68dc --- /dev/null +++ b/sglang/examples/runtime/engine/offline_batch_inference.py @@ -0,0 +1,43 @@ +""" +Usage: +python3 offline_batch_inference.py --model meta-llama/Llama-3.1-8B-Instruct +""" + +import argparse +import dataclasses + +import sglang as sgl +from sglang.srt.server_args import ServerArgs + + +def main( + server_args: ServerArgs, +): + # Sample prompts. + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + # Create a sampling params object. + sampling_params = {"temperature": 0.8, "top_p": 0.95} + + # Create an LLM. + llm = sgl.Engine(**dataclasses.asdict(server_args)) + + outputs = llm.generate(prompts, sampling_params) + # Print the outputs. + for prompt, output in zip(prompts, outputs): + print("===============================") + print(f"Prompt: {prompt}\nGenerated text: {output['text']}") + + +# The __main__ condition is necessary here because we use "spawn" to create subprocesses +# Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + args = parser.parse_args() + server_args = ServerArgs.from_cli_args(args) + main(server_args) diff --git a/sglang/examples/runtime/engine/offline_batch_inference_async.py b/sglang/examples/runtime/engine/offline_batch_inference_async.py new file mode 100644 index 0000000..578962d --- /dev/null +++ b/sglang/examples/runtime/engine/offline_batch_inference_async.py @@ -0,0 +1,65 @@ +""" +Usage: +python offline_batch_inference_async.py --model-path Qwen/Qwen2-VL-7B-Instruct + +Note: +This demo shows the usage of async generation, +which is useful to implement an online-like generation with batched inference. +""" + +import argparse +import asyncio +import dataclasses +import time + +import sglang as sgl +from sglang.srt.server_args import ServerArgs + + +class InferenceEngine: + def __init__(self, **kwargs): + self.engine = sgl.Engine(**kwargs) + + async def generate(self, prompt, sampling_params): + result = await self.engine.async_generate(prompt, sampling_params) + return result + + +async def run_server(server_args): + inference = InferenceEngine(**dataclasses.asdict(server_args)) + + # Sample prompts. + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] * 100 + + # Create a sampling params object. + sampling_params = {"temperature": 0.8, "top_p": 0.95} + + # Run the generation tasks concurrently in async mode. + tasks = [] + for prompt in prompts: + task = asyncio.create_task(inference.generate(prompt, sampling_params)) + tasks.append(task) + + # Get and print the result + for task in tasks: + await task + while True: + if not task.done(): + time.sleep(1) + else: + result = task.result() + print(f"Generated text: {result['text']}") + break + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + args = parser.parse_args() + server_args = ServerArgs.from_cli_args(args) + asyncio.run(run_server(server_args)) diff --git a/sglang/examples/runtime/engine/offline_batch_inference_eagle.py b/sglang/examples/runtime/engine/offline_batch_inference_eagle.py new file mode 100644 index 0000000..a7a89ef --- /dev/null +++ b/sglang/examples/runtime/engine/offline_batch_inference_eagle.py @@ -0,0 +1,38 @@ +import sglang as sgl + + +def main(): + # Sample prompts. + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + + # Create a sampling params object. + sampling_params = {"temperature": 0, "max_new_tokens": 30} + + # Create an LLM. + llm = sgl.Engine( + model_path="meta-llama/Llama-2-7b-chat-hf", + speculative_algorithm="EAGLE", + speculative_draft_model_path="lmsys/sglang-EAGLE-llama2-chat-7B", + speculative_num_steps=3, + speculative_eagle_topk=4, + speculative_num_draft_tokens=16, + cuda_graph_max_bs=8, + ) + + outputs = llm.generate(prompts, sampling_params) + + # Print the outputs. + for prompt, output in zip(prompts, outputs): + print("===============================") + print(f"Prompt: {prompt}\nGenerated text: {output['text']}") + + +# The __main__ condition is necessary here because we use "spawn" to create subprocesses +# Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine +if __name__ == "__main__": + main() diff --git a/sglang/examples/runtime/engine/offline_batch_inference_torchrun.py b/sglang/examples/runtime/engine/offline_batch_inference_torchrun.py new file mode 100644 index 0000000..d2185da --- /dev/null +++ b/sglang/examples/runtime/engine/offline_batch_inference_torchrun.py @@ -0,0 +1,81 @@ +import datetime +import os +import sys + +from torch.distributed.device_mesh import init_device_mesh + +from sglang.srt.entrypoints.verl_engine import VerlEngine + + +def run(): + """ + Example command: + ``` + torchrun --nproc_per_node=8 offline_batch_inference_torchrun.py + ``` + """ + + local_rank = int(os.environ["LOCAL_RANK"]) + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + + def _log(text): + t = datetime.datetime.now().strftime("%H:%M:%S") + print(f"[{t}] [rank={rank}] {text}") + + _log( + f'start {local_rank=} {rank=} {world_size=} {sys.argv=} {os.environ.get("CUDA_VISIBLE_DEVICES")}' + ) + + tp_size = 4 + dp_size = 2 + assert world_size == tp_size * dp_size + + device_mesh_kwargs = dict( + mesh_shape=(tp_size, dp_size, 1), mesh_dim_names=["tp", "dp", "pp"] + ) + device_mesh_cpu = init_device_mesh("cpu", **device_mesh_kwargs) + _log(f"{device_mesh_cpu=}") + + tp_rank = device_mesh_cpu.get_local_rank("tp") + dp_rank = device_mesh_cpu.get_local_rank("dp") + _log(f"{tp_rank=} {tp_size=} ; {dp_rank=} {dp_size=}") + + model_name, mem_fraction_static = "meta-llama/Llama-3.2-1B-Instruct", 0.1 + # model_name, mem_fraction_static = "meta-llama/Llama-3.1-70B-Instruct", 0.9 # test large models + # model_name, mem_fraction_static = "deepseek-ai/DeepSeek-V2-Lite", 0.8 + + for k in ["TORCHELASTIC_USE_AGENT_STORE"]: + if k in os.environ: + del os.environ[k] + + fragment = VerlEngine( + model_path=model_name, + mem_fraction_static=mem_fraction_static, + device_mesh_cpu=device_mesh_cpu["tp"], + base_gpu_id=dp_rank, + gpu_id_step=dp_size, + port=30000, + # for DeepSeek-V2-Lite + DP Attention + # enable_dp_attention=True, port=30000 + dp_rank * 100, + ) + _log(f"{fragment=}") + + prompt_all = [ + ["1+1=2, 1+2=3, 1+3=4, 1+4=", "9-1=8, 8-1=7, 7-1="], + ["2*1=2, 2*2=4, 2*3=", "8/2=4, 6/2="], + ] + prompt = prompt_all[dp_rank] + + output = fragment.generate( + prompt=prompt, + sampling_params=dict(max_new_tokens=16, temperature=0.0), + ) + _log(f"{prompt=} {output=}") + + fragment.shutdown() + _log(f"End script") + + +if __name__ == "__main__": + run() diff --git a/sglang/examples/runtime/engine/offline_batch_inference_vlm.py b/sglang/examples/runtime/engine/offline_batch_inference_vlm.py new file mode 100644 index 0000000..4063136 --- /dev/null +++ b/sglang/examples/runtime/engine/offline_batch_inference_vlm.py @@ -0,0 +1,52 @@ +""" +Usage: +python offline_batch_inference_vlm.py --model-path Qwen/Qwen2-VL-7B-Instruct --chat-template=qwen2-vl +""" + +import argparse +import dataclasses + +import sglang as sgl +from sglang.srt.conversation import chat_templates +from sglang.srt.server_args import ServerArgs + + +def main( + server_args: ServerArgs, +): + vlm = sgl.Engine(**dataclasses.asdict(server_args)) + + conv = chat_templates[server_args.chat_template].copy() + image_token = conv.image_token + + image_url = "https://github.com/sgl-project/sglang/blob/main/test/lang/example_image.png?raw=true" + + prompt = f"What's in this image?\n{image_token}" + + sampling_params = { + "temperature": 0.001, + "max_new_tokens": 30, + } + + output = vlm.generate( + prompt=prompt, + image_data=image_url, + sampling_params=sampling_params, + ) + + print("===============================") + print(f"Prompt: {prompt}") + print(f"Generated text: {output['text']}") + + vlm.shutdown() + + +# The __main__ condition is necessary here because we use "spawn" to create subprocesses +# Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + args = parser.parse_args() + + server_args = ServerArgs.from_cli_args(args) + main(server_args) diff --git a/sglang/examples/runtime/engine/readme.md b/sglang/examples/runtime/engine/readme.md new file mode 100644 index 0000000..d5ac936 --- /dev/null +++ b/sglang/examples/runtime/engine/readme.md @@ -0,0 +1,49 @@ +# SGLang Engine + +## Introduction +SGLang provides a direct inference engine without the need for an HTTP server. There are generally two use cases: + +1. **Offline Batch Inference** +2. **Embedding Generation** +3. **Custom Server on Top of the Engine** + +## Examples + +### [Offline Batch Inference](./offline_batch_inference.py) + +In this example, we launch an SGLang engine and feed a batch of inputs for inference. If you provide a very large batch, the engine will intelligently schedule the requests to process efficiently and prevent OOM (Out of Memory) errors. + +### [Embedding Generation](./embedding.py) + +In this example, we launch an SGLang engine and feed a batch of inputs for embedding generation. + +### [Custom Server](./custom_server.py) + +This example demonstrates how to create a custom server on top of the SGLang Engine. We use [Sanic](https://sanic.dev/en/) as an example. The server supports both non-streaming and streaming endpoints. + +#### Steps: + +1. Install Sanic: + +```bash +pip install sanic +``` + +2. Run the server: + +```bash +python custom_server +``` + +3. Send requests: + +```bash +curl -X POST http://localhost:8000/generate -H "Content-Type: application/json" -d '{"prompt": "The Transformer architecture is..."}' +curl -X POST http://localhost:8000/generate_stream -H "Content-Type: application/json" -d '{"prompt": "The Transformer architecture is..."}' --no-buffer +``` + +This will send both non-streaming and streaming requests to the server. + +### [Token-In-Token-Out for RLHF](../token_in_token_out) + +In this example, we launch an SGLang engine, feed tokens as input and generate tokens as output. diff --git a/sglang/examples/runtime/engine/save_remote_state.py b/sglang/examples/runtime/engine/save_remote_state.py new file mode 100644 index 0000000..4781269 --- /dev/null +++ b/sglang/examples/runtime/engine/save_remote_state.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Saves each worker's model state dict directly to a checkpoint, which enables a +fast load path for large tensor-parallel models where each worker only needs to +read its own shard rather than the entire checkpoint. + +Example usage: + +python save_remote_state.py \ + --model-path /path/to/load \ + --tensor-parallel-size 8 \ + --remote-model-save-url [protocol]://[host]:[port]/[model_name] \ + +Then, the model can be loaded with + +llm = Engine( + model_path="/path/to/save", + --remote-model-url [protocol]://[host]:[port]/[model_name], + tensor_parallel_size=8, +) +""" +import dataclasses +from argparse import ArgumentParser +from pathlib import Path + +from sglang import Engine, ServerArgs + +parser = ArgumentParser() +ServerArgs.add_cli_args(parser) + +parser.add_argument( + "--remote-model-save-url", + required=True, + type=str, + help="remote address to store model weights", +) + + +def main(args): + engine_args = ServerArgs.from_cli_args(args) + model_path = engine_args.model_path + if not Path(model_path).is_dir(): + raise ValueError("model path must be a local directory") + # Create LLM instance from arguments + llm = Engine(**dataclasses.asdict(engine_args)) + llm.save_remote_model(url=args.remote_model_save_url) + + +if __name__ == "__main__": + args = parser.parse_args() + main(args) diff --git a/sglang/examples/runtime/engine/save_sharded_state.py b/sglang/examples/runtime/engine/save_sharded_state.py new file mode 100644 index 0000000..80ad532 --- /dev/null +++ b/sglang/examples/runtime/engine/save_sharded_state.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +""" +Saves each worker's model state dict directly to a checkpoint, which enables a +fast load path for large tensor-parallel models where each worker only needs to +read its own shard rather than the entire checkpoint. + +Example usage: + +python save_sharded_state.py \ + --model-path /path/to/load \ + --quantization deepspeedfp \ + --tensor-parallel-size 8 \ + --output /path/to/save + +Then, the model can be loaded with + +llm = Engine( + model_path="/path/to/save", + load_format="sharded_state", + quantization="deepspeedfp", + tensor_parallel_size=8, +) +""" +import dataclasses +import os +import shutil +from argparse import ArgumentParser +from pathlib import Path + +from sglang import Engine, ServerArgs + +parser = ArgumentParser() +ServerArgs.add_cli_args(parser) + +parser.add_argument( + "--output", "-o", required=True, type=str, help="path to output checkpoint" +) +parser.add_argument( + "--file-pattern", type=str, help="string pattern of saved filenames" +) +parser.add_argument( + "--max-file-size", + type=str, + default=5 * 1024**3, + help="max size (in bytes) of each safetensors file", +) + + +def main(args): + engine_args = ServerArgs.from_cli_args(args) + model_path = engine_args.model_path + if not Path(model_path).is_dir(): + raise ValueError("model path must be a local directory") + # Create LLM instance from arguments + llm = Engine(**dataclasses.asdict(engine_args)) + Path(args.output).mkdir(exist_ok=True) + llm.save_sharded_model( + path=args.output, pattern=args.file_pattern, max_size=args.max_file_size + ) + + # Copy metadata files to output directory + for file in os.listdir(model_path): + if os.path.splitext(file)[1] not in (".bin", ".pt", ".safetensors"): + if os.path.isdir(os.path.join(model_path, file)): + shutil.copytree( + os.path.join(model_path, file), os.path.join(args.output, file) + ) + else: + shutil.copy(os.path.join(model_path, file), args.output) + + +if __name__ == "__main__": + args = parser.parse_args() + main(args) diff --git a/sglang/examples/runtime/hidden_states/hidden_states_engine.py b/sglang/examples/runtime/hidden_states/hidden_states_engine.py new file mode 100644 index 0000000..8af883a --- /dev/null +++ b/sglang/examples/runtime/hidden_states/hidden_states_engine.py @@ -0,0 +1,65 @@ +""" +Usage: +python hidden_states.py + +Note that each time you change the `return_hidden_states` parameter, +the cuda graph will be recaptured, which might lead to a performance hit. +So avoid getting hidden states and completions alternately. +""" + +import torch + +import sglang as sgl + + +def main(): + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + # Create an LLM. + llm = sgl.Engine( + model_path="Alibaba-NLP/gte-Qwen2-1.5B-instruct", + ) + + sampling_params = { + "temperature": 0.8, + "top_p": 0.95, + "max_new_tokens": 10, + } + + outputs = llm.generate( + prompts, sampling_params=sampling_params, return_hidden_states=True + ) + + llm.shutdown() + + for prompt, output in zip(prompts, outputs): + for i in range(len(output["meta_info"]["hidden_states"])): + output["meta_info"]["hidden_states"][i] = torch.tensor( + output["meta_info"]["hidden_states"][i], dtype=torch.bfloat16 + ) + print("===============================") + print( + f"Prompt: {prompt}\n" + f"Generated text: {output['text']}\n" + f"Prompt_Tokens: {output['meta_info']['prompt_tokens']}\t" + f"Completion_tokens: {output['meta_info']['completion_tokens']}" + ) + print("Hidden states: ") + hidden_states = torch.cat( + [ + i.unsqueeze(0) if len(i.shape) == 1 else i + for i in output["meta_info"]["hidden_states"] + ] + ) + print(hidden_states) + print() + + +# The __main__ condition is necessary here because we use "spawn" to create subprocesses +# Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine +if __name__ == "__main__": + main() diff --git a/sglang/examples/runtime/hidden_states/hidden_states_server.py b/sglang/examples/runtime/hidden_states/hidden_states_server.py new file mode 100644 index 0000000..96045fa --- /dev/null +++ b/sglang/examples/runtime/hidden_states/hidden_states_server.py @@ -0,0 +1,81 @@ +""" +Usage: + +python hidden_states_server.py + +Note that each time you change the `return_hidden_states` parameter, +the cuda graph will be recaptured, which might lead to a performance hit. +So avoid getting hidden states and completions alternately. +""" + +import requests +import torch + +from sglang.test.test_utils import is_in_ci +from sglang.utils import terminate_process, wait_for_server + +if is_in_ci(): + from docs.backend.patch import launch_server_cmd +else: + from sglang.utils import launch_server_cmd + + +def main(): + # Launch the server + server_process, port = launch_server_cmd( + "python -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-1.5B-instruct --host 0.0.0.0" + ) + wait_for_server(f"http://localhost:{port}") + + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + + sampling_params = { + "temperature": 0.8, + "top_p": 0.95, + "max_new_tokens": 10, + } + + json_data = { + "text": prompts, + "sampling_params": sampling_params, + "return_hidden_states": True, + } + + response = requests.post( + f"http://localhost:{port}/generate", + json=json_data, + ) + + terminate_process(server_process) + + outputs = response.json() + for prompt, output in zip(prompts, outputs): + for i in range(len(output["meta_info"]["hidden_states"])): + output["meta_info"]["hidden_states"][i] = torch.tensor( + output["meta_info"]["hidden_states"][i], dtype=torch.bfloat16 + ) + print("===============================") + print( + f"Prompt: {prompt}\n" + f"Generated text: {output['text']}\n" + f"Prompt_Tokens: {output['meta_info']['prompt_tokens']}\t" + f"Completion_tokens: {output['meta_info']['completion_tokens']}" + ) + print("Hidden states: ") + hidden_states = torch.cat( + [ + i.unsqueeze(0) if len(i.shape) == 1 else i + for i in output["meta_info"]["hidden_states"] + ] + ) + print(hidden_states) + print() + + +if __name__ == "__main__": + main() diff --git a/sglang/examples/runtime/llava_onevision/http_llama3_llava_test.py b/sglang/examples/runtime/llava_onevision/http_llama3_llava_test.py new file mode 100644 index 0000000..ed0e616 --- /dev/null +++ b/sglang/examples/runtime/llava_onevision/http_llama3_llava_test.py @@ -0,0 +1,111 @@ +""" +Usage: +# Installing latest llava-next: pip install git+https://github.com/LLaVA-VL/LLaVA-NeXT.git +# Installing latest sglang. + +# Endpoint Service CLI: +python -m sglang.launch_server --model-path lmms-lab/llama3-llava-next-8b --port=30000 + +python3 http_llama3_llava_test.py + +Output: +"Friends posing for a fun photo with a life-sized teddy bear, creating a playful and memorable moment." +""" + +import argparse +import asyncio +import copy +import json + +import aiohttp +import requests +from llava.conversation import conv_llava_llama_3 + + +async def send_request(url, data, delay=0): + await asyncio.sleep(delay) + async with aiohttp.ClientSession() as session: + async with session.post(url, json=data) as resp: + output = await resp.json() + return output + + +async def test_concurrent(args): + url = f"{args.host}:{args.port}" + + prompt = "\nPlease generate caption towards this image." + conv_template = copy.deepcopy(conv_llava_llama_3) + conv_template.append_message(role=conv_template.roles[0], message=prompt) + conv_template.append_message(role=conv_template.roles[1], message=None) + prompt_with_template = conv_template.get_prompt() + response = [] + for i in range(1): + response.append( + send_request( + url + "/generate", + { + "text": prompt_with_template, + "image_data": "https://farm4.staticflickr.com/3175/2653711032_804ff86d81_z.jpg", + "sampling_params": { + "max_new_tokens": 1024, + "temperature": 0, + "top_p": 1.0, + "presence_penalty": 2, + "frequency_penalty": 2, + "stop": "<|eot_id|>", + }, + }, + ) + ) + + rets = await asyncio.gather(*response) + for ret in rets: + print(ret["text"]) + + +def test_streaming(args): + url = f"{args.host}:{args.port}" + prompt = "\nPlease generate caption towards this image." + conv_template = copy.deepcopy(conv_llava_llama_3) + conv_template.append_message(role=conv_template.roles[0], message=prompt) + conv_template.append_message(role=conv_template.roles[1], message=None) + prompt_with_template = conv_template.get_prompt() + pload = { + "text": prompt_with_template, + "sampling_params": { + "max_new_tokens": 1024, + "temperature": 0, + "top_p": 1.0, + "presence_penalty": 2, + "frequency_penalty": 2, + "stop": "<|eot_id|>", + }, + "image_data": "https://farm4.staticflickr.com/3175/2653711032_804ff86d81_z.jpg", + "stream": True, + } + response = requests.post( + url + "/generate", + json=pload, + stream=True, + ) + + prev = 0 + for chunk in response.iter_lines(decode_unicode=False): + chunk = chunk.decode("utf-8") + if chunk and chunk.startswith("data:"): + if chunk == "data: [DONE]": + break + data = json.loads(chunk[5:].strip("\n")) + output = data["text"].strip() + print(output[prev:], end="", flush=True) + prev = len(output) + print("") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--host", type=str, default="http://127.0.0.1") + parser.add_argument("--port", type=int, default=30000) + args = parser.parse_args() + asyncio.run(test_concurrent(args)) + test_streaming(args) diff --git a/sglang/examples/runtime/llava_onevision/http_llava_onevision_test.py b/sglang/examples/runtime/llava_onevision/http_llava_onevision_test.py new file mode 100644 index 0000000..ce59f40 --- /dev/null +++ b/sglang/examples/runtime/llava_onevision/http_llava_onevision_test.py @@ -0,0 +1,264 @@ +""" +Usage: + +python3 -m sglang.launch_server --model-path lmms-lab/llava-onevision-qwen2-72b-ov --port=30000 --tp-size=8 --chat-template=chatml-llava + +python3 http_llava_onevision_test.py +""" + +import base64 +import io +import os +import sys +import time + +import numpy as np +import openai +import requests +from decord import VideoReader, cpu +from PIL import Image + +# pip install httpx==0.23.3 +# pip install decord +# pip install protobuf==3.20.0 + + +def download_video(url, cache_dir): + file_path = os.path.join(cache_dir, "jobs.mp4") + os.makedirs(cache_dir, exist_ok=True) + + response = requests.get(url) + response.raise_for_status() + + with open(file_path, "wb") as f: + f.write(response.content) + + print(f"File downloaded and saved to: {file_path}") + return file_path + + +def create_openai_client(base_url): + return openai.Client(api_key="EMPTY", base_url=base_url) + + +def image_stream_request_test(client): + print("----------------------Image Stream Request Test----------------------") + stream_request = client.chat.completions.create( + model="default", + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png" + }, + }, + { + "type": "text", + "text": "Please describe this image. Please list the benchmarks and the models.", + }, + ], + }, + ], + temperature=0.7, + max_tokens=1024, + stream=True, + ) + stream_response = "" + + for chunk in stream_request: + if chunk.choices[0].delta.content is not None: + content = chunk.choices[0].delta.content + stream_response += content + sys.stdout.write(content) + sys.stdout.flush() + + print("-" * 30) + + +def multi_image_stream_request_test(client): + print( + "----------------------Multi-Images Stream Request Test----------------------" + ) + stream_request = client.chat.completions.create( + model="default", + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png" + }, + "modalities": "multi-images", + }, + { + "type": "image_url", + "image_url": { + "url": "https://raw.githubusercontent.com/sgl-project/sglang/main/test/lang/example_image.png" + }, + "modalities": "multi-images", + }, + { + "type": "text", + "text": "I have shown you two images. Please describe the two images to me.", + }, + ], + }, + ], + temperature=0.7, + max_tokens=1024, + stream=True, + ) + stream_response = "" + + for chunk in stream_request: + if chunk.choices[0].delta.content is not None: + content = chunk.choices[0].delta.content + stream_response += content + sys.stdout.write(content) + sys.stdout.flush() + + print("-" * 30) + + +def video_stream_request_test(client, video_path): + print("------------------------Video Stream Request Test----------------------") + messages = prepare_video_messages(video_path) + + video_request = client.chat.completions.create( + model="default", + messages=messages, + temperature=0, + max_tokens=1024, + stream=True, + ) + print("-" * 30) + video_response = "" + + for chunk in video_request: + if chunk.choices[0].delta.content is not None: + content = chunk.choices[0].delta.content + video_response += content + sys.stdout.write(content) + sys.stdout.flush() + print("-" * 30) + + +def image_speed_test(client): + print("----------------------Image Speed Test----------------------") + start_time = time.time() + request = client.chat.completions.create( + model="default", + messages=[ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png" + }, + }, + { + "type": "text", + "text": "Please describe this image. Please list the benchmarks and the models.", + }, + ], + }, + ], + temperature=0, + max_tokens=1024, + ) + end_time = time.time() + response = request.choices[0].message.content + print(response) + print("-" * 30) + print_speed_test_results(request, start_time, end_time) + + +def video_speed_test(client, video_path): + print("------------------------Video Speed Test------------------------") + messages = prepare_video_messages(video_path) + + start_time = time.time() + video_request = client.chat.completions.create( + model="default", + messages=messages, + temperature=0, + max_tokens=1024, + ) + end_time = time.time() + video_response = video_request.choices[0].message.content + print(video_response) + print("-" * 30) + print_speed_test_results(video_request, start_time, end_time) + + +def prepare_video_messages(video_path): + max_frames_num = 32 + vr = VideoReader(video_path, ctx=cpu(0)) + total_frame_num = len(vr) + uniform_sampled_frames = np.linspace( + 0, total_frame_num - 1, max_frames_num, dtype=int + ) + frame_idx = uniform_sampled_frames.tolist() + frames = vr.get_batch(frame_idx).asnumpy() + + base64_frames = [] + for frame in frames: + pil_img = Image.fromarray(frame) + buff = io.BytesIO() + pil_img.save(buff, format="JPEG") + base64_str = base64.b64encode(buff.getvalue()).decode("utf-8") + base64_frames.append(base64_str) + + messages = [{"role": "user", "content": []}] + + for base64_frame in base64_frames: + frame_format = { + "type": "image_url", + "image_url": {"url": f"data:image/jpeg;base64,{base64_frame}"}, + "modalities": "video", + } + messages[0]["content"].append(frame_format) + + prompt = {"type": "text", "text": "Please describe the video in detail."} + messages[0]["content"].append(prompt) + + return messages + + +def print_speed_test_results(request, start_time, end_time): + total_tokens = request.usage.total_tokens + completion_tokens = request.usage.completion_tokens + prompt_tokens = request.usage.prompt_tokens + + print(f"Total tokens: {total_tokens}") + print(f"Completion tokens: {completion_tokens}") + print(f"Prompt tokens: {prompt_tokens}") + print(f"Time taken: {end_time - start_time} seconds") + print(f"Token per second: {total_tokens / (end_time - start_time)}") + print(f"Completion token per second: {completion_tokens / (end_time - start_time)}") + print(f"Prompt token per second: {prompt_tokens / (end_time - start_time)}") + + +def main(): + url = "https://raw.githubusercontent.com/EvolvingLMMs-Lab/sglang/dev/onevision_local/assets/jobs.mp4" + cache_dir = os.path.expanduser("~/.cache") + video_path = download_video(url, cache_dir) + + client = create_openai_client("http://127.0.0.1:30000/v1") + + image_stream_request_test(client) + multi_image_stream_request_test(client) + video_stream_request_test(client, video_path) + image_speed_test(client) + video_speed_test(client, video_path) + + +if __name__ == "__main__": + main() diff --git a/sglang/examples/runtime/llava_onevision/http_qwen_llava_test.py b/sglang/examples/runtime/llava_onevision/http_qwen_llava_test.py new file mode 100644 index 0000000..94ec4e0 --- /dev/null +++ b/sglang/examples/runtime/llava_onevision/http_qwen_llava_test.py @@ -0,0 +1,111 @@ +""" +Usage: +# Installing latest llava-next: pip install git+https://github.com/LLaVA-VL/LLaVA-NeXT.git +# Installing latest sglang. + +# Endpoint Service CLI: +python -m sglang.launch_server --model-path lmms-lab/llava-next-72b --port=30000 --tp-size=8 + +python3 http_qwen_llava_test.py + +Output: +"Two children pose with a large teddy bear, one holding a smaller stuffed bear, in a room with an American flag and potted plants." +""" + +import argparse +import asyncio +import copy +import json + +import aiohttp +import requests +from llava.conversation import conv_qwen + + +async def send_request(url, data, delay=0): + await asyncio.sleep(delay) + async with aiohttp.ClientSession() as session: + async with session.post(url, json=data) as resp: + output = await resp.json() + return output + + +async def test_concurrent(args): + url = f"{args.host}:{args.port}" + + prompt = "\nPlease generate caption towards this image." + conv_template = copy.deepcopy(conv_qwen) + conv_template.append_message(role=conv_template.roles[0], message=prompt) + conv_template.append_message(role=conv_template.roles[1], message=None) + prompt_with_template = conv_template.get_prompt() + response = [] + for i in range(1): + response.append( + send_request( + url + "/generate", + { + "text": prompt_with_template, + "image_data": "https://farm4.staticflickr.com/3175/2653711032_804ff86d81_z.jpg", + "sampling_params": { + "max_new_tokens": 1024, + "temperature": 0, + "top_p": 1.0, + "presence_penalty": 2, + "frequency_penalty": 2, + "stop": "<|im_end|>", + }, + }, + ) + ) + + rets = await asyncio.gather(*response) + for ret in rets: + print(ret["text"]) + + +def test_streaming(args): + url = f"{args.host}:{args.port}" + prompt = "\nPlease generate caption towards this image." + conv_template = copy.deepcopy(conv_qwen) + conv_template.append_message(role=conv_template.roles[0], message=prompt) + conv_template.append_message(role=conv_template.roles[1], message=None) + prompt_with_template = conv_template.get_prompt() + pload = { + "text": prompt_with_template, + "sampling_params": { + "max_new_tokens": 1024, + "temperature": 0, + "top_p": 1.0, + "presence_penalty": 2, + "frequency_penalty": 2, + "stop": "<|im_end|>", + }, + "image_data": "https://farm4.staticflickr.com/3175/2653711032_804ff86d81_z.jpg", + "stream": True, + } + response = requests.post( + url + "/generate", + json=pload, + stream=True, + ) + + prev = 0 + for chunk in response.iter_lines(decode_unicode=False): + chunk = chunk.decode("utf-8") + if chunk and chunk.startswith("data:"): + if chunk == "data: [DONE]": + break + data = json.loads(chunk[5:].strip("\n")) + output = data["text"].strip() + print(output[prev:], end="", flush=True) + prev = len(output) + print("") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--host", type=str, default="http://127.0.0.1") + parser.add_argument("--port", type=int, default=30000) + args = parser.parse_args() + asyncio.run(test_concurrent(args)) + test_streaming(args) diff --git a/sglang/examples/runtime/lora.py b/sglang/examples/runtime/lora.py new file mode 100644 index 0000000..bf3fc2d --- /dev/null +++ b/sglang/examples/runtime/lora.py @@ -0,0 +1,37 @@ +# launch server +# python -m sglang.launch_server --model mistralai/Mistral-7B-Instruct-v0.3 --lora-paths /home/ying/test_lora lora1=/home/ying/test_lora_1 lora2=/home/ying/test_lora_2 --disable-radix --disable-cuda-graph --max-loras-per-batch 4 + +# send requests +# lora_path[i] specifies the LoRA used for text[i], so make sure they have the same length +# use None to specify base-only prompt, e.x. "lora_path": [None, "/home/ying/test_lora"] +import json + +import requests + +url = "http://127.0.0.1:30000" +json_data = { + "text": [ + "prompt 1", + "prompt 2", + "prompt 3", + "prompt 4", + "prompt 5", + "prompt 6", + "prompt 7", + ], + "sampling_params": {"max_new_tokens": 32}, + "lora_path": [ + "/home/ying/test_lora", + "lora1", + "lora2", + "lora1", + "lora2", + None, + None, + ], +} +response = requests.post( + url + "/generate", + json=json_data, +) +print(json.dumps(response.json())) diff --git a/sglang/examples/runtime/multimodal_embedding.py b/sglang/examples/runtime/multimodal_embedding.py new file mode 100644 index 0000000..813c3c7 --- /dev/null +++ b/sglang/examples/runtime/multimodal_embedding.py @@ -0,0 +1,21 @@ +# launch server +# python -m sglang.launch_server --model-path Alibaba-NLP/gme-Qwen2-VL-2B-Instruct --is-embedding --chat-template gme-qwen2-vl + +import requests + +url = "http://127.0.0.1:30000" + +text_input = "Represent this image in embedding space." +image_path = "https://huggingface.co/datasets/liuhaotian/llava-bench-in-the-wild/resolve/main/images/023.jpg" + +payload = { + "model": "gme-qwen2-vl", + "input": [ + {"type": "text", "text": text_input}, + {"type": "image", "url": image_path}, + ], +} + +response = requests.post(url + "/v1/embeddings", json=payload).json() + +print("Embeddings:", [x.get("embedding") for x in response.get("data", [])]) diff --git a/sglang/examples/runtime/openai_batch_chat.py b/sglang/examples/runtime/openai_batch_chat.py new file mode 100644 index 0000000..d251ca0 --- /dev/null +++ b/sglang/examples/runtime/openai_batch_chat.py @@ -0,0 +1,94 @@ +""" +Usage: + +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +python openai_batch_chat.py + +Note: Before running this script, +you should create the input.jsonl file with the following content: +{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo-0125", "messages": [{"role": "system", "content": "You are a helpful assistant."},{"role": "user", "content": "Hello world! List 3 NBA players and tell a story"}],"max_tokens": 300}} +{"custom_id": "request-2", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo-0125", "messages": [{"role": "system", "content": "You are an assistant. "},{"role": "user", "content": "Hello world! List three capital and tell a story"}],"max_tokens": 500}} +""" + +import json +import time + +import openai + + +class OpenAIBatchProcessor: + def __init__(self): + client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="EMPTY") + + self.client = client + + def process_batch(self, input_file_path, endpoint, completion_window): + + # Upload the input file + with open(input_file_path, "rb") as file: + uploaded_file = self.client.files.create(file=file, purpose="batch") + + # Create the batch job + batch_job = self.client.batches.create( + input_file_id=uploaded_file.id, + endpoint=endpoint, + completion_window=completion_window, + ) + + # Monitor the batch job status + while batch_job.status not in ["completed", "failed", "cancelled"]: + time.sleep(3) # Wait for 3 seconds before checking the status again + print( + f"Batch job status: {batch_job.status}...trying again in 3 seconds..." + ) + batch_job = self.client.batches.retrieve(batch_job.id) + + # Check the batch job status and errors + if batch_job.status == "failed": + print(f"Batch job failed with status: {batch_job.status}") + print(f"Batch job errors: {batch_job.errors}") + return None + + # If the batch job is completed, process the results + if batch_job.status == "completed": + + # print result of batch job + print("batch", batch_job.request_counts) + + result_file_id = batch_job.output_file_id + # Retrieve the file content from the server + file_response = self.client.files.content(result_file_id) + result_content = file_response.read() # Read the content of the file + + # Save the content to a local file + result_file_name = "batch_job_chat_results.jsonl" + with open(result_file_name, "wb") as file: + file.write(result_content) # Write the binary content to the file + # Load data from the saved JSONL file + results = [] + with open(result_file_name, "r", encoding="utf-8") as file: + for line in file: + json_object = json.loads( + line.strip() + ) # Parse each line as a JSON object + results.append(json_object) + + return results + else: + print(f"Batch job failed with status: {batch_job.status}") + return None + + +# Initialize the OpenAIBatchProcessor +processor = OpenAIBatchProcessor() + +# Process the batch job +input_file_path = "input.jsonl" +endpoint = "/v1/chat/completions" +completion_window = "24h" + +# Process the batch job +results = processor.process_batch(input_file_path, endpoint, completion_window) + +# Print the results +print(results) diff --git a/sglang/examples/runtime/openai_batch_complete.py b/sglang/examples/runtime/openai_batch_complete.py new file mode 100644 index 0000000..2f5be3d --- /dev/null +++ b/sglang/examples/runtime/openai_batch_complete.py @@ -0,0 +1,93 @@ +""" +Usage: +python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 +python openai_batch_complete.py +Note: Before running this script, +you should create the input.jsonl file with the following content: +{"custom_id": "request-1", "method": "POST", "url": "/v1/completions", "body": {"model": "gpt-3.5-turbo-instruct", "prompt": "List 3 names of famous soccer player: ", "max_tokens": 200}} +{"custom_id": "request-2", "method": "POST", "url": "/v1/completions", "body": {"model": "gpt-3.5-turbo-instruct", "prompt": "List 6 names of famous basketball player: ", "max_tokens": 400}} +{"custom_id": "request-3", "method": "POST", "url": "/v1/completions", "body": {"model": "gpt-3.5-turbo-instruct", "prompt": "List 6 names of famous basketball player: ", "max_tokens": 400}} +""" + +import json +import time + +import openai + + +class OpenAIBatchProcessor: + def __init__(self): + client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="EMPTY") + + self.client = client + + def process_batch(self, input_file_path, endpoint, completion_window): + + # Upload the input file + with open(input_file_path, "rb") as file: + uploaded_file = self.client.files.create(file=file, purpose="batch") + + # Create the batch job + batch_job = self.client.batches.create( + input_file_id=uploaded_file.id, + endpoint=endpoint, + completion_window=completion_window, + ) + + # Monitor the batch job status + while batch_job.status not in ["completed", "failed", "cancelled"]: + time.sleep(3) # Wait for 3 seconds before checking the status again + print( + f"Batch job status: {batch_job.status}...trying again in 3 seconds..." + ) + batch_job = self.client.batches.retrieve(batch_job.id) + + # Check the batch job status and errors + if batch_job.status == "failed": + print(f"Batch job failed with status: {batch_job.status}") + print(f"Batch job errors: {batch_job.errors}") + return None + + # If the batch job is completed, process the results + if batch_job.status == "completed": + + # print result of batch job + print("batch", batch_job.request_counts) + + result_file_id = batch_job.output_file_id + # Retrieve the file content from the server + file_response = self.client.files.content(result_file_id) + result_content = file_response.read() # Read the content of the file + + # Save the content to a local file + result_file_name = "batch_job_complete_results.jsonl" + with open(result_file_name, "wb") as file: + file.write(result_content) # Write the binary content to the file + # Load data from the saved JSONL file + results = [] + with open(result_file_name, "r", encoding="utf-8") as file: + for line in file: + json_object = json.loads( + line.strip() + ) # Parse each line as a JSON object + results.append(json_object) + + return results + else: + print(f"Batch job failed with status: {batch_job.status}") + return None + + +# Initialize the OpenAIBatchProcessor +processor = OpenAIBatchProcessor() + +# Process the batch job +input_file_path = "input.jsonl" +endpoint = "/v1/completions" +completion_window = "24h" + +# Process the batch job +results = processor.process_batch(input_file_path, endpoint, completion_window) + +# Print the results +print(results) diff --git a/sglang/examples/runtime/openai_chat_with_response_prefill.py b/sglang/examples/runtime/openai_chat_with_response_prefill.py new file mode 100644 index 0000000..c80f97d --- /dev/null +++ b/sglang/examples/runtime/openai_chat_with_response_prefill.py @@ -0,0 +1,33 @@ +""" +Usage: +python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --port 30000 +python openai_chat.py +""" + +import openai + +client = openai.Client(base_url="http://127.0.0.1:30000/v1", api_key="EMPTY") + +response = client.chat.completions.create( + model="meta-llama/Llama-3.1-8B-Instruct", + messages=[ + {"role": "system", "content": "You are a helpful AI assistant"}, + { + "role": "user", + "content": """ +Extract the name, size, price, and color from this product description as a JSON object: + + +The SmartHome Mini is a compact smart home assistant available in black or white for only $49.99. At just 5 inches wide, it lets you control lights, thermostats, and other connected devices via voice or app—no matter where you place it in your home. This affordable little hub brings convenient hands-free control to your smart devices. + +""", + }, + { + "role": "assistant", + "content": "{\n", + }, + ], + temperature=0, +) + +print(response.choices[0].message.content) diff --git a/sglang/examples/runtime/reward_model.py b/sglang/examples/runtime/reward_model.py new file mode 100644 index 0000000..1a1177e --- /dev/null +++ b/sglang/examples/runtime/reward_model.py @@ -0,0 +1,32 @@ +# launch server +# python -m sglang.launch_server --model LxzGordon/URM-LLaMa-3.1-8B --is-embedding + +import requests + +url = "http://127.0.0.1:30000" + +PROMPT = ( + "What is the range of the numeric output of a sigmoid node in a neural network?" +) +RESPONSE1 = "The output of a sigmoid node is bounded between -1 and 1." +RESPONSE2 = "The output of a sigmoid node is bounded between 0 and 1." + +json_data = { + "conv": [ + [ + {"role": "user", "content": PROMPT}, + {"role": "assistant", "content": RESPONSE1}, + ], + [ + {"role": "user", "content": PROMPT}, + {"role": "assistant", "content": RESPONSE2}, + ], + ], +} +response = requests.post( + url + "/classify", + json=json_data, +).json() + +print(response) +print("scores:", [x["embedding"] for x in response]) diff --git a/sglang/examples/runtime/token_in_token_out/token_in_token_out_llm_engine.py b/sglang/examples/runtime/token_in_token_out/token_in_token_out_llm_engine.py new file mode 100644 index 0000000..cb1b7dd --- /dev/null +++ b/sglang/examples/runtime/token_in_token_out/token_in_token_out_llm_engine.py @@ -0,0 +1,43 @@ +""" +This example demonstrates how to provide tokenized ids to LLM as input instead of text prompt, i.e. a token-in-token-out workflow. +""" + +import sglang as sgl +from sglang.srt.hf_transformers_utils import get_tokenizer + +MODEL_PATH = "meta-llama/Llama-3.1-8B-Instruct" + + +def main(): + # Sample prompts. + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + # Create a sampling params object. + sampling_params = {"temperature": 0.8, "top_p": 0.95} + + # Tokenize inputs + tokenizer = get_tokenizer(MODEL_PATH) + token_ids_list = [tokenizer.encode(prompt) for prompt in prompts] + + # Create an LLM. + llm = sgl.Engine(model_path=MODEL_PATH, skip_tokenizer_init=True) + + outputs = llm.generate(input_ids=token_ids_list, sampling_params=sampling_params) + # Print the outputs. + for prompt, output in zip(prompts, outputs): + decode_output = tokenizer.decode(output["output_ids"]) + print("===============================") + print( + f"Prompt: {prompt}\nGenerated token ids: {output['output_ids']}\nGenerated text: {decode_output}" + ) + print() + + +# The __main__ condition is necessary here because we use "spawn" to create subprocesses +# Spawn starts a fresh program every time, if there is no __main__, it will run into infinite loop to keep spawning processes from sgl.Engine +if __name__ == "__main__": + main() diff --git a/sglang/examples/runtime/token_in_token_out/token_in_token_out_llm_server.py b/sglang/examples/runtime/token_in_token_out/token_in_token_out_llm_server.py new file mode 100644 index 0000000..00c0988 --- /dev/null +++ b/sglang/examples/runtime/token_in_token_out/token_in_token_out_llm_server.py @@ -0,0 +1,68 @@ +""" +Usage: + +python token_in_token_out_llm_server.py + +""" + +import requests + +from sglang.srt.hf_transformers_utils import get_tokenizer +from sglang.test.test_utils import is_in_ci +from sglang.utils import terminate_process, wait_for_server + +if is_in_ci(): + from docs.backend.patch import launch_server_cmd +else: + from sglang.utils import launch_server_cmd + + +MODEL_PATH = "meta-llama/Llama-3.1-8B-Instruct" + + +def main(): + # Launch the server + server_process, port = launch_server_cmd( + f"python -m sglang.launch_server --model-path {MODEL_PATH} --skip-tokenizer-init --host 0.0.0.0" + ) + wait_for_server(f"http://localhost:{port}") + + # Sample prompts. + prompts = [ + "Hello, my name is", + "The president of the United States is", + "The capital of France is", + "The future of AI is", + ] + + # Create a sampling params object. + sampling_params = {"temperature": 0.8, "top_p": 0.95} + + # Tokenize inputs + tokenizer = get_tokenizer(MODEL_PATH) + token_ids_list = [tokenizer.encode(prompt) for prompt in prompts] + + json_data = { + "input_ids": token_ids_list, + "sampling_params": sampling_params, + } + + response = requests.post( + f"http://localhost:{port}/generate", + json=json_data, + ) + + outputs = response.json() + for prompt, output in zip(prompts, outputs): + print("===============================") + decode_output = tokenizer.decode(output["output_ids"]) + print( + f"Prompt: {prompt}\nGenerated token ids: {output['output_ids']}\nGenerated text: {decode_output}" + ) + print() + + terminate_process(server_process) + + +if __name__ == "__main__": + main() diff --git a/sglang/examples/runtime/token_in_token_out/token_in_token_out_vlm_engine.py b/sglang/examples/runtime/token_in_token_out/token_in_token_out_vlm_engine.py new file mode 100644 index 0000000..3c3717d --- /dev/null +++ b/sglang/examples/runtime/token_in_token_out/token_in_token_out_vlm_engine.py @@ -0,0 +1,75 @@ +import argparse +import dataclasses +from io import BytesIO +from typing import Tuple + +import requests +from PIL import Image +from transformers import AutoProcessor + +from sglang import Engine +from sglang.lang.chat_template import get_chat_template_by_model_path +from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.server_args import ServerArgs +from sglang.test.test_utils import DEFAULT_IMAGE_URL + + +def get_input_ids( + server_args: ServerArgs, model_config: ModelConfig +) -> Tuple[list[int], list]: + chat_template = get_chat_template_by_model_path(model_config.model_path) + text = f"{chat_template.image_token}What is in this picture?" + images = [Image.open(BytesIO(requests.get(DEFAULT_IMAGE_URL).content))] + image_data = [DEFAULT_IMAGE_URL] + + processor = AutoProcessor.from_pretrained( + model_config.model_path, trust_remote_code=server_args.trust_remote_code + ) + + inputs = processor( + text=[text], + images=images, + return_tensors="pt", + ) + + return inputs.input_ids[0].tolist(), image_data + + +def token_in_out_example( + server_args: ServerArgs, +): + input_ids, image_data = get_input_ids( + server_args, + ModelConfig( + server_args.model_path, + trust_remote_code=server_args.trust_remote_code, + model_override_args=server_args.json_model_override_args, + ), + ) + backend = Engine(**dataclasses.asdict(server_args)) + + output = backend.generate( + input_ids=input_ids, + image_data=image_data, + sampling_params={ + "temperature": 0.8, + "max_new_tokens": 32, + }, + ) + + print("===============================") + print(f"Output token ids: ", output["output_ids"]) + + backend.shutdown() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + args = [ + "--model-path=Qwen/Qwen2-VL-2B", + ] + args = parser.parse_args(args=args) + server_args = ServerArgs.from_cli_args(args) + server_args.skip_tokenizer_init = True + token_in_out_example(server_args) diff --git a/sglang/examples/runtime/token_in_token_out/token_in_token_out_vlm_server.py b/sglang/examples/runtime/token_in_token_out/token_in_token_out_vlm_server.py new file mode 100644 index 0000000..01c4858 --- /dev/null +++ b/sglang/examples/runtime/token_in_token_out/token_in_token_out_vlm_server.py @@ -0,0 +1,78 @@ +""" +Usage: + +python token_in_token_out_vlm_server.py + +""" + +from io import BytesIO +from typing import Tuple + +import requests +from PIL import Image +from transformers import AutoProcessor + +from sglang.lang.chat_template import get_chat_template_by_model_path +from sglang.test.test_utils import DEFAULT_IMAGE_URL, is_in_ci +from sglang.utils import terminate_process, wait_for_server + +if is_in_ci(): + from docs.backend.patch import launch_server_cmd +else: + from sglang.utils import launch_server_cmd + + +MODEL_PATH = "Qwen/Qwen2-VL-2B" + + +def get_input_ids() -> Tuple[list[int], list]: + chat_template = get_chat_template_by_model_path(MODEL_PATH) + text = f"{chat_template.image_token}What is in this picture?" + images = [Image.open(BytesIO(requests.get(DEFAULT_IMAGE_URL).content))] + image_data = [DEFAULT_IMAGE_URL] + + processor = AutoProcessor.from_pretrained(MODEL_PATH) + + inputs = processor( + text=[text], + images=images, + return_tensors="pt", + ) + + return inputs.input_ids[0].tolist(), image_data + + +def main(): + # Launch the server + server_process, port = launch_server_cmd( + f"python -m sglang.launch_server --model-path {MODEL_PATH} --skip-tokenizer-init --host 0.0.0.0" + ) + wait_for_server(f"http://localhost:{port}") + + input_ids, image_data = get_input_ids() + + sampling_params = { + "temperature": 0.8, + "max_new_tokens": 32, + } + + json_data = { + "input_ids": input_ids, + "image_data": image_data, + "sampling_params": sampling_params, + } + + response = requests.post( + f"http://localhost:{port}/generate", + json=json_data, + ) + + output = response.json() + print("===============================") + print(f"Output token ids: ", output["output_ids"]) + + terminate_process(server_process) + + +if __name__ == "__main__": + main() diff --git a/sglang/examples/runtime/vertex_predict.py b/sglang/examples/runtime/vertex_predict.py new file mode 100644 index 0000000..58a41b1 --- /dev/null +++ b/sglang/examples/runtime/vertex_predict.py @@ -0,0 +1,66 @@ +""" +Usage: +python -m sglang.launch_server --model meta-llama/Llama-2-7b-hf --port 30000 +python vertex_predict.py + +This example shows the request and response formats of the prediction route for +Google Cloud Vertex AI Online Predictions. + +Vertex AI SDK for Python is recommended for deploying models to Vertex AI +instead of a local server. After deploying the model to a Vertex AI Online +Prediction Endpoint, send requests via the Python SDK: + +response = endpoint.predict( + instances=[ + {"text": "The capital of France is"}, + {"text": "What is a car?"}, + ], + parameters={"sampling_params": {"max_new_tokens": 16}}, +) +print(response.predictions) + +More details about get online predictions from Vertex AI can be found at +https://cloud.google.com/vertex-ai/docs/predictions/get-online-predictions. +""" + +from dataclasses import dataclass +from typing import List, Optional + +import requests + + +@dataclass +class VertexPrediction: + predictions: List + + +class LocalVertexEndpoint: + def __init__(self) -> None: + self.base_url = "http://127.0.0.1:30000" + + def predict(self, instances: List[dict], parameters: Optional[dict] = None): + response = requests.post( + self.base_url + "/vertex_generate", + json={ + "instances": instances, + "parameters": parameters, + }, + ) + return VertexPrediction(predictions=response.json()["predictions"]) + + +endpoint = LocalVertexEndpoint() + +# Predict with a single prompt. +response = endpoint.predict(instances=[{"text": "The capital of France is"}]) +print(response.predictions) + +# Predict with multiple prompts and parameters. +response = endpoint.predict( + instances=[ + {"text": "The capital of France is"}, + {"text": "What is a car?"}, + ], + parameters={"sampling_params": {"max_new_tokens": 16}}, +) +print(response.predictions) diff --git a/sglang/python/pyproject.toml b/sglang/python/pyproject.toml new file mode 100644 index 0000000..a1a0dcd --- /dev/null +++ b/sglang/python/pyproject.toml @@ -0,0 +1,134 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "sglang" +version = "0.4.4.post3" +description = "SGLang is yet another fast serving framework for large language models and vision language models." +readme = "README.md" +requires-python = ">=3.8" +license = { file = "LICENSE" } +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: Apache Software License", +] +dependencies = ["aiohttp", "requests", "tqdm", "numpy", "IPython", "setproctitle"] + +[project.optional-dependencies] +runtime_common = [ + "compressed-tensors", + "datasets", + "decord", + "fastapi", + "hf_transfer", + "huggingface_hub", + "interegular", + "llguidance>=0.6.15", + "modelscope", + "ninja", + "orjson", + "packaging", + "pillow", + "prometheus-client>=0.20.0", + "psutil", + "pydantic", + "python-multipart", + "pyzmq>=25.1.2", + "soundfile==0.13.1", + "torchao>=0.7.0", + "transformers==4.50.0", + "uvicorn", + "uvloop", + "compressed-tensors", + "xgrammar==0.1.17", +] + +srt = [ + "sglang[runtime_common]", + "sgl-kernel==0.0.5.post4", + "flashinfer_python==0.2.3", + "torch==2.5.1", + "cuda-python", + "outlines>=0.0.44,<=0.1.11", + "partial_json_parser", + "einops", +] + +# HIP (Heterogeneous-computing Interface for Portability) for AMD +# => base docker rocm/vllm-dev:20250114, not from public vllm whl +srt_hip = [ + "sglang[runtime_common]", + "torch", + "vllm==0.6.7.dev2", + "outlines==0.1.11" +] + +# xpu is not enabled in public vllm and torch whl, +# need to follow https://docs.vllm.ai/en/latest/getting_started/xpu-installation.htmlinstall vllm +srt_xpu = ["sglang[runtime_common]", "outlines>=0.0.44,<=0.1.11"] + +# For Intel Gaudi(device : hpu) follow the installation guide +# https://docs.vllm.ai/en/latest/getting_started/gaudi-installation.html +srt_hpu = ["sglang[runtime_common]", "outlines>=0.0.44,<=0.1.11"] + +# CPU: currently, there are no pre-built vllm wheels for CPU. +# To install vllm for CPU, please follow the instruction here: +# https://docs.vllm.ai/en/latest/getting_started/installation/cpu/index.html +srt_cpu = ["sglang[runtime_common]", "outlines>=0.0.44,<=0.1.11", "torch"] + +openai = ["openai>=1.0", "tiktoken"] +anthropic = ["anthropic>=0.20.0"] +litellm = ["litellm>=1.0.0"] +torch_memory_saver = ["torch_memory_saver>=0.0.4"] +test = [ + "jsonlines", + "matplotlib", + "pandas", + "sentence_transformers", + "accelerate", + "peft", +] +all = ["sglang[srt]", "sglang[openai]", "sglang[anthropic]", "sglang[litellm]"] +all_hip = ["sglang[srt_hip]", "sglang[openai]", "sglang[anthropic]", "sglang[litellm]"] +all_xpu = ["sglang[srt_xpu]", "sglang[openai]", "sglang[anthropic]", "sglang[litellm]"] +all_hpu = ["sglang[srt_hpu]", "sglang[openai]", "sglang[anthropic]", "sglang[litellm]"] +all_cpu = ["sglang[srt_cpu]", "sglang[openai]", "sglang[anthropic]", "sglang[litellm]"] + +dev = ["sglang[all]", "sglang[test]"] +dev_hip = ["sglang[all_hip]", "sglang[test]"] +dev_xpu = ["sglang[all_xpu]", "sglang[test]"] +dev_hpu = ["sglang[all_hpu]", "sglang[test]"] +dev_cpu = ["sglang[all_cpu]", "sglang[test]"] + +[project.urls] +"Homepage" = "https://github.com/sgl-project/sglang" +"Bug Tracker" = "https://github.com/sgl-project/sglang/issues" + +[tool.setuptools.package-data] +"sglang" = [ + "srt/layers/moe/fused_moe_triton/configs/*.json", + "srt/layers/quantization/configs/*.json", +] + +[tool.setuptools.packages.find] +exclude = [ + "assets*", + "benchmark*", + "docs*", + "dist*", + "playground*", + "scripts*", + "tests*", +] + +[tool.wheel] +exclude = [ + "assets*", + "benchmark*", + "docs*", + "dist*", + "playground*", + "scripts*", + "tests*", +] diff --git a/sglang/python/sglang/README.md b/sglang/python/sglang/README.md new file mode 100644 index 0000000..0234973 --- /dev/null +++ b/sglang/python/sglang/README.md @@ -0,0 +1,16 @@ +# Code Structures + +- `lang`: The frontend language. +- `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 throughput 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 the local server. +- `llama3_eval.py`: Evaluation of Llama 3 using the Meta Llama dataset. +- `utils.py`: Common utilities. +- `version.py`: Version info. diff --git a/sglang/python/sglang/__init__.py b/sglang/python/sglang/__init__.py new file mode 100644 index 0000000..db0cf26 --- /dev/null +++ b/sglang/python/sglang/__init__.py @@ -0,0 +1,78 @@ +# SGLang public APIs + +# Frontend Language APIs +from sglang.api import ( + Engine, + Runtime, + assistant, + assistant_begin, + assistant_end, + flush_cache, + function, + gen, + gen_int, + gen_string, + get_server_info, + image, + select, + set_default_backend, + system, + system_begin, + system_end, + user, + user_begin, + user_end, + video, +) +from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint +from sglang.lang.choices import ( + greedy_token_selection, + token_length_normalized, + unconditional_likelihood_normalized, +) +from sglang.utils import LazyImport + +ServerArgs = LazyImport("sglang.srt.server_args", "ServerArgs") +Anthropic = LazyImport("sglang.lang.backend.anthropic", "Anthropic") +LiteLLM = LazyImport("sglang.lang.backend.litellm", "LiteLLM") +OpenAI = LazyImport("sglang.lang.backend.openai", "OpenAI") +VertexAI = LazyImport("sglang.lang.backend.vertexai", "VertexAI") + +# Other configs +from sglang.global_config import global_config +from sglang.version import __version__ + +__all__ = [ + "Engine", + "Runtime", + "assistant", + "assistant_begin", + "assistant_end", + "flush_cache", + "function", + "gen", + "gen_int", + "gen_string", + "get_server_info", + "image", + "select", + "set_default_backend", + "system", + "system_begin", + "system_end", + "user", + "user_begin", + "user_end", + "video", + "RuntimeEndpoint", + "greedy_token_selection", + "token_length_normalized", + "unconditional_likelihood_normalized", + "ServerArgs", + "Anthropic", + "LiteLLM", + "OpenAI", + "VertexAI", + "global_config", + "__version__", +] diff --git a/sglang/python/sglang/api.py b/sglang/python/sglang/api.py new file mode 100644 index 0000000..50319c4 --- /dev/null +++ b/sglang/python/sglang/api.py @@ -0,0 +1,279 @@ +"""Public APIs of the language.""" + +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.ir import ( + SglExpr, + SglExprList, + SglFunction, + SglGen, + SglImage, + SglRoleBegin, + SglRoleEnd, + SglSelect, + SglVideo, +) + + +def function( + func: Optional[Callable] = None, num_api_spec_tokens: Optional[int] = None +): + if func: + return SglFunction(func, num_api_spec_tokens=num_api_spec_tokens) + + def decorator(func): + return SglFunction(func, num_api_spec_tokens=num_api_spec_tokens) + + return decorator + + +def Runtime(*args, **kwargs): + # Avoid importing unnecessary dependency + from sglang.lang.backend.runtime_endpoint import Runtime + + return Runtime(*args, **kwargs) + + +def Engine(*args, **kwargs): + # Avoid importing unnecessary dependency + from sglang.srt.entrypoints.engine import Engine + + return Engine(*args, **kwargs) + + +def set_default_backend(backend: BaseBackend): + global_config.default_backend = backend + + +def flush_cache(backend: Optional[BaseBackend] = None): + backend = backend or global_config.default_backend + if backend is None: + return False + + # If backend is Runtime + if hasattr(backend, "endpoint"): + backend = backend.endpoint + return backend.flush_cache() + + +def get_server_info(backend: Optional[BaseBackend] = None): + backend = backend or global_config.default_backend + if backend is None: + return None + + # If backend is Runtime + if hasattr(backend, "endpoint"): + backend = backend.endpoint + return backend.get_server_info() + + +def gen( + name: Optional[str] = None, + max_tokens: Optional[int] = None, + min_tokens: Optional[int] = None, + n: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + min_p: Optional[float] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + ignore_eos: Optional[bool] = None, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, + dtype: Optional[Union[type, str]] = None, + choices: Optional[List[str]] = None, + choices_method: Optional[ChoicesSamplingMethod] = None, + regex: Optional[str] = None, + json_schema: Optional[str] = None, +): + """Call the model to generate. See the meaning of the arguments in docs/backend/sampling_params.md""" + + if choices: + return SglSelect( + name, + choices, + 0.0 if temperature is None else temperature, + token_length_normalized if choices_method is None else choices_method, + ) + + # check regex is valid + if regex is not None: + try: + re.compile(regex) + except re.error as e: + raise e + + return SglGen( + name, + max_tokens, + min_tokens, + n, + stop, + stop_token_ids, + temperature, + top_p, + top_k, + min_p, + frequency_penalty, + presence_penalty, + ignore_eos, + return_logprob, + logprob_start_len, + top_logprobs_num, + return_text_in_logprobs, + dtype, + regex, + json_schema, + ) + + +def gen_int( + name: Optional[str] = None, + max_tokens: Optional[int] = None, + n: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + min_p: Optional[float] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + ignore_eos: Optional[bool] = None, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, +): + return SglGen( + name, + max_tokens, + None, + n, + stop, + stop_token_ids, + temperature, + top_p, + top_k, + min_p, + frequency_penalty, + presence_penalty, + ignore_eos, + return_logprob, + logprob_start_len, + top_logprobs_num, + return_text_in_logprobs, + int, + None, + ) + + +def gen_string( + name: Optional[str] = None, + max_tokens: Optional[int] = None, + n: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + min_p: Optional[float] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + ignore_eos: Optional[bool] = None, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, +): + return SglGen( + name, + max_tokens, + None, + n, + stop, + stop_token_ids, + temperature, + top_p, + top_k, + min_p, + frequency_penalty, + presence_penalty, + ignore_eos, + return_logprob, + logprob_start_len, + top_logprobs_num, + return_text_in_logprobs, + str, + None, + ) + + +def image(expr: SglExpr): + return SglImage(expr) + + +def video(path: str, num_frames: int): + return SglVideo(path, num_frames) + + +def select( + name: Optional[str] = None, + choices: Optional[List[str]] = None, + temperature: float = 0.0, + choices_method: ChoicesSamplingMethod = token_length_normalized, +): + assert choices is not None + return SglSelect(name, choices, temperature, choices_method) + + +def _role_common(name: str, expr: Optional[SglExpr] = None): + if expr is None: + return SglExprList([SglRoleBegin(name), SglRoleEnd(name)]) + else: + return SglExprList([SglRoleBegin(name), expr, SglRoleEnd(name)]) + + +def system(expr: Optional[SglExpr] = None): + return _role_common("system", expr) + + +def user(expr: Optional[SglExpr] = None): + return _role_common("user", expr) + + +def assistant(expr: Optional[SglExpr] = None): + return _role_common("assistant", expr) + + +def system_begin(): + return SglRoleBegin("system") + + +def system_end(): + return SglRoleEnd("system") + + +def user_begin(): + return SglRoleBegin("user") + + +def user_end(): + return SglRoleEnd("user") + + +def assistant_begin(): + return SglRoleBegin("assistant") + + +def assistant_end(): + return SglRoleEnd("assistant") diff --git a/sglang/python/sglang/bench_offline_throughput.py b/sglang/python/sglang/bench_offline_throughput.py new file mode 100644 index 0000000..69bbc3e --- /dev/null +++ b/sglang/python/sglang/bench_offline_throughput.py @@ -0,0 +1,424 @@ +""" +Benchmark the throughput in the offline mode. +It accepts server arguments (the same as launch_server.py) and benchmark arguments (the same as bench_serving.py). + +# Usage +## Sharegpt dataset with default args +python -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --num-prompts 10 + +## Random dataset with default args +python -m sglang.bench_offline_throughput --model-path meta-llama/Meta-Llama-3.1-8B-Instruct --dataset-name random --random-input 1024 --random-output 1024 +""" + +import argparse +import dataclasses +import json +import logging +import os +import random +import time +from typing import Dict, List, Optional, Tuple + +import numpy as np + +from sglang.bench_serving import ( + get_dataset, + get_tokenizer, + sample_random_requests, + set_ulimit, +) +from sglang.lang.backend.runtime_endpoint import Runtime +from sglang.srt.entrypoints.engine import Engine +from sglang.srt.server_args import ServerArgs + + +@dataclasses.dataclass +class BenchArgs: + backend: str = "engine" + result_filename: str = "" + dataset_name: str = "sharegpt" + dataset_path: str = "" + num_prompts: int = 1000 + sharegpt_output_len: Optional[int] = None + sharegpt_context_len: Optional[int] = None + random_input_len: int = 1024 + random_output_len: int = 1024 + random_range_ratio: float = 0.0 + gsp_num_groups: int = 64 + gsp_prompts_per_group: int = 16 + gsp_system_prompt_len: int = 2048 + gsp_question_len: int = 128 + gsp_output_len: int = 256 + seed: int = 1 + disable_ignore_eos: bool = False + extra_request_body: Optional[str] = None + apply_chat_template: bool = False + profile: bool = False + skip_warmup: bool = False + do_not_exit: bool = False + prompt_suffix: str = "" + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument("--backend", type=str, default=BenchArgs.backend) + parser.add_argument( + "--result-filename", type=str, default=BenchArgs.result_filename + ) + parser.add_argument( + "--dataset-name", + type=str, + default="sharegpt", + choices=["sharegpt", "random", "generated-shared-prefix"], + help="Name of the dataset to benchmark on.", + ) + parser.add_argument( + "--dataset-path", type=str, default="", help="Path to the dataset." + ) + parser.add_argument( + "--num-prompts", + type=int, + default=BenchArgs.num_prompts, + help="Number of prompts to process. Default is 1000.", + ) + parser.add_argument( + "--sharegpt-output-len", + type=int, + default=BenchArgs.sharegpt_output_len, + help="Output length for each request. Overrides the output length from the ShareGPT dataset.", + ) + parser.add_argument( + "--sharegpt-context-len", + type=int, + default=BenchArgs.sharegpt_context_len, + help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.", + ) + parser.add_argument( + "--random-input-len", + type=int, + default=BenchArgs.random_input_len, + help="Number of input tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-output-len", + type=int, + default=BenchArgs.random_output_len, + help="Number of output tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-range-ratio", + type=float, + default=BenchArgs.random_range_ratio, + help="Range of sampled ratio of input/output length, " + "used only for random dataset.", + ) + parser.add_argument( + "--gsp-num-groups", + type=int, + default=BenchArgs.gsp_num_groups, + help="Number of groups with shared prefix, used" + "only for generate-shared-prefix", + ) + parser.add_argument( + "--gsp-prompts-per-group", + type=int, + default=BenchArgs.gsp_prompts_per_group, + help="Number of prompts per group of shared prefix, used" + "only for generate-shared-prefix", + ) + parser.add_argument( + "--gsp-system-prompt-len", + type=int, + default=BenchArgs.gsp_system_prompt_len, + help="System prompt length, used" "only for generate-shared-prefix", + ) + parser.add_argument( + "--gsp-question-len", + type=int, + default=BenchArgs.gsp_question_len, + help="Question length, used" "only for generate-shared-prefix", + ) + parser.add_argument( + "--gsp-output-len", + type=int, + default=BenchArgs.gsp_output_len, + help="Target length in tokens for outputs in generated-shared-prefix dataset", + ) + parser.add_argument("--seed", type=int, default=1, help="The random seed.") + parser.add_argument( + "--disable-ignore-eos", + action="store_true", + help="Disable ignore EOS token", + ) + parser.add_argument( + "--extra-request-body", + metavar='{"key1": "value1", "key2": "value2"}', + type=str, + default=BenchArgs.extra_request_body, + help="Append given JSON object to the request payload. You can use this to specify" + "additional generate params like sampling params.", + ) + parser.add_argument( + "--apply-chat-template", + action="store_true", + help="Apply chat template", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Use Torch Profiler. The endpoint must be launched with " + "SGLANG_TORCH_PROFILER_DIR to enable profiler.", + ) + parser.add_argument( + "--skip-warmup", + action="store_true", + help="Skip the warmup batches.", + ) + parser.add_argument( + "--do-not-exit", + action="store_true", + help="Do not exit the program. This is useful for nsys profile with --duration and --delay.", + ) + parser.add_argument( + "--prompt-suffix", + type=str, + default="", + help="Suffix applied to the end of all user prompts, followed by assistant prompt suffix.", + ) + + @classmethod + def from_cli_args(cls, args: argparse.Namespace): + attrs = [attr.name for attr in dataclasses.fields(cls)] + return cls(**{attr: getattr(args, attr) for attr in attrs}) + + +def throughput_test_once( + backend_name: str, + backend, + reqs: List[Tuple[str, int, int]], + ignore_eos: bool, + extra_request_body: Dict, + profile: bool, +): + measurement_results = { + "backend": backend_name, + "successful_requests": len(reqs), + "total_latency": -1, + "total_input_tokens": sum(r[1] for r in reqs), + "total_output_tokens": -1, + "request_throughput": -1, + "input_throughput": -1, + "output_throughput": -1, + "total_throughput": -1, + } + + prompt = [r[0] for r in reqs] + sampling_params = [ + { + "temperature": 0, + "max_new_tokens": r[2], + "ignore_eos": ignore_eos, + **extra_request_body, + } + for r in reqs + ] + + if profile: + assert ( + "SGLANG_TORCH_PROFILER_DIR" in os.environ + ), "Please set SGLANG_TORCH_PROFILER_DIR." + os.makedirs(os.environ["SGLANG_TORCH_PROFILER_DIR"], exist_ok=True) + backend.start_profile() + + st = time.perf_counter() + gen_out = backend.generate(prompt=prompt, sampling_params=sampling_params) + latency = time.perf_counter() - st + + if profile: + backend.stop_profile() + monitor_trace_file(os.getenv("SGLANG_TORCH_PROFILER_DIR")) + + if backend_name == "runtime": + gen_out = json.loads(gen_out) + + server_info = backend.get_server_info() + + measurement_results["total_latency"] = latency + measurement_results["total_output_tokens"] = sum( + o["meta_info"]["completion_tokens"] for o in gen_out + ) + measurement_results["request_throughput"] = ( + measurement_results["successful_requests"] / latency + ) + measurement_results["input_throughput"] = ( + measurement_results["total_input_tokens"] / latency + ) + measurement_results["output_throughput"] = ( + measurement_results["total_output_tokens"] / latency + ) + measurement_results["total_throughput"] = ( + measurement_results["total_input_tokens"] + + measurement_results["total_output_tokens"] + ) / latency + measurement_results["last_gen_throughput"] = server_info["last_gen_throughput"] + + return measurement_results + + +def monitor_trace_file(directory, interval=1): + + print(f"Monitoring {directory} for new trace files...") + + known_files = set(os.listdir(directory)) + + while True: + flag = False + time.sleep(interval) + current_files = set(os.listdir(directory)) + + new_files = current_files - known_files + for new_file in new_files: + new_file_path = os.path.join(directory, new_file) + print(f"New file detected: {new_file}") + + previous_size = 0 + while True: + try: + current_size = os.path.getsize(new_file_path) + except FileNotFoundError: + print(f"File {new_file} is no longer accessible.") + break + + if current_size > previous_size: + previous_size = current_size + else: + flag = True + break + + time.sleep(interval) + if flag: + break + + +def throughput_test( + server_args: ServerArgs, + bench_args: BenchArgs, +): + if bench_args.backend == "engine": + backend = Engine(**dataclasses.asdict(server_args)) + if not backend: + raise ValueError("Please provide valid engine arguments") + elif bench_args.backend == "runtime": + backend = Runtime(**dataclasses.asdict(server_args)) + else: + raise ValueError('Please set backend to either "engine" or "runtime"') + + tokenizer_id = server_args.tokenizer_path or server_args.model_path + tokenizer = get_tokenizer(tokenizer_id) + + # Set global environmnets + set_ulimit() + random.seed(bench_args.seed) + np.random.seed(bench_args.seed) + + # Parse args + extra_request_body = {} + if bench_args.extra_request_body: + extra_request_body = json.loads(args.extra_request_body) + + # Read dataset + input_requests = get_dataset(bench_args, tokenizer) + + warmup_requests = sample_random_requests( + input_len=256, + output_len=16, + num_prompts=min(bench_args.num_prompts, 16), + range_ratio=1.0, + tokenizer=tokenizer, + dataset_path=bench_args.dataset_path, + ) + + # Warm up + if not bench_args.skip_warmup: + logging.info("\nWarmup...") + throughput_test_once( + backend_name=bench_args.backend, + backend=backend, + reqs=warmup_requests, + ignore_eos=not bench_args.disable_ignore_eos, + extra_request_body=extra_request_body, + profile=False, + ) + time.sleep(0.5) + + logging.info("\nBenchmark...") + result = throughput_test_once( + backend_name=bench_args.backend, + backend=backend, + reqs=input_requests, + ignore_eos=not bench_args.disable_ignore_eos, + extra_request_body=extra_request_body, + profile=bench_args.profile, + ) + backend.shutdown() + + if bench_args.result_filename: + with open(bench_args.result_filename, "a") as fout: + fout.write(json.dumps(result) + "\n") + + print( + "\n{s:{c}^{n}}".format(s=" Offline Throughput Benchmark Result ", n=50, c="=") + ) + print("{:<40} {:<10}".format("Backend:", result["backend"])) + print("{:<40} {:<10}".format("Successful requests:", result["successful_requests"])) + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", result["total_latency"])) + print("{:<40} {:<10}".format("Total input tokens:", result["total_input_tokens"])) + print( + "{:<40} {:<10}".format("Total generated tokens:", result["total_output_tokens"]) + ) + print( + "{:<40} {:<10.2f}".format( + "Last generation throughput (tok/s):", result["last_gen_throughput"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Request throughput (req/s):", result["request_throughput"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Input token throughput (tok/s):", result["input_throughput"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Output token throughput (tok/s):", result["output_throughput"] + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Total token throughput (tok/s):", result["total_throughput"] + ) + ) + print("=" * 50) + + return result + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + BenchArgs.add_cli_args(parser) + args = parser.parse_args() + server_args = ServerArgs.from_cli_args(args) + bench_args = BenchArgs.from_cli_args(args) + + logging.basicConfig( + level=getattr(logging, server_args.log_level.upper()), + format="%(message)s", + ) + + throughput_test(server_args, bench_args) + + while bench_args.do_not_exit: + pass diff --git a/sglang/python/sglang/bench_one_batch.py b/sglang/python/sglang/bench_one_batch.py new file mode 100644 index 0000000..e1a69ae --- /dev/null +++ b/sglang/python/sglang/bench_one_batch.py @@ -0,0 +1,531 @@ +""" +Benchmark the latency of running a single static batch without a server. + +This script does not launch a server and uses the low-level APIs. +It accepts server arguments (the same as launch_server.py) and benchmark arguments (e.g., batch size, input lengths). + +# Usage (latency test) +## with dummy weights: +python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --load-format dummy +## sweep through multiple data points and store (append) the results in a jsonl file: +python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 12 14 --input-len 256 512 --output-len 32 256 --run-name test_run +## run with profiling: +python -m sglang.bench_one_batch --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 12 14 --input-len 256 512 --profile +# Usage (correctness test): +python -m sglang.bench_one_batch --model-path TinyLlama/TinyLlama-1.1B-Chat-v0.4 --correct + +## Reference output (of the correctness test above, can be gpu dependent): +input_ids=[[1, 450, 7483, 310, 3444, 338], [1, 450, 7483, 310, 278, 3303, 13187, 290, 338], [1, 20628, 338, 263, 6575, 1460, 2462, 322, 306, 763]] + +prefill logits (first half): tensor([[-10.0312, -9.5000, 0.8931, ..., -4.9414, -3.2422, -3.3633], + [-10.0312, -9.5000, 0.8931, ..., -4.9414, -3.2422, -3.3633], + [ -9.1875, -10.2500, 2.7129, ..., -4.3359, -4.0664, -4.1328]], + device='cuda:0') + +prefill logits (final): tensor([[-8.3125, -7.1172, 3.3457, ..., -4.9570, -4.1328, -3.4141], + [-8.9141, -9.0156, 4.1445, ..., -4.9922, -4.4961, -4.0781], + [-9.6328, -9.0547, 4.0195, ..., -5.3047, -4.7148, -4.4570]], + device='cuda:0') + +========== Prompt 0 ========== + The capital of France is Paris. +The capital of the United States is Washington, D.C. + + +========== Prompt 1 ========== + The capital of the United Kindom is London. +The capital of the United Kingdom is London. +The capital of the + +========== Prompt 2 ========== + Today is a sunny day and I like to go for a walk in the park. +I'm going to the park +""" + +import argparse +import dataclasses +import itertools +import json +import logging +import multiprocessing +import os +import time +from typing import Tuple + +import numpy as np +import torch +import torch.distributed as dist + +from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.entrypoints.engine import _set_envs_and_config +from sglang.srt.hf_transformers_utils import get_tokenizer +from sglang.srt.managers.schedule_batch import Req, ScheduleBatch +from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_executor.model_runner import ModelRunner +from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.srt.server_args import PortArgs, ServerArgs +from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.utils import ( + configure_logger, + get_bool_env_var, + kill_process_tree, + set_gpu_proc_affinity, + suppress_other_loggers, +) + + +@dataclasses.dataclass +class BenchArgs: + run_name: str = "default" + batch_size: Tuple[int] = (1,) + input_len: Tuple[int] = (1024,) + output_len: Tuple[int] = (16,) + result_filename: str = "result.jsonl" + correctness_test: bool = False + # This is only used for correctness test + cut_len: int = 4 + profile: bool = False + profile_filename_prefix: str = "profile" + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument("--run-name", type=str, default=BenchArgs.run_name) + parser.add_argument( + "--batch-size", type=int, nargs="+", default=BenchArgs.batch_size + ) + parser.add_argument( + "--input-len", type=int, nargs="+", default=BenchArgs.input_len + ) + parser.add_argument( + "--output-len", type=int, nargs="+", default=BenchArgs.output_len + ) + parser.add_argument( + "--result-filename", type=str, default=BenchArgs.result_filename + ) + parser.add_argument("--correctness-test", action="store_true") + parser.add_argument("--cut-len", type=int, default=BenchArgs.cut_len) + parser.add_argument( + "--profile", action="store_true", help="Use Torch Profiler." + ) + parser.add_argument( + "--profile-filename-prefix", + type=str, + default=BenchArgs.profile_filename_prefix, + help="Prefix of the profiling file names. The full profiling result file(s) be " + '"[profile_filename_prefix]_batch[batch_size]_input[input_len]_output[output_len].trace.json.gz"', + ) + + @classmethod + def from_cli_args(cls, args: argparse.Namespace): + # use the default value's type to cast the args into correct types. + attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)] + return cls( + **{attr: attr_type(getattr(args, attr)) for attr, attr_type in attrs} + ) + + +def load_model(server_args, port_args, tp_rank): + suppress_other_loggers() + rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None + + model_config = ModelConfig( + server_args.model_path, + trust_remote_code=server_args.trust_remote_code, + revision=server_args.revision, + context_length=server_args.context_length, + model_override_args=server_args.json_model_override_args, + is_embedding=server_args.is_embedding, + dtype=server_args.dtype, + quantization=server_args.quantization, + ) + model_runner = ModelRunner( + model_config=model_config, + mem_fraction_static=server_args.mem_fraction_static, + gpu_id=tp_rank, + tp_rank=tp_rank, + tp_size=server_args.tp_size, + nccl_port=port_args.nccl_port, + server_args=server_args, + ) + rank_print(f"max_total_num_tokens={model_runner.max_total_num_tokens}") + tokenizer = get_tokenizer( + server_args.tokenizer_path, + tokenizer_mode=server_args.tokenizer_mode, + trust_remote_code=server_args.trust_remote_code, + ) + if server_args.tp_size > 1: + dist.barrier() + return model_runner, tokenizer + + +def prepare_inputs_for_correctness_test(bench_args, tokenizer): + prompts = [ + "The capital of France is", + "The capital of the United Kindom is", + "Today is a sunny day and I like", + ] + input_ids = [tokenizer.encode(p) for p in prompts] + sampling_params = SamplingParams( + temperature=0, + max_new_tokens=BenchArgs.output_len, + ) + + reqs = [] + for i in range(len(prompts)): + assert len(input_ids[i]) > bench_args.cut_len + + tmp_input_ids = input_ids[i][: bench_args.cut_len] + req = Req( + rid=i, + origin_input_text=prompts[i], + origin_input_ids=tmp_input_ids, + sampling_params=sampling_params, + ) + req.prefix_indices = [] + req.fill_ids = req.origin_input_ids + req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + reqs.append(req) + + return input_ids, reqs + + +def prepare_extend_inputs_for_correctness_test( + bench_args, input_ids, reqs, model_runner +): + for i in range(len(reqs)): + req = reqs[i] + req.fill_ids += input_ids[i][bench_args.cut_len :] + req.prefix_indices = model_runner.req_to_token_pool.req_to_token[ + i, : bench_args.cut_len + ] + req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + return reqs + + +def prepare_synthetic_inputs_for_latency_test(batch_size, input_len): + input_ids = np.ones((batch_size, input_len), dtype=np.int32) + sampling_params = SamplingParams( + temperature=0, + max_new_tokens=BenchArgs.output_len, + ) + + reqs = [] + for i in range(len(input_ids)): + req = Req( + rid=i, + origin_input_text="", + origin_input_ids=list(input_ids[i]), + sampling_params=sampling_params, + ) + req.prefix_indices = [] + req.fill_ids = req.origin_input_ids + req.extend_input_len = len(req.fill_ids) - len(req.prefix_indices) + reqs.append(req) + + return reqs + + +@torch.no_grad +def extend(reqs, model_runner): + batch = ScheduleBatch.init_new( + reqs=reqs, + req_to_token_pool=model_runner.req_to_token_pool, + token_to_kv_pool_allocator=model_runner.token_to_kv_pool_allocator, + tree_cache=None, + model_config=model_runner.model_config, + enable_overlap=False, + spec_algorithm=SpeculativeAlgorithm.NONE, + enable_custom_logit_processor=False, + ) + batch.prepare_for_extend() + model_worker_batch = batch.get_model_worker_batch() + forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner) + logits_output = model_runner.forward(forward_batch) + next_token_ids = model_runner.sample(logits_output, forward_batch) + return next_token_ids, logits_output.next_token_logits, batch + + +@torch.no_grad +def decode(input_token_ids, batch, model_runner): + batch.output_ids = input_token_ids + batch.prepare_for_decode() + model_worker_batch = batch.get_model_worker_batch() + forward_batch = ForwardBatch.init_new(model_worker_batch, model_runner) + logits_output = model_runner.forward(forward_batch) + next_token_ids = model_runner.sample(logits_output, forward_batch) + return next_token_ids, logits_output.next_token_logits + + +def correctness_test( + server_args, + port_args, + bench_args, + tp_rank, +): + # Configure the logger + configure_logger(server_args, prefix=f" TP{tp_rank}") + rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None + + # Load the model + model_runner, tokenizer = load_model(server_args, port_args, tp_rank) + + # Prepare inputs + input_ids, reqs = prepare_inputs_for_correctness_test(bench_args, tokenizer) + rank_print(f"\n{input_ids=}\n") + + if bench_args.cut_len > 0: + # Prefill + next_token_ids, next_token_logits, batch = extend(reqs, model_runner) + rank_print(f"prefill logits (first half): {next_token_logits} \n") + + # Prepare extend inputs + reqs = prepare_extend_inputs_for_correctness_test( + bench_args, input_ids, reqs, model_runner + ) + + # Extend (prefill w/ KV cache) + next_token_ids, next_token_logits, batch = extend(reqs, model_runner) + rank_print(f"prefill logits (final): {next_token_logits} \n") + + # Decode + output_ids = [input_ids[i] + [next_token_ids[i]] for i in range(len(input_ids))] + for _ in range(bench_args.output_len[0] - 1): + next_token_ids, _ = decode(next_token_ids, batch, model_runner) + next_token_ids_list = next_token_ids.tolist() + for i in range(len(reqs)): + output_ids[i].append(next_token_ids_list[i]) + + # Print output texts + for i in range(len(reqs)): + rank_print(f"========== Prompt {i} ==========") + rank_print(tokenizer.decode(output_ids[i]), "\n") + + +def synchronize(device): + torch.get_device_module(device).synchronize() + + +def latency_test_run_once( + run_name, + model_runner, + rank_print, + reqs, + batch_size, + input_len, + output_len, + device, + profile, + profile_filename_prefix, +): + max_batch_size = model_runner.max_total_num_tokens // (input_len + output_len) + if batch_size > max_batch_size: + rank_print( + f"skipping ({batch_size}, {input_len}, {output_len}) due to max batch size limit" + ) + return + + # Clear the pools. + model_runner.req_to_token_pool.clear() + model_runner.token_to_kv_pool_allocator.clear() + + measurement_results = { + "run_name": run_name, + "batch_size": batch_size, + "input_len": input_len, + "output_len": output_len, + } + + tot_latency = 0 + + profiler = None + if profile: + profiler = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + with_stack=True, + ) + profiler.start() + + # Prefill + synchronize(device) + tic = time.time() + next_token_ids, _, batch = extend(reqs, model_runner) + synchronize(device) + prefill_latency = time.time() - tic + tot_latency += prefill_latency + throughput = input_len * batch_size / prefill_latency + rank_print( + f"Prefill. latency: {prefill_latency:6.5f} s, throughput: {throughput:9.2f} token/s" + ) + measurement_results["prefill_latency"] = prefill_latency + measurement_results["prefill_throughput"] = throughput + + # Decode + decode_latencies = [] + for i in range(output_len - 1): + synchronize(device) + tic = time.time() + next_token_ids, _ = decode(next_token_ids, batch, model_runner) + synchronize(device) + latency = time.time() - tic + tot_latency += latency + throughput = batch_size / latency + decode_latencies.append(latency) + if i < 5: + rank_print( + f"Decode. latency: {latency:6.5f} s, throughput: {throughput:9.2f} token/s" + ) + + if profile: + profiler.stop() + profile_filename = f"{profile_filename_prefix}_batch{batch_size}_input{input_len}_output{output_len}.trace.json.gz" + parent_dir = os.path.dirname(os.path.abspath(profile_filename)) + os.makedirs(parent_dir, exist_ok=True) + profiler.export_chrome_trace(profile_filename) + rank_print(f"torch profiler chrome trace saved to {profile_filename}") + + # Record decode timing from 2nd output + if output_len > 1: + med_decode_latency = np.median(decode_latencies) + med_decode_throughput = batch_size / med_decode_latency + rank_print( + f"Decode. median latency: {med_decode_latency:6.5f} s, median throughput: {med_decode_throughput:9.2f} token/s" + ) + measurement_results["median_decode_latency"] = med_decode_latency + measurement_results["median_decode_throughput"] = med_decode_throughput + + throughput = (input_len + output_len) * batch_size / tot_latency + rank_print( + f"Total. latency: {tot_latency:6.3f} s, throughput: {throughput:9.2f} token/s" + ) + measurement_results["total_latency"] = tot_latency + measurement_results["overall_throughput"] = throughput + return measurement_results + + +def latency_test( + server_args, + port_args, + bench_args, + tp_rank, +): + # Set CPU affinity + if get_bool_env_var("SGLANG_SET_CPU_AFFINITY"): + set_gpu_proc_affinity(server_args.tp_size, server_args.nnodes, tp_rank) + + # Configure the logger + configure_logger(server_args, prefix=f" TP{tp_rank}") + rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None + + # Load the model + model_runner, tokenizer = load_model(server_args, port_args, tp_rank) + + # Prepare inputs for warm up + reqs = prepare_synthetic_inputs_for_latency_test( + bench_args.batch_size[0], bench_args.input_len[0] + ) + + # Warm up + rank_print("Warmup ...") + latency_test_run_once( + bench_args.run_name, + model_runner, + rank_print, + reqs, + bench_args.batch_size[0], + bench_args.input_len[0], + 8, # shorter decoding to speed up the warmup + server_args.device, + profile=False, + profile_filename_prefix="", # not used + ) + + rank_print("Benchmark ...") + + # Run the sweep + result_list = [] + for bs, il, ol in itertools.product( + bench_args.batch_size, bench_args.input_len, bench_args.output_len + ): + reqs = prepare_synthetic_inputs_for_latency_test(bs, il) + ret = latency_test_run_once( + bench_args.run_name, + model_runner, + rank_print, + reqs, + bs, + il, + ol, + server_args.device, + bench_args.profile if tp_rank == 0 else None, + bench_args.profile_filename_prefix, + ) + if ret is not None: + result_list.append(ret) + + # Write results in jsonlines format on rank 0. + if tp_rank == 0 and bench_args.result_filename: + with open(bench_args.result_filename, "a") as fout: + for result in result_list: + fout.write(json.dumps(result) + "\n") + + +def main(server_args, bench_args): + _set_envs_and_config(server_args) + + if server_args.model_path: + if bench_args.correctness_test: + work_func = correctness_test + else: + work_func = latency_test + else: + raise ValueError( + "Provide --model-path for running the tests or " + "provide --result-filename for plotting the results" + ) + + port_args = PortArgs.init_new(server_args) + + if server_args.tp_size == 1: + work_func(server_args, port_args, bench_args, 0) + else: + workers = [] + for tp_rank in range(server_args.tp_size): + proc = multiprocessing.Process( + target=work_func, + args=( + server_args, + port_args, + bench_args, + tp_rank, + ), + ) + proc.start() + workers.append(proc) + + for proc in workers: + proc.join() + + proc.terminate() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + BenchArgs.add_cli_args(parser) + args = parser.parse_args() + server_args = ServerArgs.from_cli_args(args) + bench_args = BenchArgs.from_cli_args(args) + + logging.basicConfig( + level=getattr(logging, server_args.log_level.upper()), + format="%(message)s", + ) + + try: + main(server_args, bench_args) + finally: + if server_args.tp_size != 1: + kill_process_tree(os.getpid(), include_parent=False) diff --git a/sglang/python/sglang/bench_one_batch_server.py b/sglang/python/sglang/bench_one_batch_server.py new file mode 100644 index 0000000..1fc4ff5 --- /dev/null +++ b/sglang/python/sglang/bench_one_batch_server.py @@ -0,0 +1,192 @@ +""" +Benchmark the latency of running a single batch with a server. + +This script launches a server and uses the HTTP interface. +It accepts server arguments (the same as launch_server.py) and benchmark arguments (e.g., batch size, input lengths). + +Usage: +python3 -m sglang.bench_one_batch_server --model meta-llama/Meta-Llama-3.1-8B --batch-size 1 16 64 --input-len 1024 --output-len 8 + +python3 -m sglang.bench_one_batch_server --model None --base-url http://localhost:30000 --batch-size 16 --input-len 1024 --output-len 8 +""" + +import argparse +import dataclasses +import itertools +import json +import multiprocessing +import os +import time +from typing import Tuple + +import numpy as np +import requests + +from sglang.srt.entrypoints.http_server import launch_server +from sglang.srt.server_args import ServerArgs +from sglang.srt.utils import kill_process_tree + + +@dataclasses.dataclass +class BenchArgs: + run_name: str = "default" + batch_size: Tuple[int] = (1,) + input_len: Tuple[int] = (1024,) + output_len: Tuple[int] = (16,) + result_filename: str = "result.jsonl" + base_url: str = "" + skip_warmup: bool = False + + @staticmethod + def add_cli_args(parser: argparse.ArgumentParser): + parser.add_argument("--run-name", type=str, default=BenchArgs.run_name) + parser.add_argument( + "--batch-size", type=int, nargs="+", default=BenchArgs.batch_size + ) + parser.add_argument( + "--input-len", type=int, nargs="+", default=BenchArgs.input_len + ) + parser.add_argument( + "--output-len", type=int, nargs="+", default=BenchArgs.output_len + ) + parser.add_argument( + "--result-filename", type=str, default=BenchArgs.result_filename + ) + parser.add_argument("--base-url", type=str, default=BenchArgs.base_url) + parser.add_argument("--skip-warmup", action="store_true") + + @classmethod + def from_cli_args(cls, args: argparse.Namespace): + # use the default value's type to cast the args into correct types. + attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)] + return cls( + **{attr: attr_type(getattr(args, attr)) for attr, attr_type in attrs} + ) + + +def launch_server_internal(server_args): + try: + launch_server(server_args) + except Exception as e: + raise e + finally: + kill_process_tree(os.getpid(), include_parent=False) + + +def launch_server_process(server_args: ServerArgs): + proc = multiprocessing.Process(target=launch_server_internal, args=(server_args,)) + proc.start() + base_url = f"http://{server_args.host}:{server_args.port}" + timeout = 600 + + start_time = time.time() + while time.time() - start_time < timeout: + try: + headers = { + "Content-Type": "application/json; charset=utf-8", + } + response = requests.get(f"{base_url}/v1/models", headers=headers) + if response.status_code == 200: + return proc, base_url + except requests.RequestException: + pass + time.sleep(10) + raise TimeoutError("Server failed to start within the timeout period.") + + +def run_one_case( + url: str, + batch_size: int, + input_len: int, + output_len: int, + run_name: str, + result_filename: str, +): + input_ids = [ + [int(x) for x in np.random.randint(0, high=16384, size=(input_len,))] + for _ in range(batch_size) + ] + + tic = time.time() + response = requests.post( + url + "/generate", + json={ + "input_ids": input_ids, + "sampling_params": { + "temperature": 0, + "max_new_tokens": output_len, + "ignore_eos": True, + }, + }, + ) + latency = time.time() - tic + + _ = response.json() + output_throughput = batch_size * output_len / latency + overall_throughput = batch_size * (input_len + output_len) / latency + + print(f"batch size: {batch_size}") + print(f"latency: {latency:.2f} s") + print(f"output throughput: {output_throughput:.2f} token/s") + print(f"(input + output) throughput: {overall_throughput:.2f} token/s") + + if result_filename: + with open(result_filename, "a") as fout: + res = { + "run_name": run_name, + "batch_size": batch_size, + "input_len": input_len, + "output_len": output_len, + "latency": round(latency, 4), + "output_throughput": round(output_throughput, 2), + "overall_throughput": round(overall_throughput, 2), + } + fout.write(json.dumps(res) + "\n") + + +def run_benchmark(server_args: ServerArgs, bench_args: BenchArgs): + if bench_args.base_url: + proc, base_url = None, bench_args.base_url + else: + proc, base_url = launch_server_process(server_args) + + # warmup + if not bench_args.skip_warmup: + run_one_case( + base_url, + batch_size=16, + input_len=1024, + output_len=16, + run_name="", + result_filename="", + ) + + # benchmark + try: + for bs, il, ol in itertools.product( + bench_args.batch_size, bench_args.input_len, bench_args.output_len + ): + run_one_case( + base_url, + bs, + il, + ol, + bench_args.run_name, + bench_args.result_filename, + ) + finally: + if proc: + kill_process_tree(proc.pid) + + print(f"\nResults are saved to {bench_args.result_filename}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + BenchArgs.add_cli_args(parser) + args = parser.parse_args() + server_args = ServerArgs.from_cli_args(args) + bench_args = BenchArgs.from_cli_args(args) + + run_benchmark(server_args, bench_args) diff --git a/sglang/python/sglang/bench_serving.py b/sglang/python/sglang/bench_serving.py new file mode 100644 index 0000000..95cd6a3 --- /dev/null +++ b/sglang/python/sglang/bench_serving.py @@ -0,0 +1,1579 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/backend_request_func.py +# Adapted from https://github.com/vllm-project/vllm/blob/6366efc67b0aedd2c1721c14385370e50b297fb3/benchmarks/benchmark_serving.py + +""" +Benchmark online serving with dynamic requests. + +Usage: +python3 -m sglang.bench_serving --backend sglang --num-prompt 10 + +python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 3000 --random-input 1024 --random-output 1024 --random-range-ratio 0.5 +""" + +import argparse +import asyncio +import json +import os +import pickle +import random +import resource +import sys +import time +import traceback +import warnings +from argparse import ArgumentParser +from dataclasses import dataclass, field +from datetime import datetime +from pathlib import Path +from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union + +import aiohttp +import numpy as np +import requests +from tqdm.asyncio import tqdm +from transformers import ( + AutoTokenizer, + PreTrainedTokenizer, + PreTrainedTokenizerBase, + PreTrainedTokenizerFast, +) + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60) +ASSISTANT_SUFFIX = "Assistant:" + +global args + + +@dataclass +class RequestFuncInput: + prompt: str + api_url: str + prompt_len: int + output_len: int + model: str + lora_name: str + extra_request_body: Dict[str, Any] + + +@dataclass +class RequestFuncOutput: + generated_text: str = "" + success: bool = False + latency: float = 0.0 + ttft: float = 0.0 # Time to first token + itl: List[float] = field(default_factory=list) # List of inter-token latencies + prompt_len: int = 0 + error: str = "" + output_len: int = 0 + + +def remove_prefix(text: str, prefix: str) -> str: + return text[len(prefix) :] if text.startswith(prefix) else text + + +def remove_suffix(text: str, suffix: str) -> str: + return text[: -len(suffix)] if text.endswith(suffix) else text + + +def get_auth_headers() -> Dict[str, str]: + api_key = os.environ.get("OPENAI_API_KEY") + if api_key: + return {"Authorization": f"Bearer {api_key}"} + else: + return {} + + +# trt llm does not support ignore_eos +# https://github.com/triton-inference-server/tensorrtllm_backend/issues/505 +async def async_request_trt_llm( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith("generate_stream") + + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + payload = { + "accumulate_tokens": True, + "text_input": request_func_input.prompt, + "temperature": 0.000001, + "top_p": 1.0, + "max_tokens": request_func_input.output_len, + "stream": True, + "min_length": request_func_input.output_len, + "end_id": 1048576, + **request_func_input.extra_request_body, + } + if args.disable_ignore_eos: + del payload["min_length"] + del payload["end_id"] + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data:") + + data = json.loads(chunk) + output.generated_text += data["text_output"] + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = timestamp - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + + output.latency = most_recent_timestamp - st + output.success = True + output.output_len = request_func_input.output_len + + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +# set ignore_eos True by default +async def async_request_openai_completions( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith( + "completions" + ), "OpenAI Completions API URL must end with 'completions'." + + prompt = request_func_input.prompt + + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + payload = { + "model": request_func_input.model, + "prompt": prompt, + "temperature": 0.0, + "best_of": 1, + "max_tokens": request_func_input.output_len, + "stream": not args.disable_stream, + "ignore_eos": not args.disable_ignore_eos, + **request_func_input.extra_request_body, + } + headers = get_auth_headers() + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + output_len = request_func_input.output_len + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if data["choices"][0]["text"]: + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + generated_text += data["choices"][0]["text"] + output_len = (data.get("usage") or {}).get( + "completion_tokens", output_len + ) + + output.generated_text = generated_text + output.success = True + output.latency = latency + output.output_len = output_len + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_truss( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + + prompt = request_func_input.prompt + + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + payload = { + "model": request_func_input.model, + "prompt": prompt, + "temperature": 0.0, + "best_of": 1, + "max_tokens": request_func_input.output_len, + "stream": not args.disable_stream, + "ignore_eos": not args.disable_ignore_eos, + **request_func_input.extra_request_body, + } + headers = get_auth_headers() + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if data["choices"][0]["delta"]["content"]: + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - most_recent_timestamp) + + most_recent_timestamp = timestamp + generated_text += data["choices"][0]["delta"]["content"] + + output.generated_text = generated_text + output.success = True + output.latency = latency + output.output_len = request_func_input.output_len + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_sglang_generate( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + prompt = request_func_input.prompt + + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + payload = { + "text": prompt, + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": request_func_input.output_len, + "ignore_eos": not args.disable_ignore_eos, + }, + "stream": not args.disable_stream, + "lora_path": request_func_input.lora_name, + "return_logprob": args.return_logprob, + "logprob_start_len": -1, + **request_func_input.extra_request_body, + } + headers = get_auth_headers() + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + output_len = request_func_input.output_len + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + last_output_len = 0 + try: + async with session.post( + url=api_url, json=payload, headers=headers + ) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + # print(chunk_bytes) + + chunk = remove_prefix(chunk_bytes.decode("utf-8"), "data: ") + latency = time.perf_counter() - st + if chunk == "[DONE]": + pass + else: + data = json.loads(chunk) + + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if data["text"]: + timestamp = time.perf_counter() + generated_text = data["text"] + output_len = data["meta_info"]["completion_tokens"] + + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + num_new_tokens = output_len - last_output_len + if num_new_tokens == 0: + continue + adjust_itl = ( + timestamp - most_recent_timestamp + ) / num_new_tokens + output.itl.extend([adjust_itl] * num_new_tokens) + + most_recent_timestamp = timestamp + last_output_len = output_len + + output.generated_text = generated_text + output.success = True + output.latency = latency + output.output_len = output_len + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + print(f"{output.error=}") + + if pbar: + pbar.update(1) + return output + + +async def async_request_gserver( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + raise NotImplementedError() + + +async def async_request_profile(api_url: str) -> RequestFuncOutput: + async with aiohttp.ClientSession(timeout=AIOHTTP_TIMEOUT) as session: + output = RequestFuncOutput() + try: + async with session.post(url=api_url) as response: + if response.status == 200: + output.success = True + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + return output + + +def get_model(pretrained_model_name_or_path: str) -> str: + if os.getenv("SGLANG_USE_MODELSCOPE", "false").lower() == "true": + import huggingface_hub.constants + from modelscope import snapshot_download + + model_path = snapshot_download( + model_id=pretrained_model_name_or_path, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"], + ) + + return model_path + return pretrained_model_name_or_path + + +def get_tokenizer( + pretrained_model_name_or_path: str, +) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: + if pretrained_model_name_or_path.endswith( + ".json" + ) or pretrained_model_name_or_path.endswith(".model"): + from sglang.srt.hf_transformers_utils import get_tokenizer + + return get_tokenizer(pretrained_model_name_or_path) + + if pretrained_model_name_or_path is not None and not os.path.exists( + pretrained_model_name_or_path + ): + pretrained_model_name_or_path = get_model(pretrained_model_name_or_path) + return AutoTokenizer.from_pretrained( + pretrained_model_name_or_path, trust_remote_code=True + ) + + +def get_dataset(args, tokenizer): + if args.dataset_name == "sharegpt": + input_requests = sample_sharegpt_requests( + dataset_path=args.dataset_path, + num_requests=args.num_prompts, + tokenizer=tokenizer, + fixed_output_len=args.sharegpt_output_len, + context_len=args.sharegpt_context_len, + prompt_suffix=args.prompt_suffix, + apply_chat_template=args.apply_chat_template, + ) + elif args.dataset_name == "random": + input_requests = sample_random_requests( + input_len=args.random_input_len, + output_len=args.random_output_len, + num_prompts=args.num_prompts, + range_ratio=args.random_range_ratio, + tokenizer=tokenizer, + dataset_path=args.dataset_path, + ) + elif args.dataset_name == "generated-shared-prefix": + input_requests = sample_generated_shared_prefix_requests( + num_groups=args.gsp_num_groups, + prompts_per_group=args.gsp_prompts_per_group, + system_prompt_len=args.gsp_system_prompt_len, + question_len=args.gsp_question_len, + output_len=args.gsp_output_len, + tokenizer=tokenizer, + args=args, + ) + else: + raise ValueError(f"Unknown dataset: {args.dataset_name}") + return input_requests + + +ASYNC_REQUEST_FUNCS = { + "sglang": async_request_sglang_generate, + "sglang-native": async_request_sglang_generate, + "sglang-oai": async_request_openai_completions, + "vllm": async_request_openai_completions, + "lmdeploy": async_request_openai_completions, + "trt": async_request_trt_llm, + "gserver": async_request_gserver, + "truss": async_request_truss, +} + + +@dataclass +class BenchmarkMetrics: + completed: int + total_input: int + total_output: int + total_output_retokenized: int + request_throughput: float + input_throughput: float + output_throughput: float + output_throughput_retokenized: float + total_throughput: float + total_throughput_retokenized: float + mean_ttft_ms: float + median_ttft_ms: float + std_ttft_ms: float + p99_ttft_ms: float + mean_tpot_ms: float + median_tpot_ms: float + std_tpot_ms: float + p99_tpot_ms: float + mean_itl_ms: float + median_itl_ms: float + std_itl_ms: float + p95_itl_ms: float + p99_itl_ms: float + max_itl_ms: float + mean_e2e_latency_ms: float + median_e2e_latency_ms: float + std_e2e_latency_ms: float + p99_e2e_latency_ms: float + concurrency: float + + +SHAREGPT_URL = "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json" + + +def download_and_cache_file(url: str, filename: Optional[str] = None): + """Read and cache a file from a url.""" + if filename is None: + filename = os.path.join("/tmp", url.split("/")[-1]) + + # Check if the cache file already exists + if os.path.exists(filename): + return filename + + print(f"Downloading from {url} to {filename}") + + # Stream the response to show the progress bar + response = requests.get(url, stream=True) + response.raise_for_status() # Check for request errors + + # Total size of the file in bytes + total_size = int(response.headers.get("content-length", 0)) + chunk_size = 1024 # Download in chunks of 1KB + + # Use tqdm to display the progress bar + with open(filename, "wb") as f, tqdm( + desc=filename, + total=total_size, + unit="B", + unit_scale=True, + unit_divisor=1024, + ) as bar: + for chunk in response.iter_content(chunk_size=chunk_size): + f.write(chunk) + bar.update(len(chunk)) + + return filename + + +def sample_sharegpt_requests( + dataset_path: str, + num_requests: int, + tokenizer: PreTrainedTokenizerBase, + fixed_output_len: Optional[int] = None, + context_len: Optional[int] = None, + prompt_suffix: Optional[str] = "", + apply_chat_template=False, +) -> List[Tuple[str, int, int]]: + if fixed_output_len is not None and fixed_output_len < 4: + raise ValueError("output_len too small") + + # Download sharegpt if necessary + if not os.path.isfile(dataset_path) and dataset_path == "": + dataset_path = download_and_cache_file(SHAREGPT_URL) + + # Load the dataset. + with open(dataset_path) as f: + dataset = json.load(f) + + # Filter out the conversations with less than 2 turns. + dataset = [ + data + for data in dataset + if len(data.get("conversations", data.get("conversation", []))) >= 2 + ] + # Only keep the first two turns of each conversation. + dataset = [ + ( + data.get("conversations", data.get("conversation", []))[0]["value"], + data.get("conversations", data.get("conversation", []))[1]["value"], + ) + for data in dataset + ] + + # Shuffle the dataset. + random.shuffle(dataset) + + # Filter out sequences that are too long or too short + filtered_dataset: List[Tuple[str, int, int]] = [] + for i in range(len(dataset)): + if len(filtered_dataset) == num_requests: + break + + # Tokenize the prompts and completions. + prompt = dataset[i][0] + if prompt_suffix: + prompt = ( + remove_suffix(prompt, ASSISTANT_SUFFIX) + + prompt_suffix + + ASSISTANT_SUFFIX + ) + + if apply_chat_template: + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=False, + ) + prompt = prompt.replace(tokenizer.bos_token, "") + + prompt_token_ids = tokenizer.encode(prompt) + completion = dataset[i][1] + completion_token_ids = tokenizer.encode(completion) + prompt_len = len(prompt_token_ids) + output_len = ( + len(completion_token_ids) if fixed_output_len is None else fixed_output_len + ) + + if prompt_len < 2 or output_len < 2: + # Prune too short sequences. + continue + + if context_len and prompt_len + output_len > context_len: + # Prune too long sequences. + continue + + filtered_dataset.append((prompt, prompt_len, output_len)) + + print(f"#Input tokens: {np.sum([x[1] for x in filtered_dataset])}") + print(f"#Output tokens: {np.sum([x[2] for x in filtered_dataset])}") + return filtered_dataset + + +def sample_random_requests( + input_len: int, + output_len: int, + num_prompts: int, + range_ratio: float, + tokenizer: PreTrainedTokenizerBase, + dataset_path: str, +) -> List[Tuple[str, int, int]]: + + input_lens = np.random.randint( + max(int(input_len * range_ratio), 1), + input_len + 1, + size=num_prompts, + ) + output_lens = np.random.randint( + int(output_len * range_ratio), + output_len + 1, + size=num_prompts, + ) + + if True: + # Sample token ids from ShareGPT and repeat/truncate them to satisfy the input_lens + + # Download sharegpt if necessary + if not os.path.isfile(dataset_path): + dataset_path = download_and_cache_file(SHAREGPT_URL) + + # Load the dataset. + with open(dataset_path) as f: + dataset = json.load(f) + # Filter out the conversations with less than 2 turns. + dataset = [ + data + for data in dataset + if len(data.get("conversations", data.get("conversation", []))) >= 2 + ] + # Only keep the first two turns of each conversation. + dataset = [ + ( + data.get("conversations", data.get("conversation", []))[0]["value"], + data.get("conversations", data.get("conversation", []))[1]["value"], + ) + for data in dataset + ] + # Shuffle the dataset. + random.shuffle(dataset) + + # Filter out sequences that are too long or too short + input_requests: List[Tuple[str, int, int]] = [] + for data in dataset: + i = len(input_requests) + if i == num_prompts: + break + + # Tokenize the prompts and completions. + prompt = data[0] + prompt_token_ids = tokenizer.encode(prompt) + prompt_len = len(prompt_token_ids) + + # Skip empty prompt + if prompt_len == 0: + continue + + if prompt_len > input_lens[i]: + input_ids = prompt_token_ids[: input_lens[i]] + else: + ratio = (input_lens[i] + prompt_len - 1) // prompt_len + input_ids = (prompt_token_ids * ratio)[: input_lens[i]] + prompt = tokenizer.decode(input_ids) + input_requests.append((prompt, int(input_lens[i]), int(output_lens[i]))) + else: + # Sample token ids from random integers. This can cause some NaN issues. + offsets = np.random.randint(0, tokenizer.vocab_size, size=num_prompts) + input_requests = [] + for i in range(num_prompts): + prompt = tokenizer.decode( + [ + (offsets[i] + i + j) % tokenizer.vocab_size + for j in range(input_lens[i]) + ] + ) + input_requests.append((prompt, int(input_lens[i]), int(output_lens[i]))) + + print(f"#Input tokens: {np.sum(input_lens)}") + print(f"#Output tokens: {np.sum(output_lens)}") + return input_requests + + +def gen_prompt(tokenizer, token_num): + """Generate a random prompt of specified token length using tokenizer vocabulary.""" + all_available_tokens = list(tokenizer.get_vocab().values()) + selected_tokens = random.choices(all_available_tokens, k=token_num) + return tokenizer.decode(selected_tokens) + + +def get_gen_prefix_cache_path(args, tokenizer): + """Create cache directory under ~/.cache/sglang/benchmark""" + cache_dir = Path.home() / ".cache" / "sglang" / "benchmark" + + # Create a unique cache filename based on the generation parameters + cache_key = ( + f"gen_shared_prefix_{args.gsp_num_groups}_{args.gsp_prompts_per_group}_" + f"{args.gsp_system_prompt_len}_{args.gsp_question_len}_{args.gsp_output_len}_" + f"{tokenizer.__class__.__name__}.pkl" + ) + return cache_dir / cache_key + + +def sample_generated_shared_prefix_requests( + num_groups: int, + prompts_per_group: int, + system_prompt_len: int, + question_len: int, + output_len: int, + tokenizer: PreTrainedTokenizerBase, + args: argparse.Namespace, +) -> List[Tuple[str, int, int]]: + """Generate benchmark requests with shared system prompts using random tokens and caching.""" + cache_path = get_gen_prefix_cache_path(args, tokenizer) + + # Try to load from cache first + if cache_path.exists(): + print(f"\nLoading cached generated input data from {cache_path}") + with open(cache_path, "rb") as f: + return pickle.load(f) + + print("\nGenerating new input data...") + + # Generate system prompts for each group + system_prompts = [] + for _ in range(num_groups): + system_prompt = gen_prompt(tokenizer, system_prompt_len) + system_prompts.append(system_prompt) + + # Generate questions + questions = [] + for _ in range(num_groups * prompts_per_group): + question = gen_prompt(tokenizer, question_len) + questions.append(question) + + # Combine system prompts with questions + input_requests = [] + total_input_tokens = 0 + total_output_tokens = 0 + + for group_idx in tqdm(range(num_groups), desc="Generating system prompt"): + system_prompt = system_prompts[group_idx] + for prompt_idx in tqdm( + range(prompts_per_group), desc="Generating questions", leave=False + ): + question = questions[group_idx * prompts_per_group + prompt_idx] + full_prompt = f"{system_prompt}\n\n{question}" + prompt_len = len(tokenizer.encode(full_prompt)) + + input_requests.append((full_prompt, prompt_len, output_len)) + total_input_tokens += prompt_len + total_output_tokens += output_len + + # Shuffle questions + random.shuffle(input_requests) + + # Print statistics + print(f"\nGenerated shared prefix dataset statistics:") + print(f"Number of groups: {num_groups}") + print(f"Prompts per group: {prompts_per_group}") + print(f"Total prompts: {len(input_requests)}") + print(f"Total input tokens: {total_input_tokens}") + print(f"Total output tokens: {total_output_tokens}") + print( + f"Average system prompt length: {sum(len(tokenizer.encode(sp)) for sp in system_prompts) / len(system_prompts):.1f} tokens" + ) + print( + f"Average question length: {sum(len(tokenizer.encode(q)) for q in questions) / len(questions):.1f} tokens\n" + ) + + # Save to cache + cache_path.parent.mkdir(parents=True, exist_ok=True) + print(f"Caching generated input data to {cache_path}") + with open(cache_path, "wb") as f: + pickle.dump(input_requests, f) + + return input_requests + + +async def get_request( + input_requests: List[Tuple[str, int, int]], + request_rate: float, +) -> AsyncGenerator[Tuple[str, int, int], None]: + input_requests = iter(input_requests) + for request in input_requests: + yield request + + if request_rate == float("inf"): + # If the request rate is infinity, then we don't need to wait. + continue + + # Sample the request interval from the exponential distribution. + interval = np.random.exponential(1.0 / request_rate) + # The next request will be sent after the interval. + await asyncio.sleep(interval) + + +def calculate_metrics( + input_requests: List[Tuple[str, int, int]], + outputs: List[RequestFuncOutput], + dur_s: float, + tokenizer: PreTrainedTokenizerBase, + backend: str, +) -> Tuple[BenchmarkMetrics, List[int]]: + output_lens: List[int] = [] + retokenized_output_lens: List[int] = [] + total_input = 0 + completed = 0 + itls: List[float] = [] + tpots: List[float] = [] + ttfts: List[float] = [] + e2e_latencies: List[float] = [] + for i in range(len(outputs)): + if outputs[i].success: + output_len = outputs[i].output_len + output_lens.append(output_len) + retokenized_output_len = len( + tokenizer.encode(outputs[i].generated_text, add_special_tokens=False) + ) + retokenized_output_lens.append(retokenized_output_len) + total_input += input_requests[i][1] + if output_len > 1: + tpots.append((outputs[i].latency - outputs[i].ttft) / (output_len - 1)) + itls += outputs[i].itl + ttfts.append(outputs[i].ttft) + + e2e_latencies.append(outputs[i].latency) + + completed += 1 + else: + output_lens.append(0) + retokenized_output_lens.append(0) + + if completed == 0: + warnings.warn( + "All requests failed. This is likely due to a misconfiguration " + "on the benchmark arguments.", + stacklevel=2, + ) + metrics = BenchmarkMetrics( + completed=completed, + total_input=total_input, + total_output=sum(output_lens), + total_output_retokenized=sum(retokenized_output_lens), + request_throughput=completed / dur_s, + input_throughput=total_input / dur_s, + output_throughput=sum(output_lens) / dur_s, + output_throughput_retokenized=sum(retokenized_output_lens) / dur_s, + total_throughput=(total_input + sum(output_lens)) / dur_s, + total_throughput_retokenized=(total_input + sum(retokenized_output_lens)) + / dur_s, + mean_ttft_ms=np.mean(ttfts or 0) + * 1000, # ttfts is empty if streaming is not supported by backend + median_ttft_ms=np.median(ttfts or 0) * 1000, + std_ttft_ms=np.std(ttfts or 0) * 1000, + p99_ttft_ms=np.percentile(ttfts or 0, 99) * 1000, + mean_tpot_ms=np.mean(tpots or 0) * 1000, + median_tpot_ms=np.median(tpots or 0) * 1000, + std_tpot_ms=np.std(tpots or 0) * 1000, + p99_tpot_ms=np.percentile(tpots or 0, 99) * 1000, + mean_itl_ms=np.mean(itls or 0) * 1000, + median_itl_ms=np.median(itls or 0) * 1000, + std_itl_ms=np.std(itls or 0) * 1000, + p95_itl_ms=np.percentile(itls or 0, 95) * 1000, + p99_itl_ms=np.percentile(itls or 0, 99) * 1000, + max_itl_ms=np.max(itls or 0) * 1000, + mean_e2e_latency_ms=np.mean(e2e_latencies) * 1000, + median_e2e_latency_ms=np.median(e2e_latencies) * 1000, + std_e2e_latency_ms=np.std(e2e_latencies) * 1000, + p99_e2e_latency_ms=np.percentile(e2e_latencies, 99) * 1000, + concurrency=np.sum(e2e_latencies) / dur_s, + ) + + return metrics, output_lens + + +async def benchmark( + backend: str, + api_url: str, + base_url: str, + model_id: str, + tokenizer: PreTrainedTokenizerBase, + input_requests: List[Tuple[str, int, int]], + request_rate: float, + max_concurrency: Optional[int], + disable_tqdm: bool, + lora_names: List[str], + extra_request_body: Dict[str, Any], + profile: bool, + pd_seperated: bool = False, +): + if backend in ASYNC_REQUEST_FUNCS: + request_func = ASYNC_REQUEST_FUNCS[backend] + else: + raise ValueError(f"Unknown backend: {backend}") + + # Limit concurrency + # From https://github.com/vllm-project/vllm/pull/9390 + semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None + + async def limited_request_func(request_func_input, pbar): + if semaphore is None: + return await request_func(request_func_input=request_func_input, pbar=pbar) + async with semaphore: + return await request_func(request_func_input=request_func_input, pbar=pbar) + + # Warmup + print("Starting initial single prompt test run...") + test_prompt, test_prompt_len, test_output_len = input_requests[0] + if lora_names != None and len(lora_names) != 0: + lora_name = lora_names[0] + else: + lora_name = None + + test_input = RequestFuncInput( + model=model_id, + prompt=test_prompt, + api_url=api_url, + prompt_len=test_prompt_len, + output_len=min(test_output_len, 32), + lora_name=lora_name, + extra_request_body=extra_request_body, + ) + test_output = await request_func(request_func_input=test_input) + if not test_output.success: + raise ValueError( + "Initial test run failed - Please make sure benchmark arguments " + f"are correctly specified. Error: {test_output.error}" + ) + else: + print("Initial test run completed. Starting main benchmark run...") + + # Flush cache + if "sglang" in backend: + requests.post(base_url + "/flush_cache", headers=get_auth_headers()) + + time.sleep(1.0) + + # Start profiler + if profile: + print("Starting profiler...") + profile_output = await async_request_profile( + api_url=base_url + "/start_profile" + ) + if profile_output.success: + print("Profiler started") + + pbar = None if disable_tqdm else tqdm(total=len(input_requests)) + + # Run all requests + benchmark_start_time = time.perf_counter() + tasks: List[asyncio.Task] = [] + async for request in get_request(input_requests, request_rate): + prompt, prompt_len, output_len = request + if lora_names != None and len(lora_names) != 0: + idx = random.randint(0, len(lora_names) - 1) + lora_name = lora_names[idx] + else: + lora_name = None + + request_func_input = RequestFuncInput( + model=model_id, + prompt=prompt, + api_url=api_url, + prompt_len=prompt_len, + output_len=output_len, + lora_name=lora_name, + extra_request_body=extra_request_body, + ) + tasks.append( + asyncio.create_task( + limited_request_func(request_func_input=request_func_input, pbar=pbar) + ) + ) + outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) + + # Stop profiler + if profile: + print("Stopping profiler...") + profile_output = await async_request_profile(api_url=base_url + "/stop_profile") + if profile_output.success: + print("Profiler stopped") + + if pbar is not None: + pbar.close() + + if "sglang" in backend: + server_info = requests.get(base_url + "/get_server_info") + if pd_seperated: + accept_length = server_info.json()["decode"][0].get( + "avg_spec_accept_length", None + ) + else: + accept_length = server_info.json().get("avg_spec_accept_length", None) + else: + accept_length = None + + # Compute metrics and print results + benchmark_duration = time.perf_counter() - benchmark_start_time + metrics, output_lens = calculate_metrics( + input_requests=input_requests, + outputs=outputs, + dur_s=benchmark_duration, + tokenizer=tokenizer, + backend=backend, + ) + + print("\n{s:{c}^{n}}".format(s=" Serving Benchmark Result ", n=50, c="=")) + print("{:<40} {:<10}".format("Backend:", backend)) + print("{:<40} {:<10}".format("Traffic request rate:", request_rate)) + print( + "{:<40} {:<10}".format( + "Max reqeuest concurrency:", + max_concurrency if max_concurrency else "not set", + ) + ) + print("{:<40} {:<10}".format("Successful requests:", metrics.completed)) + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", benchmark_duration)) + print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) + print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output)) + print( + "{:<40} {:<10}".format( + "Total generated tokens (retokenized):", metrics.total_output_retokenized + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Request throughput (req/s):", metrics.request_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Input token throughput (tok/s):", metrics.input_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Output token throughput (tok/s):", metrics.output_throughput + ) + ) + print( + "{:<40} {:<10.2f}".format( + "Total token throughput (tok/s):", metrics.total_throughput + ) + ) + print("{:<40} {:<10.2f}".format("Concurrency:", metrics.concurrency)) + if accept_length: + print("{:<40} {:<10.2f}".format("Accept length:", accept_length)) + print("{s:{c}^{n}}".format(s="End-to-End Latency", n=50, c="-")) + print( + "{:<40} {:<10.2f}".format("Mean E2E Latency (ms):", metrics.mean_e2e_latency_ms) + ) + print( + "{:<40} {:<10.2f}".format( + "Median E2E Latency (ms):", metrics.median_e2e_latency_ms + ) + ) + print("{s:{c}^{n}}".format(s="Time to First Token", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean TTFT (ms):", metrics.mean_ttft_ms)) + print("{:<40} {:<10.2f}".format("Median TTFT (ms):", metrics.median_ttft_ms)) + print("{:<40} {:<10.2f}".format("P99 TTFT (ms):", metrics.p99_ttft_ms)) + print("{s:{c}^{n}}".format(s="Inter-Token Latency", n=50, c="-")) + print("{:<40} {:<10.2f}".format("Mean ITL (ms):", metrics.mean_itl_ms)) + print("{:<40} {:<10.2f}".format("Median ITL (ms):", metrics.median_itl_ms)) + print("{:<40} {:<10.2f}".format("P95 ITL (ms):", metrics.p95_itl_ms)) + print("{:<40} {:<10.2f}".format("P99 ITL (ms):", metrics.p99_itl_ms)) + print("{:<40} {:<10.2f}".format("Max ITL (ms):", metrics.max_itl_ms)) + print("=" * 50) + + if ( + metrics.median_ttft_ms is not None + and metrics.mean_itl_ms is not None + and metrics.output_throughput is not None + ): + result = { + # Arguments + "backend": args.backend, + "dataset_name": args.dataset_name, + "request_rate": request_rate, + "max_concurrency": max_concurrency, + "sharegpt_output_len": args.sharegpt_output_len, + "random_input_len": args.random_input_len, + "random_output_len": args.random_output_len, + "random_range_ratio": args.random_range_ratio, + # Results + "duration": benchmark_duration, + "completed": metrics.completed, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "total_output_tokens_retokenized": metrics.total_output_retokenized, + "request_throughput": metrics.request_throughput, + "input_throughput": metrics.input_throughput, + "output_throughput": metrics.output_throughput, + "mean_e2e_latency_ms": metrics.mean_e2e_latency_ms, + "median_e2e_latency_ms": metrics.median_e2e_latency_ms, + "std_e2e_latency_ms": metrics.std_e2e_latency_ms, + "p99_e2e_latency_ms": metrics.p99_e2e_latency_ms, + "mean_ttft_ms": metrics.mean_ttft_ms, + "median_ttft_ms": metrics.median_ttft_ms, + "std_ttft_ms": metrics.std_ttft_ms, + "p99_ttft_ms": metrics.p99_ttft_ms, + "mean_tpot_ms": metrics.mean_tpot_ms, + "median_tpot_ms": metrics.median_tpot_ms, + "std_tpot_ms": metrics.std_tpot_ms, + "p99_tpot_ms": metrics.p99_tpot_ms, + "mean_itl_ms": metrics.mean_itl_ms, + "median_itl_ms": metrics.median_itl_ms, + "std_itl_ms": metrics.std_itl_ms, + "p95_itl_ms": metrics.p95_itl_ms, + "p99_itl_ms": metrics.p99_itl_ms, + "concurrency": metrics.concurrency, + "accept_length": accept_length, + } + else: + print(f"Error running benchmark for request rate: {request_rate}") + print("-" * 30) + + # Determine output file name + if args.output_file: + output_file_name = args.output_file + else: + now = datetime.now().strftime("%m%d") + if args.dataset_name == "random": + output_file_name = f"{args.backend}_{now}_{args.num_prompts}_{args.random_input_len}_{args.random_output_len}.jsonl" + else: + output_file_name = f"{args.backend}_{now}_{args.num_prompts}_sharegpt.jsonl" + + # Append results to a JSONL file + with open(output_file_name, "a") as file: + file.write(json.dumps(result) + "\n") + + result.update( + { + "input_lens": [output.prompt_len for output in outputs], + "output_lens": output_lens, + "ttfts": [output.ttft for output in outputs], + "itls": [output.itl for output in outputs], + "generated_texts": [output.generated_text for output in outputs], + "errors": [output.error for output in outputs], + } + ) + return result + + +def check_chat_template(model_path): + try: + tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) + return "chat_template" in tokenizer.init_kwargs + except Exception as e: + print(f"Fail to load tokenizer config with error={e}") + return False + + +def set_global_args(args_: argparse.Namespace): + """Set the global args.""" + global args + args = args_ + + +def run_benchmark(args_: argparse.Namespace): + global args + args = args_ + + # Set default value for max_concurrency if not present + if not hasattr(args, "max_concurrency"): + args.max_concurrency = None + + print(f"benchmark_args={args}") + + # Set global environments + set_ulimit() + random.seed(args.seed) + np.random.seed(args.seed) + + extra_request_body = {} + if args.extra_request_body: + extra_request_body = json.loads(args.extra_request_body) + + # Set url + if args.port is None: + args.port = { + "sglang": 30000, + "sglang-native": 30000, + "sglang-oai": 30000, + "lmdeploy": 23333, + "vllm": 8000, + "trt": 8000, + "gserver": 9988, + "truss": 8080, + }.get(args.backend, 30000) + + model_url = ( + f"{args.base_url}/v1/models" + if args.base_url + else f"http://{args.host}:{args.port}/v1/models" + ) + + if args.backend in ["sglang", "sglang-native"]: + api_url = ( + f"{args.base_url}/generate" + if args.base_url + else f"http://{args.host}:{args.port}/generate" + ) + elif args.backend in ["sglang-oai", "vllm", "lmdeploy"]: + api_url = ( + f"{args.base_url}/v1/completions" + if args.base_url + else f"http://{args.host}:{args.port}/v1/completions" + ) + elif args.backend == "trt": + api_url = ( + f"{args.base_url}/v2/models/ensemble/generate_stream" + if args.base_url + else f"http://{args.host}:{args.port}/v2/models/ensemble/generate_stream" + ) + if args.model is None: + print("Please provide a model using `--model` when using `trt` backend.") + sys.exit(1) + elif args.backend == "gserver": + api_url = args.base_url if args.base_url else f"{args.host}:{args.port}" + args.model = args.model or "default" + elif args.backend == "truss": + api_url = ( + f"{args.base_url}/v1/models/model:predict" + if args.base_url + else f"http://{args.host}:{args.port}/v1/models/model:predict" + ) + base_url = ( + f"http://{args.host}:{args.port}" if args.base_url is None else args.base_url + ) + + # Get model name + if args.model is None: + if args.backend == "truss": + print( + "Please provide a model with `--model` when using truss backend. e.g. --model meta-llama/Llama-3.1-8B-Instruct" + ) + sys.exit(1) + try: + response = requests.get(model_url, headers=get_auth_headers()) + model_list = response.json().get("data", []) + args.model = model_list[0]["id"] if model_list else None + except Exception as e: + print(f"Failed to fetch model from {model_url}. Error: {e}") + print( + "Please specify the correct host and port using `--host` and `--port`." + ) + sys.exit(1) + + if args.model is None: + print("No model specified or found. Please provide a model using `--model`.") + sys.exit(1) + + if not check_chat_template(args.model): + print( + "\nWARNING It is recommended to use the `Chat` or `Instruct` model for benchmarking.\n" + "Because when the tokenizer counts the output tokens, if there is gibberish, it might count incorrectly.\n" + ) + + print(f"{args}\n") + + # Read dataset + backend = args.backend + model_id = args.model + tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model + tokenizer = get_tokenizer(tokenizer_id) + input_requests = get_dataset(args, tokenizer) + + return asyncio.run( + benchmark( + backend=backend, + api_url=api_url, + base_url=base_url, + model_id=model_id, + tokenizer=tokenizer, + input_requests=input_requests, + request_rate=args.request_rate, + max_concurrency=args.max_concurrency, + disable_tqdm=args.disable_tqdm, + lora_names=args.lora_name, + extra_request_body=extra_request_body, + profile=args.profile, + pd_seperated=args.pd_seperated, + ) + ) + + +def set_ulimit(target_soft_limit=65535): + resource_type = resource.RLIMIT_NOFILE + current_soft, current_hard = resource.getrlimit(resource_type) + + if current_soft < target_soft_limit: + try: + resource.setrlimit(resource_type, (target_soft_limit, current_hard)) + except ValueError as e: + print(f"Fail to set RLIMIT_NOFILE: {e}") + + +class LoRAPathAction(argparse.Action): + def __call__(self, parser, namespace, values, option_string=None): + setattr(namespace, self.dest, []) + for lora_name in values: + getattr(namespace, self.dest).append(lora_name) + + +if __name__ == "__main__": + parser = ArgumentParser(description="Benchmark the online serving throughput.") + parser.add_argument( + "--backend", + type=str, + choices=list(ASYNC_REQUEST_FUNCS.keys()), + default="sglang", + help="Must specify a backend, depending on the LLM Inference Engine.", + ) + parser.add_argument( + "--base-url", + type=str, + default=None, + help="Server or API base url if not using http host and port.", + ) + parser.add_argument( + "--host", type=str, default="0.0.0.0", help="Default host is 0.0.0.0." + ) + parser.add_argument( + "--port", + type=int, + help="If not set, the default port is configured according to its default value for different LLM Inference Engines.", + ) + parser.add_argument( + "--dataset-name", + type=str, + default="sharegpt", + choices=["sharegpt", "random", "generated-shared-prefix"], + help="Name of the dataset to benchmark on.", + ) + parser.add_argument( + "--dataset-path", type=str, default="", help="Path to the dataset." + ) + parser.add_argument( + "--model", + type=str, + help="Name or path of the model. If not set, the default model will request /v1/models for conf.", + ) + parser.add_argument( + "--tokenizer", + type=str, + help="Name or path of the tokenizer. If not set, using the model conf.", + ) + parser.add_argument( + "--num-prompts", + type=int, + default=1000, + help="Number of prompts to process. Default is 1000.", + ) + parser.add_argument( + "--sharegpt-output-len", + type=int, + default=None, + help="Output length for each request. Overrides the output length from the ShareGPT dataset.", + ) + parser.add_argument( + "--sharegpt-context-len", + type=int, + default=None, + help="The context length of the model for the ShareGPT dataset. Requests longer than the context length will be dropped.", + ) + parser.add_argument( + "--random-input-len", + type=int, + default=1024, + help="Number of input tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-output-len", + default=1024, + type=int, + help="Number of output tokens per request, used only for random dataset.", + ) + parser.add_argument( + "--random-range-ratio", + type=float, + default=0.0, + help="Range of sampled ratio of input/output length, " + "used only for random dataset.", + ) + parser.add_argument( + "--request-rate", + type=float, + default=float("inf"), + help="Number of requests per second. If this is inf, then all the requests are sent at time 0. " + "Otherwise, we use Poisson process to synthesize the request arrival times. Default is inf.", + ) + parser.add_argument( + "--max-concurrency", + type=int, + default=None, + help="Maximum number of concurrent requests. This can be used " + "to help simulate an environment where a higher level component " + "is enforcing a maximum number of concurrent requests. While the " + "--request-rate argument controls the rate at which requests are " + "initiated, this argument will control how many are actually allowed " + "to execute at a time. This means that when used in combination, the " + "actual request rate may be lower than specified with --request-rate, " + "if the server is not processing requests fast enough to keep up.", + ) + parser.add_argument("--output-file", type=str, help="Output JSONL file name.") + parser.add_argument( + "--disable-tqdm", + action="store_true", + help="Specify to disable tqdm progress bar.", + ) + parser.add_argument( + "--disable-stream", + action="store_true", + help="Disable streaming mode.", + ) + parser.add_argument( + "--return-logprob", + action="store_true", + help="Return logprob.", + ) + parser.add_argument("--seed", type=int, default=1, help="The random seed.") + parser.add_argument( + "--disable-ignore-eos", + action="store_true", + help="Disable ignoring EOS.", + ) + parser.add_argument( + "--extra-request-body", + metavar='{"key1": "value1", "key2": "value2"}', + type=str, + help="Append given JSON object to the request payload. You can use this to specify" + "additional generate params like sampling params.", + ) + parser.add_argument( + "--apply-chat-template", + action="store_true", + help="Apply chat template", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Use Torch Profiler. The endpoint must be launched with " + "SGLANG_TORCH_PROFILER_DIR to enable profiler.", + ) + parser.add_argument( + "--lora-name", + type=str, + nargs="*", + default=None, + action=LoRAPathAction, + help="The names of LoRA adapters. You can provide a list of names in the format {name} {name} {name}...", + ) + parser.add_argument( + "--prompt-suffix", + type=str, + default="", + help="Suffix applied to the end of all user prompts, followed by assistant prompt suffix.", + ) + parser.add_argument( + "--pd-seperated", + action="store_true", + help="Benchmark PD disaggregation server", + ) + + group = parser.add_argument_group("generated-shared-prefix dataset arguments") + group.add_argument( + "--gsp-num-groups", + type=int, + default=64, + help="Number of system prompt groups for generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-prompts-per-group", + type=int, + default=16, + help="Number of prompts per system prompt group for generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-system-prompt-len", + type=int, + default=2048, + help="Target length in tokens for system prompts in generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-question-len", + type=int, + default=128, + help="Target length in tokens for questions in generated-shared-prefix dataset", + ) + group.add_argument( + "--gsp-output-len", + type=int, + default=256, + help="Target length in tokens for outputs in generated-shared-prefix dataset", + ) + args = parser.parse_args() + run_benchmark(args) diff --git a/sglang/python/sglang/check_env.py b/sglang/python/sglang/check_env.py new file mode 100644 index 0000000..d525da9 --- /dev/null +++ b/sglang/python/sglang/check_env.py @@ -0,0 +1,305 @@ +"""Check environment configurations and dependency versions.""" + +import importlib.metadata +import os +import resource +import subprocess +import sys +from collections import OrderedDict, defaultdict + +import torch + +from sglang.srt.utils import is_hip + + +def is_cuda_v2(): + return torch.version.cuda is not None + + +# List of packages to check versions +PACKAGE_LIST = [ + "sglang", + "sgl_kernel", + "flashinfer", + "triton", + "transformers", + "torchao", + "numpy", + "aiohttp", + "fastapi", + "hf_transfer", + "huggingface_hub", + "interegular", + "modelscope", + "orjson", + "outlines", + "packaging", + "psutil", + "pydantic", + "multipart", + "zmq", + "torchao", + "uvicorn", + "uvloop", + "vllm", + "xgrammar", + "openai", + "tiktoken", + "anthropic", + "litellm", + "decord", +] + + +def get_package_versions(packages): + """ + Get versions of specified packages. + """ + versions = {} + for package in packages: + package_name = package.split("==")[0].split(">=")[0].split("<=")[0] + try: + version = importlib.metadata.version(package_name) + versions[package_name] = version + except ModuleNotFoundError: + versions[package_name] = "Module Not Found" + return versions + + +def get_cuda_info(): + """ + Get CUDA-related information if available. + """ + if is_cuda_v2(): + cuda_info = {"CUDA available": torch.cuda.is_available()} + + if cuda_info["CUDA available"]: + cuda_info.update(_get_gpu_info()) + cuda_info.update(_get_cuda_version_info()) + + return cuda_info + elif is_hip(): + cuda_info = {"ROCM available": torch.cuda.is_available()} + + if cuda_info["ROCM available"]: + cuda_info.update(_get_gpu_info()) + cuda_info.update(_get_cuda_version_info()) + + return cuda_info + + +def _get_gpu_info(): + """ + Get information about available GPUs. + """ + devices = defaultdict(list) + capabilities = defaultdict(list) + for k in range(torch.cuda.device_count()): + devices[torch.cuda.get_device_name(k)].append(str(k)) + capability = torch.cuda.get_device_capability(k) + capabilities[f"{capability[0]}.{capability[1]}"].append(str(k)) + + gpu_info = {} + for name, device_ids in devices.items(): + gpu_info[f"GPU {','.join(device_ids)}"] = name + + if len(capabilities) == 1: + # All GPUs have the same compute capability + cap, gpu_ids = list(capabilities.items())[0] + gpu_info[f"GPU {','.join(gpu_ids)} Compute Capability"] = cap + else: + # GPUs have different compute capabilities + for cap, gpu_ids in capabilities.items(): + gpu_info[f"GPU {','.join(gpu_ids)} Compute Capability"] = cap + + return gpu_info + + +def _get_cuda_version_info(): + """ + Get CUDA version information. + """ + if is_cuda_v2(): + from torch.utils.cpp_extension import CUDA_HOME + + cuda_info = {"CUDA_HOME": CUDA_HOME} + + if CUDA_HOME and os.path.isdir(CUDA_HOME): + cuda_info.update(_get_nvcc_info()) + cuda_info.update(_get_cuda_driver_version()) + + return cuda_info + elif is_hip(): + from torch.utils.cpp_extension import ROCM_HOME as ROCM_HOME + + cuda_info = {"ROCM_HOME": ROCM_HOME} + + if ROCM_HOME and os.path.isdir(ROCM_HOME): + cuda_info.update(_get_nvcc_info()) + cuda_info.update(_get_cuda_driver_version()) + + return cuda_info + else: + cuda_info = {"CUDA_HOME": ""} + return cuda_info + + +def _get_nvcc_info(): + """ + Get NVCC version information. + """ + if is_cuda_v2(): + from torch.utils.cpp_extension import CUDA_HOME + + try: + nvcc = os.path.join(CUDA_HOME, "bin/nvcc") + nvcc_output = ( + subprocess.check_output(f'"{nvcc}" -V', shell=True) + .decode("utf-8") + .strip() + ) + return { + "NVCC": nvcc_output[ + nvcc_output.rfind("Cuda compilation tools") : nvcc_output.rfind( + "Build" + ) + ].strip() + } + except subprocess.SubprocessError: + return {"NVCC": "Not Available"} + elif is_hip(): + from torch.utils.cpp_extension import ROCM_HOME + + try: + hipcc = os.path.join(ROCM_HOME, "bin/hipcc") + hipcc_output = ( + subprocess.check_output(f'"{hipcc}" --version', shell=True) + .decode("utf-8") + .strip() + ) + return { + "HIPCC": hipcc_output[ + hipcc_output.rfind("HIP version") : hipcc_output.rfind("AMD clang") + ].strip() + } + except subprocess.SubprocessError: + return {"HIPCC": "Not Available"} + else: + return {"NVCC": "Not Available"} + + +def _get_cuda_driver_version(): + """ + Get CUDA driver version. + """ + versions = set() + if is_cuda_v2(): + try: + output = subprocess.check_output( + [ + "nvidia-smi", + "--query-gpu=driver_version", + "--format=csv,noheader,nounits", + ] + ) + versions = set(output.decode().strip().split("\n")) + if len(versions) == 1: + return {"CUDA Driver Version": versions.pop()} + else: + return {"CUDA Driver Versions": ", ".join(sorted(versions))} + except subprocess.SubprocessError: + return {"CUDA Driver Version": "Not Available"} + elif is_hip(): + try: + output = subprocess.check_output( + [ + "rocm-smi", + "--showdriverversion", + "--csv", + ] + ) + versions = set(output.decode().strip().split("\n")) + versions.discard("name, value") + ver = versions.pop() + ver = ver.replace('"Driver version", ', "").replace('"', "") + + return {"ROCM Driver Version": ver} + except subprocess.SubprocessError: + return {"ROCM Driver Version": "Not Available"} + else: + return {"CUDA Driver Version": "Not Available"} + + +def get_gpu_topology(): + """ + Get GPU topology information. + """ + if is_cuda_v2(): + try: + result = subprocess.run( + ["nvidia-smi", "topo", "-m"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=True, + ) + return "\n" + result.stdout if result.returncode == 0 else None + except subprocess.SubprocessError: + return None + elif is_hip(): + try: + result = subprocess.run( + ["rocm-smi", "--showtopotype"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + check=True, + ) + return "\n" + result.stdout if result.returncode == 0 else None + except subprocess.SubprocessError: + return None + else: + return None + + +def get_hypervisor_vendor(): + try: + output = subprocess.check_output(["lscpu"], text=True) + for line in output.split("\n"): + if "Hypervisor vendor:" in line: + return line.split(":")[1].strip() + return None + except: + return None + + +def check_env(): + """ + Check and print environment information. + """ + env_info = OrderedDict() + env_info["Python"] = sys.version.replace("\n", "") + env_info.update(get_cuda_info()) + env_info["PyTorch"] = torch.__version__ + env_info.update(get_package_versions(PACKAGE_LIST)) + + gpu_topo = get_gpu_topology() + if gpu_topo: + if is_cuda_v2(): + env_info["NVIDIA Topology"] = gpu_topo + elif is_hip(): + env_info["AMD Topology"] = gpu_topo + + hypervisor_vendor = get_hypervisor_vendor() + if hypervisor_vendor: + env_info["Hypervisor vendor"] = hypervisor_vendor + + ulimit_soft, _ = resource.getrlimit(resource.RLIMIT_NOFILE) + env_info["ulimit soft"] = ulimit_soft + + for k, v in env_info.items(): + print(f"{k}: {v}") + + +if __name__ == "__main__": + check_env() diff --git a/sglang/python/sglang/global_config.py b/sglang/python/sglang/global_config.py new file mode 100644 index 0000000..c331e64 --- /dev/null +++ b/sglang/python/sglang/global_config.py @@ -0,0 +1,49 @@ +"""Global configurations""" + +import os + + +class GlobalConfig: + """ + Store some global constants. + + See also python/sglang/srt/managers/schedule_batch.py::global_server_args_dict, which stores + many global runtime arguments as well. + """ + + def __init__(self): + # Verbosity level + # 0: do not output anything + # 2: output final text after every run + self.verbosity = 0 + + # Default backend of the language + self.default_backend = None + + # Runtime constants: New generation token ratio estimation + self.default_init_new_token_ratio = float( + os.environ.get("SGLANG_INIT_NEW_TOKEN_RATIO", 0.7) + ) + self.default_min_new_token_ratio_factor = float( + os.environ.get("SGLANG_MIN_NEW_TOKEN_RATIO_FACTOR", 0.14) + ) + self.default_new_token_ratio_decay_steps = float( + os.environ.get("SGLANG_NEW_TOKEN_RATIO_DECAY_STEPS", 600) + ) + + # Runtime constants: others + self.retract_decode_steps = 20 + self.flashinfer_workspace_size = os.environ.get( + "FLASHINFER_WORKSPACE_SIZE", 384 * 1024 * 1024 + ) + + # Output tokenization configs + self.skip_special_tokens_in_output = True + self.spaces_between_special_tokens_in_out = True + + # Language frontend interpreter optimization configs + self.enable_precache_with_tracing = True + self.enable_parallel_encoding = True + + +global_config = GlobalConfig() diff --git a/sglang/python/sglang/lang/__init__.py b/sglang/python/sglang/lang/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sglang/python/sglang/lang/backend/__init__.py b/sglang/python/sglang/lang/backend/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/sglang/python/sglang/lang/backend/anthropic.py b/sglang/python/sglang/lang/backend/anthropic.py new file mode 100644 index 0000000..5a36bd9 --- /dev/null +++ b/sglang/python/sglang/lang/backend/anthropic.py @@ -0,0 +1,77 @@ +from typing import List, Optional, Union + +import numpy as np + +from sglang.lang.backend.base_backend import BaseBackend +from sglang.lang.chat_template import get_chat_template +from sglang.lang.interpreter import StreamExecutor +from sglang.lang.ir import SglSamplingParams + +try: + import anthropic +except ImportError as e: + anthropic = e + + +class Anthropic(BaseBackend): + def __init__(self, model_name, *args, **kwargs): + super().__init__() + + if isinstance(anthropic, Exception): + raise anthropic + + self.model_name = model_name + self.chat_template = get_chat_template("claude") + self.client = anthropic.Anthropic(*args, **kwargs) + + def get_chat_template(self): + return self.chat_template + + def generate( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + if s.messages_: + messages = s.messages_ + else: + messages = [{"role": "user", "content": s.text_}] + + if messages and messages[0]["role"] == "system": + system = messages.pop(0)["content"] + else: + system = "" + + ret = self.client.messages.create( + model=self.model_name, + system=system, + messages=messages, + **sampling_params.to_anthropic_kwargs(), + ) + comp = ret.content[0].text + + return comp, {} + + def generate_stream( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + if s.messages_: + messages = s.messages_ + else: + messages = [{"role": "user", "content": s.text_}] + + if messages and messages[0]["role"] == "system": + system = messages.pop(0)["content"] + else: + system = "" + + with self.client.messages.stream( + model=self.model_name, + system=system, + messages=messages, + **sampling_params.to_anthropic_kwargs(), + ) as stream: + for text in stream.text_stream: + yield text, {} diff --git a/sglang/python/sglang/lang/backend/base_backend.py b/sglang/python/sglang/lang/backend/base_backend.py new file mode 100644 index 0000000..725c0a9 --- /dev/null +++ b/sglang/python/sglang/lang/backend/base_backend.py @@ -0,0 +1,82 @@ +from typing import Callable, List, Optional, Union + +from sglang.lang.chat_template import get_chat_template +from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod +from sglang.lang.interpreter import StreamExecutor +from sglang.lang.ir import SglSamplingParams + + +class BaseBackend: + def __init__(self) -> None: + self.support_concate_and_append = False + self.chat_template = get_chat_template("default") + + def get_model_name(self): + raise NotImplementedError() + + def get_chat_template(self): + return self.chat_template + + def cache_prefix(self, prefix_str: str): + pass + + def uncache_prefix(self, rid: str): + pass + + def end_request(self, rid: Union[str, List[str]]): + pass + + def begin_program(self, s: StreamExecutor): + pass + + def end_program(self, s: Union[StreamExecutor, List[StreamExecutor]]): + pass + + def commit_lazy_operations(self, s: StreamExecutor): + pass + + def fork_program( + self, + src: StreamExecutor, + dst: List[StreamExecutor], + position_ids_offset: Optional[List[int]] = None, + ): + pass + + def fill_image(self, s: StreamExecutor): + pass + + def generate( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + raise NotImplementedError() + + def generate_stream( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + raise NotImplementedError() + + def select( + self, + s: StreamExecutor, + choices: List[str], + temperature: float, + choices_method: Optional[ChoicesSamplingMethod] = None, + ) -> ChoicesDecision: + raise NotImplementedError() + + def concatenate_and_append(self, src_rids: List[str], dst_rid: str): + raise NotImplementedError() + + def shutdown(self): + pass + + def flush_cache(self): + pass + + def get_server_info(self): + pass diff --git a/sglang/python/sglang/lang/backend/litellm.py b/sglang/python/sglang/lang/backend/litellm.py new file mode 100644 index 0000000..5803b54 --- /dev/null +++ b/sglang/python/sglang/lang/backend/litellm.py @@ -0,0 +1,90 @@ +from typing import Mapping, Optional + +from sglang.lang.backend.base_backend import BaseBackend +from sglang.lang.chat_template import get_chat_template_by_model_path +from sglang.lang.interpreter import StreamExecutor +from sglang.lang.ir import SglSamplingParams + +try: + import litellm +except ImportError as e: + litellm = e + litellm.num_retries = 1 + + +class LiteLLM(BaseBackend): + def __init__( + self, + model_name, + chat_template=None, + api_key=None, + organization: Optional[str] = None, + base_url: Optional[str] = None, + timeout: Optional[float] = 600, + max_retries: Optional[int] = litellm.num_retries, + default_headers: Optional[Mapping[str, str]] = None, + ): + super().__init__() + + if isinstance(litellm, Exception): + raise litellm + + self.model_name = model_name + + self.chat_template = chat_template or get_chat_template_by_model_path( + model_name + ) + + self.client_params = { + "api_key": api_key, + "organization": organization, + "base_url": base_url, + "timeout": timeout, + "max_retries": max_retries, + "default_headers": default_headers, + } + + def get_chat_template(self): + return self.chat_template + + def generate( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + if s.messages_: + messages = s.messages_ + else: + messages = [{"role": "user", "content": s.text_}] + + ret = litellm.completion( + model=self.model_name, + messages=messages, + **self.client_params, + **sampling_params.to_litellm_kwargs(), + ) + comp = ret.choices[0].message.content + + return comp, {} + + def generate_stream( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + if s.messages_: + messages = s.messages_ + else: + messages = [{"role": "user", "content": s.text_}] + + ret = litellm.completion( + model=self.model_name, + messages=messages, + stream=True, + **self.client_params, + **sampling_params.to_litellm_kwargs(), + ) + for chunk in ret: + text = chunk.choices[0].delta.content + if text is not None: + yield text, {} diff --git a/sglang/python/sglang/lang/backend/openai.py b/sglang/python/sglang/lang/backend/openai.py new file mode 100644 index 0000000..c6e31f7 --- /dev/null +++ b/sglang/python/sglang/lang/backend/openai.py @@ -0,0 +1,471 @@ +import dataclasses +import logging +import time +import warnings +from typing import Callable, List, Optional, Union + +import numpy as np + +from sglang.lang.backend.base_backend import BaseBackend +from sglang.lang.chat_template import ChatTemplate, get_chat_template_by_model_path +from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod +from sglang.lang.interpreter import StreamExecutor +from sglang.lang.ir import SglSamplingParams + +try: + import openai + import tiktoken +except ImportError as e: + openai = tiktoken = e + + +logger = logging.getLogger(__name__) + + +def create_logit_bias_int(tokenizer): + """Get logit bias for integer numbers.""" + int_token_ids = [] + + tokens = tokenizer._mergeable_ranks + for token, token_id in tokens.items(): + s = tokenizer.decode([token_id]) + if all([c.isdigit() for c in s]) or s in [" "]: + int_token_ids.append(token_id) + if len(int_token_ids) >= 300: # OpenAI API limit + break + special_tokens = tokenizer._special_tokens + mask = {t: 100 for t in int_token_ids[:299]} + mask[special_tokens["<|endoftext|>"]] = 100 + return mask + + +INSTRUCT_MODEL_NAMES = [ + "gpt-3.5-turbo-instruct", +] + + +@dataclasses.dataclass +class TokenUsage: + prompt_tokens: int + completion_tokens: int + + def reset(self): + self.prompt_tokens = self.completion_tokens = 0 + + +class OpenAI(BaseBackend): + def __init__( + self, + model_name: str, + is_chat_model: Optional[bool] = None, + chat_template: Optional[ChatTemplate] = None, + is_azure: bool = False, + *args, + **kwargs, + ): + super().__init__() + + if isinstance(openai, Exception): + raise openai + + if is_azure: + self.client = openai.AzureOpenAI(*args, **kwargs) + else: + self.client = openai.OpenAI(*args, **kwargs) + + self.model_name = model_name + try: + self.tokenizer = tiktoken.encoding_for_model(model_name) + except KeyError: + self.tokenizer = tiktoken.get_encoding("cl100k_base") + self.logit_bias_int = create_logit_bias_int(self.tokenizer) + + self.chat_template = chat_template or get_chat_template_by_model_path( + model_name + ) + + if is_chat_model is not None: + self.is_chat_model = is_chat_model + else: + if model_name in INSTRUCT_MODEL_NAMES: + self.is_chat_model = False + else: + self.is_chat_model = True + + self.chat_prefix = self.chat_template.role_prefix_and_suffix["assistant"][0] + + # Usage + self.token_usage = TokenUsage(0, 0) + + # API speculative execution + # TODO(ying): This does not support multi-threading (run_batch) + self.spec_kwargs = {} + self.spec_format = [] + self.spec_max_num_tries = 3 + + def get_chat_template(self): + return self.chat_template + + def _prepare_spec_execution( + self, + sampling_params: SglSamplingParams, + num_api_spec_tokens: int, + spec_var_name: str, + ): + if "max_tokens" not in self.spec_kwargs: + self.spec_kwargs["max_tokens"] = num_api_spec_tokens + else: + assert self.spec_kwargs["max_tokens"] == num_api_spec_tokens + + params = sampling_params.to_openai_kwargs() + for key, value in params.items(): + if key in ["stop"]: + continue + if key in ["max_tokens"]: + warnings.warn( + "The parameter max_tokens will be overwritten by speculated number of tokens." + ) + continue + if key not in self.spec_kwargs: + self.spec_kwargs[key] = value + else: + assert ( + value == self.spec_kwargs[key] + ), "sampling parameters should be consistent if turn on api speculative execution." + self.spec_format.append( + {"text": "", "stop": params["stop"], "name": spec_var_name} + ) + return "", {} + + def generate( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + spec_var_name: str = None, + ): + if sampling_params.dtype is None: + if self.is_chat_model: + if s.num_api_spec_tokens is None: + if not s.text_.endswith(self.chat_prefix): + raise RuntimeError( + "This use case is not supported if api speculative execution is off. " + "For OpenAI chat models, sgl.gen must be right after sgl.assistant. " + "Example of adding api speculative execution: @function(num_api_spec_tokens=128)." + ) + prompt = s.messages_ + else: + return self._prepare_spec_execution( + sampling_params, s.num_api_spec_tokens, spec_var_name + ) + else: + prompt = s.text_ + + kwargs = sampling_params.to_openai_kwargs() + if self.model_name.startswith("o1") or self.model_name.startswith("o3"): + kwargs.pop("max_tokens", None) + else: + kwargs.pop("max_completion_tokens", None) + + comp = openai_completion( + client=self.client, + token_usage=self.token_usage, + is_chat=self.is_chat_model, + model=self.model_name, + prompt=prompt, + **kwargs, + ) + # Keep the returned list (or string) as is. + elif sampling_params.dtype in [str, "str", "string"]: + assert ( + not self.is_chat_model + ), "constrained type not supported on chat model" + kwargs = sampling_params.to_openai_kwargs() + kwargs.pop("stop") + comp = openai_completion( + client=self.client, + token_usage=self.token_usage, + is_chat=self.is_chat_model, + model=self.model_name, + prompt=s.text_ + '"', + stop='"', + **kwargs, + ) + # Wrap each element in quotes if we have a list. + if isinstance(comp, list): + comp = ['"' + x + '"' for x in comp] + else: + comp = '"' + comp + '"' + elif sampling_params.dtype in [int, "int"]: + assert ( + not self.is_chat_model + ), "constrained type not supported on chat model" + kwargs = sampling_params.to_openai_kwargs() + kwargs.pop("stop") + comp = openai_completion( + client=self.client, + token_usage=self.token_usage, + is_chat=self.is_chat_model, + model=self.model_name, + prompt=s.text_, + logit_bias=self.logit_bias_int, + stop=[" "], + **kwargs, + ) + # Leave as a list if that's what is returned. + else: + raise ValueError(f"Unknown dtype: {sampling_params.dtype}") + + return comp, {} + + def spec_fill(self, value: str): + assert self.is_chat_model + self.spec_format.append({"text": value, "stop": None, "name": None}) + + def spec_pattern_match(self, comp): + for i, term in enumerate(self.spec_format): + text = term["text"] + if text != "": + if comp.startswith(text): + comp = comp[len(text) :] + else: + return False + else: + pos = comp.find(term["stop"]) + if pos != -1: + term["text"] = comp[:pos] + comp = comp[pos:] + else: + if i == len(self.spec_format) - 1: + term["text"] = comp + else: + return False + return True + + def role_end_generate( + self, + s: StreamExecutor, + ): + if s.num_api_spec_tokens is None or not s.text_.endswith(self.chat_prefix): + return + + comp = "" + if not all(x["name"] is None for x in self.spec_format): + # TODO(ying): throw errors or warnings + for i in range(self.spec_max_num_tries): + comp = openai_completion( + client=self.client, + token_usage=self.token_usage, + is_chat=self.is_chat_model, + model=self.model_name, + prompt=s.messages_, + **self.spec_kwargs, + ) + # Use a string for pattern matching. + comp_for_match = comp[0] if isinstance(comp, list) else comp + if self.spec_pattern_match(comp_for_match): + break + + for term in self.spec_format: + s.text_ += term["text"] + name = term["name"] + if name is not None: + s.variables[name] = term["text"] + s.meta_info[name] = {} + s.variable_event[name].set() + + self.spec_kwargs = {} + self.spec_format = [] + + def generate_stream( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + if sampling_params.dtype is None: + if self.is_chat_model: + if not s.text_.endswith(self.chat_prefix): + raise RuntimeError( + "This use case is not supported. " + "For OpenAI chat models, sgl.gen must be right after sgl.assistant" + ) + prompt = s.messages_ + else: + prompt = s.text_ + + kwargs = sampling_params.to_openai_kwargs() + generator = openai_completion_stream( + client=self.client, + token_usage=self.token_usage, + is_chat=self.is_chat_model, + model=self.model_name, + prompt=prompt, + **kwargs, + ) + return generator + else: + raise ValueError(f"Unknown dtype: {sampling_params.dtype}") + + def select( + self, + s: StreamExecutor, + choices: List[str], + temperature: float, + choices_method: ChoicesSamplingMethod, + ) -> ChoicesDecision: + """Note: `choices_method` is not used by the OpenAI backend.""" + if self.is_chat_model: + raise NotImplementedError( + "select/choices is not supported for chat models. " + "Please try to use a non-chat model such as gpt-3.5-turbo-instruct" + ) + + n_choices = len(choices) + token_ids = [self.tokenizer.encode(x) for x in choices] + scores = [0] * n_choices + valid = [len(x) > 0 for x in token_ids] + prompt_tokens = self.tokenizer.encode(s.text_) + + max_len = max([len(x) for x in token_ids]) + for step in range(max_len): + # Build logit bias + logit_bias = {} + for i in range(n_choices): + if valid[i]: + logit_bias[token_ids[i][step]] = 100 + + # Call API + ret = self.client.completions.create( + model=self.model_name, + prompt=prompt_tokens, + logit_bias=logit_bias, + max_tokens=1, + temperature=temperature, + ) + ret_str = ret.choices[0].text + ret_token = self.tokenizer.encode(ret_str)[0] + self.token_usage.prompt_tokens += ret.usage.prompt_tokens + self.token_usage.completion_tokens = ret.usage.completion_tokens + + # TODO: + # 1. return logits as the scores + # 2. compute logits of the full choice + # 3. consider chunk-based decoding + + # Update valid + hit = False + for i in range(n_choices): + if valid[i]: + if step == len(token_ids[i]) - 1: + valid[i] = False + + if ret_token == token_ids[i][step]: + scores[i] += 1 + hit = True + else: + valid[i] = False + assert hit + + if np.sum(valid) <= 1: + break + + prompt_tokens.append(ret_token) + + return ChoicesDecision( + decision=choices[np.argmax(scores)], + meta_info={"scores": scores}, + ) + + +def openai_completion( + client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs +) -> Union[str, List[str]]: + # if "ebnf" is in kwargs, warn and remove + if "ebnf" in kwargs: + warnings.warn("EBNF is not officially supported by OpenAI endpoints. Ignoring.") + del kwargs["ebnf"] + + for attempt in range(retries): + try: + if is_chat: + if "stop" in kwargs and kwargs["stop"] is None: + kwargs.pop("stop") + ret = client.chat.completions.create(messages=prompt, **kwargs) + if len(ret.choices) == 1: + comp = ret.choices[0].message.content + else: + comp = [c.message.content for c in ret.choices] + else: + ret = client.completions.create(prompt=prompt, **kwargs) + if isinstance(prompt, (list, tuple)): + comp = [c.text for c in ret.choices] + else: + comp = ret.choices[0].text + if len(ret.choices) > 1: + comp = [c.text for c in ret.choices] + + token_usage.prompt_tokens += ret.usage.prompt_tokens + token_usage.completion_tokens += ret.usage.completion_tokens + break + except (openai.APIError, openai.APIConnectionError, openai.RateLimitError) as e: + logger.error(f"OpenAI Error: {e}. Waiting 5 seconds...") + time.sleep(5) + if attempt == retries - 1: + raise e + except Exception as e: + logger.error(f"RuntimeError {e}.") + raise e + + return comp + + +def openai_completion_stream( + client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs +): + # if "ebnf" is in kwargs, warn and remove + if "ebnf" in kwargs: + warnings.warn("EBNF is not officially supported by OpenAI endpoints. Ignoring.") + del kwargs["ebnf"] + + for attempt in range(retries): + try: + if is_chat: + if "stop" in kwargs and kwargs["stop"] is None: + kwargs.pop("stop") + generator = client.chat.completions.create( + messages=prompt, + stream=True, + stream_options={"include_usage": True}, + **kwargs, + ) + for ret in generator: + if len(ret.choices) == 0: + continue + try: + content = ret.choices[0].delta.content + except IndexError: + content = None + yield content or "", {} + else: + generator = client.completions.create( + prompt=prompt, + stream=True, + stream_options={"include_usage": True}, + **kwargs, + ) + for ret in generator: + if len(ret.choices) == 0: + continue + content = ret.choices[0].text + yield content or "", {} + + token_usage.prompt_tokens += ret.usage.prompt_tokens + token_usage.completion_tokens += ret.usage.completion_tokens + break + except (openai.APIError, openai.APIConnectionError, openai.RateLimitError) as e: + logger.error(f"OpenAI Error: {e}. Waiting 5 seconds...") + time.sleep(5) + if attempt == retries - 1: + raise e + except Exception as e: + logger.error(f"RuntimeError {e}.") + raise e diff --git a/sglang/python/sglang/lang/backend/runtime_endpoint.py b/sglang/python/sglang/lang/backend/runtime_endpoint.py new file mode 100644 index 0000000..1cd3d52 --- /dev/null +++ b/sglang/python/sglang/lang/backend/runtime_endpoint.py @@ -0,0 +1,500 @@ +import atexit +import json +import multiprocessing +import warnings +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.interpreter import StreamExecutor +from sglang.lang.ir import ( + REGEX_BOOL, + REGEX_FLOAT, + REGEX_INT, + REGEX_STR, + SglSamplingParams, +) +from sglang.utils import http_request + + +class RuntimeEndpoint(BaseBackend): + def __init__( + self, + base_url: str, + api_key: Optional[str] = None, + verify: Optional[str] = None, + chat_template_name: Optional[str] = None, + ): + super().__init__() + self.support_concate_and_append = True + + self.base_url = base_url + self.api_key = api_key + self.verify = verify + + res = http_request( + self.base_url + "/get_model_info", + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + self.model_info = res.json() + + if chat_template_name: + self.chat_template = get_chat_template(chat_template_name) + else: + self.chat_template = get_chat_template_by_model_path( + self.model_info["model_path"] + ) + + def get_model_name(self): + return self.model_info["model_path"] + + def flush_cache(self): + res = http_request( + self.base_url + "/flush_cache", + api_key=self.api_key, + verify=self.verify, + method="POST", + ) + self._assert_success(res) + + def get_server_info(self): + res = http_request( + self.base_url + "/get_server_info", + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + return res.json() + + def get_chat_template(self): + return self.chat_template + + def cache_prefix(self, prefix_str: str): + res = http_request( + self.base_url + "/generate", + json={"text": prefix_str, "sampling_params": {"max_new_tokens": 0}}, + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + + def commit_lazy_operations(self, s: StreamExecutor): + data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}} + self._add_images(s, data) + res = http_request( + self.base_url + "/generate", + json=data, + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + + def fill_image(self, s: StreamExecutor): + data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}} + self._add_images(s, data) + res = http_request( + self.base_url + "/generate", + json=data, + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + + def _handle_dtype_to_regex(self, sampling_params: SglSamplingParams): + if sampling_params.dtype is None: + return + + if sampling_params.stop == (): + sampling_params.stop = [] + + dtype_regex = None + if sampling_params.dtype in ["int", int]: + + dtype_regex = REGEX_INT + sampling_params.stop.extend([" ", "\n"]) + elif sampling_params.dtype in ["float", float]: + + dtype_regex = REGEX_FLOAT + sampling_params.stop.extend([" ", "\n"]) + elif sampling_params.dtype in ["str", str]: + + dtype_regex = REGEX_STR + elif sampling_params.dtype in ["bool", bool]: + + dtype_regex = REGEX_BOOL + else: + raise RuntimeError(f"Invalid dtype: {sampling_params.dtype}") + + if dtype_regex is not None and sampling_params.regex is not None: + warnings.warn( + f"Both dtype and regex are set. Only dtype will be used. dtype: {sampling_params.dtype}, regex: {sampling_params.regex}" + ) + + sampling_params.regex = dtype_regex + + def generate( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + self._handle_dtype_to_regex(sampling_params) + data = { + "text": s.text_, + "sampling_params": { + "skip_special_tokens": global_config.skip_special_tokens_in_output, + "spaces_between_special_tokens": global_config.spaces_between_special_tokens_in_out, + **sampling_params.to_srt_kwargs(), + }, + } + + for item in [ + "return_logprob", + "logprob_start_len", + "top_logprobs_num", + "return_text_in_logprobs", + ]: + value = getattr(sampling_params, item, None) + if value is not None: + data[item] = value + + self._add_images(s, data) + + res = http_request( + self.base_url + "/generate", + json=data, + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + + obj = res.json() + comp = obj["text"] + return comp, obj["meta_info"] + + def generate_stream( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + self._handle_dtype_to_regex(sampling_params) + + data = { + "text": s.text_, + "sampling_params": { + "skip_special_tokens": global_config.skip_special_tokens_in_output, + "spaces_between_special_tokens": global_config.spaces_between_special_tokens_in_out, + **sampling_params.to_srt_kwargs(), + }, + } + + for item in [ + "return_logprob", + "logprob_start_len", + "top_logprobs_num", + "return_text_in_logprobs", + ]: + value = getattr(sampling_params, item, None) + if value is not None: + data[item] = value + + data["stream"] = True + self._add_images(s, data) + + res = http_request( + self.base_url + "/generate", + json=data, + stream=True, + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + pos = 0 + + for chunk in res.iter_lines(decode_unicode=False): + chunk = chunk.decode("utf-8") + if chunk and chunk.startswith("data:"): + if chunk == "data: [DONE]": + break + data = json.loads(chunk[5:].strip("\n")) + chunk_text = data["text"][pos:] + meta_info = data["meta_info"] + pos += len(chunk_text) + yield chunk_text, meta_info + + def select( + self, + s: StreamExecutor, + choices: List[str], + temperature: float, + choices_method: ChoicesSamplingMethod, + ) -> ChoicesDecision: + assert temperature <= 1e-5 + + # Cache common prefix + data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}} + obj = self._generate_http_request(s, data) + prompt_len = obj["meta_info"]["prompt_tokens"] + logprob_start_len = max(prompt_len - 2, 0) # For token healing + + # Compute logprob + data = { + "text": [s.text_ + c for c in choices], + "sampling_params": { + "max_new_tokens": 0, + "temperature": 0, + }, + "return_logprob": True, + "return_text_in_logprobs": True, + "logprob_start_len": logprob_start_len, + } + obj = self._generate_http_request(s, data) + + input_token_logprobs = [r["meta_info"]["input_token_logprobs"] for r in obj] + output_token_logprobs = [r["meta_info"]["output_token_logprobs"] for r in obj] + normalized_prompt_logprobs = [ + compute_normalized_prompt_logprobs(r["meta_info"]["input_token_logprobs"]) + for r in obj + ] + + # Remove extra token if no token healing occurred + for i in range(len(input_token_logprobs)): + healed_token_str = input_token_logprobs[i][0][-1] + if s.text_.endswith(healed_token_str): + healed_token_logprob = input_token_logprobs[i][0][0] + normalized_prompt_logprobs[i] = ( + normalized_prompt_logprobs[i] * len(input_token_logprobs[i]) + - healed_token_logprob + ) / (len(input_token_logprobs[i]) - 1) + input_token_logprobs[i] = input_token_logprobs[i][1:] + + # Compute unconditional logprobs if required + if choices_method.requires_unconditional_logprobs: + input_ids = [[el[1] for el in subl] for subl in input_token_logprobs] + data = { + "input_ids": input_ids, + "sampling_params": {"max_new_tokens": 0}, + "return_logprob": True, + } + obj = self._generate_http_request(s, data) + unconditional_token_logprobs = [ + r["meta_info"]["input_token_logprobs"] for r in obj + ] + else: + unconditional_token_logprobs = None + + return choices_method( + choices=choices, + normalized_prompt_logprobs=normalized_prompt_logprobs, + input_token_logprobs=input_token_logprobs, + output_token_logprobs=output_token_logprobs, + unconditional_token_logprobs=unconditional_token_logprobs, + ) + + def concatenate_and_append(self, src_rids: List[str], dst_rid: str): + res = http_request( + self.base_url + "/concate_and_append_request", + json={"src_rids": src_rids, "dst_rid": dst_rid}, + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + + def _generate_http_request(self, s: StreamExecutor, data): + self._add_images(s, data) + res = http_request( + self.base_url + "/generate", + json=data, + api_key=self.api_key, + verify=self.verify, + ) + self._assert_success(res) + return res.json() + + def _add_images(self, s: StreamExecutor, data): + if s.images_: + assert len(s.images_) == 1, "Only support one image." + data["image_data"] = s.images_[0][1] + + def _assert_success(self, res): + if res.status_code != 200: + raise RuntimeError(res.json()) + + +def compute_normalized_prompt_logprobs(input_logprobs): + values = [x[0] for x in input_logprobs if x[0]] + return sum(values) / len(values) + + +class Runtime: + """ + A wrapper for the HTTP server. + This is used for launching the server in a python program without + using the command line interface. + + It is mainly used for the frontend language. + You should use the Engine class if you want to do normal offline processing without the frontend language. + """ + + def __init__( + self, + log_level: str = "error", + *args, + **kwargs, + ): + """See the arguments in server_args.py::ServerArgs""" + # We delay the import of any `sglang.srt` components in `sglang.lang`, so users can run + # client code without installing SRT server and its dependency if they want. + from sglang.srt.entrypoints.http_server import launch_server + from sglang.srt.server_args import ServerArgs + from sglang.srt.utils import is_port_available + + self.server_args = ServerArgs(*args, log_level=log_level, **kwargs) + + # Pre-allocate ports + for port in range(self.server_args.port, 40000): + if is_port_available(port): + break + self.server_args.port = port + + self.url = self.server_args.url() + self.generate_url = self.url + "/generate" + + # NOTE: We store pid instead of proc to fix some issues during __delete__ + self.pid = None + pipe_reader, pipe_writer = multiprocessing.Pipe(duplex=False) + + proc = multiprocessing.Process( + target=launch_server, + args=(self.server_args, pipe_writer), + ) + proc.start() + pipe_writer.close() + self.pid = proc.pid + + # Before python program terminates, call shutdown implicitly. Therefore, users don't have to explicitly call .shutdown() + atexit.register(self.shutdown) + + # TODO: remove this pipe_writer mechanism and use `/health_generate` instead. + try: + init_state = pipe_reader.recv() + except EOFError: + init_state = "" + + if init_state != "ready": + self.shutdown() + raise RuntimeError( + "Initialization failed. Please see the error messages above." + ) + + self.endpoint = RuntimeEndpoint(self.url) + + def shutdown(self): + from sglang.srt.utils import kill_process_tree + + if self.pid is not None: + kill_process_tree(self.pid) + self.pid = None + + def cache_prefix(self, prefix: str): + self.endpoint.cache_prefix(prefix) + + def get_tokenizer(self): + from sglang.srt.hf_transformers_utils import get_tokenizer + + return get_tokenizer( + self.server_args.tokenizer_path, + tokenizer_mode=self.server_args.tokenizer_mode, + trust_remote_code=self.server_args.trust_remote_code, + revision=self.server_args.revision, + ) + + async def async_generate( + self, + prompt: str, + sampling_params: Optional[Dict] = None, + ): + if self.server_args.skip_tokenizer_init: + json_data = { + "input_ids": prompt, + "sampling_params": sampling_params, + "stream": True, + } + else: + json_data = { + "text": prompt, + "sampling_params": sampling_params, + "stream": True, + } + pos = 0 + + timeout = aiohttp.ClientTimeout(total=3 * 3600) + async with aiohttp.ClientSession(timeout=timeout, trust_env=True) as session: + async with session.post(self.generate_url, json=json_data) as response: + async for chunk, _ in response.content.iter_chunks(): + chunk = chunk.decode("utf-8") + if chunk and chunk.startswith("data:"): + if chunk == "data: [DONE]\n\n": + break + data = json.loads(chunk[5:].strip("\n")) + if "text" in data: + cur = data["text"][pos:] + if cur: + yield cur + pos += len(cur) + else: + yield data + + add_request = async_generate + + def generate( + self, + prompt: Union[str, List[str]], + sampling_params: Optional[Dict] = None, + return_logprob: Optional[Union[List[bool], bool]] = False, + logprob_start_len: Optional[Union[List[int], int]] = None, + top_logprobs_num: Optional[Union[List[int], int]] = None, + lora_path: Optional[List[Optional[str]]] = None, + ): + json_data = { + "text": prompt, + "sampling_params": sampling_params, + "return_logprob": return_logprob, + "logprob_start_len": logprob_start_len, + "top_logprobs_num": top_logprobs_num, + "lora_path": lora_path, + } + assert not isinstance(lora_path, list) or len(lora_path) == len(prompt) + response = requests.post( + self.url + "/generate", + json=json_data, + ) + return json.dumps(response.json()) + + def encode( + self, + prompt: Union[str, List[str], List[Dict], List[List[Dict]]], + ): + json_data = {"text": prompt} + response = requests.post(self.url + "/encode", json=json_data) + return json.dumps(response.json()) + + async def get_server_info(self): + async with aiohttp.ClientSession() as session: + async with session.get(f"{self.url}/get_server_info") as response: + if response.status == 200: + return await response.json() + else: + error_data = await response.json() + raise RuntimeError( + f"Failed to get server info. {error_data['error']['message']}" + ) + + def __del__(self): + self.shutdown() diff --git a/sglang/python/sglang/lang/backend/vertexai.py b/sglang/python/sglang/lang/backend/vertexai.py new file mode 100644 index 0000000..c27733b --- /dev/null +++ b/sglang/python/sglang/lang/backend/vertexai.py @@ -0,0 +1,149 @@ +import os +import warnings +from typing import Optional + +from sglang.lang.backend.base_backend import BaseBackend +from sglang.lang.chat_template import get_chat_template +from sglang.lang.interpreter import StreamExecutor +from sglang.lang.ir import SglSamplingParams + +try: + import vertexai + from vertexai.preview.generative_models import ( + GenerationConfig, + GenerativeModel, + Image, + ) +except ImportError as e: + GenerativeModel = e + + +class VertexAI(BaseBackend): + def __init__(self, model_name, safety_settings=None): + super().__init__() + + if isinstance(GenerativeModel, Exception): + raise GenerativeModel + + project_id = os.environ["GCP_PROJECT_ID"] + location = os.environ.get("GCP_LOCATION") + vertexai.init(project=project_id, location=location) + + self.model_name = model_name + self.chat_template = get_chat_template("default") + self.safety_settings = safety_settings + + def get_chat_template(self): + return self.chat_template + + def generate( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + if s.messages_: + prompt = self.messages_to_vertexai_input(s.messages_) + else: + # single-turn + prompt = ( + self.text_to_vertexai_input(s.text_, s.cur_images) + if s.cur_images + else s.text_ + ) + ret = GenerativeModel(self.model_name).generate_content( + prompt, + generation_config=GenerationConfig(**sampling_params.to_vertexai_kwargs()), + safety_settings=self.safety_settings, + ) + + comp = ret.text + + return comp, {} + + def generate_stream( + self, + s: StreamExecutor, + sampling_params: SglSamplingParams, + ): + if s.messages_: + prompt = self.messages_to_vertexai_input(s.messages_) + else: + # single-turn + prompt = ( + self.text_to_vertexai_input(s.text_, s.cur_images) + if s.cur_images + else s.text_ + ) + generator = GenerativeModel(self.model_name).generate_content( + prompt, + stream=True, + generation_config=GenerationConfig(**sampling_params.to_vertexai_kwargs()), + safety_settings=self.safety_settings, + ) + for ret in generator: + yield ret.text, {} + + def text_to_vertexai_input(self, text, images): + input = [] + # split with image token + text_segs = text.split(self.chat_template.image_token) + for image_path, image_base64_data in images: + text_seg = text_segs.pop(0) + if text_seg != "": + input.append(text_seg) + input.append(Image.from_bytes(image_base64_data)) + text_seg = text_segs.pop(0) + if text_seg != "": + input.append(text_seg) + return input + + def messages_to_vertexai_input(self, messages): + vertexai_message = [] + # from openai message format to vertexai message format + for msg in messages: + if isinstance(msg["content"], str): + text = msg["content"] + else: + text = msg["content"][0]["text"] + + if msg["role"] == "system": + warnings.warn("Warning: system prompt is not supported in VertexAI.") + vertexai_message.append( + { + "role": "user", + "parts": [{"text": "System prompt: " + text}], + } + ) + vertexai_message.append( + { + "role": "model", + "parts": [{"text": "Understood."}], + } + ) + continue + if msg["role"] == "user": + vertexai_msg = { + "role": "user", + "parts": [{"text": text}], + } + elif msg["role"] == "assistant": + vertexai_msg = { + "role": "model", + "parts": [{"text": text}], + } + + # images + if isinstance(msg["content"], list) and len(msg["content"]) > 1: + for image in msg["content"][1:]: + assert image["type"] == "image_url" + vertexai_msg["parts"].append( + { + "inline_data": { + "data": image["image_url"]["url"].split(",")[1], + "mime_type": "image/jpeg", + } + } + ) + + vertexai_message.append(vertexai_msg) + return vertexai_message diff --git a/sglang/python/sglang/lang/chat_template.py b/sglang/python/sglang/lang/chat_template.py new file mode 100644 index 0000000..8554d28 --- /dev/null +++ b/sglang/python/sglang/lang/chat_template.py @@ -0,0 +1,555 @@ +from dataclasses import dataclass +from enum import Enum, auto +from typing import Callable, Dict, List, Tuple + + +class ChatTemplateStyle(Enum): + PLAIN = auto() + LLAMA2 = auto() + + +@dataclass +class ChatTemplate: + name: str + default_system_prompt: str + role_prefix_and_suffix: Dict[str, Tuple[str, str]] + stop_str: List[str] = () + image_token: str = "" + audio_token: str = ""), + }, + image_token=" \n", + ) +) + +register_chat_template( + ChatTemplate( + name="llama-2-chat", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ("<>\n", "\n<>\n\n"), + "user": ("[INST] ", " [/INST]"), + "assistant": ("", " "), + }, + style=ChatTemplateStyle.LLAMA2, + ) +) + +register_chat_template( + ChatTemplate( + name="llama-3-instruct", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ( + "<|start_header_id|>system<|end_header_id|>\n\n", + "<|eot_id|>", + ), + "user": ( + "<|start_header_id|>user<|end_header_id|>\n\n", + "<|eot_id|>", + ), + "assistant": ( + "<|start_header_id|>assistant<|end_header_id|>\n\n", + "<|eot_id|>", + ), + }, + stop_str=("<|eot_id|>",), + image_token="<|image|>", + ) +) + +# https://huggingface.co/openbmb/MiniCPM-V-2_6 +register_chat_template( + ChatTemplate( + name="minicpmv", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ("", " "), + "user": ("user:", " "), + "assistant": ("assistant:", ""), + }, + stop_str=("<|im_end|>", "<|endoftext|>"), + image_token="(./)", + ) +) + +register_chat_template( + ChatTemplate( + name="janus-pro", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ( + "", + "", + ), + "User": ( + "<｜User｜>", + "", + ), + "assistant": ( + "<｜Assistant｜>", + "<｜end▁of▁sentence｜>", + ), + }, + stop_str=("<｜end▁of▁sentence｜>",), + image_token="\n", + ) +) + +# https://huggingface.co/openbmb/MiniCPM-o-2_6 +register_chat_template( + ChatTemplate( + name="minicpmo", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ("", " "), + "user": ("user:", " "), + "assistant": ("assistant:", ""), + }, + stop_str=("<|im_end|>", "<|endoftext|>"), + image_token="(./)", + audio_token="()", + ) +) + +# The difference between "llama-3-instruct-llava" and "llama-3-instruct" is that llava uses a different image_token. +register_chat_template( + ChatTemplate( + name="llama-3-instruct-llava", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ( + "<|start_header_id|>system<|end_header_id|>\n\n", + "<|eot_id|>", + ), + "user": ( + "<|start_header_id|>user<|end_header_id|>\n\n", + "<|eot_id|>", + ), + "assistant": ( + "<|start_header_id|>assistant<|end_header_id|>\n\n", + "<|eot_id|>", + ), + }, + stop_str=("<|eot_id|>",), + image_token="\n", + ) +) + +# Reference: https://modelscope.cn/models/01ai/Yi-1.5-34B-Chat/file/view/master?fileName=tokenizer_config.json&status=1 +register_chat_template( + ChatTemplate( + name="yi-1.5", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ("", ""), + "user": ("<|im_start|>user\n", "<|im_end|>\n<|im_start|>assistant\n"), + "assistant": ("", "<|im_end|>\n"), + }, + style=ChatTemplateStyle.PLAIN, + stop_str=("<|im_end|>",), + ) +) + +# Reference: https://github.com/01-ai/Yi/tree/main/VL#major-difference-with-llava +register_chat_template( + ChatTemplate( + name="yi-vl", + default_system_prompt=( + "This is a chat between an inquisitive human and an AI assistant. Assume the role of the AI assistant. Read all the images carefully, and respond to the human's questions with informative, helpful, detailed and polite answers." + "这是一个好奇的人类和一个人工智能助手之间的对话。假设你扮演这个AI助手的角色。仔细阅读所有的图像，并对人类的问题做出信息丰富、有帮助、详细的和礼貌的回答。" + ), + role_prefix_and_suffix={ + "system": ("", "\n\n"), + "user": ("### Human:", "\n"), + "assistant": ("### Assistant:", "\n"), + }, + image_token=" \n", + ) +) + +register_chat_template( + ChatTemplate( + name="gemma-it", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ("", ""), + "user": ("user\n", "\n"), + "assistant": ("model\n", "\n"), + }, + style=ChatTemplateStyle.PLAIN, + ) +) + +register_chat_template( + ChatTemplate( + name="dbrx-instruct", + default_system_prompt="You are DBRX, created by Databricks. You were last updated in December 2023. You answer questions based on information available up to that point.\nYOU PROVIDE SHORT RESPONSES TO SHORT QUESTIONS OR STATEMENTS, but provide thorough responses to more complex and open-ended questions.\nYou assist with various tasks, from writing to coding (using markdown for code blocks — remember to use ``` with code, JSON, and tables).\n(You do not have real-time data access or code execution capabilities. You avoid stereotyping and provide balanced perspectives on controversial topics. You do not provide song lyrics, poems, or news articles and do not divulge details of your training data.)\nThis is your system prompt, guiding your responses. Do not reference it, just respond to the user. If you find yourself talking about this message, stop. You should be responding appropriately and usually that means not mentioning this.\nYOU DO NOT MENTION ANY OF THIS INFORMATION ABOUT YOURSELF UNLESS THE INFORMATION IS DIRECTLY PERTINENT TO THE USER'S QUERY.", + role_prefix_and_suffix={ + "system": ("<|im_start|>system\n", "<|im_end|>"), + "user": ("\n<|im_start|>user\n", "<|im_end|>"), + "assistant": ("\n<|im_start|>assistant\n", "<|im_end|>"), + }, + stop_str=("<|im_end|>",), + ) +) + +register_chat_template( + ChatTemplate( + name="c4ai-command-r", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ( + "<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>", + "<|END_OF_TURN_TOKEN|>", + ), + "user": ("<|START_OF_TURN_TOKEN|><|USER_TOKEN|>", "<|END_OF_TURN_TOKEN|>"), + "assistant": ( + "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>", + "<|END_OF_TURN_TOKEN|>", + ), + }, + style=ChatTemplateStyle.PLAIN, + ) +) + +register_chat_template( + ChatTemplate( + name="granite-3-instruct", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ( + "<|start_of_role|>system<|end_of_role|>", + "<|end_of_text|>", + ), + "user": ( + "<|start_of_role|>user<|end_of_role|>", + "<|end_of_text|>", + ), + "assistant": ( + "<|start_of_role|>assistant<|end_of_role|>", + "<|end_of_text|>", + ), + }, + stop_str=("<|end_of_text|>",), + ) +) + +register_chat_template( + ChatTemplate( + name="deepseek-v3", + default_system_prompt=None, + role_prefix_and_suffix={ + "system": ( + "", + "", + ), + "user": ( + "<｜User｜>", + "", + ), + "assistant": ( + "<｜Assistant｜>", + "<｜end▁of▁sentence｜>", + ), + }, + stop_str=("<｜end▁of▁sentence｜>",), + ) +) + + +@register_chat_template_matching_function +def match_deepseek(model_path: str): + if ( + "deepseek-v3" in model_path.lower() or "deepseek-r1" in model_path.lower() + ) and "base" not in model_path.lower(): + return get_chat_template("deepseek-v3") + + +@register_chat_template_matching_function +def match_deepseek_janus_pro(model_path: str): + if "janus" in model_path.lower(): + return get_chat_template("janus-pro") + + +@register_chat_template_matching_function +def match_dbrx(model_path: str): + if "dbrx" in model_path.lower() and "instruct" in model_path.lower(): + return get_chat_template("dbrx-instruct") + + +@register_chat_template_matching_function +def match_vicuna(model_path: str): + if "vicuna" in model_path.lower(): + return get_chat_template("vicuna_v1.1") + if "llava-v1.5" in model_path.lower(): + return get_chat_template("vicuna_v1.1") + if "llava-next-video-7b" in model_path.lower(): + return get_chat_template("vicuna_v1.1") + + +@register_chat_template_matching_function +def match_llama2_chat(model_path: str): + model_path = model_path.lower() + if "llama-2" in model_path and "chat" in model_path: + return get_chat_template("llama-2-chat") + if ( + "mistral" in model_path or "mixtral" in model_path + ) and "instruct" in model_path: + return get_chat_template("llama-2-chat") + if "codellama" in model_path and "instruct" in model_path: + return get_chat_template("llama-2-chat") + + +@register_chat_template_matching_function +def match_llama3_instruct(model_path: str): + model_path = model_path.lower() + if "llama-3" in model_path and "instruct" in model_path: + return get_chat_template("llama-3-instruct") + + +@register_chat_template_matching_function +def match_chat_ml(model_path: str): + # import pdb;pdb.set_trace() + model_path = model_path.lower() + if "tinyllama" in model_path: + return get_chat_template("chatml") + # Now the suffix for qwen2 chat model is "instruct" + if "qwen" in model_path and "vl" in model_path: + return get_chat_template("qwen2-vl") + if "qwen" in model_path: + if "vl" in model_path: + return get_chat_template("qwen2-vl") + if ("chat" in model_path or "instruct" in model_path) and ( + "llava" not in model_path + ): + return get_chat_template("qwen") + if ( + "llava-v1.6-34b" in model_path + or "llava-v1.6-yi-34b" in model_path + or "llava-next-video-34b" in model_path + or "llava-onevision-qwen2" in model_path + ): + return get_chat_template("chatml-llava") + + +@register_chat_template_matching_function +def match_chat_yi(model_path: str): + model_path = model_path.lower() + if "yi-vl" in model_path and "llava" not in model_path: + return get_chat_template("yi-vl") + elif "yi-1.5" in model_path and "chat" in model_path: + return get_chat_template("yi-1.5") + + +@register_chat_template_matching_function +def match_gemma_it(model_path: str): + model_path = model_path.lower() + if "gemma" in model_path and "it" in model_path: + return get_chat_template("gemma-it") + + +@register_chat_template_matching_function +def match_openbmb_minicpm(model_path: str): + model_path = model_path.lower() + if "minicpm-v" in model_path: + return get_chat_template("minicpmv") + elif "minicpm-o" in model_path: + return get_chat_template("minicpmo") + + +@register_chat_template_matching_function +def match_c4ai_command_r(model_path: str): + model_path = model_path.lower() + if "c4ai-command-r" in model_path: + return get_chat_template("c4ai-command-r") + + +@register_chat_template_matching_function +def match_granite_instruct(model_path: str): + model_path = model_path.lower() + # When future versions of Granite are released, this code may + # need to be updated. For now, assume that the Granite 3.0 + # template works across the board. + if "granite" in model_path and "instruct" in model_path: + return get_chat_template("granite-3-instruct") + + +@register_chat_template_matching_function +def match_gemma3_instruct(model_path: str): + model_path = model_path.lower() + if "gemma-3" in model_path and "1b" not in model_path: + # gemma-3-1b-it is completion model + return get_chat_template("gemma-it") + + +if __name__ == "__main__": + messages = [ + {"role": "system", "content": None}, # None means default + # {"role": "system", "content": "You are a helpful, respectful and honest assistant."}, + {"role": "user", "content": "Hello!"}, + {"role": "assistant", "content": "Hi!"}, + {"role": "user", "content": "What can you do?"}, + {"role": "assistant", "content": "I can chat with you."}, + ] + + template = get_chat_template("llama-2-chat") + print(template.get_prompt(messages)) diff --git a/sglang/python/sglang/lang/choices.py b/sglang/python/sglang/lang/choices.py new file mode 100644 index 0000000..e52c6b3 --- /dev/null +++ b/sglang/python/sglang/lang/choices.py @@ -0,0 +1,164 @@ +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Any, Dict, List, Optional + +import numpy as np + + +@dataclass +class ChoicesDecision: + decision: str + meta_info: Optional[Dict[str, Any]] = None + + +class ChoicesSamplingMethod(ABC): + + @property + def requires_unconditional_logprobs(self) -> bool: + return False + + @abstractmethod + def __call__( + self, + *, + choices: List[str], + normalized_prompt_logprobs: List[float], + input_token_logprobs: List[List[Any]], + output_token_logprobs: List[List[Any]], + unconditional_token_logprobs: Optional[List[List[Any]]] = None, + ) -> ChoicesDecision: ... + + +class TokenLengthNormalized(ChoicesSamplingMethod): + + def __call__( + self, + *, + choices: List[str], + normalized_prompt_logprobs: List[float], + input_token_logprobs: List[List[Any]], + output_token_logprobs: List[List[Any]], + unconditional_token_logprobs: Optional[List[List[Any]]] = None, + ) -> ChoicesDecision: + """Select the option with the highest token length normalized prompt logprob.""" + best_choice = choices[np.argmax(normalized_prompt_logprobs)] + meta_info = { + "normalized_prompt_logprobs": normalized_prompt_logprobs, + "input_token_logprobs": input_token_logprobs, + "output_token_logprobs": output_token_logprobs, + } + return ChoicesDecision(decision=best_choice, meta_info=meta_info) + + +token_length_normalized = TokenLengthNormalized() + + +class GreedyTokenSelection(ChoicesSamplingMethod): + + def __call__( + self, + *, + choices: List[str], + normalized_prompt_logprobs: List[float], + input_token_logprobs: List[List[Any]], + output_token_logprobs: List[List[Any]], + unconditional_token_logprobs: Optional[List[List[Any]]] = None, + ) -> ChoicesDecision: + """Select the option based on greedy logprob selection. For overlapping options + where one option is a subset of a longer option, extend the shorter option using + its average logprob for comparison against the longer option.""" + + num_options = len(choices) + max_tokens = max(len(option) for option in input_token_logprobs) + logprob_matrix = self._build_logprob_matrix( + input_token_logprobs, max_tokens, num_options + ) + remaining = self._greedy_selection(logprob_matrix, num_options, max_tokens) + + best_choice = choices[remaining[0]] + meta_info = { + "normalized_prompt_logprobs": normalized_prompt_logprobs, + "input_token_logprobs": input_token_logprobs, + "output_token_logprobs": output_token_logprobs, + "greedy_logprob_matrix": logprob_matrix.tolist(), + } + return ChoicesDecision(decision=best_choice, meta_info=meta_info) + + def _build_logprob_matrix(self, input_token_logprobs, max_tokens, num_options): + logprob_matrix = np.zeros((num_options, max_tokens)) + for i, option in enumerate(input_token_logprobs): + actual_logprobs = [token[0] for token in option] + avg_logprob = np.mean(actual_logprobs) + logprob_matrix[i, : len(option)] = actual_logprobs + if len(option) < max_tokens: + logprob_matrix[i, len(option) :] = avg_logprob + return logprob_matrix + + def _greedy_selection(self, logprob_matrix, num_options, max_tokens): + remaining = np.arange(num_options) + for j in range(max_tokens): + max_logprob = np.max(logprob_matrix[remaining, j]) + remaining = remaining[logprob_matrix[remaining, j] == max_logprob] + if len(remaining) == 1: + break + return remaining + + +greedy_token_selection = GreedyTokenSelection() + + +class UnconditionalLikelihoodNormalized(ChoicesSamplingMethod): + + @property + def requires_unconditional_logprobs(self) -> bool: + return True + + def __call__( + self, + *, + choices: List[str], + normalized_prompt_logprobs: List[float], + input_token_logprobs: List[List[Any]], + output_token_logprobs: List[List[Any]], + unconditional_token_logprobs: Optional[List[List[Any]]] = None, + ) -> ChoicesDecision: + """Select the option with the highest average token logprob once normalized by + the unconditional token logprobs. + + The first unconditional token logprob is assumed to be None. If so, it is + replaced with 0 for the purposes of normalization.""" + + if unconditional_token_logprobs is None: + raise ValueError( + "Unconditional token logprobs are required for this method." + ) + + normalized_unconditional_prompt_logprobs = self._normalize_logprobs( + input_token_logprobs, unconditional_token_logprobs + ) + + best_choice = choices[np.argmax(normalized_unconditional_prompt_logprobs)] + meta_info = { + "normalized_prompt_logprobs": normalized_prompt_logprobs, + "input_token_logprobs": input_token_logprobs, + "output_token_logprobs": output_token_logprobs, + "unconditional_token_logprobs": unconditional_token_logprobs, + "normalized_unconditional_prompt_logprobs": normalized_unconditional_prompt_logprobs, + } + return ChoicesDecision(decision=best_choice, meta_info=meta_info) + + def _normalize_logprobs(self, input_token_logprobs, unconditional_token_logprobs): + normalized_unconditional_prompt_logprobs = [] + for inputs, unconditionals in zip( + input_token_logprobs, unconditional_token_logprobs + ): + inputs_logprobs = np.array([token[0] for token in inputs]) + unconditionals_logprobs = np.array([token[0] for token in unconditionals]) + unconditionals_logprobs[0] = unconditionals_logprobs[0] or 0 + normalized_unconditional_prompt_logprobs.append( + float(np.mean(inputs_logprobs - unconditionals_logprobs)) + ) + return normalized_unconditional_prompt_logprobs + + +unconditional_likelihood_normalized = UnconditionalLikelihoodNormalized() diff --git a/sglang/python/sglang/lang/compiler.py b/sglang/python/sglang/lang/compiler.py new file mode 100644 index 0000000..5e1b411 --- /dev/null +++ b/sglang/python/sglang/lang/compiler.py @@ -0,0 +1,237 @@ +import multiprocessing +from concurrent.futures import ThreadPoolExecutor +from queue import Queue +from typing import List, Union + +from sglang.global_config import global_config +from sglang.lang.interpreter import ProgramState, StreamExecutor, cache_program +from sglang.lang.ir import ( + SglArgument, + SglConstantText, + SglExpr, + SglSamplingParams, + SglVariable, +) + + +def compile_func(function, backend): + tracer = function.trace(backend=backend) + compiler = CompiledFunction(tracer, function) + return compiler + + +class CompiledFunction: + def __init__(self, tracer, function): + self.function = function + + self.last_node = CompGraphNode(tracer.last_node) + self.expr_to_node = {} + self.build_graph(tracer) + self.topological_sort() + + def build_graph(self, tracer): + self.nodes = [self.last_node] + self.expr_to_node[tracer.last_node] = self.nodes[-1] + + rename_pid = {} + + visited = set([tracer.last_node]) + head = 0 + while head < len(self.nodes): + cur_node = self.nodes[head] + + # add prev node + prev_node = cur_node.expr.prev_node + if prev_node is not None: + if prev_node not in visited: + visited.add(prev_node) + self.nodes.append(CompGraphNode(prev_node)) + self.expr_to_node[prev_node] = self.nodes[-1] + cur_node.prev_node = self.expr_to_node[prev_node] + self.expr_to_node[prev_node].add_next_node(cur_node) + + # add source node + if isinstance(cur_node.expr, SglVariable): + if cur_node.expr.name in tracer.variables: + source = tracer.variables[cur_node.expr.name].source + else: + source = cur_node.expr.source + if source not in visited: + visited.add(source) + self.nodes.append(CompGraphNode(source)) + self.expr_to_node[source] = self.nodes[-1] + cur_node.source_node = self.expr_to_node[source] + self.expr_to_node[source].add_next_node(cur_node) + head += 1 + + # rename pid + if cur_node.expr.pid not in rename_pid: + rename_pid[cur_node.expr.pid] = len(rename_pid) + cur_node.expr.pid = rename_pid[cur_node.expr.pid] + + def topological_sort(self): + prevd = {} + cand = Queue() + for x in self.nodes: + prevd[x] = (x.prev_node is not None) + (x.source_node is not None) + if prevd[x] == 0: + cand.put(x) + new_list = [] + while cand.qsize() > 0: + head = cand.get() + new_list.append(head) + for x in head.next_nodes: + prevd[x] -= 1 + if prevd[x] == 0: + cand.put(x) + self.nodes = new_list + + def print_graph( + self, + ): + for node in self.nodes: + print(node) + + def run_internal( + self, + backend, + kwargs, + default_sampling_para, + ): + stream_executor_ids = set([x.expr.pid for x in self.nodes]) + stream_executors = {} + for x in stream_executor_ids: + arguments = kwargs if x == self.last_node.expr.pid else {} + stream_executors[x] = StreamExecutor( + backend, arguments, default_sampling_para, None, False + ) + for node in self.nodes: + se_id = node.expr.pid + expr = node.expr + if isinstance(expr, SglVariable): + # Make a copy for SglVariable + expr = SglVariable(expr.name, expr.source) + expr.source_stream_executor = stream_executors[ + node.source_node.expr.pid + ] + elif isinstance(expr, SglArgument): + # Substitute SglArgument + expr = kwargs[expr.name] + stream_executors[se_id].submit(expr) + for stream_executor in stream_executors.values(): + stream_executor.end() + return ProgramState(stream_executors[self.last_node.expr.pid]) + + def run( + self, + *, + max_new_tokens: int = 128, + stop: Union[str, List[str]] = (), + temperature: float = 1.0, + top_p: float = 1.0, + top_k: int = -1, + min_p: float = 0.0, + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, + backend=None, + **kwargs, + ): + backend = backend or global_config.default_backend + + kwargs.update(self.function.bind_arguments) + + default_sampling_para = SglSamplingParams( + max_new_tokens=max_new_tokens, + stop=stop, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + ) + + return self.run_internal(backend, kwargs, default_sampling_para) + + def run_batch( + self, + batch_kwargs, + *, + max_new_tokens: int = 128, + stop: Union[str, List[str]] = (), + temperature: float = 1.0, + top_p: float = 1.0, + top_k: int = -1, + min_p: float = 0.0, + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, + backend=None, + num_threads: Union[str, int] = "auto", + ): + assert isinstance(batch_kwargs, (list, tuple)) + if len(batch_kwargs) == 0: + return [] + assert isinstance(batch_kwargs[0], dict) + + backend = backend or global_config.default_backend + + default_sampling_para = SglSamplingParams( + max_new_tokens=max_new_tokens, + stop=stop, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + ) + + # Extract prefix by tracing and cache it + if len(batch_kwargs) > 1: + cache_program(self.function, backend) + + # Run all programs + if num_threads == "auto": + num_threads = multiprocessing.cpu_count() + num_threads = min(num_threads, len(batch_kwargs)) + + if num_threads == 1: + rets = [] + for arguments in batch_kwargs: + rets.append( + self.run_internal(backend, arguments, default_sampling_para) + ) + else: + with ThreadPoolExecutor(num_threads) as executor: + futures = [] + for arguments in batch_kwargs: + futures.append( + executor.submit( + self.run_internal, backend, arguments, default_sampling_para + ) + ) + rets = [f.result() for f in futures] + rets[-1].sync() + + return rets + + +class CompGraphNode: + def __init__( + self, expr: SglExpr, prev_node=None, next_nodes=None, source_node=None + ): + self.expr = expr + self.next_nodes = next_nodes or [] + self.prev_node = prev_node + self.source_node = source_node + + def add_next_node(self, other): + self.next_nodes.append(other) + + def __repr__(self): + re = f"stream {self.expr.pid:2d}: " + re += f"%{self.expr.node_id} = " + if self.prev_node is not None: + re += f"%{self.prev_node.expr.node_id} + " + re += repr(self.expr) + return re diff --git a/sglang/python/sglang/lang/interpreter.py b/sglang/python/sglang/lang/interpreter.py new file mode 100644 index 0000000..1ac9164 --- /dev/null +++ b/sglang/python/sglang/lang/interpreter.py @@ -0,0 +1,1021 @@ +"""The interpreter that executes SGL programs""" + +import asyncio +import contextvars +import copy +import multiprocessing +import queue +import threading +import uuid +import warnings +from concurrent.futures import ThreadPoolExecutor +from contextlib import contextmanager +from typing import Any, Callable, Dict, List, Optional + +import tqdm + +from sglang.global_config import global_config +from sglang.lang.ir import ( + SglCommitLazy, + SglConcateAndAppend, + SglConstantText, + SglExpr, + SglExprList, + SglGen, + SglImage, + SglRoleBegin, + SglRoleEnd, + SglSelect, + SglVariable, + SglVarScopeBegin, + SglVarScopeEnd, + SglVideo, +) +from sglang.utils import ( + encode_image_base64, + encode_video_base64, + get_exception_traceback, +) + + +def run_internal(state, program, func_args, func_kwargs, sync): + try: + state.ret_value = program.func(state, *func_args, **func_kwargs) + except Exception as e: + raise e + finally: + state.stream_executor.end() + + if sync: + state.stream_executor.sync() + + if global_config.verbosity >= 2: + print(state.text()) + + +def run_program( + program, + backend, + func_args, + func_kwargs, + default_sampling_para, + stream, + sync=False, + use_thread=True, +): + if hasattr(backend, "endpoint"): + backend = backend.endpoint + assert backend is not None, "Please specify a backend" + func_kwargs.update(program.bind_arguments) + stream_executor = StreamExecutor( + backend, + func_kwargs, + default_sampling_para, + chat_template=None, + stream=stream, + num_api_spec_tokens=program.num_api_spec_tokens, + use_thread=use_thread, + ) + state = ProgramState(stream_executor) + + if stream: + t = threading.Thread( + target=run_internal, args=(state, program, func_args, func_kwargs, sync) + ) + t.start() + return state + else: + run_internal(state, program, func_args, func_kwargs, sync) + return state + + +def run_program_batch( + program, + backend, + batch_arguments, + default_sampling_para, + num_threads, + progress_bar, + generator_style=False, +): + if hasattr(backend, "endpoint"): + backend = backend.endpoint + + # Pre-cache the common prefix for a batch. The prefix is extracted by tracing the program. + if global_config.enable_precache_with_tracing and len(batch_arguments) > 1: + cache_program(program, backend) + + # Run all programs + if num_threads == "auto": + num_threads = max(96, multiprocessing.cpu_count() * 16) + num_threads = min(num_threads, len(batch_arguments)) + + if generator_style: + return _run_program_batch_generator( + program, + backend, + batch_arguments, + default_sampling_para, + num_threads, + progress_bar, + ) + + # Original code path when generator_style=False + if num_threads == 1: + rets = [] + if progress_bar: + for arguments in tqdm.tqdm(batch_arguments): + rets.append( + run_program( + program, + backend, + (), + arguments, + default_sampling_para, + False, + True, + ) + ) + else: + for arguments in batch_arguments: + rets.append( + run_program( + program, + backend, + (), + arguments, + default_sampling_para, + False, + True, + ) + ) + else: + if progress_bar: + pbar = tqdm.tqdm(total=len(batch_arguments)) + + with ThreadPoolExecutor(num_threads) as executor: + futures = [] + for arguments in batch_arguments: + futures.append( + executor.submit( + run_program, + program, + backend, + (), + arguments, + default_sampling_para, + False, + True, + ) + ) + if progress_bar: + futures[-1].add_done_callback(lambda _: pbar.update()) + + rets = [f.result() for f in futures] + rets[-1].sync() + + if progress_bar: + pbar.close() + + return rets + + +def _run_program_batch_generator( + program, + backend, + batch_arguments, + default_sampling_para, + num_threads, + progress_bar, +): + """Helper function that yields results one by one using chunking to avoid overwhelming ThreadPoolExecutor.""" + if num_threads == 1: + iterator = tqdm.tqdm(batch_arguments) if progress_bar else batch_arguments + for arguments in iterator: + yield run_program( + program, + backend, + (), + arguments, + default_sampling_para, + False, + True, + ) + else: + pbar = tqdm.tqdm(total=len(batch_arguments)) if progress_bar else None + + # Process in chunks to avoid overwhelming ThreadPoolExecutor + # Otherwise, ThreadPoolExecutor.submit will block after adding certain number of tasks + # so we will never reach "yield" until all tasks are done + chunk_size = 200 + + with ThreadPoolExecutor(num_threads) as executor: + for chunk_start in range(0, len(batch_arguments), chunk_size): + chunk_end = min(chunk_start + chunk_size, len(batch_arguments)) + chunk_futures = [] + + # Submit chunk of tasks + for i in range(chunk_start, chunk_end): + future = executor.submit( + run_program, + program, + backend, + (), + batch_arguments[i], + default_sampling_para, + False, + True, + ) + if pbar: + future.add_done_callback(lambda _: pbar.update()) + chunk_futures.append(future) + + # Yield results from this chunk as they complete + for future in chunk_futures: + yield future.result() + + if pbar: + pbar.close() + + +def cache_program(program, backend): + from sglang.lang.tracer import extract_prefix_by_tracing + + prefix = extract_prefix_by_tracing(program, backend) + if prefix and len(prefix) > 64: + backend.cache_prefix(prefix) + + +class StreamExecutor: + """A stream executor that executes SGL expressions in a background thread.""" + + def __init__( + self, + backend, + arguments, + default_sampling_para, + chat_template, + stream, + num_api_spec_tokens=None, + use_thread=True, + ): + from sglang.lang.backend.base_backend import BaseBackend + + self.sid = uuid.uuid4().hex + self.backend: BaseBackend = backend + self.arguments: Dict[str, Any] = arguments + self.default_sampling_para = default_sampling_para + self.stream = stream + + self.variables = {} # Dict[name: str -> value: str] + self.variable_event = {} # Dict[name: str -> event: threading.Event] + self.meta_info = {} # Dict[name: str -> info: str] + self.is_finished = False + self.error_ = None + + # For completion + self.text_ = "" # The full text + + # For chat + self.messages_ = [] # The messages in the OpenAI API format + self.chat_template = chat_template or self.backend.get_chat_template() + self.cur_role = None + self.cur_role_begin_pos = None + + # For vision + self.images_ = [] + self.cur_images = [] + + # For fork/join + self.fork_start_text_pos = None + + # For speculative execution + self.num_api_spec_tokens = num_api_spec_tokens + self.speculated_text = "" + + # Worker thread + self.use_thread = use_thread + if self.use_thread: + self.queue = queue.Queue() + + def _run_worker_in_context(): + self._thread_worker_func() + + self.worker = threading.Thread( + target=contextvars.copy_context().run, args=(_run_worker_in_context,) + ) + self.worker.start() + + # For streaming + if stream: + self.stream_text_event = threading.Event() + self.stream_var_event = {} + else: + self.stream_text_event = None + self.stream_var_event = None + + def submit(self, expr: SglExpr): + self._init_var_event(expr) + + if self.use_thread: + self.queue.put(expr) + else: + self._execute(expr) + + def sync(self): + if self.use_thread: + self.queue.join() + + def get_var(self, name): + if name in self.variable_event: + self.variable_event[name].wait() + return self.variables[name] + + def set_var(self, name, value): + self.variables[name] = value + + def get_meta_info(self, name, timeout=None): + if name in self.variable_event: + got = self.variable_event[name].wait(timeout) + if not got: + raise TimeoutError(f"Timeout while waiting for event '{name}'") + ret = self.meta_info.get(name, None) + return ret + + def fork( + self, + size: int = 1, + position_ids_offset: Optional[List[int]] = None, + ): + if size > 1 and str(self.text_): + self.submit(SglCommitLazy()) + + self.sync() + size = int(size) + + exes = [ + StreamExecutor( + self.backend, + self.arguments, + self.default_sampling_para, + self.chat_template, + self.stream, + ) + for _ in range(size) + ] + for i in range(size): + exes[i].variables = dict(self.variables) + exes[i].text_ = str(self.text_) + exes[i].messages_ = list(self.messages_) + exes[i].cur_role = self.cur_role + exes[i].cur_role_begin_pos = self.cur_role_begin_pos + exes[i].fork_start_text_pos = len(self.text_) + exes[i].images_ = list(self.images_) + + # TODO(ying): handle API speculative execution + + return exes + + def text(self): + self.sync() + return self.text_ + + def messages(self): + self.sync() + return self.messages_ + + def error(self): + self.sync() + return self.error_ + + def end(self): + if self.use_thread: + if self.worker.is_alive(): + self.queue.put(None) + self.backend.end_program(self) + + def _thread_worker_func(self): + error = None + + while True: + expr = self.queue.get() + if expr is None: + self.queue.task_done() + break + + try: + self._execute(expr) + except Exception as e: + warnings.warn(f"Error in stream_executor: {get_exception_traceback()}") + error = e + break + self.queue.task_done() + if self.stream_text_event: + self.stream_text_event.set() + + # Clean the queue and events + if error is not None: + try: + while True: + self.queue.task_done() + self.queue.get_nowait() + except queue.Empty: + pass + for name in self.variable_event: + self.variable_event[name].set() + if self.stream_var_event: + for name in self.stream_var_event: + self.stream_var_event[name].set() + self.error_ = error + + if self.stream_text_event: + self.stream_text_event.set() + + self.is_finished = True + + def _execute(self, other): + if isinstance(other, str): + other = SglConstantText(other) + + assert isinstance(other, SglExpr), f"{other}" + + if isinstance(other, SglConstantText): + self._execute_fill(other.value) + elif isinstance(other, SglGen): + self._execute_gen(other) + elif isinstance(other, SglSelect): + self._execute_select(other) + elif isinstance(other, SglExprList): + for x in other.expr_list: + self._execute(x) + elif isinstance(other, SglRoleBegin): + self._execute_role_begin(other) + elif isinstance(other, SglRoleEnd): + self._execute_role_end(other) + elif isinstance(other, SglImage): + self._execute_image(other) + elif isinstance(other, SglVideo): + self._execute_video(other) + elif isinstance(other, SglVariable): + self._execute_variable(other) + elif isinstance(other, SglVarScopeBegin): + self._execute_var_scope_begin(other) + elif isinstance(other, SglVarScopeEnd): + self._execute_var_scope_end(other) + elif isinstance(other, SglCommitLazy): + self._execute_commit_lazy_operations(other) + elif isinstance(other, SglConcateAndAppend): + if ( + global_config.enable_parallel_encoding + and self.backend.support_concate_and_append + ): + self._execute_concatenate_and_append_kv_cache(other) + else: + self._execute_concatenate_and_append_text(other) + else: + raise ValueError(f"Unknown type: {type(other)}") + + def _execute_fill(self, value: str, prefix=False): + value = str(value) + + if ( + self.cur_role == "assistant" + and self.num_api_spec_tokens is not None + and self.backend.is_chat_model + and not prefix + ): + self.backend.spec_fill(value) + return + + if self.speculated_text.startswith(value): + self.speculated_text = self.speculated_text[len(value) :] + else: + self.speculated_text = "" + + self.text_ += value + + def _execute_image(self, expr: SglImage): + path = expr.path + + base64_data = encode_image_base64(path) + + self.images_.append((path, base64_data)) + self.cur_images.append((path, base64_data)) + self.text_ += self.chat_template.image_token + + def _execute_video(self, expr: SglVideo): + path = expr.path + num_frames = expr.num_frames + + base64_data = encode_video_base64(path, num_frames) + + self.images_.append((path, base64_data)) + self.cur_images.append((path, base64_data)) + self.text_ += self.chat_template.image_token + + def _spec_gen(self, sampling_params): + stop = sampling_params.stop + max_new_tokens = sampling_params.max_new_tokens + meta_info = {} + + def regen(): + nonlocal meta_info + + sampling_params.max_new_tokens = max( + sampling_params.max_new_tokens, self.num_api_spec_tokens + ) + sampling_params.stop = None + self.speculated_text, meta_info = self.backend.generate( + self, sampling_params=sampling_params + ) + + def find_stop(): + if isinstance(stop, str): + return self.speculated_text.find(stop) + elif isinstance(stop, (tuple, list)): + pos = -1 + for stop_str in stop: + stop_pos = self.speculated_text.find(stop_str) + if stop_pos != -1 and (pos == -1 or stop_pos < pos): + pos = stop_pos + return pos + else: + raise Exception("Wrong type of stop in sampling parameters.") + + if stop is None: + if len(self.speculated_text) < max_new_tokens: + regen() + comp = self.speculated_text[:max_new_tokens] + self.speculated_text = self.speculated_text[max_new_tokens:] + elif isinstance(stop, (str, list, tuple)): + if self.speculated_text == "": + regen() + stop_pos = find_stop() + if stop_pos == -1: + stop_pos = min( + sampling_params.max_new_tokens, + len(self.speculated_text), + ) + comp = self.speculated_text[:stop_pos] + self.speculated_text = self.speculated_text[stop_pos:] + else: + raise ValueError("Wrong type of stop in sampling parameters.") + + return comp, meta_info + + def _execute_gen(self, expr: SglGen): + sampling_params = self._resolve_sampling_params(expr.sampling_params) + name = expr.name + if not self.stream: + if self.num_api_spec_tokens is None: + comp, meta_info = self.backend.generate( + self, + sampling_params=sampling_params, + ) + + else: + if self.backend.is_chat_model: + # Speculative execution on models with only chat interface. + # Store the calls into a temporary list. + # They will be lazily executed later. + comp, meta_info = self.backend.generate( + self, + sampling_params=sampling_params, + spec_var_name=name, + ) + return + + else: # Speculative execution on models with completion interface + comp, meta_info = self._spec_gen(sampling_params) + if isinstance(comp, list): + self.text_ += comp[0] + else: + assert isinstance(comp, str) + self.text_ += comp + + self.variables[name] = comp + self.meta_info[name] = meta_info + self.variable_event[name].set() + else: + assert ( + self.num_api_spec_tokens is None + ), "stream is not supported with api speculative execution" + generator = self.backend.generate_stream( + self, sampling_params=sampling_params + ) + + self.variables[name] = "" + self.stream_var_event[name].set() + + for comp, meta_info in generator: + self.text_ += comp + self.variables[name] += comp + self.meta_info[name] = meta_info + self.stream_var_event[name].set() + self.stream_text_event.set() + + self.variable_event[name].set() + self.stream_var_event[name].set() + + def _execute_select(self, expr: SglSelect): + choices_decision = self.backend.select( + self, expr.choices, expr.temperature, expr.choices_method + ) + if expr.name is not None: + name = expr.name + self.variables[name] = choices_decision.decision + self.meta_info[name] = choices_decision.meta_info + self.variable_event[name].set() + if self.stream_var_event: + self.stream_var_event[name].set() + self.text_ += choices_decision.decision + + def _execute_variable(self, expr: SglVariable): + src_executor = expr.source_stream_executor + value = src_executor.get_var(expr.name) + self._execute_fill(value) + + def _execute_role_begin(self, expr: SglRoleBegin): + assert self.cur_role is None, "Nested roles are not allowed." + + if len(self.messages_) == 0 and expr.role != "system": + # Insert the default system message + default_system = self.chat_template.default_system_prompt + if default_system: + self._execute_role_begin(SglRoleBegin("system")) + self._execute_fill(default_system) + self._execute_role_end(SglRoleEnd("system")) + + self.cur_role = expr.role + + prefix, _ = self.chat_template.get_prefix_and_suffix(expr.role, self.messages_) + + self._execute_fill(prefix, prefix=True) + self.cur_role_begin_pos = len(self.text_) + + def _execute_role_end(self, expr: SglRoleEnd): + if ( + self.cur_role == "assistant" + and self.num_api_spec_tokens is not None + and self.backend.is_chat_model + ): + # Execute the stored lazy generation calls + self.backend.role_end_generate(self) + self.cur_role = None + + new_text = self.text_[self.cur_role_begin_pos :].lstrip() + + _, suffix = self.chat_template.get_prefix_and_suffix(expr.role, self.messages_) + self._execute_fill(suffix) + + if self.cur_images: + # OpenAI vision API format + last_msg = { + "role": expr.role, + "content": [{"type": "text", "text": new_text}], + } + for image_path, image_base64_data in self.cur_images: + last_msg["content"].append( + { + "type": "image_url", + "image_url": { + "url": f"data:image/jpeg;base64,{image_base64_data}" + }, + } + ) + self.messages_.append(last_msg) + self.cur_images = [] + else: + # OpenAI chat API format + self.messages_.append({"role": expr.role, "content": new_text}) + + def _execute_var_scope_begin(self, expr: SglVarScopeBegin): + self.variables[expr.name] = int(len(self.text_)) + + def _execute_var_scope_end(self, expr: SglVarScopeEnd): + self.variables[expr.name] = self.text_[self.variables[expr.name] :] + self.variable_event[expr.name].set() + + def _execute_commit_lazy_operations(self, expr: SglCommitLazy): + self.backend.commit_lazy_operations(self) + + def _execute_concatenate_and_append_text(self, expr: SglConcateAndAppend): + new_text = "" + for s in expr.states: + exe = s.stream_executor + exe.sync() + new_text += exe.text_[exe.fork_start_text_pos :] + + self._execute_fill(new_text) + + def _execute_concatenate_and_append_kv_cache(self, expr: SglConcateAndAppend): + self_len = len(self.text_) + + for i, s in enumerate(expr.states): + exe = s.stream_executor + exe.submit(SglCommitLazy()) + + for i, s in enumerate(expr.states): + exe = s.stream_executor + exe.sync() + assert exe.fork_start_text_pos == self_len + self.text_ += exe.text_[exe.fork_start_text_pos :] + + src_rids = [state.stream_executor.sid for state in expr.states] + self.backend.concatenate_and_append(src_rids, self.sid) + + def _init_var_event(self, expr): + if isinstance(expr, (SglGen, SglSelect, SglVarScopeBegin)): + self.variable_event[expr.name] = threading.Event() + if self.stream: + self.stream_var_event[expr.name] = threading.Event() + elif isinstance(expr, SglExprList): + for e in expr.expr_list: + self._init_var_event(e) + + def _resolve_sampling_params(self, sampling_params): + """ + Construct sampling param based on default + override values + + The default values of sampling are populated in `default_sampling_para` via sgl.function.run(...sampling_args) + , and `sampling_params` contains the override values from sgl.gen(). + + Here we use default_sampling_para as the base and override the values if they exist in `sampling_params`. + It also extends the stop tokens based on the chat template. + """ + + # deepcopy is required because the dict has lists inside + clone = copy.deepcopy(self.default_sampling_para) + + for item in [ + "max_new_tokens", + "min_new_tokens", + "n", + "stop", + "stop_token_ids", + "temperature", + "top_p", + "top_k", + "min_p", + "frequency_penalty", + "presence_penalty", + "ignore_eos", + "return_logprob", + "logprob_start_len", + "top_logprobs_num", + "return_text_in_logprobs", + "dtype", + "regex", + "json_schema", + ]: + value = getattr(sampling_params, item, None) + if value is not None: + setattr(clone, item, value) + + if self.chat_template.stop_str: + if clone.stop == (): + clone.stop = [] + elif isinstance(clone.stop, str): + clone.stop = [clone.stop] + clone.stop += self.chat_template.stop_str + + return clone + + def __del__(self): + self.end() + + +class ProgramState: + """The state of an SGL program.""" + + def __init__(self, stream_executor: StreamExecutor): + self.stream_executor = stream_executor + + def _role_common(self, name: str, expr: Optional[SglExpr] = None): + if expr is not None: + role_expr = SglExprList([SglRoleBegin(name), expr, SglRoleEnd(name)]) + self.stream_executor.submit(role_expr) + return role_expr + else: + + @contextmanager + def role_scope(): + self.stream_executor.submit(SglRoleBegin(name)) + yield + self.stream_executor.submit(SglRoleEnd(name)) + + return role_scope() + + def system(self, expr: Optional[SglExpr] = None): + return self._role_common("system", expr) + + def user(self, expr: Optional[SglExpr] = None): + return self._role_common("user", expr) + + def assistant(self, expr: Optional[SglExpr] = None): + return self._role_common("assistant", expr) + + @contextmanager + def var_scope(self, name: str): + self.stream_executor.submit(SglVarScopeBegin(name)) + yield + self.stream_executor.submit(SglVarScopeEnd(name)) + + def fork( + self, + size: int = 1, + position_ids_offset: Optional[List[int]] = None, + ): + stream_executors = self.stream_executor.fork(size, position_ids_offset) + states = [ProgramState(x) for x in stream_executors] + state_group = ProgramStateGroup(states, self) + return state_group + + @contextmanager + def copy(self, position_ids_offset: Optional[List[int]] = None): + state_group = self.fork(1, position_ids_offset) + try: + yield state_group[0] + finally: + state_group.join() + + def text(self): + return self.stream_executor.text() + + def messages(self): + return self.stream_executor.messages() + + def sync(self): + return self.stream_executor.sync() + + def error(self): + return self.stream_executor.error() + + def text_iter(self, var_name: Optional[str] = None): + if self.stream_executor.stream: + prev = 0 + if var_name is None: + event = self.stream_executor.stream_text_event + while True: + event.wait() + event.clear() + out = str(self.stream_executor.text_[prev:]) + prev += len(out) + if out: + yield out + if self.stream_executor.is_finished: + break + else: + event = None + while not event: + if var_name in self.stream_executor.stream_var_event: + event = self.stream_executor.stream_var_event[var_name] + if self.stream_executor.is_finished: + yield "" + return + + while True: + event.wait() + event.clear() + out = str(self.stream_executor.variables[var_name][prev:]) + prev += len(out) + if out: + yield out + if self.stream_executor.variable_event[var_name].is_set(): + break + else: + if var_name is None: + yield self.text() + else: + yield self.get_var(var_name) + + async def text_async_iter( + self, var_name: Optional[str] = None, return_meta_data: bool = False + ): + loop = asyncio.get_running_loop() + + if self.stream_executor.stream: + prev = 0 + if var_name is None: + event = self.stream_executor.stream_text_event + while True: + await loop.run_in_executor(None, event.wait) + event.clear() + out = str(self.stream_executor.text_[prev:]) + prev += len(out) + if out: + yield out + if self.stream_executor.is_finished: + break + else: + event = None + while not event: + if var_name in self.stream_executor.stream_var_event: + event = self.stream_executor.stream_var_event[var_name] + if self.stream_executor.is_finished: + yield "" + return + + while True: + await loop.run_in_executor(None, event.wait) + event.clear() + out = str(self.stream_executor.variables[var_name][prev:]) + prev += len(out) + if out: + if return_meta_data: + yield out, self.stream_executor.meta_info[var_name] + else: + yield out + if self.stream_executor.variable_event[var_name].is_set(): + break + else: + if var_name is None: + yield self.text() + else: + yield self.get_var(var_name) + + def get_var(self, name): + return self.stream_executor.get_var(name) + + def set_var(self, name, value): + return self.stream_executor.set_var(name, value) + + def get_meta_info(self, name): + return self.stream_executor.get_meta_info(name) + + def __iadd__(self, other): + if other is None: + raise ValueError("Tried to append None to state.") + self.stream_executor.submit(other) + return self + + def __getitem__(self, name): + return self.get_var(name) + + def __setitem__(self, name, value): + self.set_var(name, value) + + def __contains__(self, name): + return name in self.stream_executor.variables + + def __del__(self): + self.stream_executor.end() + + def __repr__(self) -> str: + return f"ProgramState({self.text()})" + + +class ProgramStateGroup: + def __init__( + self, states: List[ProgramState], src_state: Optional[ProgramState] = None + ): + self.states = states + self.src_state = src_state + + def join(self, mode: str = "gather_variable"): + if mode == "gather_variable": + # Copy variables back + src_vars = self.src_state.stream_executor.variables + src_var_set = set(src_vars.keys()) + for child_state in self.states: + child_state.stream_executor.sync() + child_vars = child_state.stream_executor.variables + new_vars = set(child_vars.keys()) - src_var_set + + for k in new_vars: + if k in src_vars: + src_vars[k].append(child_vars[k]) + else: + src_vars[k] = [child_vars[k]] + elif mode == "concate_and_append": + # Concatenate and append KV cache + self.src_state += SglConcateAndAppend(self.states) + # Need a sync here. Otherwise, `states` can be deleted. + self.src_state.stream_executor.sync() + else: + raise ValueError(f"Invalid join mode: {mode}") + + for s in self.states: + s.stream_executor.end() + + def __getitem__(self, i: int): + return self.states[i] + + def __setitem__(self, i: int, value): + assert self.states[i] == value + + def __iadd__(self, other): + if isinstance(other, Callable): + # lambda function + for i in range(len(self.states)): + self.states[i] += other(i) + elif isinstance(other, SglExpr): + for i in range(len(self.states)): + self.states[i] += other + elif isinstance(other, (list, tuple)): + for i in range(len(self.states)): + self.states[i] += other[i] + else: + raise ValueError(f"Invalid value: {other}") + + return self diff --git a/sglang/python/sglang/lang/ir.py b/sglang/python/sglang/lang/ir.py new file mode 100644 index 0000000..f1775bc --- /dev/null +++ b/sglang/python/sglang/lang/ir.py @@ -0,0 +1,608 @@ +"""The intermediate representation.""" + +import dataclasses +import inspect +import warnings +from typing import List, Optional, Union + +from sglang.global_config import global_config +from sglang.lang.choices import ChoicesSamplingMethod + +REGEX_INT = r"[-+]?[0-9]+[ \n]*" +REGEX_FLOAT = r"[-+]?[0-9]*\.?[0-9]+[ \n]*" +REGEX_BOOL = r"(True|False)" +REGEX_STR = r"\"[\w\d\s]*\"" # bugs with regex r"\".*\"" in interegular pkg + + +@dataclasses.dataclass +class SglSamplingParams: + max_new_tokens: int = 128 + min_new_tokens: int = 0 + n: int = 1 + stop: Union[str, List[str]] = () + stop_token_ids: Optional[List[int]] = () + temperature: float = 1.0 + top_p: float = 1.0 + top_k: int = -1 # -1 means disable + min_p: float = 0.0 + frequency_penalty: float = 0.0 + presence_penalty: float = 0.0 + ignore_eos: bool = False + return_logprob: Optional[bool] = None + logprob_start_len: Optional[int] = (None,) + top_logprobs_num: Optional[int] = (None,) + return_text_in_logprobs: Optional[bool] = (None,) + json_schema: Optional[str] = None + + # for constrained generation, not included in to_xxx_kwargs + dtype: Optional[str] = None + regex: Optional[str] = None + + def clone(self): + return SglSamplingParams( + self.max_new_tokens, + self.min_new_tokens, + self.n, + self.stop, + self.stop_token_ids, + self.temperature, + self.top_p, + self.top_k, + self.min_p, + self.frequency_penalty, + self.presence_penalty, + self.ignore_eos, + self.return_logprob, + self.logprob_start_len, + self.top_logprobs_num, + self.return_text_in_logprobs, + self.json_schema, + ) + + def to_openai_kwargs(self): + # OpenAI does not support top_k, so we drop it here + if self.regex is not None: + warnings.warn("Regular expression is not supported in the OpenAI backend.") + return { + "max_tokens": self.max_new_tokens, + "max_completion_tokens": self.max_new_tokens, + "n": self.n, + "stop": self.stop or None, + "temperature": self.temperature, + "top_p": self.top_p, + "frequency_penalty": self.frequency_penalty, + "presence_penalty": self.presence_penalty, + } + + def to_vertexai_kwargs(self): + if self.regex is not None: + warnings.warn( + "Regular expression is not supported in the VertexAI backend." + ) + return { + "candidate_count": 1, + "max_output_tokens": self.max_new_tokens, + "stop_sequences": self.stop, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": self.top_k if self.top_k > 0 else None, + } + + def to_anthropic_kwargs(self): + # Anthropic does not support frequency_penalty or presence_penalty, so we drop it here + if self.regex is not None: + warnings.warn( + "Regular expression is not supported in the Anthropic backend." + ) + return { + "max_tokens": self.max_new_tokens, + "stop_sequences": ( + self.stop if isinstance(self.stop, (list, tuple)) else [self.stop] + ), + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": self.top_k, + } + + def to_litellm_kwargs(self): + if self.regex is not None: + warnings.warn("Regular expression is not supported in the LiteLLM backend.") + return { + "max_tokens": self.max_new_tokens, + "stop": self.stop or None, + "temperature": self.temperature, + "top_p": self.top_p, + "frequency_penalty": self.frequency_penalty, + "presence_penalty": self.presence_penalty, + } + + def to_srt_kwargs(self): + return { + "max_new_tokens": self.max_new_tokens, + "min_new_tokens": self.min_new_tokens, + "n": self.n, + "stop": self.stop, + "stop_token_ids": self.stop_token_ids, + "temperature": self.temperature, + "top_p": self.top_p, + "top_k": self.top_k, + "min_p": self.min_p, + "frequency_penalty": self.frequency_penalty, + "presence_penalty": self.presence_penalty, + "ignore_eos": self.ignore_eos, + "regex": self.regex, + "json_schema": self.json_schema, + } + + +class SglFunction: + def __init__(self, func, num_api_spec_tokens=None, bind_arguments=None): + self.func = func + self.num_api_spec_tokens = num_api_spec_tokens + self.bind_arguments = bind_arguments or {} + self.pin_prefix_rid = None + + # Parse arguments + argspec = inspect.getfullargspec(func) + assert argspec.args[0] == "s", 'The first argument must be "s"' + self.arg_names = argspec.args[1:] + self.arg_defaults = argspec.defaults if argspec.defaults is not None else [] + + def bind(self, **kwargs): + assert all(key in self.arg_names for key in kwargs) + + new_bind_dict = {**self.bind_arguments, **kwargs} + return SglFunction(self.func, bind_arguments=new_bind_dict) + + def run( + self, + *args, + max_new_tokens: int = 128, + n: int = 1, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: float = 1.0, + top_p: float = 1.0, + top_k: int = -1, + min_p: float = 0.0, + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, + ignore_eos: bool = False, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, + stream: bool = False, + backend=None, + use_thread: bool = True, + **kwargs, + ): + from sglang.lang.interpreter import run_program + + # avoid using [] as the default arg: https://nikos7am.com/posts/mutable-default-arguments/ + if stop is None: + stop = [] + if stop_token_ids is None: + stop_token_ids = [] + + default_sampling_para = SglSamplingParams( + max_new_tokens=max_new_tokens, + n=n, + stop=stop, + stop_token_ids=stop_token_ids, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + ignore_eos=ignore_eos, + return_logprob=return_logprob, + logprob_start_len=logprob_start_len, + top_logprobs_num=top_logprobs_num, + return_text_in_logprobs=return_text_in_logprobs, + ) + backend = backend or global_config.default_backend + return run_program( + self, + backend, + args, + kwargs, + default_sampling_para, + stream, + use_thread=use_thread, + ) + + def run_batch( + self, + batch_kwargs, + *, + max_new_tokens: int = 128, + n: int = 1, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: float = 1.0, + top_p: float = 1.0, + top_k: int = -1, + min_p: float = 0.0, + frequency_penalty: float = 0.0, + presence_penalty: float = 0.0, + ignore_eos: bool = False, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, + backend=None, + num_threads: Union[str, int] = "auto", + progress_bar: bool = False, + generator_style: bool = False, + ): + from sglang.lang.interpreter import run_program_batch + + if stop is None: + stop = [] + if stop_token_ids is None: + stop_token_ids = [] + + assert isinstance(batch_kwargs, (list, tuple)) + if len(batch_kwargs) == 0: + return [] + if not isinstance(batch_kwargs[0], dict): + num_programs = len(batch_kwargs) + # change the list of argument values to dict of arg_name -> arg_value + batch_kwargs = [ + {self.arg_names[i]: v for i, v in enumerate(arg_values)} + for arg_values in batch_kwargs + if isinstance(arg_values, (list, tuple)) + and len(self.arg_names) - len(self.arg_defaults) + <= len(arg_values) + <= len(self.arg_names) + ] + # Ensure to raise an exception if the number of arguments mismatch + if len(batch_kwargs) != num_programs: + raise Exception("Given arguments mismatch the SGL function signature") + + default_sampling_para = SglSamplingParams( + max_new_tokens=max_new_tokens, + n=n, + stop=stop, + stop_token_ids=stop_token_ids, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + ignore_eos=ignore_eos, + return_logprob=return_logprob, + logprob_start_len=logprob_start_len, + top_logprobs_num=top_logprobs_num, + return_text_in_logprobs=return_text_in_logprobs, + ) + backend = backend or global_config.default_backend + return run_program_batch( + self, + backend, + batch_kwargs, + default_sampling_para, + num_threads, + progress_bar, + generator_style=generator_style, + ) + + def trace(self, *, backend=None, **kwargs): + from sglang.lang.tracer import trace_program + + backend = backend or global_config.default_backend + return trace_program(self, kwargs, backend) + + def cache(self, backend=None): + from sglang.lang.interpreter import cache_program + + backend = backend or global_config.default_backend + return cache_program(self, backend) + + def compile(self, *, backend=None): + from sglang.lang.compiler import compile_func + + return compile_func(self, backend) + + def __call__(self, *args, **kwargs): + from sglang.lang.tracer import TracingScope + + tracing_scope = TracingScope.get_current_scope() + if tracing_scope is None: + return self.run(*args, **kwargs) + else: + kwargs["backend"] = tracing_scope.tracer_state.backend + return self.trace(*args, **kwargs) + + +class SglExpr: + node_ct = 0 + + def __init__(self): + self.node_id = SglExpr.node_ct + self.prev_node = None + self.pid = None + SglExpr.node_ct += 1 + + def __add__(self, other): + if isinstance(other, str): + other = SglConstantText(other) + assert isinstance(other, SglExpr) + + return self.concatenate_ir(self, other) + + def __radd__(self, other): + if isinstance(other, str): + other = SglConstantText(other) + assert isinstance(other, SglExpr), f"{other}" + + return self.concatenate_ir(other, self) + + def concatenate_ir(self, a, b): + if isinstance(a, SglExprList): + if isinstance(b, SglExprList): + return SglExprList(a.expr_list + b.expr_list) + else: + return SglExprList(a.expr_list + [b]) + elif isinstance(b, SglExprList): + return SglExprList([a] + b.expr_list) + + return SglExprList([a, b]) + + def print_graph_dfs(self): + ret = [""] + visited = set() + + def dfs_print(x): + if x is None or x in visited: + return + visited.add(x) + + # Print dependency + if x.prev_node is not None: + dfs_print(x.prev_node) + + if isinstance(x, SglExprList): + for y in x.expr_list: + dfs_print(y) + # elif isinstance(x, SglRole): + # dfs_print(x.expr) + elif isinstance(x, SglVariable): + dfs_print(x.source) + + # Print the node itself + if isinstance(x, (SglFork, SglGetForkItem)): + ret[0] += f"%{x.node_id} = {x}\n" + else: + if x.prev_node is not None: + ret[0] += ( + f"%{x.node_id} = %{x.prev_node.node_id} + " + str(x) + "\n" + ) + else: + ret[0] += f"%{x.node_id} = " + str(x) + "\n" + + dfs_print(self) + return ret[0] + + +class SglExprList(SglExpr): + def __init__(self, expr_list: List[SglExpr]): + super().__init__() + self.expr_list = expr_list + + def __repr__(self): + return f"ExprList({self.expr_list})" + + +class SglArgument(SglExpr): + def __init__(self, name: str, value: str): + super().__init__() + self.name = name + self.value = value + + def __repr__(self): + return f"Argument(name={self.name}, value={repr(self.value)})" + + def __len__(self): + return len(self.value) + + def __getitem__(self, i): + return self.value[i] + + def __int__(self): + return self.value + + def __bool__(self): + return self.value + + def __format__(self, *args): + raise TypeError( + "Cannot put argument inside a f-string. " + "This is not compatible with the tracer. " + ) + + +class SglImage(SglExpr): + def __init__(self, path: str): + self.path = path + + def __repr__(self) -> str: + return f"SglImage({self.path})" + + +class SglVideo(SglExpr): + def __init__(self, path: str, num_frames: int): + self.path = path + self.num_frames = num_frames + + def __repr__(self) -> str: + return f"SglVideo({self.path}, {self.num_frames})" + + +class SglGen(SglExpr): + def __init__( + self, + name: Optional[str] = None, + max_new_tokens: Optional[int] = None, + min_new_tokens: Optional[int] = None, + n: Optional[int] = None, + stop: Optional[Union[str, List[str]]] = None, + stop_token_ids: Optional[List[int]] = None, + temperature: Optional[float] = None, + top_p: Optional[float] = None, + top_k: Optional[int] = None, + min_p: Optional[float] = None, + frequency_penalty: Optional[float] = None, + presence_penalty: Optional[float] = None, + ignore_eos: Optional[bool] = None, + return_logprob: Optional[bool] = None, + logprob_start_len: Optional[int] = None, + top_logprobs_num: Optional[int] = None, + return_text_in_logprobs: Optional[bool] = None, + dtype: Optional[type] = None, + regex: Optional[str] = None, + json_schema: Optional[str] = None, + ): + """Call the model to generate. See the meaning of the arguments in docs/backend/sampling_params.md""" + super().__init__() + self.name = name + self.sampling_params = SglSamplingParams( + max_new_tokens=max_new_tokens, + min_new_tokens=min_new_tokens, + n=n, + stop=stop, + stop_token_ids=stop_token_ids, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + frequency_penalty=frequency_penalty, + presence_penalty=presence_penalty, + ignore_eos=ignore_eos, + return_logprob=return_logprob, + logprob_start_len=logprob_start_len, + top_logprobs_num=top_logprobs_num, + return_text_in_logprobs=return_text_in_logprobs, + dtype=dtype, + regex=regex, + json_schema=json_schema, + ) + + def __repr__(self): + return f"Gen('{self.name}')" + + +class SglConstantText(SglExpr): + def __init__(self, value: str): + super().__init__() + self.value = value + + def __repr__(self): + return f"Constant({repr(self.value)})" + + +class SglRoleBegin(SglExpr): + def __init__(self, role: str): + super().__init__() + self.role = role + + def __repr__(self): + return f"RoleBegin({self.role})" + + +class SglRoleEnd(SglExpr): + def __init__(self, role: str): + super().__init__() + self.role = role + + def __repr__(self): + return f"RoleEnd({self.role})" + + +class SglSelect(SglExpr): + + def __init__( + self, + name: str, + choices: List[str], + temperature: float, + choices_method: ChoicesSamplingMethod, + ): + super().__init__() + self.name = name + self.choices = choices + self.temperature = temperature + self.choices_method = choices_method + + def __repr__(self): + return f"Select({self.name}, choices={self.choices}, choices_method={self.choices_method})" + + +class SglFork(SglExpr): + def __init__(self, number: int, position_ids_offset=None): + super().__init__() + self.number = number + self.position_ids_offset = position_ids_offset + + def __repr__(self): + return ( + f"Fork(%{self.prev_node.node_id}, number={self.number}, " + f"position_ids_offset={self.position_ids_offset})" + ) + + +class SglGetForkItem(SglExpr): + def __init__(self, index: int): + super().__init__() + self.index = index + + def __repr__(self): + return f"GetForkItem(%{self.prev_node.node_id}, index={self.index})" + + +class SglVariable(SglExpr): + def __init__(self, name: str, source): + super().__init__() + self.name = name + self.source = source + + def __repr__(self): + return f"Variable('{self.name}', source=%{self.source.node_id})" + + +class SglVarScopeBegin(SglExpr): + def __init__(self, name: str): + super().__init__() + self.name = name + + def __repr__(self): + return f"VarScopeBegin('{self.name}')" + + +class SglVarScopeEnd(SglExpr): + def __init__(self, name: str): + super().__init__() + self.name = name + + def __repr__(self): + return f"VarScopeEnd('{self.name}')" + + +class SglConcateAndAppend(SglExpr): + def __init__(self, states): + super().__init__() + self.states = states + + def __repr__(self): + return f"ConcatenateAndAppend('{self.states}')" + + +class SglCommitLazy(SglExpr): + def __init__(self): + super().__init__() + + def __repr__(self): + return "CommitLazy()" diff --git a/sglang/python/sglang/lang/tracer.py b/sglang/python/sglang/lang/tracer.py new file mode 100644 index 0000000..69f035b --- /dev/null +++ b/sglang/python/sglang/lang/tracer.py @@ -0,0 +1,283 @@ +"""Tracing a program.""" + +import uuid +from typing import Any, Callable, Dict, List, Optional, Union + +from sglang.global_config import global_config +from sglang.lang.backend.base_backend import BaseBackend +from sglang.lang.interpreter import ProgramState, ProgramStateGroup +from sglang.lang.ir import ( + SglArgument, + SglCommitLazy, + SglConcateAndAppend, + SglConstantText, + SglExpr, + SglExprList, + SglFork, + SglFunction, + SglGen, + SglGetForkItem, + SglRoleBegin, + SglRoleEnd, + SglSelect, + SglVariable, + SglVarScopeBegin, + SglVarScopeEnd, +) + + +class StopTracing(Exception): + pass + + +def extract_prefix_by_tracing(program, backend): + # Create dummy arguments + dummy_arguments = {name: SglArgument(name, None) for name in program.arg_names} + arguments = dummy_arguments + arguments.update(program.bind_arguments) + + # Trace + tracer = TracerProgramState(backend, arguments, only_trace_prefix=True) + try: + with TracingScope(tracer): + tracer.ret_value = program.func(tracer, **arguments) + except (StopTracing, TypeError, AttributeError): + # Some exceptions may not be catched + pass + + # Run and cache prefix + prefix = "" + for expr in tracer.flatten_nodes(): + if isinstance(expr, SglConstantText): + prefix += expr.value + else: + break + return prefix + + +def trace_program(program, arguments, backend): + # Create dummy backend + if backend is None: + backend = BaseBackend() + + # Create dummy arguments + dummy_arguments = { + name: SglArgument(name, None) + for name in program.arg_names + if name not in arguments + } + arguments.update(dummy_arguments) + arguments.update(program.bind_arguments) + + # Trace + tracer = TracerProgramState(backend, arguments, only_trace_prefix=False) + with TracingScope(tracer): + tracer.ret_value = program.func(tracer, **arguments) + return tracer + + +class TracerProgramState(ProgramState): + def __init__(self, backend, arguments, only_trace_prefix): + self.pid = uuid.uuid4().hex + self.backend = backend + self.arguments: Dict[str, Any] = arguments + self.only_trace_prefix = only_trace_prefix + + if hasattr(backend, "endpoint"): + self.backend = backend.endpoint + + self.nodes = [] + self.last_node = None + self.variables = {} + self.ret_value = None + + # For completion + + # For chat + self.messages_ = [] + self.cur_role = None + self.chat_template = self.backend.get_chat_template() + + # For multi states + self.child_states = [] + + cur_scope = TracingScope.get_current_scope() + if cur_scope is not None: + cur_scope.add_child_state(self) + + ################################## + ########### Public API ########### + ################################## + + def fork(self, size: int = 1, position_ids_offset: Optional[List[int]] = None): + assert size >= 1 + + if self.only_trace_prefix: + raise StopTracing() + + fork_node = SglFork(size) + fork_node.prev_node = self.last_node + + states = [ + TracerProgramState(self.backend, self.arguments, self.only_trace_prefix) + for _ in range(size) + ] + + for i in range(size): + node = SglGetForkItem(i) + node.prev_node = fork_node + states[i].last_node = node + states[i].variables = dict(self.variables) + states[i].messages_ = list(self.messages_) + states[i].cur_role = self.cur_role + states[i].chat_template = self.chat_template + + state_group = ProgramStateGroup(states, self) + + return state_group + + ################################## + ########## Internal API ########## + ################################## + + def _append_node(self, other: SglExpr): + self.nodes.append(other) + other.prev_node = self.last_node + self.last_node = other + + def _execute(self, other: SglExpr): + if isinstance(other, str): + other = SglConstantText(other) + + other.pid = self.pid + + if isinstance(other, SglConstantText): + self._execute_fill(other) + elif isinstance(other, SglGen): + self._execute_gen(other) + elif isinstance(other, SglSelect): + self._execute_select(other) + elif isinstance(other, SglExprList): + for x in other.expr_list: + self._execute(x) + elif isinstance(other, SglRoleBegin): + self._execute_role_begin(other) + elif isinstance(other, SglRoleEnd): + self._execute_role_end(other) + elif isinstance(other, SglVarScopeBegin): + self._execute_var_scope_begin(other) + elif isinstance(other, SglVarScopeEnd): + self._execute_var_scope_end(other) + else: + if self.only_trace_prefix: + raise StopTracing() + else: + self._append_node(other) + + return self + + def __iadd__(self, other): + self._execute(other) + return self + + def _execute_fill(self, expr: SglConstantText): + if isinstance(expr, str): + expr = SglConstantText(expr) + self._append_node(expr) + + def _execute_gen(self, expr: SglGen): + name = expr.name if expr.name is not None else "gen_" + str(len(self.variables)) + new_node = SglVariable(name, source=expr) + self.variables[name] = new_node + self._append_node(expr) + + def _execute_select(self, expr: SglSelect): + name = ( + expr.name if expr.name is not None else "select_" + str(len(self.variables)) + ) + new_node = SglVariable(name, source=expr) + self.variables[name] = new_node + self._append_node(expr) + + def _execute_role_begin(self, expr: SglRoleBegin): + assert self.cur_role is None, "Nested roles are not allowed." + + if len(self.messages_) == 0 and expr.role != "system": + # Insert default system message + default_system = self.chat_template.default_system_prompt + if default_system: + self._execute_role_begin(SglRoleBegin("system")) + self._execute_fill(default_system) + self._execute_role_end(SglRoleEnd("system")) + + self.cur_role = expr.role + + prefix, suffix = self.chat_template.get_prefix_and_suffix( + expr.role, self.messages_ + ) + + self._execute_fill(prefix) + + def _execute_role_end(self, expr: SglRoleEnd): + prefix, suffix = self.chat_template.get_prefix_and_suffix( + expr.role, self.messages_ + ) + + self._execute_fill(suffix) + + self.messages_.append({"role": expr.role, "content": ""}) + + self.cur_role = None + + def _execute_var_scope_end(self, expr: SglVarScopeEnd): + new_node = SglVariable(name, source=self.last_node) + self.variables[name] = new_node + + def get_var(self, name): + ret = self.arguments.get(name, None) + if ret is not None: + return ret + + v = self.variables[name] + return SglVariable(v.name, v.source) + + def flatten_nodes(self): + def traverse(cur): + if isinstance(cur, SglExprList): + for child in cur.expr_list: + traverse(child) + else: + ret.append(cur) + + ret = [] + for x in self.nodes: + traverse(x) + return ret + + def __del__(self): + pass + + +class TracingScope: + cur_scope = None + + def __init__(self, tracer_state: TracerProgramState): + self.tracer_state = tracer_state + self.last_scope = TracingScope.cur_scope + + def __enter__(self): + TracingScope.cur_scope = self + return self + + def __exit__(self, exc_type, exc_value, traceback): + TracingScope.cur_scope = self.last_scope + + @staticmethod + def get_current_scope(): + return TracingScope.cur_scope + + def add_child_state(self, state: TracerProgramState): + cur_scope = self + while cur_scope is not None: + cur_scope.tracer_state.child_states.append(state) + cur_scope = cur_scope.last_scope diff --git a/sglang/python/sglang/launch_server.py b/sglang/python/sglang/launch_server.py new file mode 100644 index 0000000..caae7b0 --- /dev/null +++ b/sglang/python/sglang/launch_server.py @@ -0,0 +1,16 @@ +"""Launch the inference server.""" + +import os +import sys + +from sglang.srt.entrypoints.http_server import launch_server +from sglang.srt.server_args import prepare_server_args +from sglang.srt.utils import kill_process_tree + +if __name__ == "__main__": + server_args = prepare_server_args(sys.argv[1:]) + + try: + launch_server(server_args) + finally: + kill_process_tree(os.getpid(), include_parent=False) diff --git a/sglang/python/sglang/llama3_eval.py b/sglang/python/sglang/llama3_eval.py new file mode 100644 index 0000000..35bd4a7 --- /dev/null +++ b/sglang/python/sglang/llama3_eval.py @@ -0,0 +1,316 @@ +# 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 +import transformers +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") == 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/sglang/python/sglang/srt/_custom_ops.py b/sglang/python/sglang/srt/_custom_ops.py new file mode 100644 index 0000000..d0bc512 --- /dev/null +++ b/sglang/python/sglang/srt/_custom_ops.py @@ -0,0 +1,150 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/_custom_ops.py +import logging +import os +from typing import List, Tuple + +import torch +import torch.library + +from sglang.srt.utils import get_bool_env_var, is_hip, is_hpu + +logger = logging.getLogger(__name__) +use_vllm_custom_allreduce = get_bool_env_var( + "USE_VLLM_CUSTOM_ALLREDUCE", default="false" +) + +if not is_hpu(): + # ROCm does not use vllm custom allreduce + if use_vllm_custom_allreduce and not is_hip(): + try: + import vllm._C + except ImportError as e: + logger.warning("Failed to import from vllm._C with %r", e) + else: + try: + import sgl_kernel + except ImportError as e: + logger.warning("Failed to import from custom_ar with %r", e) + + +if use_vllm_custom_allreduce and not is_hip(): + # vLLM custom allreduce + def init_custom_ar( + ipc_tensors: List[torch.Tensor], + rank_data: torch.Tensor, + rank: int, + full_nvlink: bool, + ) -> int: + return torch.ops._C_custom_ar.init_custom_ar( + ipc_tensors, rank_data, rank, full_nvlink + ) + + def all_reduce( + fa: int, + inp: torch.Tensor, + out: torch.Tensor, + reg_buffer: int, + reg_buffer_sz_bytes: int, + ) -> None: + torch.ops._C_custom_ar.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes) + + def dispose(fa: int) -> None: + torch.ops._C_custom_ar.dispose(fa) + + def meta_size() -> int: + return torch.ops._C_custom_ar.meta_size() + + def register_buffer(fa: int, ipc_tensors: List[int]) -> None: + return torch.ops._C_custom_ar.register_buffer(fa, ipc_tensors) + + def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[int], List[int]]: + return torch.ops._C_custom_ar.get_graph_buffer_ipc_meta(fa) + + def register_graph_buffers( + fa: int, handles: List[List[int]], offsets: List[List[int]] + ) -> None: + torch.ops._C_custom_ar.register_graph_buffers(fa, handles, offsets) + +else: + if is_hip(): + # ROCM custom allreduce + + def init_custom_ar( + meta: torch.Tensor, + rank_data: torch.Tensor, + handles: List[str], + offsets: List[int], + rank: int, + full_nvlink: bool, + ) -> int: + return sgl_kernel.allreduce.init_custom_ar( + meta, rank_data, handles, offsets, rank, full_nvlink + ) + + def all_reduce_reg(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None: + sgl_kernel.allreduce.all_reduce_reg(fa, inp, out) + + def all_reduce_unreg( + fa: int, inp: torch.Tensor, reg_buffer: torch.Tensor, out: torch.Tensor + ) -> None: + sgl_kernel.allreduce.all_reduce_unreg(fa, inp, reg_buffer, out) + + def dispose(fa: int) -> None: + sgl_kernel.allreduce.dispose(fa) + + def meta_size() -> int: + return sgl_kernel.allreduce.meta_size() + + def register_buffer( + fa: int, t: torch.Tensor, handles: List[str], offsets: List[int] + ) -> None: + return sgl_kernel.allreduce.register_buffer(fa, t, handles, offsets) + + def get_graph_buffer_ipc_meta(fa: int) -> Tuple[torch.Tensor, List[int]]: + return sgl_kernel.allreduce.get_graph_buffer_ipc_meta(fa) + + def register_graph_buffers( + fa: int, handles: List[str], offsets: List[List[int]] + ) -> None: + sgl_kernel.allreduce.register_graph_buffers(fa, handles, offsets) + + def allocate_meta_buffer(size: int) -> torch.Tensor: + return sgl_kernel.allreduce.allocate_meta_buffer(size) + + def get_meta_buffer_ipc_handle(inp: torch.Tensor) -> torch.Tensor: + return sgl_kernel.allreduce.get_meta_buffer_ipc_handle(inp) + + else: + # TRTLLM custom allreduce + def init_custom_ar( + rank_id: int, + world_size: int, + rank_data_base: torch.Tensor, + buffers: List[int], + tmp_result_buffers: List[int], + barrier_in: List[int], + barrier_out: List[int], + ) -> int: + return sgl_kernel.init_custom_reduce( + rank_id, + world_size, + rank_data_base, + buffers, + tmp_result_buffers, + barrier_in, + barrier_out, + ) + + def all_reduce(fa: int, inp: torch.Tensor, out: torch.Tensor) -> None: + sgl_kernel.custom_reduce(fa, inp, out) + + def dispose(fa: int) -> None: + sgl_kernel.custom_dispose(fa) + + def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[int], List[int]]: + return sgl_kernel.get_graph_buffer_ipc_meta(fa) + + def register_graph_buffers( + fa: int, handles: List[List[int]], offsets: List[List[int]] + ) -> None: + sgl_kernel.register_graph_buffers(fa, handles, offsets) diff --git a/sglang/python/sglang/srt/aio_rwlock.py b/sglang/python/sglang/srt/aio_rwlock.py new file mode 100644 index 0000000..deda1fe --- /dev/null +++ b/sglang/python/sglang/srt/aio_rwlock.py @@ -0,0 +1,100 @@ +import asyncio + + +class RWLock: + def __init__(self): + # Protects internal state + self._lock = asyncio.Lock() + + # Condition variable used to wait for state changes + self._cond = asyncio.Condition(self._lock) + + # Number of readers currently holding the lock + self._readers = 0 + + # Whether a writer is currently holding the lock + self._writer_active = False + + # How many writers are queued waiting for a turn + self._waiting_writers = 0 + + @property + def reader_lock(self): + """ + A context manager for acquiring a shared (reader) lock. + + Example: + async with rwlock.reader_lock: + # read-only access + """ + return _ReaderLock(self) + + @property + def writer_lock(self): + """ + A context manager for acquiring an exclusive (writer) lock. + + Example: + async with rwlock.writer_lock: + # exclusive access + """ + return _WriterLock(self) + + async def acquire_reader(self): + async with self._lock: + # Wait until there is no active writer or waiting writer + # to ensure fairness. + while self._writer_active or self._waiting_writers > 0: + await self._cond.wait() + self._readers += 1 + + async def release_reader(self): + async with self._lock: + self._readers -= 1 + # If this was the last reader, wake up anyone waiting + # (potentially a writer or new readers). + if self._readers == 0: + self._cond.notify_all() + + async def acquire_writer(self): + async with self._lock: + # Increment the count of writers waiting + self._waiting_writers += 1 + try: + # Wait while either a writer is active or readers are present + while self._writer_active or self._readers > 0: + await self._cond.wait() + self._writer_active = True + finally: + # Decrement waiting writers only after we've acquired the writer lock + self._waiting_writers -= 1 + + async def release_writer(self): + async with self._lock: + self._writer_active = False + # Wake up anyone waiting (readers or writers) + self._cond.notify_all() + + +class _ReaderLock: + def __init__(self, rwlock: RWLock): + self._rwlock = rwlock + + async def __aenter__(self): + await self._rwlock.acquire_reader() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self._rwlock.release_reader() + + +class _WriterLock: + def __init__(self, rwlock: RWLock): + self._rwlock = rwlock + + async def __aenter__(self): + await self._rwlock.acquire_writer() + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + await self._rwlock.release_writer() diff --git a/sglang/python/sglang/srt/code_completion_parser.py b/sglang/python/sglang/srt/code_completion_parser.py new file mode 100644 index 0000000..94a98b0 --- /dev/null +++ b/sglang/python/sglang/srt/code_completion_parser.py @@ -0,0 +1,174 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Completion templates.""" + + +import dataclasses +import json +import logging +import os +from enum import auto + +from sglang.srt.openai_api.protocol import ChatCompletionRequest + +logger = logging.getLogger(__name__) +completion_template_name = None + + +class FimPosition: + """Postion of fim middle token.""" + + MIDDLE = auto() + END = auto() + + +@dataclasses.dataclass +class CompletionTemplate: + """A class that manages completion prompt templates. only for code completion currently.""" + + # The name of this template + name: str + + # the fim begin token + fim_begin_token: str + + # The fim middle token + fim_middle_token: str + + # The fim end token + fim_end_token: str + + # The position of the fim middle token + fim_position: FimPosition + + +# A global registry for all completion templates +completion_templates: dict[str, CompletionTemplate] = {} + + +def load_completion_template_for_openai_api(completion_template_arg): + global completion_template_name + + logger.info( + f"Use completion template for the OpenAI-compatible API server: {completion_template_arg}" + ) + + if not completion_template_exists(completion_template_arg): + if not os.path.exists(completion_template_arg): + raise RuntimeError( + f"Completion template {completion_template_arg} is not a built-in template name " + "or a valid completion template file path." + ) + + assert completion_template_arg.endswith( + ".json" + ), "unrecognized format of completion template file" + with open(completion_template_arg, "r") as filep: + template = json.load(filep) + try: + fim_position = FimPosition[template["fim_position"]] + except KeyError: + raise ValueError( + f"Unknown fim position: {template['fim_position']}" + ) from None + register_completion_template( + CompletionTemplate( + name=template["name"], + fim_begin_token=template["fim_begin_token"], + fim_middle_token=template["fim_middle_token"], + fim_end_token=template["fim_end_token"], + fim_position=fim_position, + ), + override=True, + ) + completion_template_name = template["name"] + else: + completion_template_name = completion_template_arg + + +def register_completion_template(template: CompletionTemplate, override: bool = False): + """Register a new completion template.""" + if not override: + assert ( + template.name not in completion_templates + ), f"{template.name} has been registered." + + completion_templates[template.name] = template + + +def completion_template_exists(template_name: str) -> bool: + return template_name in completion_templates + + +def is_completion_template_defined() -> bool: + global completion_template_name + return completion_template_name != None + + +def generate_completion_prompt_from_request(request: ChatCompletionRequest) -> str: + global completion_template_name + if request.suffix == "": + return request.prompt + + return generate_completion_prompt( + request.prompt, request.suffix, completion_template_name + ) + + +def generate_completion_prompt(prompt: str, suffix: str, template_name: str) -> str: + + completion_template = completion_templates[template_name] + fim_begin_token = completion_template.fim_begin_token + fim_middle_token = completion_template.fim_middle_token + fim_end_token = completion_template.fim_end_token + fim_position = completion_template.fim_position + + if fim_position == FimPosition.MIDDLE: + prompt = f"{fim_begin_token}{prompt}{fim_middle_token}{suffix}{fim_end_token}" + elif fim_position == FimPosition.END: + prompt = f"{fim_begin_token}{prompt}{fim_end_token}{suffix}{fim_middle_token}" + + return prompt + + +register_completion_template( + CompletionTemplate( + name="deepseek_coder", + fim_begin_token="<｜fim▁begin｜>", + fim_middle_token="<｜fim▁hole｜>", + fim_end_token="<｜fim▁end｜>", + fim_position=FimPosition.MIDDLE, + ) +) + + +register_completion_template( + CompletionTemplate( + name="star_coder", + fim_begin_token="", + fim_middle_token="", + fim_end_token="", + fim_position=FimPosition.END, + ) +) + +register_completion_template( + CompletionTemplate( + name="qwen_coder", + fim_begin_token="<|fim_prefix|>", + fim_middle_token="<|fim_middle|>", + fim_end_token="<|fim_suffix|>", + fim_position=FimPosition.END, + ) +) diff --git a/sglang/python/sglang/srt/configs/__init__.py b/sglang/python/sglang/srt/configs/__init__.py new file mode 100644 index 0000000..1e8370b --- /dev/null +++ b/sglang/python/sglang/srt/configs/__init__.py @@ -0,0 +1,13 @@ +from sglang.srt.configs.chatglm import ChatGLMConfig +from sglang.srt.configs.dbrx import DbrxConfig +from sglang.srt.configs.deepseekvl2 import DeepseekVL2Config +from sglang.srt.configs.exaone import ExaoneConfig +from sglang.srt.configs.janus_pro import MultiModalityConfig + +__all__ = [ + "ExaoneConfig", + "ChatGLMConfig", + "DbrxConfig", + "DeepseekVL2Config", + "MultiModalityConfig", +] diff --git a/sglang/python/sglang/srt/configs/chatglm.py b/sglang/python/sglang/srt/configs/chatglm.py new file mode 100644 index 0000000..9370c21 --- /dev/null +++ b/sglang/python/sglang/srt/configs/chatglm.py @@ -0,0 +1,78 @@ +# Adapted from +# https://github.com/THUDM/ChatGLM2-6B +# https://github.com/vllm-project/vllm/blob/main/vllm/transformers_utils/configs/chatglm.py + +# ChatGLM2 and ChatGLM3 share the same config. +# ChatGLM4 is officially supported by Huggingface +# transformers >= 4.46.0 is required +# https://huggingface.co/docs/transformers/en/model_doc/glm +from transformers import PretrainedConfig + + +class ChatGLMConfig(PretrainedConfig): + model_type = "chatglm" + attribute_map = { + "num_hidden_layers": "num_layers", + "n_head_kv": "multi_query_group_num", + } + + def __init__( + self, + num_layers=28, + padded_vocab_size=65024, + hidden_size=4096, + ffn_hidden_size=13696, + kv_channels=128, + num_attention_heads=32, + seq_length=2048, + hidden_dropout=0.0, + attention_dropout=0.0, + layernorm_epsilon=1e-5, + rmsnorm=True, + apply_residual_connection_post_layernorm=False, + post_layer_norm=True, + add_bias_linear=False, + add_qkv_bias=False, + interleaved_qkv=False, + bias_dropout_fusion=True, + multi_query_attention=False, + multi_query_group_num=1, + apply_query_key_layer_scaling=True, + attention_softmax_in_fp32=True, + fp32_residual_connection=False, + quantization_bit=0, + pre_seq_len=None, + prefix_projection=False, + **kwargs + ): + self.num_layers = num_layers + self.vocab_size = padded_vocab_size + self.padded_vocab_size = padded_vocab_size + self.hidden_size = hidden_size + self.ffn_hidden_size = ffn_hidden_size + self.kv_channels = kv_channels + self.num_attention_heads = num_attention_heads + self.seq_length = seq_length + # It is to be compatible with long lora. + self.max_position_embeddings = seq_length + self.hidden_dropout = hidden_dropout + self.attention_dropout = attention_dropout + self.layernorm_epsilon = layernorm_epsilon + self.rmsnorm = rmsnorm + self.apply_residual_connection_post_layernorm = ( + apply_residual_connection_post_layernorm + ) + self.post_layer_norm = post_layer_norm + self.add_bias_linear = add_bias_linear + self.add_qkv_bias = add_qkv_bias + self.bias_dropout_fusion = bias_dropout_fusion + self.multi_query_attention = multi_query_attention + self.multi_query_group_num = multi_query_group_num + self.apply_query_key_layer_scaling = apply_query_key_layer_scaling + self.attention_softmax_in_fp32 = attention_softmax_in_fp32 + self.fp32_residual_connection = fp32_residual_connection + self.quantization_bit = quantization_bit + self.pre_seq_len = pre_seq_len + self.prefix_projection = prefix_projection + self.interleaved_qkv = interleaved_qkv + super().__init__(**kwargs) diff --git a/sglang/python/sglang/srt/configs/dbrx.py b/sglang/python/sglang/srt/configs/dbrx.py new file mode 100644 index 0000000..75ccbde --- /dev/null +++ b/sglang/python/sglang/srt/configs/dbrx.py @@ -0,0 +1,279 @@ +# Adapted from +# https://huggingface.co/databricks/dbrx-base/blob/main/configuration_dbrx.py +# https://github.com/vllm-project/vllm/blob/main/vllm/transformers_utils/configs/dbrx.py +"""Dbrx configuration.""" + +from typing import Any, Optional + +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + +DBRX_PRETRAINED_CONFIG_ARCHIVE_MAP = {} # type: ignore + + +class DbrxAttentionConfig(PretrainedConfig): + """Configuration class for Dbrx Attention. + + [`DbrxAttention`] class. It is used to instantiate attention layers + according to the specified arguments, defining the layers architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + attn_pdrop (`float`, *optional*, defaults to 0.0): + The dropout probability for the attention layers. + clip_qkv (`float`, *optional*, defaults to None): + If not `None`, clip the queries, keys, and values in the attention layer to this value. + kv_n_heads (Optional[int]): For grouped_query_attention only, allow user to specify number of kv heads. + rope_theta (float): The base frequency for rope. + """ + + def __init__( + self, + attn_pdrop: float = 0, + clip_qkv: Optional[float] = None, + kv_n_heads: int = 1, + rope_theta: float = 10000.0, + **kwargs: Any, + ): + super().__init__(**kwargs) + self.attn_pdrop = attn_pdrop + self.clip_qkv = clip_qkv + self.kv_n_heads = kv_n_heads + self.rope_theta = rope_theta + + for k in ["model_type"]: + if k in kwargs: + kwargs.pop(k) + if len(kwargs) != 0: + raise ValueError(f"Found unknown {kwargs=}") + + @classmethod + def from_pretrained( + cls, pretrained_model_name_or_path: str, **kwargs: Any + ) -> "PretrainedConfig": + cls._set_token_in_kwargs(kwargs) + + config_dict, kwargs = cls.get_config_dict( + pretrained_model_name_or_path, **kwargs + ) + + if config_dict.get("model_type") == "dbrx": + config_dict = config_dict["attn_config"] + + if ( + "model_type" in config_dict + and hasattr(cls, "model_type") + and config_dict["model_type"] != cls.model_type + ): + logger.warning( + "You are using a model of type %s to instantiate a model of " + "type %s. This is not supported for all configurations of " + "models and can yield errors.", + config_dict["model_type"], + cls.model_type, + ) + + return cls.from_dict(config_dict, **kwargs) + + +class DbrxFFNConfig(PretrainedConfig): + """Configuration class for Dbrx FFN. + + [`DbrxFFN`] class. It is used to instantiate feedforward layers according to + the specified arguments, defining the layers architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + Args: + ffn_act_fn (dict, optional): A dict specifying activation function for the FFN. + The dict should have a key 'name' with the value being the name of + the activation function along with any additional keyword arguments. + ffn_hidden_size (int, optional): The hidden size of the feedforward network. + moe_num_experts (int, optional): The number of experts in the mixture of experts layer. + moe_top_k (int, optional): The number of experts to use in the mixture of experts layer. + moe_jitter_eps (float, optional): The jitter epsilon for the mixture of experts layer. + moe_loss_weight (float, optional): The loss weight for the mixture of experts layer. + moe_normalize_expert_weights (float, optional): The normalization factor for the expert weights. + uniform_expert_assignment (bool, optional): Whether to use uniform expert assignment. + This should only be used for benchmarking purposes. + """ + + def __init__( + self, + ffn_act_fn: Optional[dict] = None, + ffn_hidden_size: int = 3584, + moe_num_experts: int = 4, + moe_top_k: int = 1, + moe_jitter_eps: Optional[float] = None, + moe_loss_weight: float = 0.01, + moe_normalize_expert_weights: Optional[float] = 1, + uniform_expert_assignment: bool = False, + **kwargs: Any, + ): + super().__init__() + if ffn_act_fn is None: + ffn_act_fn = {"name": "silu"} + self.ffn_act_fn = ffn_act_fn + self.ffn_hidden_size = ffn_hidden_size + self.moe_num_experts = moe_num_experts + self.moe_top_k = moe_top_k + self.moe_jitter_eps = moe_jitter_eps + self.moe_loss_weight = moe_loss_weight + self.moe_normalize_expert_weights = moe_normalize_expert_weights + self.uniform_expert_assignment = uniform_expert_assignment + + for k in ["model_type"]: + if k in kwargs: + kwargs.pop(k) + if len(kwargs) != 0: + raise ValueError(f"Found unknown {kwargs=}") + + @classmethod + def from_pretrained( + cls, pretrained_model_name_or_path: str, **kwargs: Any + ) -> "PretrainedConfig": + cls._set_token_in_kwargs(kwargs) + + config_dict, kwargs = cls.get_config_dict( + pretrained_model_name_or_path, **kwargs + ) + + if config_dict.get("model_type") == "dbrx": + config_dict = config_dict["ffn_config"] + + if ( + "model_type" in config_dict + and hasattr(cls, "model_type") + and config_dict["model_type"] != cls.model_type + ): + logger.warning( + "You are using a model of type %s to instantiate a model of " + "type %s. This is not supported for all " + "configurations of models and can yield errors.", + config_dict["model_type"], + cls.model_type, + ) + + return cls.from_dict(config_dict, **kwargs) + + +class DbrxConfig(PretrainedConfig): + """Configuration class for Dbrx. + + [`DbrxModel`]. It is used to instantiate a Dbrx model according to the + specified arguments, defining the model architecture. + + Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the + documentation from [`PretrainedConfig`] for more information. + + + Args: + d_model (`int`, *optional*, defaults to 6144): + Dimensionality of the embeddings and hidden states. + n_heads (`int`, *optional*, defaults to 48): + Number of attention heads for each attention layer in the Transformer encoder. + n_layers (`int`, *optional*, defaults to 40): + Number of hidden layers in the Transformer encoder. + max_seq_len (`int`, *optional*, defaults to 32768): + The maximum sequence length of the model. + vocab_size (`int`, *optional*, defaults to 100352): + Vocabulary size of the Dbrx model. Defines the maximum number of different tokens that can be represented by + the `inputs_ids` passed when calling [`DbrxModel`]. + resid_pdrop (`float`, *optional*, defaults to 0.0): + The dropout probability applied to the attention output before combining with residual. + emb_pdrop (`float`, *optional*, defaults to 0.0): + The dropout probability for the embedding layer. + attn_config (`dict`, *optional*): + A dictionary used to configure the model's attention module. + ffn_config (`dict`, *optional*): + A dictionary used to configure the model's FFN module. + use_cache (`bool`, *optional*, defaults to `False`): + Whether or not the model should return the last key/values attentions (not used by all models). + initializer_range (`float`, *optional*, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + output_router_logits (`bool`, *optional*, defaults to `False`): + Whether or not the router logits should be returned by the model. Enabling this will also + allow the model to output the auxiliary loss. See [here]() for more details + router_aux_loss_coef (`float`, *optional*, defaults to 0.001): + The aux loss factor for the total loss. + + + Example: + ```python + >>> from transformers import DbrxConfig, DbrxModel + + >>> # Initializing a Dbrx configuration + >>> configuration = DbrxConfig() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = DbrxModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "dbrx" + attribute_map = { + "num_attention_heads": "n_heads", + "hidden_size": "d_model", + "num_hidden_layers": "n_layers", + "max_position_embeddings": "max_seq_len", + } + + def __init__( + self, + d_model: int = 2048, + n_heads: int = 16, + n_layers: int = 24, + max_seq_len: int = 2048, + vocab_size: int = 32000, + resid_pdrop: float = 0.0, + emb_pdrop: float = 0.0, + attn_config: Optional[DbrxAttentionConfig] = None, + ffn_config: Optional[DbrxFFNConfig] = None, + use_cache: bool = True, + initializer_range: float = 0.02, + output_router_logits: bool = False, + router_aux_loss_coef: float = 0.05, + **kwargs: Any, + ): + if attn_config is None: + self.attn_config = DbrxAttentionConfig() + elif isinstance(attn_config, dict): + self.attn_config = DbrxAttentionConfig(**attn_config) + else: + self.attn_config = attn_config + + if ffn_config is None: + self.ffn_config = DbrxFFNConfig() + elif isinstance(ffn_config, dict): + self.ffn_config = DbrxFFNConfig(**ffn_config) + else: + self.ffn_config = ffn_config + + self.d_model = d_model + self.n_heads = n_heads + self.n_layers = n_layers + self.max_seq_len = max_seq_len + self.vocab_size = vocab_size + self.resid_pdrop = resid_pdrop + self.emb_pdrop = emb_pdrop + self.use_cache = use_cache + self.initializer_range = initializer_range + self.output_router_logits = output_router_logits + self.router_aux_loss_coef = router_aux_loss_coef + + tie_word_embeddings = kwargs.pop("tie_word_embeddings", False) + if tie_word_embeddings: + raise ValueError("tie_word_embeddings is not supported for Dbrx models.") + + super().__init__( + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) diff --git a/sglang/python/sglang/srt/configs/deepseekvl2.py b/sglang/python/sglang/srt/configs/deepseekvl2.py new file mode 100644 index 0000000..a045492 --- /dev/null +++ b/sglang/python/sglang/srt/configs/deepseekvl2.py @@ -0,0 +1,676 @@ +import math +import os +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import torch +from PIL import Image, ImageOps +from transformers import ( + AutoProcessor, + LlamaTokenizerFast, + PretrainedConfig, + ProcessorMixin, +) + + +def select_best_resolution(image_size, candidate_resolutions): + # used for cropping + original_width, original_height = image_size + best_fit = None + max_effective_resolution = 0 + min_wasted_resolution = float("inf") + + for width, height in candidate_resolutions: + scale = min(width / original_width, height / original_height) + downscaled_width, downscaled_height = int(original_width * scale), int( + original_height * scale + ) + effective_resolution = min( + downscaled_width * downscaled_height, original_width * original_height + ) + wasted_resolution = (width * height) - effective_resolution + + if effective_resolution > max_effective_resolution or ( + effective_resolution == max_effective_resolution + and wasted_resolution < min_wasted_resolution + ): + max_effective_resolution = effective_resolution + min_wasted_resolution = wasted_resolution + best_fit = (width, height) + + return best_fit + + +class DictOutput(object): + def keys(self): + return self.__dict__.keys() + + def __getitem__(self, item): + return self.__dict__[item] + + def __setitem__(self, key, value): + self.__dict__[key] = value + + +@dataclass +class VLChatProcessorOutput(DictOutput): + input_ids: torch.LongTensor + target_ids: torch.LongTensor + images: torch.Tensor + images_seq_mask: torch.BoolTensor + images_spatial_crop: torch.LongTensor + + def __len__(self): + return len(self.input_ids) + + +class ImageTransform(object): + def __init__( + self, + mean: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5), + std: Optional[Tuple[float, float, float]] = (0.5, 0.5, 0.5), + normalize: bool = True, + ): + self.mean = mean + self.std = std + self.normalize = normalize + + # only load torchvision.transforms when needed + try: + import torchvision.transforms as T + + # FIXME: add version check for gguf + except ImportError as err: + raise ImportError( + "Please install torchvision via `pip install torchvision` to use Deepseek-VL2." + ) from err + + transform_pipelines = [T.ToTensor()] + + if normalize: + transform_pipelines.append(T.Normalize(mean, std)) + + self.transform = T.Compose(transform_pipelines) + + def __call__(self, pil_img: Image.Image): + x = self.transform(pil_img) + return x + + +class DeepseekVLV2Processor(ProcessorMixin): + tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast") + attributes = ["tokenizer"] + + def __init__( + self, + tokenizer: LlamaTokenizerFast, + candidate_resolutions: Tuple[Tuple[int, int]], + patch_size: int, + downsample_ratio: int, + image_mean: Tuple[float, float, float] = (0.5, 0.5, 0.5), + image_std: Tuple[float, float, float] = (0.5, 0.5, 0.5), + normalize: bool = True, + image_token: str = "", + pad_token: str = "<｜▁pad▁｜>", + add_special_token: bool = False, + sft_format: str = "deepseek", + mask_prompt: bool = True, + ignore_id: int = -100, + **kwargs, + ): + + self.candidate_resolutions = candidate_resolutions + self.image_size = candidate_resolutions[0][0] + self.patch_size = patch_size + self.image_mean = image_mean + self.image_std = image_std + self.normalize = normalize + self.downsample_ratio = downsample_ratio + + self.image_transform = ImageTransform( + mean=image_mean, std=image_std, normalize=normalize + ) + self.tokenizer = tokenizer + # must set this，padding side with make a difference in batch inference + self.tokenizer.padding_side = "left" + + # add the pad_token as special token to use 'tokenizer.pad_token' and 'tokenizer.pad_token_id' + if tokenizer.pad_token is None: + self.tokenizer.add_special_tokens({"pad_token": pad_token}) + + # add image token + image_token_id = self.tokenizer.vocab.get(image_token) + if image_token_id is None: + special_tokens = [image_token] + special_tokens_dict = {"additional_special_tokens": special_tokens} + self.tokenizer.add_special_tokens(special_tokens_dict) + self.image_token_id = self.tokenizer.vocab.get(image_token) + + # add five special tokens for grounding-related tasks + # <|ref|>, <|/ref|>, <|det|>, <|/det|>, <|grounding|> + special_tokens = ["<|ref|>", "<|/ref|>", "<|det|>", "<|/det|>", "<|grounding|>"] + special_tokens_dict = {"additional_special_tokens": special_tokens} + self.tokenizer.add_special_tokens(special_tokens_dict) + + # add special tokens for SFT data + special_tokens = ["<|User|>", "<|Assistant|>"] + special_tokens_dict = {"additional_special_tokens": special_tokens} + self.tokenizer.add_special_tokens(special_tokens_dict) + + self.image_token = image_token + self.pad_token = pad_token + self.add_special_token = add_special_token + self.sft_format = sft_format + self.mask_prompt = mask_prompt + self.ignore_id = ignore_id + + super().__init__( + tokenizer, + **kwargs, + ) + + def format_messages_v2(self, messages, pil_images, max_req_input_len=-1): + """play the role of format_messages_v2 and get_images_info in the last version""" + tokenized_data = [] + masked_tokenized_data = [] # labels + images_list = [] + images_seq_mask = [] + images_spatial_crop = [] + + image_index = 0 + image_token_cnt = messages.count(self.image_token) + tokenized_str, images, seq_mask, spatial_crop = self.tokenize_with_images( + messages, + pil_images[image_index : image_index + image_token_cnt], + bos=False, + eos=True, + cropping=len(pil_images) <= 2, + max_req_input_len=max_req_input_len, + ) + + image_index = image_token_cnt + tokenized_data += tokenized_str + if self.mask_prompt: + masked_tokenized_data += [self.ignore_id] * len(tokenized_str) + else: + masked_tokenized_data += tokenized_str + images_list += images + images_seq_mask += seq_mask + images_spatial_crop += spatial_crop + + assert len(tokenized_data) == len( + images_seq_mask + ), f"format_messages_v2: tokenized_str's length {len(tokenized_str)} is not equal to imags_seq_mask's length {len(images_seq_mask)}" + + return ( + tokenized_data, + masked_tokenized_data, + images_list, + images_seq_mask, + images_spatial_crop, + ) + + @property + def bos_id(self): + return self.tokenizer.bos_token_id + + @property + def eos_id(self): + return self.tokenizer.eos_token_id + + @property + def pad_id(self): + return self.tokenizer.pad_token_id + + def encode(self, text: str, bos: bool = True, eos: bool = False): + t = self.tokenizer.encode(text, add_special_tokens=False) + + if bos: + t = [self.bos_id] + t + if eos: + t = t + [self.eos_id] + + return t + + def decode(self, t: List[int], **kwargs) -> str: + return self.tokenizer.decode(t, **kwargs) + + def process_one( + self, + prompt: str = None, + conversations: List[Dict[str, str]] = None, + images: List[Image.Image] = None, + apply_sft_format: bool = False, + inference_mode: bool = True, + system_prompt: str = "", + max_req_input_len: int = -1, + **kwargs, + ): + """ + + Args: + prompt (str): the formatted prompt; + conversations (List[Dict]): conversations with a list of messages; + images (List[ImageType]): the list of images; + apply_sft_format (bool): if prompt is not None, then apply the SFT format to prompt; + if conversations is not None, then it will always apply the SFT format to conversations; + inference_mode (bool): if True, then remove the last eos token; + system_prompt (str): the system prompt; + **kwargs: + + Returns: + outputs (BaseProcessorOutput): the output of the processor, + - input_ids (torch.LongTensor): [N + image tokens] + - target_ids (torch.LongTensor): [N + image tokens] + - images (torch.FloatTensor): [n_images, 3, H, W] + - image_id (int): the id of the image token + - num_image_tokens (List[int]): the number of image tokens + """ + + assert ( + prompt is None or conversations is None + ), "prompt and conversations cannot be used at the same time." + + ( + tokenized_str, + masked_tokenized_str, + images_list, + images_seq_mask, + images_spatial_crop, + ) = self.format_messages_v2(conversations, images, max_req_input_len) + + assert ( + len(tokenized_str) == len(images_seq_mask) == len(masked_tokenized_str) + ), ( + f"tokenized_str's length {len(tokenized_str)}, input_ids' length {len(masked_tokenized_str)}, " + f"imags_seq_mask's length {len(images_seq_mask)}, are not equal" + ) + + input_ids = torch.LongTensor(tokenized_str) + target_ids = torch.LongTensor(masked_tokenized_str) + images_seq_mask = torch.tensor(images_seq_mask, dtype=torch.bool) + + # set input_ids < 0 | input_ids == self.image_token_id as ignore_id + target_ids[(input_ids < 0) | (input_ids == self.image_token_id)] = ( + self.ignore_id + ) + input_ids[input_ids < 0] = self.pad_id + + if inference_mode: + assert input_ids[-1] == self.eos_id + input_ids = input_ids[:-1] + target_ids = target_ids[:-1] + images_seq_mask = images_seq_mask[:-1] + + if len(images_list) == 0: + images = torch.zeros((1, 3, self.image_size, self.image_size)) + images_spatial_crop = torch.zeros((1, 2), dtype=torch.long) + else: + images = torch.stack(images_list, dim=0) + images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long) + + prepare = VLChatProcessorOutput( + input_ids=input_ids, + target_ids=target_ids, + images=images, + images_seq_mask=images_seq_mask, + images_spatial_crop=images_spatial_crop, + ) + + return prepare + + def __call__( + self, + *, + prompt: str = None, + conversations: List[Dict[str, str]] = None, + images: List[Image.Image] = None, + apply_sft_format: bool = False, + inference_mode: bool = True, + system_prompt: str = "", + max_req_input_len: int = -1, + **kwargs, + ): + prepare = self.process_one( + prompt=prompt, + conversations=conversations, + images=images, + apply_sft_format=apply_sft_format, + inference_mode=inference_mode, + system_prompt=system_prompt, + max_req_input_len=max_req_input_len, + ) + + return prepare + + def find_all_indices(self, messages, target_value): + indices = [] + for index, item in enumerate(messages): + if item == target_value: + indices.append(index) + return indices + + def tokenize_with_images( + self, + conversation: str, + images: List[Image.Image], + bos: bool = True, + eos: bool = True, + cropping: bool = True, + max_req_input_len: int = -1, + ): + """Tokenize text with tags.""" + images_list, images_seq_mask, images_spatial_crop = [], [], [] + text_splits = conversation.split(self.image_token) + tokenized_str = [] + for text_sep, image in zip(text_splits, images): + """encode text_sep""" + tokenized_sep = self.encode(text_sep, bos=False, eos=False) + tokenized_str += tokenized_sep + images_seq_mask += [False] * len(tokenized_sep) + + """select best resolution for anyres""" + if cropping: + best_width, best_height = select_best_resolution( + image.size, self.candidate_resolutions + ) + else: + best_width, best_height = self.image_size, self.image_size + # print(image.size, (best_width, best_height)) # check the select_best_resolutions func + + """process the global view""" + global_view = ImageOps.pad( + image, + (self.image_size, self.image_size), + color=tuple(int(x * 255) for x in self.image_transform.mean), + ) + images_list.append(self.image_transform(global_view)) + + """process the local views""" + local_view = ImageOps.pad( + image, + (best_width, best_height), + color=tuple(int(x * 255) for x in self.image_transform.mean), + ) + for i in range(0, best_height, self.image_size): + for j in range(0, best_width, self.image_size): + images_list.append( + self.image_transform( + local_view.crop( + (j, i, j + self.image_size, i + self.image_size) + ) + ) + ) + + """record height / width crop num""" + num_width_tiles, num_height_tiles = ( + best_width // self.image_size, + best_height // self.image_size, + ) + images_spatial_crop.append([num_width_tiles, num_height_tiles]) + + """add image tokens""" + h = w = math.ceil( + (self.image_size // self.patch_size) / self.downsample_ratio + ) + # global views tokens h * (w + 1), 1 is for line seperator + tokenized_image = [self.image_token_id] * h * (w + 1) + # add a seperator between global and local views + tokenized_image += [self.image_token_id] + # local views tokens, (num_height_tiles * h) * (num_width_tiles * w + 1) + tokenized_image += ( + [self.image_token_id] + * (num_height_tiles * h) + * (num_width_tiles * w + 1) + ) + + tokenized_str += tokenized_image + images_seq_mask += [True] * len(tokenized_image) + # print(width_crop_num, height_crop_num, len(tokenized_image)) # test the correctness of the number of image-related tokens + + """process the last text split""" + tokenized_sep = self.encode(text_splits[-1], bos=False, eos=False) + # deal with video, limit with request len + if max_req_input_len > -1: + if max_req_input_len < len(tokenized_sep) + len(tokenized_str) - 1: + rest = max_req_input_len - len(tokenized_sep) - 1 - 1024 + tokenized_str = tokenized_str[:rest] + images_seq_mask = images_seq_mask[:rest] + tokenized_str += tokenized_sep + images_seq_mask += [False] * len(tokenized_sep) + + """add the bos and eos tokens""" + if bos: + tokenized_str = [self.bos_id] + tokenized_str + images_seq_mask = [False] + images_seq_mask + if eos: + tokenized_str = tokenized_str + [self.eos_id] + images_seq_mask = images_seq_mask + [False] + + assert len(tokenized_str) == len( + images_seq_mask + ), f"tokenize_with_images func: tokenized_str's length {len(tokenized_str)} is not equal to imags_seq_mask's length {len(images_seq_mask)}" + + return tokenized_str, images_list, images_seq_mask, images_spatial_crop + + +class DeepseekVL2VisionEncoderConfig(PretrainedConfig): + model_type: str = "vision" + + model_name: str = "siglip_large_patch16_384" + image_size: int = 384 + patch_size: int = 16 + width: int = 1024 + layers: int = 24 + heads: int = 16 + mlp_ratio: int = 4 + global_pool: str = "map" + ignore_head: bool = True + class_token: bool = False + num_classes: int = 0 + use_checkpoint: bool = False + weight_init: str = "skip" + deterministic: bool = False + num_recomputing_layers: int = 0 + + def __init__( + self, + model_name: str = "siglip_large_patch16_384", + image_size: int = 384, + patch_size: int = 16, + width: int = 1024, + layers: int = 24, + heads: int = 16, + mlp_ratio: int = 4, + global_pool: str = "map", + ignore_head: bool = True, + class_token: bool = False, + num_classes: int = 0, + use_checkpoint: bool = False, + **kwargs, + ): + self.model_name = model_name + self.image_size = image_size + self.patch_size = patch_size + self.width = width + self.layers = layers + self.heads = heads + self.mlp_ratio = mlp_ratio + self.global_pool = global_pool + self.ignore_head = ignore_head + self.class_token = class_token + self.num_classes = num_classes + self.use_checkpoint = use_checkpoint + + super().__init__(**kwargs) + + +class DeepseekVL2MlpProjectorConfig(PretrainedConfig): + model_type = "mlp_projector" + projector_type: str = "downsample_mlp_gelu" + input_dim: int = 1152 + n_embed: int = 2048 + depth: int = 2 + mlp_ratio: int = 1 + downsample_ratio: int = 2 + token_pooling: bool = False + + def __init__( + self, + projector_type: str = "downsample_mlp_gelu", + input_dim: int = 1152, + n_embed: int = 2048, + depth: int = 2, + mlp_ratio: int = 1, + downsample_ratio: int = 2, + **kwargs, + ): + self.projector_type = projector_type + self.input_dim = input_dim + self.n_embed = n_embed + self.depth = depth + self.mlp_ratio = mlp_ratio + self.downsample_ratio = downsample_ratio + + super().__init__(**kwargs) + + +class DeepseekV2Config(PretrainedConfig): + + model_type = "deepseek_v2" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + vocab_size=102400, + hidden_size=4096, + intermediate_size=11008, + moe_intermediate_size=1407, + num_hidden_layers=30, + num_attention_heads=32, + num_key_value_heads=32, + n_shared_experts=None, + n_routed_experts=None, + ep_size=1, + routed_scaling_factor=1.0, + kv_lora_rank=512, + q_lora_rank=1536, + qk_rope_head_dim=64, + v_head_dim=128, + qk_nope_head_dim=128, + topk_method="gready", + n_group=None, + topk_group=None, + num_experts_per_tok=None, + moe_layer_freq=1, + first_k_dense_replace=0, + norm_topk_prob=False, + scoring_func="softmax", + aux_loss_alpha=0.001, + seq_aux=True, + hidden_act="silu", + max_position_embeddings=2048, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + pad_token_id=None, + bos_token_id=100000, + eos_token_id=100001, + pretraining_tp=1, + tie_word_embeddings=False, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + use_mla=True, + **kwargs, + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.moe_intermediate_size = moe_intermediate_size + self.num_hidden_layers = num_hidden_layers + self.num_attention_heads = num_attention_heads + self.n_shared_experts = n_shared_experts + self.n_routed_experts = n_routed_experts + self.ep_size = ep_size + self.routed_scaling_factor = routed_scaling_factor + self.kv_lora_rank = kv_lora_rank + self.q_lora_rank = q_lora_rank + self.qk_rope_head_dim = qk_rope_head_dim + self.v_head_dim = v_head_dim + self.qk_nope_head_dim = qk_nope_head_dim + self.topk_method = topk_method + self.n_group = n_group + self.topk_group = topk_group + self.num_experts_per_tok = num_experts_per_tok + self.moe_layer_freq = moe_layer_freq + self.first_k_dense_replace = first_k_dense_replace + self.norm_topk_prob = norm_topk_prob + self.scoring_func = scoring_func + self.aux_loss_alpha = aux_loss_alpha + self.seq_aux = seq_aux + # for backward compatibility + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + + self.num_key_value_heads = num_key_value_heads + self.hidden_act = hidden_act + self.initializer_range = initializer_range + self.rms_norm_eps = float(rms_norm_eps) + self.pretraining_tp = pretraining_tp + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + self.attention_bias = attention_bias + self.attention_dropout = attention_dropout + self.use_mla = use_mla + + super().__init__( + pad_token_id=pad_token_id, + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs, + ) + + +class DeepseekVL2Config(PretrainedConfig): + model_type = "deepseek_vl_v2" + vision_config: DeepseekVL2VisionEncoderConfig + projector_config: DeepseekVL2MlpProjectorConfig + language_config: DeepseekV2Config + + tile_tag: str = "2D" + global_view_pos: str = "head" + candidate_resolutions: Tuple[Tuple[int, int]] = ((384, 384),) + + def __init__( + self, + tile_tag: str = "tile_tag", + global_view_pos: str = "head", + candidate_resolutions: Tuple[Tuple[int, int]] = ((384, 384),), + **kwargs, + ): + super().__init__(**kwargs) + + vision_config = kwargs.get("vision_config", {}) + self.vision_config = DeepseekVL2VisionEncoderConfig(**vision_config) + + projector_config = kwargs.get("projector_config", {}) + self.projector_config = DeepseekVL2MlpProjectorConfig(**projector_config) + + language_config = kwargs.get("language_config", {}) + if isinstance(language_config, DeepseekV2Config): + self.language_config = language_config + else: + self.language_config = DeepseekV2Config(**language_config) + + self.tile_tag = tile_tag + self.global_view_pos = global_view_pos + self.candidate_resolutions = candidate_resolutions + self.architectures = ["DeepseekVL2ForCausalLM"] + + +AutoProcessor.register(DeepseekVL2Config, DeepseekVLV2Processor) diff --git a/sglang/python/sglang/srt/configs/device_config.py b/sglang/python/sglang/srt/configs/device_config.py new file mode 100644 index 0000000..d95e848 --- /dev/null +++ b/sglang/python/sglang/srt/configs/device_config.py @@ -0,0 +1,17 @@ +import logging +from typing import Optional + +import torch + +logger = logging.getLogger(__name__) + + +class DeviceConfig: + device: Optional[torch.device] + + def __init__(self, device: str = "cuda") -> None: + if device in ["cuda", "xpu", "hpu", "cpu"]: + self.device_type = device + else: + raise RuntimeError(f"Not supported device type: {device}") + self.device = torch.device(self.device_type) diff --git a/sglang/python/sglang/srt/configs/exaone.py b/sglang/python/sglang/srt/configs/exaone.py new file mode 100644 index 0000000..7b0a2d2 --- /dev/null +++ b/sglang/python/sglang/srt/configs/exaone.py @@ -0,0 +1,195 @@ +# coding=utf-8 +# Copyright 2024 The LG AI Research EXAONE Lab. All rights reserved. +# Copyright 2024 The LG CNS AI Engineering Team. +# Copyright 2023-2024 SGLang Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" EXAONE model configuration """ +from typing import Any, Dict + +from transformers.configuration_utils import PretrainedConfig +from transformers.utils import logging + +logger = logging.get_logger(__name__) + +EXAONE_PRETRAINED_CONFIG_ARCHIVE_MAP: Dict[str, Any] = {} + + +# ruff: noqa: E501 +class ExaoneConfig(PretrainedConfig): + r""" + This is the configuration class to store the configuration of a :class:`~transformers.ExaoneModel`. It is used to + instantiate a EXAONE model according to the specified arguments, defining the model architecture. Instantiating a + configuration with the defaults will yield a similar configuration to that of the Exaone + + Configuration objects inherit from :class:`~transformers.PretrainedConfig` and can be used to control the model + outputs. Read the documentation from :class:`~transformers.PretrainedConfig` for more information. + + + Args: + vocab_size (:obj:`int`, `optional`, defaults to 102400): + Vocabulary size of the EXAONE model. Defines the number of different tokens that can be represented by the + :obj:`inputs_ids` passed when calling :class:`~transformers.ExaoneModel`. Vocabulary size of the model. + Defines the different tokens that can be represented by the `inputs_ids` passed to the forward method of + :class:`~transformers.EXAONEModel`. + max_position_embeddings (:obj:`int`, `optional`, defaults to 2048): + The maximum sequence length that this model might ever be used with. Typically set this to something large + just in case (e.g., 512 or 1024 or 2048). + hidden_size (:obj:`int`, `optional`, defaults to 2048): + Dimensionality of the encoder layers and the pooler layer. + num_layers (:obj:`int`, `optional`, defaults to 32): + Number of hidden layers in the Transformer encoder. + num_attention_heads (:obj:`int`, `optional`, defaults to 32): + Number of attention heads for each attention layer in the Transformer decoder. + num_key_value_heads (:obj:`int`, `optional`): + This is the number of key_value heads that should be used to implement Grouped Query Attention. If + `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if + `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When + converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed + by meanpooling all the original heads within that group. For more details checkout [this + paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to + `num_attention_heads`. + intermediate_size (:obj:`int`, `optional`, defaults to `hidden_size * 4`): + Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder. + activation_function (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"silu"`): + The non-linear activation function (function or string) in the decoder. + rope_theta (:obj:`float`, `optional`, defaults to 10000.0): + The base period of the RoPE embeddings. + rope_scaling (:obj:`Dict`, `optional`): + Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type + and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value + accordingly. + Expected contents: + `rope_type` (:obj:`str`): + The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope', + 'llama3'], with 'default' being the original RoPE implementation. + `factor` (:obj:`float`, `optional`): + Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In + most scaling types, a `factor` of x will enable the model to handle sequences of length x * + original maximum pre-trained length. + `original_max_position_embeddings` (:obj:`int`, `optional`): + Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during + pretraining. + `attention_factor` (:obj:`float`, `optional`): + Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention + computation. If unspecified, it defaults to value recommended by the implementation, using the + `factor` field to infer the suggested value. + `beta_fast` (:obj:`float`, `optional`): + Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear + ramp function. If unspecified, it defaults to 32. + `beta_slow` (:obj:`float`, `optional`): + Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear + ramp function. If unspecified, it defaults to 1. + `short_factor` (:obj:`List[float]`, `optional`): + Only used with 'longrope'. The scaling factor to be applied to short contexts (< + `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden + size divided by the number of attention heads divided by 2 + `long_factor` (:obj:`List[float]`, `optional`): + Only used with 'longrope'. The scaling factor to be applied to long contexts (< + `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden + size divided by the number of attention heads divided by 2 + `low_freq_factor` (:obj:`float`, `optional`): + Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE + `high_freq_factor` (:obj:`float`, `optional`): + Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE + embed_dropout (:obj:`float`, `optional`, defaults to 0.0): + The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler. + attention_dropout (:obj:`float`, `optional`, defaults to 0.0): + The dropout ratio for the attention probabilities. + layer_norm_epsilon (:obj:`float`, `optional`, defaults to 1e-5): + The epsilon used by the layer normalization layers. + initializer_range (:obj:`float`, `optional`, defaults to 0.02): + The standard deviation of the truncated_normal_initializer for initializing all weight matrices. + use_cache (:obj:`bool`, `optional`, defaults to :obj:`True`): + Whether or not the model should return the last key/values attentions (not used by all models). Only + relevant if ``configs.is_decoder=True``. + bos_token_id (:obj:`int`, `optional`, defaults to 0): + Beginning of stream token id. + eos_token_id (:obj:`int`, `optional`, defaults to 2): + End of stream token id. + tie_word_embeddings (:obj:`bool`, `optional`, defaults to :obj:`True`): + Whether to tie weight embeddings + gradient_checkpointing (:obj:`bool`, `optional`, defaults to :obj:`False`): + If True, use gradient checkpointing to save memory at the expense of slower backward pass. + + Example:: + + >>> from transformers import EXAONEModel, ExaoneConfig + + >>> # Initializing a EXAONE configuration + >>> configuration = ExaoneConfig() + + >>> # Initializing a model from configuration + >>> model = EXAONEModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.configs + """ + + model_type = "exaone" + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = {"num_hidden_layers": "num_layers"} + + def __init__( + self, + vocab_size=102400, + max_position_embeddings=2048, + hidden_size=2048, + num_layers=32, + num_attention_heads=32, + num_key_value_heads=None, + intermediate_size=None, + activation_function="silu", + rope_theta=10000.0, + rope_scaling=None, + embed_dropout=0.0, + attention_dropout=0.0, + layer_norm_epsilon=1e-5, + initializer_range=0.02, + use_cache=True, + bos_token_id=0, + eos_token_id=2, + tie_word_embeddings=True, + **kwargs + ): + self.vocab_size = vocab_size + self.max_position_embeddings = max_position_embeddings + self.hidden_size = hidden_size + self.num_layers = num_layers + self.num_attention_heads = num_attention_heads + self.num_hidden_layers = num_layers + if num_key_value_heads is None: + num_key_value_heads = num_attention_heads + self.num_key_value_heads = num_key_value_heads + if intermediate_size: + self.intermediate_size = intermediate_size + else: + self.intermediate_size = hidden_size * 4 + self.activation_function = activation_function + self.embed_dropout = embed_dropout + self.attention_dropout = attention_dropout + self.layer_norm_epsilon = layer_norm_epsilon + self.initializer_range = initializer_range + self.use_cache = use_cache + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + + self.bos_token_id = bos_token_id + self.eos_token_id = eos_token_id + + super().__init__( + bos_token_id=bos_token_id, + eos_token_id=eos_token_id, + tie_word_embeddings=tie_word_embeddings, + **kwargs + ) diff --git a/sglang/python/sglang/srt/configs/janus_pro.py b/sglang/python/sglang/srt/configs/janus_pro.py new file mode 100644 index 0000000..ad254ed --- /dev/null +++ b/sglang/python/sglang/srt/configs/janus_pro.py @@ -0,0 +1,628 @@ +# Adapted from: +# https://github.com/deepseek-ai/Janus/tree/main/janus/models + +from dataclasses import dataclass +from typing import Dict, List, Tuple, Union + +import numpy as np +import PIL +import torch +from PIL.Image import Image +from transformers import ( + BaseImageProcessor, + BatchFeature, + LlamaConfig, + LlamaTokenizerFast, + PretrainedConfig, + ProcessorMixin, +) +from transformers.image_utils import to_numpy_array + +from sglang.srt.configs.utils import register_image_processor, register_processor +from sglang.srt.mm_utils import expand2square + + +class DictToObject(dict): + def __init__(self, dictionary): + super(self).__init__(dictionary) + + for key, value in dictionary.items(): + if isinstance(value, dict): + value = DictToObject(value) + setattr(self, key, value) + + +class VisionConfig(PretrainedConfig): + model_type = "vision" + cls: str = "" + params = {} + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + self.cls = kwargs.get("cls", "") + if not isinstance(self.cls, str): + self.cls = self.cls.__name__ + + self.params = kwargs.get("params", {}) + + +class GenAlignerConfig(PretrainedConfig): + model_type = "gen_aligner" + cls: str = "" + params = {} + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + self.cls = kwargs.get("cls", "") + if not isinstance(self.cls, str): + self.cls = self.cls.__name__ + + self.params = kwargs.get("params", {}) + + +class GenHeadConfig(PretrainedConfig): + model_type = "gen_head" + cls: str = "" + params = {} + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + self.cls = kwargs.get("cls", "") + if not isinstance(self.cls, str): + self.cls = self.cls.__name__ + + self.params = kwargs.get("params", {}) + + +class AlignerConfig(PretrainedConfig): + model_type = "aligner" + cls: str = "" + params = {} + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + self.cls = kwargs.get("cls", "") + if not isinstance(self.cls, str): + self.cls = self.cls.__name__ + + self.params = kwargs.get("params", {}) + + +class GenVisionConfig(PretrainedConfig): + model_type = "gen_vision" + cls: str = "" + params = {} + + def __init__(self, **kwargs): + super().__init__(**kwargs) + + self.cls = kwargs.get("cls", "") + if not isinstance(self.cls, str): + self.cls = self.cls.__name__ + + self.params = kwargs.get("params", {}) + + +@dataclass +class SigLIPVisionCfg: + width: int = 1152 + layers: Union[Tuple[int, int, int, int], int] = 27 + heads: int = 16 + patch_size: int = 14 + image_size: Union[Tuple[int, int], int] = 336 + global_pool: str = "map" + mlp_ratio: float = 3.7362 + class_token: bool = False + num_classes: int = 0 + use_checkpoint: bool = False + + +class MultiModalityConfig(PretrainedConfig): + model_type = "multi_modality" + vision_config: VisionConfig + aligner_config: AlignerConfig + + gen_vision_config: GenVisionConfig + gen_aligner_config: GenAlignerConfig + gen_head_config: GenHeadConfig + + language_config: LlamaConfig + + def __init__(self, **kwargs): + super().__init__(**kwargs) + vision_config = kwargs.get("vision_config", {}) + self.vision_config = VisionConfig(**vision_config) + + aligner_config = kwargs.get("aligner_config", {}) + self.aligner_config = AlignerConfig(**aligner_config) + + gen_vision_config = kwargs.get("gen_vision_config", {}) + self.gen_vision_config = GenVisionConfig(**gen_vision_config) + + gen_aligner_config = kwargs.get("gen_aligner_config", {}) + self.gen_aligner_config = GenAlignerConfig(**gen_aligner_config) + + gen_head_config = kwargs.get("gen_head_config", {}) + self.gen_head_config = GenHeadConfig(**gen_head_config) + + language_config = kwargs.get("language_config", {}) + if isinstance(language_config, LlamaConfig): + self.language_config = language_config + else: + self.language_config = LlamaConfig(**language_config) + + +class VLMImageProcessor(BaseImageProcessor): + model_input_names = ["pixel_values"] + + def __init__( + self, + image_size: int, + min_size: int = 14, + image_mean: Union[Tuple[float, float, float], List[float]] = ( + 0.48145466, + 0.4578275, + 0.40821073, + ), + image_std: Union[Tuple[float, float, float], List[float]] = ( + 0.26862954, + 0.26130258, + 0.27577711, + ), + rescale_factor: float = 1.0 / 255.0, + do_normalize: bool = True, + **kwargs, + ): + super().__init__(**kwargs) + + self.image_size = image_size + self.rescale_factor = rescale_factor + self.image_mean = image_mean + self.image_std = image_std + self.min_size = min_size + self.do_normalize = do_normalize + + if image_mean is None: + self.background_color = (127, 127, 127) + else: + self.background_color = tuple([int(x * 255) for x in image_mean]) + + def resize(self, pil_img: Image) -> np.ndarray: + """ + + Args: + pil_img (PIL.Image): [H, W, 3] in PIL.Image in RGB + + Returns: + x (np.ndarray): [3, self.image_size, self.image_size] + """ + + width, height = pil_img.size + max_size = max(width, height) + + size = [ + max(int(height / max_size * self.image_size), self.min_size), + max(int(width / max_size * self.image_size), self.min_size), + ] + + if width <= 0 or height <= 0 or size[0] <= 0 or size[1] <= 0: + # print(f"orig size = {pil_img.size}, new size = {size}") + raise ValueError("Invalid size!") + + def resize( + pil_img, size, interpolation=PIL.Image.Resampling.BICUBIC, antialias=True + ): + if isinstance(size, int): + w, h = pil_img.size + if (w <= h and w == size) or (h <= w and h == size): + return pil_img + if w < h: + ow = size + oh = int(size * h / w) + else: + oh = size + ow = int(size * w / h) + size = (ow, oh) + else: + size = (size[1], size[0]) + + return pil_img.resize( + size, resample=interpolation, reducing_gap=None if antialias else 3.0 + ) + + pil_img = resize( + pil_img, size, interpolation=PIL.Image.Resampling.BICUBIC, antialias=True + ) + + pil_img = expand2square(pil_img, self.background_color) + x = to_numpy_array(pil_img) + + # [H, W, 3] -> [3, H, W] + x = np.transpose(x, (2, 0, 1)) + + return x + + def preprocess(self, images, return_tensors: str = "pt", **kwargs) -> BatchFeature: + # resize and pad to [self.image_size, self.image_size] + # then convert from [H, W, 3] to [3, H, W] + if not isinstance(images, list): + images = [images] + images: List[np.ndarray] = [self.resize(image) for image in images] + images = [image[:3, ...] for image in images] + + # rescale from [0, 255] -> [0, 1] + images = [ + self.rescale( + image=image, + scale=self.rescale_factor, + input_data_format="channels_first", + ) + for image in images + ] + + # normalize + if self.do_normalize: + images = [ + self.normalize( + image=image, + mean=self.image_mean, + std=self.image_std, + input_data_format="channels_first", + ) + for image in images + ] + data = {"pixel_values": images} + return BatchFeature(data=data, tensor_type=return_tensors) + + @property + def default_shape(self): + return [3, self.image_size, self.image_size] + + +class DictOutput(object): + def keys(self): + return self.__dict__.keys() + + def __getitem__(self, item): + return self.__dict__[item] + + def __setitem__(self, key, value): + self.__dict__[key] = value + + +@dataclass +class VLChatProcessorOutput(DictOutput): + sft_format: str + input_ids: torch.Tensor + pixel_values: torch.Tensor + num_image_tokens: torch.IntTensor + + def __len__(self): + return len(self.input_ids) + + +@dataclass +class BatchedVLChatProcessorOutput(DictOutput): + sft_format: List[str] + input_ids: torch.Tensor + pixel_values: torch.Tensor + attention_mask: torch.Tensor + images_seq_mask: torch.BoolTensor + images_emb_mask: torch.BoolTensor + + +# FIXME: had to place Official Processor here, since image_processor module would not be imported in all threads, +# hence AutoProcessor registration would not be affective in some cases +class VLChatProcessor(ProcessorMixin): + image_processor_class = "AutoImageProcessor" + tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast") + + attributes = ["image_processor", "tokenizer"] + + def __init__( + self, + image_processor: VLMImageProcessor, + tokenizer: LlamaTokenizerFast, + image_tag: str = "", + image_start_tag: str = "", + image_end_tag: str = "", + pad_tag: str = "<｜▁pad▁｜>", + num_image_tokens: int = 576, + add_special_token: bool = False, + sft_format: str = "deepseek", + mask_prompt: bool = True, + ignore_id: int = -100, + **kwargs, + ): + self.image_processor = image_processor + self.tokenizer = tokenizer + + image_id = self.tokenizer.vocab.get(image_tag) + if image_id is None: + special_tokens = [image_tag] + special_tokens_dict = {"additional_special_tokens": special_tokens} + self.tokenizer.add_special_tokens(special_tokens_dict) + # print(f"Add image tag = {image_tag} to the tokenizer") + + self.image_tag = image_tag + self.image_start_tag = image_start_tag + self.image_end_tag = image_end_tag + self.pad_tag = pad_tag + + self.num_image_tokens = num_image_tokens + self.add_special_token = add_special_token + self.sft_format = sft_format + self.ignore_id = ignore_id + + super().__init__( + image_processor, + tokenizer, + **kwargs, + ) + + @property + def image_token(self): + return self.image_tag + + @property + def image_id(self) -> int: + image_id = self.tokenizer.vocab.get(self.image_tag) + return image_id + + @property + def image_start_id(self): + image_start_id = self.tokenizer.vocab.get(self.image_start_tag) + return image_start_id + + @property + def image_end_id(self): + image_end_id = self.tokenizer.vocab.get(self.image_end_tag) + return image_end_id + + @property + def image_start_token(self): + return self.image_start_tag + + @property + def image_end_token(self): + return self.image_end_tag + + @property + def pad_id(self): + pad_id = self.tokenizer.vocab.get(self.pad_tag) + return pad_id + + def add_image_token( + self, + image_indices: List[int], + input_ids: torch.LongTensor, + ): + """ + + Args: + image_indices (List[int]): [index_0, index_1, ..., index_j] + input_ids (torch.LongTensor): [N] + + Returns: + input_ids (torch.LongTensor): [N + image tokens] + num_image_tokens (torch.IntTensor): [n_images] + """ + + input_slices = [] + + start = 0 + for index in image_indices: + if self.add_special_token: + end = index + 1 + else: + end = index + + # original text tokens + input_slices.append(input_ids[start:end]) + + # add boi, image tokens, eoi and set the mask as False + input_slices.append(self.image_start_id * torch.ones((1), dtype=torch.long)) + input_slices.append( + self.image_id * torch.ones((self.num_image_tokens,), dtype=torch.long) + ) + input_slices.append(self.image_end_id * torch.ones((1), dtype=torch.long)) + start = index + 1 + + # the left part + input_slices.append(input_ids[start:]) + + # concat all slices + input_ids = torch.cat(input_slices, dim=0) + num_image_tokens = torch.IntTensor([self.num_image_tokens] * len(image_indices)) + + return input_ids, num_image_tokens + + def process_one( + self, + prompt: str = None, + images: List[Image] = None, + **kwargs, + ): + """ + + Args: + prompt (str): the formatted prompt; + images (List[ImageType]): the list of images; + **kwargs: + + Returns: + outputs (BaseProcessorOutput): the output of the processor, + - input_ids (torch.LongTensor): [N + image tokens] + - target_ids (torch.LongTensor): [N + image tokens] + - images (torch.FloatTensor): [n_images, 3, H, W] + - image_id (int): the id of the image token + - num_image_tokens (List[int]): the number of image tokens + """ + + sft_format = prompt + # tokenize + input_ids = self.tokenizer.encode(sft_format) + input_ids = torch.LongTensor(input_ids) + + # add image tokens to the input_ids + image_token_mask: torch.Tensor = (input_ids == self.image_id).to(torch.bool) + image_indices = image_token_mask.nonzero() + input_ids, num_image_tokens = self.add_image_token( + image_indices=image_indices, + input_ids=input_ids, + ) + + # load images + images_outputs = self.image_processor(images, return_tensors="pt") + + prepare = VLChatProcessorOutput( + sft_format=sft_format, + input_ids=input_ids, + pixel_values=images_outputs.pixel_values, + num_image_tokens=num_image_tokens, + ) + + return prepare + + def __call__( + self, + *, + prompt: str = None, + conversations: List[Dict[str, str]] = None, + images: List[Image] = None, + force_batchify: bool = True, + **kwargs, + ): + """ + + Args: + prompt (str): the formatted prompt; + conversations (List[Dict]): conversations with a list of messages; + images (List[ImageType]): the list of images; + force_batchify (bool): force batchify the inputs; + **kwargs: + + Returns: + outputs (BaseProcessorOutput): the output of the processor, + - input_ids (torch.LongTensor): [N + image tokens] + - images (torch.FloatTensor): [n_images, 3, H, W] + - image_id (int): the id of the image token + - num_image_tokens (List[int]): the number of image tokens + """ + + prepare = self.process_one( + prompt=prompt, conversations=conversations, images=images + ) + + if force_batchify: + prepare = self.batchify([prepare]) + + return prepare + + def batchify( + self, prepare_list: List[VLChatProcessorOutput] + ) -> BatchedVLChatProcessorOutput: + """ + Preprocesses the inputs for multimodal inference. + + Args: + prepare_list (List[VLChatProcessorOutput]): A list of VLChatProcessorOutput. + + Returns: + BatchedVLChatProcessorOutput: A dictionary of the inputs to use for multimodal inference. + """ + + batch_size = len(prepare_list) + sft_format = [] + n_images = [] + seq_lens = [] + for prepare in prepare_list: + n_images.append(len(prepare.num_image_tokens)) + seq_lens.append(len(prepare)) + + input_token_max_len = max(seq_lens) + max_n_images = max(1, max(n_images)) + + batched_input_ids = torch.full( + (batch_size, input_token_max_len), self.pad_id + ).long() # FIXME + batched_attention_mask = torch.zeros((batch_size, input_token_max_len)).long() + batched_pixel_values = torch.zeros( + (batch_size, max_n_images, *self.image_processor.default_shape) + ).float() + batched_images_seq_mask = torch.zeros((batch_size, input_token_max_len)).bool() + batched_images_emb_mask = torch.zeros( + (batch_size, max_n_images, self.num_image_tokens) + ).bool() + + for i, prepare in enumerate(prepare_list): + input_ids = prepare.input_ids + seq_len = len(prepare) + n_image = len(prepare.num_image_tokens) + # left-padding + batched_attention_mask[i, -seq_len:] = 1 + batched_input_ids[i, -seq_len:] = torch.LongTensor(input_ids) + batched_images_seq_mask[i, -seq_len:] = input_ids == self.image_id + + if n_image > 0: + batched_pixel_values[i, :n_image] = prepare.pixel_values + for j, n_image_tokens in enumerate(prepare.num_image_tokens): + batched_images_emb_mask[i, j, :n_image_tokens] = True + + sft_format.append(prepare.sft_format) + + batched_prepares = BatchedVLChatProcessorOutput( + input_ids=batched_input_ids, + attention_mask=batched_attention_mask, + pixel_values=batched_pixel_values, + images_seq_mask=batched_images_seq_mask, + images_emb_mask=batched_images_emb_mask, + sft_format=sft_format, + ) + + return batched_prepares + + +class VLMImageProcessorConfig(PretrainedConfig): + model_type = "deepseek_vlm" + image_size: int + min_size: int + image_mean: Union[Tuple[float, float, float], List[float]] + image_std: Union[Tuple[float, float, float], List[float]] + rescale_factor: float + do_normalize: bool + + def __init__( + self, + image_size: int, + min_size: int = 14, + image_mean: Union[Tuple[float, float, float], List[float]] = ( + 0.48145466, + 0.4578275, + 0.40821073, + ), + image_std: Union[Tuple[float, float, float], List[float]] = ( + 0.26862954, + 0.26130258, + 0.27577711, + ), + rescale_factor: float = 1.0 / 255.0, + do_normalize: bool = True, + **kwargs, + ): + self.image_size = image_size + self.min_size = min_size + self.image_mean = image_mean + self.image_std = image_std + self.rescale_factor = rescale_factor + self.do_normalize = do_normalize + + super().__init__(**kwargs) + + +register_processor(MultiModalityConfig, VLChatProcessor) +register_image_processor(MultiModalityConfig, VLMImageProcessor) diff --git a/sglang/python/sglang/srt/configs/load_config.py b/sglang/python/sglang/srt/configs/load_config.py new file mode 100644 index 0000000..be9a40b --- /dev/null +++ b/sglang/python/sglang/srt/configs/load_config.py @@ -0,0 +1,89 @@ +# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py +import enum +import json +import logging +from dataclasses import dataclass, field +from typing import List, Optional, Union + +from sglang.srt.utils import is_hip + +logger = logging.getLogger(__name__) + + +class LoadFormat(str, enum.Enum): + AUTO = "auto" + PT = "pt" + SAFETENSORS = "safetensors" + NPCACHE = "npcache" + DUMMY = "dummy" + SHARDED_STATE = "sharded_state" + GGUF = "gguf" + BITSANDBYTES = "bitsandbytes" + MISTRAL = "mistral" + LAYERED = "layered" + JAX = "jax" + REMOTE = "remote" + + +@dataclass +class LoadConfig: + """ + download_dir: Directory to download and load the weights, default to the + default cache directory of huggingface. + load_format: The format of the model weights to load: + "auto" will try to load the weights in the safetensors format and + fall back to the pytorch bin format if safetensors format is + not available. + "pt" will load the weights in the pytorch bin format. + "safetensors" will load the weights in the safetensors format. + "npcache" will load the weights in pytorch format and store + a numpy cache to speed up the loading. + "dummy" will initialize the weights with random values, which is + mainly for profiling. + "bitsandbytes" will load nf4 type weights. + ignore_patterns: The list of patterns to ignore when loading the model. + Default to "original/**/*" to avoid repeated loading of llama's + checkpoints. + decryption_key_file: If set, decrypts the output files with a password read + from this file (after PBKDF2). + """ + + load_format: Union[str, LoadFormat] = LoadFormat.AUTO + download_dir: Optional[str] = None + model_loader_extra_config: Optional[Union[str, dict]] = field(default_factory=dict) + ignore_patterns: Optional[Union[List[str], str]] = None + decryption_key_file: Optional[str] = None + + def __post_init__(self): + model_loader_extra_config = self.model_loader_extra_config or {} + if isinstance(model_loader_extra_config, str): + self.model_loader_extra_config = json.loads(model_loader_extra_config) + self._verify_load_format() + + if self.ignore_patterns is not None and len(self.ignore_patterns) > 0: + logger.info( + "Ignoring the following patterns when downloading weights: %s", + self.ignore_patterns, + ) + else: + self.ignore_patterns = ["original/**/*"] + + def _verify_load_format(self) -> None: + if not isinstance(self.load_format, str): + return + + load_format = self.load_format.lower() + self.load_format = LoadFormat(load_format) + + rocm_not_supported_load_format: List[str] = [] + if is_hip() and load_format in rocm_not_supported_load_format: + rocm_supported_load_format = [ + f + for f in LoadFormat.__members__ + if (f not in rocm_not_supported_load_format) + ] + raise ValueError( + f"load format '{load_format}' is not supported in ROCm. " + f"Supported load formats are " + f"{rocm_supported_load_format}" + ) diff --git a/sglang/python/sglang/srt/configs/model_config.py b/sglang/python/sglang/srt/configs/model_config.py new file mode 100644 index 0000000..8bc8d12 --- /dev/null +++ b/sglang/python/sglang/srt/configs/model_config.py @@ -0,0 +1,512 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import json +import logging +import math +from enum import IntEnum, auto +from typing import List, Optional, Set, Union + +import torch +from transformers import PretrainedConfig + +from sglang.srt.hf_transformers_utils import get_config, get_context_length +from sglang.srt.layers.quantization import QUANTIZATION_METHODS +from sglang.srt.utils import get_bool_env_var, is_hip + +logger = logging.getLogger(__name__) + + +class AttentionArch(IntEnum): + MLA = auto() + MHA = auto() + + +class ModelConfig: + def __init__( + self, + model_path: str, + trust_remote_code: bool = True, + revision: Optional[str] = None, + context_length: Optional[int] = None, + model_override_args: Optional[str] = None, + is_embedding: Optional[bool] = None, + dtype: str = "auto", + quantization: Optional[str] = None, + override_config_file: Optional[str] = None, + ) -> None: + self.model_path = model_path + self.revision = revision + self.quantization = quantization + + # Parse args + self.maybe_pull_model_tokenizer_from_remote() + self.model_override_args = json.loads(model_override_args) + kwargs = {} + if override_config_file and override_config_file.strip(): + kwargs["_configuration_file"] = override_config_file.strip() + + self.hf_config = get_config( + self.model_path, + trust_remote_code=trust_remote_code, + revision=revision, + model_override_args=self.model_override_args, + **kwargs, + ) + self.hf_text_config = get_hf_text_config(self.hf_config) + + # Check model type + self.is_generation = is_generation_model( + self.hf_config.architectures, is_embedding + ) + self.is_multimodal = is_multimodal_model(self.hf_config.architectures) + self.is_multimodal_gen = is_multimodal_gen_model(self.hf_config.architectures) + self.is_image_gen = is_image_gen_model(self.hf_config.architectures) + self.is_audio_model = is_audio_model(self.hf_config.architectures) + self.is_encoder_decoder = is_encoder_decoder_model(self.hf_config.architectures) + self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype) + + # Derive context length + derived_context_len = get_context_length(self.hf_text_config) + if context_length is not None: + if context_length > derived_context_len: + if get_bool_env_var( + "SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN", default="True" + ): + logger.warning( + f"Warning: User-specified context_length ({context_length}) is greater than the derived context_length ({derived_context_len}). " + f"This may lead to incorrect model outputs or CUDA errors." + ) + self.context_len = context_length + else: + raise ValueError( + f"User-specified context_length ({context_length}) is greater than the derived context_length ({derived_context_len}). " + f"This may lead to incorrect model outputs or CUDA errors. Note that the derived context_length may differ from max_position_embeddings in the model's config. " + f"To allow overriding this maximum, set the env var SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1" + ) + else: + self.context_len = context_length + else: + self.context_len = derived_context_len + + # Unify the config keys for hf_text_config + self.head_dim = getattr( + self.hf_text_config, + "head_dim", + self.hf_text_config.hidden_size // self.hf_text_config.num_attention_heads, + ) + + # FIXME: temporary special judge for MLA architecture + if ( + "DeepseekV2ForCausalLM" in self.hf_config.architectures + or "DeepseekV3ForCausalLM" in self.hf_config.architectures + or "DeepseekV3ForCausalLMNextN" in self.hf_config.architectures + ): + self.head_dim = 256 + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_config.kv_lora_rank + self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim + self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim + self.v_head_dim = self.hf_config.v_head_dim + + # Handle rope scaling with yarn + self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim) + if self.hf_config.rope_scaling: + mscale_all_dim = self.hf_config.rope_scaling.get( + "mscale_all_dim", False + ) + scaling_factor = self.hf_config.rope_scaling["factor"] + mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim)) + self.scaling = self.scaling * mscale * mscale + + elif "MiniCPM3ForCausalLM" in self.hf_config.architectures: + self.head_dim = 128 + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_config.kv_lora_rank + self.qk_rope_head_dim = self.hf_config.qk_rope_head_dim + elif "DeepseekVL2ForCausalLM" in self.hf_config.architectures: + self.head_dim = 256 + self.attention_arch = AttentionArch.MLA + self.kv_lora_rank = self.hf_text_config.kv_lora_rank + self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim + else: + self.attention_arch = AttentionArch.MHA + + self.num_attention_heads = self.hf_text_config.num_attention_heads + self.num_key_value_heads = getattr( + self.hf_text_config, "num_key_value_heads", None + ) + + # for Dbrx and MPT models + if self.hf_config.model_type in ["dbrx", "mpt"]: + self.num_key_value_heads = getattr( + self.hf_config.attn_config, "kv_n_heads", None + ) + + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + self.hidden_size = self.hf_text_config.hidden_size + self.num_hidden_layers = self.hf_text_config.num_hidden_layers + self.vocab_size = self.hf_text_config.vocab_size + + # Verify quantization + self._verify_quantization() + + # Cache attributes + self.hf_eos_token_id = self.get_hf_eos_token_id() + self.image_token_id = getattr(self.hf_config, "image_token_id", None) + + # adapted from https://github.com/vllm-project/vllm/blob/main/vllm/config.py#L289 + def get_total_num_kv_heads(self) -> int: + """Returns the total number of KV heads.""" + # For GPTBigCode & Falcon: + # NOTE: for falcon, when new_decoder_architecture is True, the + # multi_query flag is ignored and we use n_head_kv for the number of + # KV heads. + falcon_model_types = ["falcon", "RefinedWeb", "RefinedWebModel"] + new_decoder_arch_falcon = ( + self.hf_config.model_type in falcon_model_types + and getattr(self.hf_config, "new_decoder_architecture", False) + ) + if not new_decoder_arch_falcon and getattr( + self.hf_text_config, "multi_query", False + ): + # Multi-query attention, only one KV head. + # Currently, tensor parallelism is not supported in this case. + return 1 + + # For DBRX and MPT + if self.hf_config.model_type in ["mpt"]: + if "kv_n_heads" in self.hf_config.attn_config: + return self.hf_config.attn_config["kv_n_heads"] + return self.hf_config.num_attention_heads + if self.hf_config.model_type in ["dbrx"]: + return getattr( + self.hf_config.attn_config, + "kv_n_heads", + self.hf_config.num_attention_heads, + ) + + attributes = [ + # For Falcon: + "n_head_kv", + "num_kv_heads", + # For LLaMA-2: + "num_key_value_heads", + # For ChatGLM: + "multi_query_group_num", + ] + for attr in attributes: + num_kv_heads = getattr(self.hf_text_config, attr, None) + if num_kv_heads is not None: + return num_kv_heads + + # For non-grouped-query attention models, the number of KV heads is + # equal to the number of attention heads. + return self.hf_text_config.num_attention_heads + + def get_num_kv_heads(self, tensor_parallel_size) -> int: + """Returns the number of KV heads per GPU.""" + total_num_kv_heads = self.get_total_num_kv_heads() + # If tensor parallelism is used, we divide the number of KV heads by + # the tensor parallel size. We will replicate the KV heads in the + # case where the number of KV heads is smaller than the tensor + # parallel size so each GPU has at least one KV head. + return max(1, total_num_kv_heads // tensor_parallel_size) + + # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py + def _parse_quant_hf_config(self): + quant_cfg = getattr(self.hf_config, "quantization_config", None) + if quant_cfg is None: + # compressed-tensors uses a "compression_config" key + quant_cfg = getattr(self.hf_config, "compression_config", None) + return quant_cfg + + # adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py + def _verify_quantization(self) -> None: + supported_quantization = [*QUANTIZATION_METHODS] + rocm_supported_quantization = [ + "awq", + "gptq", + "fp8", + "compressed_tensors", + "compressed-tensors", + "fbgemm_fp8", + "w8a8_fp8", + ] + optimized_quantization_methods = [ + "fp8", + "marlin", + "modelopt", + "gptq_marlin_24", + "gptq_marlin", + "awq_marlin", + "fbgemm_fp8", + "compressed_tensors", + "compressed-tensors", + "experts_int8", + "w8a8_int8", + "w8a8_fp8", + ] + compatible_quantization_methods = { + "w8a8_int8": ["compressed-tensors", "compressed_tensors"], + "w8a8_fp8": ["compressed-tensors", "compressed_tensors"], + } + if self.quantization is not None: + self.quantization = self.quantization.lower() + + # Parse quantization method from the HF model config, if available. + quant_cfg = self._parse_quant_hf_config() + + if quant_cfg is not None: + quant_method = quant_cfg.get("quant_method", "").lower() + + # Detect which checkpoint is it + for _, method in QUANTIZATION_METHODS.items(): + quantization_override = method.override_quantization_method( + quant_cfg, self.quantization + ) + if quantization_override: + quant_method = quantization_override + self.quantization = quantization_override + break + + # Verify quantization configurations. + if self.quantization is None: + self.quantization = quant_method + elif self.quantization != quant_method: + if ( + self.quantization not in compatible_quantization_methods + or quant_method + not in compatible_quantization_methods[self.quantization] + ): + raise ValueError( + "Quantization method specified in the model config " + f"({quant_method}) does not match the quantization " + f"method specified in the `quantization` argument " + f"({self.quantization})." + ) + + if self.quantization is not None: + if self.quantization not in supported_quantization: + raise ValueError( + f"Unknown quantization method: {self.quantization}. Must " + f"be one of {supported_quantization}." + ) + if is_hip() and self.quantization not in rocm_supported_quantization: + raise ValueError( + f"{self.quantization} quantization is currently not " + f"supported in ROCm." + ) + if self.quantization not in optimized_quantization_methods: + logger.warning( + "%s quantization is not fully " + "optimized yet. The speed can be slower than " + "non-quantized models.", + self.quantization, + ) + + def get_hf_eos_token_id(self) -> Optional[Set[int]]: + eos_ids = getattr(self.hf_config, "eos_token_id", None) + if eos_ids: + # it can be either int or list of int + eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids) + return eos_ids + + def maybe_pull_model_tokenizer_from_remote(self) -> None: + """ + Pull the model config files to a temporary + directory in case of remote. + + Args: + model: The model name or path. + + """ + from sglang.srt.connector import create_remote_connector + from sglang.srt.utils import is_remote_url + + if is_remote_url(self.model_path): + logger.info("Pulling model configs from remote...") + # BaseConnector implements __del__() to clean up the local dir. + # Since config files need to exist all the time, so we DO NOT use + # with statement to avoid closing the client. + client = create_remote_connector(self.model_path) + if is_remote_url(self.model_path): + client.pull_files(allow_pattern=["*config.json"]) + self.model_weights = self.model_path + self.model_path = client.get_local_dir() + + +def get_hf_text_config(config: PretrainedConfig): + """Get the "sub" config relevant to llm for multi modal models. + No op for pure text models. + """ + class_name = config.architectures[0] + if class_name.startswith("Llava") and class_name.endswith("ForCausalLM"): + # We support non-hf version of llava models, so we do not want to + # read the wrong values from the unused default text_config. + # NOTE(HandH1998): We set `torch_dtype` of config to `torch.float16` for the weights, as + # `torch.float16` is default used for image features in `python/sglang/srt/models/llava.py`. + setattr(config, "torch_dtype", torch.float16) + return config + + if hasattr(config, "text_config"): + # The code operates under the assumption that text_config should have + # `num_attention_heads` (among others). Assert here to fail early + # if transformers config doesn't align with this assumption. + assert hasattr(config.text_config, "num_attention_heads") + return config.text_config + if hasattr(config, "language_config"): + return config.language_config + else: + return config + + +# adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py +_STR_DTYPE_TO_TORCH_DTYPE = { + "half": torch.float16, + "float16": torch.float16, + "float": torch.float32, + "float32": torch.float32, + "bfloat16": torch.bfloat16, +} + + +# adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py +def _get_and_verify_dtype( + config: PretrainedConfig, + dtype: Union[str, torch.dtype], +) -> torch.dtype: + # NOTE: getattr(config, "torch_dtype", torch.float32) is not correct + # because config.torch_dtype can be None. + config_dtype = getattr(config, "torch_dtype", None) + if config_dtype is None: + config_dtype = torch.float32 + + if isinstance(dtype, str): + dtype = dtype.lower() + if dtype == "auto": + if config_dtype == torch.float32: + if config.model_type.startswith("gemma"): + if config.model_type == "gemma": + gemma_version = "" + else: + gemma_version = config.model_type[5] + logger.info( + f"For Gemma {gemma_version}, we downcast float32 to bfloat16 instead " + "of float16 by default. Please specify `dtype` if you " + "want to use float16." + ) + torch_dtype = torch.bfloat16 + else: + # Following the common practice, we use float16 for float32 + # models. + torch_dtype = torch.float16 + else: + torch_dtype = config_dtype + else: + if dtype not in _STR_DTYPE_TO_TORCH_DTYPE: + raise ValueError(f"Unknown dtype: {dtype}") + torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype] + elif isinstance(dtype, torch.dtype): + torch_dtype = dtype + else: + raise ValueError(f"Unknown dtype: {dtype}") + + # Verify the dtype. + if torch_dtype != config_dtype: + if torch_dtype == torch.float32: + # Upcasting to float32 is allowed. + logger.info("Upcasting %s to %s.", config_dtype, torch_dtype) + pass + elif config_dtype == torch.float32: + # Downcasting from float32 to float16 or bfloat16 is allowed. + logger.info("Downcasting %s to %s.", config_dtype, torch_dtype) + pass + else: + # Casting between float16 and bfloat16 is allowed with a warning. + logger.warning("Casting %s to %s.", config_dtype, torch_dtype) + + return torch_dtype + + +def is_generation_model(model_architectures: List[str], is_embedding: bool = False): + # We have two ways to determine whether a model is a generative model. + # 1. Check the model architecture + # 2. check the `is_embedding` server args + + if ( + "LlamaEmbeddingModel" in model_architectures + or "MistralModel" in model_architectures + or "LlamaForSequenceClassification" in model_architectures + or "LlamaForSequenceClassificationWithNormal_Weights" in model_architectures + or "InternLM2ForRewardModel" in model_architectures + or "Qwen2ForRewardModel" in model_architectures + or "Qwen2ForSequenceClassification" in model_architectures + or "CLIPModel" in model_architectures + ): + return False + else: + return not is_embedding + + +multimodal_model_archs = [ + "DeepseekVL2ForCausalLM", + "Gemma3ForConditionalGeneration", + "Grok1VForCausalLM", + "Grok1AForCausalLM", + "LlavaLlamaForCausalLM", + "LlavaMistralForCausalLM", + "LlavaQwenForCausalLM", + "LlavaVidForCausalLM", + "MiniCPMO", + "MiniCPMV", + "MultiModalityCausalLM", + "MllamaForConditionalGeneration", + "Qwen2VLForConditionalGeneration", + "Qwen2_5_VLForConditionalGeneration", + "CLIPModel", +] + + +def is_multimodal_model(model_architectures: List[str]): + if any( + multi_model_arch in model_architectures + for multi_model_arch in multimodal_model_archs + ): + return True + else: + return False + + +def is_multimodal_gen_model(model_architectures: List[str]): + return False + + +def is_image_gen_model(model_architectures: List[str]): + return False + + +def is_audio_model(model_architectures: List[str]): + return False + + +def is_encoder_decoder_model(model_architectures: List[str]): + return "MllamaForConditionalGeneration" in model_architectures + + +def yarn_get_mscale(scale: float = 1, mscale: float = 1) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 diff --git a/sglang/python/sglang/srt/configs/utils.py b/sglang/python/sglang/srt/configs/utils.py new file mode 100644 index 0000000..8403d95 --- /dev/null +++ b/sglang/python/sglang/srt/configs/utils.py @@ -0,0 +1,25 @@ +from typing import Type + +from transformers import ( + AutoImageProcessor, + AutoProcessor, + BaseImageProcessor, + PretrainedConfig, + ProcessorMixin, +) + + +def register_image_processor( + config: Type[PretrainedConfig], image_processor: Type[BaseImageProcessor] +): + """ + register customized hf image processor while removing hf impl + """ + AutoImageProcessor.register(config, None, image_processor, None, exist_ok=True) + + +def register_processor(config: Type[PretrainedConfig], processor: Type[ProcessorMixin]): + """ + register customized hf processor while removing hf impl + """ + AutoProcessor.register(config, processor, exist_ok=True) diff --git a/sglang/python/sglang/srt/connector/__init__.py b/sglang/python/sglang/srt/connector/__init__.py new file mode 100644 index 0000000..829644c --- /dev/null +++ b/sglang/python/sglang/srt/connector/__init__.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 + +import enum +import logging + +from sglang.srt.connector.base_connector import ( + BaseConnector, + BaseFileConnector, + BaseKVConnector, +) +from sglang.srt.connector.redis import RedisConnector +from sglang.srt.connector.s3 import S3Connector +from sglang.srt.utils import parse_connector_type + +logger = logging.getLogger(__name__) + + +class ConnectorType(str, enum.Enum): + FS = "filesystem" + KV = "KV" + + +def create_remote_connector(url, device="cpu") -> BaseConnector: + connector_type = parse_connector_type(url) + if connector_type == "redis": + return RedisConnector(url) + elif connector_type == "s3": + return S3Connector(url) + else: + raise ValueError(f"Invalid connector type: {url}") + + +def get_connector_type(client: BaseConnector) -> ConnectorType: + if isinstance(client, BaseKVConnector): + return ConnectorType.KV + if isinstance(client, BaseFileConnector): + return ConnectorType.FS + + raise ValueError(f"Invalid connector type: {client}") + + +__all__ = [ + "BaseConnector", + "BaseFileConnector", + "BaseKVConnector", + "RedisConnector", + "S3Connector", + "ConnectorType", + "create_remote_connector", + "get_connector_type", +] diff --git a/sglang/python/sglang/srt/connector/base_connector.py b/sglang/python/sglang/srt/connector/base_connector.py new file mode 100644 index 0000000..a9c00d0 --- /dev/null +++ b/sglang/python/sglang/srt/connector/base_connector.py @@ -0,0 +1,112 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os +import shutil +import signal +import tempfile +from abc import ABC, abstractmethod +from typing import Generator, List, Optional, Tuple + +import torch + + +class BaseConnector(ABC): + """ + For fs connector such as s3: + :/// + + For kv connector such as redis: + ://://keys/ + ://files/ + """ + + def __init__(self, url: str, device: torch.device = "cpu"): + self.url = url + self.device = device + self.closed = False + self.local_dir = tempfile.mkdtemp() + for sig in (signal.SIGINT, signal.SIGTERM): + existing_handler = signal.getsignal(sig) + signal.signal(sig, self._close_by_signal(existing_handler)) + + def get_local_dir(self): + return self.local_dir + + @abstractmethod + def weight_iterator( + self, rank: int = 0 + ) -> Generator[Tuple[str, torch.Tensor], None, None]: + raise NotImplementedError() + + @abstractmethod + def pull_files( + self, + allow_pattern: Optional[List[str]] = None, + ignore_pattern: Optional[List[str]] = None, + ) -> None: + raise NotImplementedError() + + def close(self): + if self.closed: + return + + self.closed = True + if os.path.exists(self.local_dir): + shutil.rmtree(self.local_dir) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + self.close() + + def __del__(self): + self.close() + + def _close_by_signal(self, existing_handler=None): + + def new_handler(signum, frame): + self.close() + if existing_handler: + existing_handler(signum, frame) + + return new_handler + + +class BaseKVConnector(BaseConnector): + + @abstractmethod + def get(self, key: str) -> Optional[torch.Tensor]: + raise NotImplementedError() + + @abstractmethod + def getstr(self, key: str) -> Optional[str]: + raise NotImplementedError() + + @abstractmethod + def set(self, key: str, obj: torch.Tensor) -> None: + raise NotImplementedError() + + @abstractmethod + def setstr(self, key: str, obj: str) -> None: + raise NotImplementedError() + + @abstractmethod + def list(self, prefix: str) -> List[str]: + raise NotImplementedError() + + +class BaseFileConnector(BaseConnector): + """ + List full file names from remote fs path and filter by allow pattern. + + Args: + allow_pattern: A list of patterns of which files to pull. + + Returns: + list[str]: List of full paths allowed by the pattern + """ + + @abstractmethod + def glob(self, allow_pattern: str) -> List[str]: + raise NotImplementedError() diff --git a/sglang/python/sglang/srt/connector/redis.py b/sglang/python/sglang/srt/connector/redis.py new file mode 100644 index 0000000..761594f --- /dev/null +++ b/sglang/python/sglang/srt/connector/redis.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 + +import logging +from typing import Generator, List, Optional, Tuple +from urllib.parse import urlparse + +import torch + +from sglang.srt.connector import BaseKVConnector +from sglang.srt.connector.serde import create_serde +from sglang.srt.connector.utils import pull_files_from_db + +logger = logging.getLogger(__name__) + + +class RedisConnector(BaseKVConnector): + + def __init__(self, url: str, device: torch.device = "cpu"): + import redis + + super().__init__(url, device) + parsed_url = urlparse(url) + self.connection = redis.Redis(host=parsed_url.hostname, port=parsed_url.port) + self.model_name = parsed_url.path.lstrip("/") + # TODO: more serde options + self.s, self.d = create_serde("safe") + + def get(self, key: str) -> Optional[torch.Tensor]: + val = self.connection.get(key) + + if val is None: + logger.error("Key %s not found", key) + return None + + return self.d.from_bytes(val) + + def getstr(self, key: str) -> Optional[str]: + val = self.connection.get(key) + if val is None: + logger.error("Key %s not found", key) + return None + + return val.decode("utf-8") + + def set(self, key: str, tensor: torch.Tensor) -> None: + assert tensor is not None + self.connection.set(key, self.s.to_bytes(tensor)) + + def setstr(self, key: str, obj: str) -> None: + self.connection.set(key, obj) + + def list(self, prefix: str) -> List[str]: + cursor = 0 + all_keys: List[bytes] = [] + + while True: + ret: Tuple[int, List[bytes]] = self.connection.scan( + cursor=cursor, match=f"{prefix}*" + ) # type: ignore + cursor, keys = ret + all_keys.extend(keys) + if cursor == 0: + break + + return [key.decode("utf-8") for key in all_keys] + + def weight_iterator( + self, rank: int = 0 + ) -> Generator[Tuple[str, bytes], None, None]: + keys = self.list(f"{self.model_name}/keys/rank_{rank}/") + for key in keys: + val = self.get(key) + key = key.removeprefix(f"{self.model_name}/keys/rank_{rank}/") + yield key, val + + def pull_files( + self, + allow_pattern: Optional[List[str]] = None, + ignore_pattern: Optional[List[str]] = None, + ) -> None: + pull_files_from_db(self, self.model_name, allow_pattern, ignore_pattern) + + def close(self): + self.connection.close() + super().close() diff --git a/sglang/python/sglang/srt/connector/s3.py b/sglang/python/sglang/srt/connector/s3.py new file mode 100644 index 0000000..7bef8f5 --- /dev/null +++ b/sglang/python/sglang/srt/connector/s3.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: Apache-2.0 + +import fnmatch +import os +from pathlib import Path +from typing import Generator, Optional, Tuple + +import torch + +from sglang.srt.connector import BaseFileConnector + + +def _filter_allow(paths: list[str], patterns: list[str]) -> list[str]: + return [ + path + for path in paths + if any(fnmatch.fnmatch(path, pattern) for pattern in patterns) + ] + + +def _filter_ignore(paths: list[str], patterns: list[str]) -> list[str]: + return [ + path + for path in paths + if not any(fnmatch.fnmatch(path, pattern) for pattern in patterns) + ] + + +def list_files( + s3, + path: str, + allow_pattern: Optional[list[str]] = None, + ignore_pattern: Optional[list[str]] = None, +) -> tuple[str, str, list[str]]: + """ + List files from S3 path and filter by pattern. + + Args: + s3: S3 client to use. + path: The S3 path to list from. + allow_pattern: A list of patterns of which files to pull. + ignore_pattern: A list of patterns of which files not to pull. + + Returns: + tuple[str, str, list[str]]: A tuple where: + - The first element is the bucket name + - The second element is string represent the bucket + and the prefix as a dir like string + - The third element is a list of files allowed or + disallowed by pattern + """ + parts = path.removeprefix("s3://").split("/") + prefix = "/".join(parts[1:]) + bucket_name = parts[0] + + objects = s3.list_objects_v2(Bucket=bucket_name, Prefix=prefix) + paths = [obj["Key"] for obj in objects.get("Contents", [])] + + paths = _filter_ignore(paths, ["*/"]) + if allow_pattern is not None: + paths = _filter_allow(paths, allow_pattern) + + if ignore_pattern is not None: + paths = _filter_ignore(paths, ignore_pattern) + + return bucket_name, prefix, paths + + +class S3Connector(BaseFileConnector): + + def __init__(self, url: str) -> None: + import boto3 + + super().__init__(url) + self.client = boto3.client("s3") + + def glob(self, allow_pattern: Optional[list[str]] = None) -> list[str]: + bucket_name, _, paths = list_files( + self.client, path=self.url, allow_pattern=allow_pattern + ) + return [f"s3://{bucket_name}/{path}" for path in paths] + + def pull_files( + self, + allow_pattern: Optional[list[str]] = None, + ignore_pattern: Optional[list[str]] = None, + ) -> None: + """ + Pull files from S3 storage into the temporary directory. + + Args: + s3_model_path: The S3 path of the model. + allow_pattern: A list of patterns of which files to pull. + ignore_pattern: A list of patterns of which files not to pull. + + """ + bucket_name, base_dir, files = list_files( + self.client, self.url, allow_pattern, ignore_pattern + ) + if len(files) == 0: + return + + for file in files: + destination_file = os.path.join(self.local_dir, file.removeprefix(base_dir)) + local_dir = Path(destination_file).parent + os.makedirs(local_dir, exist_ok=True) + self.client.download_file(bucket_name, file, destination_file) + + def weight_iterator( + self, rank: int = 0 + ) -> Generator[Tuple[str, torch.Tensor], None, None]: + from sglang.srt.model_loader.weight_utils import ( + runai_safetensors_weights_iterator, + ) + + # only support safetensor files now + hf_weights_files = self.glob(allow_pattern=["*.safetensors"]) + return runai_safetensors_weights_iterator(hf_weights_files) + + def close(self): + self.client.close() + super().close() diff --git a/sglang/python/sglang/srt/connector/serde/__init__.py b/sglang/python/sglang/srt/connector/serde/__init__.py new file mode 100644 index 0000000..394dba0 --- /dev/null +++ b/sglang/python/sglang/srt/connector/serde/__init__.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 + +# inspired by LMCache +from typing import Optional, Tuple + +import torch + +from sglang.srt.connector.serde.safe_serde import SafeDeserializer, SafeSerializer +from sglang.srt.connector.serde.serde import Deserializer, Serializer + + +def create_serde(serde_type: str) -> Tuple[Serializer, Deserializer]: + s: Optional[Serializer] = None + d: Optional[Deserializer] = None + + if serde_type == "safe": + s = SafeSerializer() + d = SafeDeserializer(torch.uint8) + else: + raise ValueError(f"Unknown serde type: {serde_type}") + + return s, d + + +__all__ = [ + "Serializer", + "Deserializer", + "SafeSerializer", + "SafeDeserializer", + "create_serde", +] diff --git a/sglang/python/sglang/srt/connector/serde/safe_serde.py b/sglang/python/sglang/srt/connector/serde/safe_serde.py new file mode 100644 index 0000000..0163af9 --- /dev/null +++ b/sglang/python/sglang/srt/connector/serde/safe_serde.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 + +from typing import Union + +import torch +from safetensors.torch import load, save + +from sglang.srt.connector.serde.serde import Deserializer, Serializer + + +class SafeSerializer(Serializer): + + def __init__(self): + super().__init__() + + def to_bytes(self, t: torch.Tensor) -> bytes: + return save({"tensor_bytes": t.cpu().contiguous()}) + + +class SafeDeserializer(Deserializer): + + def __init__(self, dtype): + super().__init__(dtype) + + def from_bytes_normal(self, b: Union[bytearray, bytes]) -> torch.Tensor: + return load(bytes(b))["tensor_bytes"].to(dtype=self.dtype) + + def from_bytes(self, b: Union[bytearray, bytes]) -> torch.Tensor: + return self.from_bytes_normal(b) diff --git a/sglang/python/sglang/srt/connector/serde/serde.py b/sglang/python/sglang/srt/connector/serde/serde.py new file mode 100644 index 0000000..3d6f804 --- /dev/null +++ b/sglang/python/sglang/srt/connector/serde/serde.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 + +import abc +from abc import ABC, abstractmethod + +import torch + + +class Serializer(ABC): + + @abstractmethod + def to_bytes(self, t: torch.Tensor) -> bytes: + """ + Serialize a pytorch tensor to bytes. The serialized bytes should contain + both the data and the metadata (shape, dtype, etc.) of the tensor. + + Input: + t: the input pytorch tensor, can be on any device, in any shape, + with any dtype + + Returns: + bytes: the serialized bytes + """ + raise NotImplementedError + + +class Deserializer(metaclass=abc.ABCMeta): + + def __init__(self, dtype): + self.dtype = dtype + + @abstractmethod + def from_bytes(self, bs: bytes) -> torch.Tensor: + """ + Deserialize a pytorch tensor from bytes. + + Input: + bytes: a stream of bytes + + Output: + torch.Tensor: the deserialized pytorch tensor + """ + raise NotImplementedError diff --git a/sglang/python/sglang/srt/connector/utils.py b/sglang/python/sglang/srt/connector/utils.py new file mode 100644 index 0000000..6cd12da --- /dev/null +++ b/sglang/python/sglang/srt/connector/utils.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os +from pathlib import Path +from typing import Optional +from urllib.parse import urlparse + +from sglang.srt.connector import BaseConnector + + +def parse_model_name(url: str) -> str: + """ + Parse the model name from the url. + Only used for db connector + """ + parsed_url = urlparse(url) + return parsed_url.path.lstrip("/") + + +def pull_files_from_db( + connector: BaseConnector, + model_name: str, + allow_pattern: Optional[list[str]] = None, + ignore_pattern: Optional[list[str]] = None, +) -> None: + prefix = f"{model_name}/files/" + local_dir = connector.get_local_dir() + files = connector.list(prefix) + + for file in files: + destination_file = os.path.join(local_dir, file.removeprefix(prefix)) + local_dir = Path(destination_file).parent + os.makedirs(local_dir, exist_ok=True) + with open(destination_file, "wb") as f: + f.write(connector.getstr(file).encode("utf-8")) diff --git a/sglang/python/sglang/srt/constrained/base_grammar_backend.py b/sglang/python/sglang/srt/constrained/base_grammar_backend.py new file mode 100644 index 0000000..5d1b0da --- /dev/null +++ b/sglang/python/sglang/srt/constrained/base_grammar_backend.py @@ -0,0 +1,194 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""The baseclass of a backend for grammar-guided constrained decoding.""" + +import logging +from abc import ABC, abstractmethod +from concurrent.futures import Future, ThreadPoolExecutor +from dataclasses import dataclass +from threading import Event, Lock +from typing import Dict, List, Optional, Tuple + +import torch + +from sglang.srt.server_args import ServerArgs + +logger = logging.getLogger(__name__) + + +class BaseGrammarObject(ABC): + @abstractmethod + def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]: + """ + Try to jump forward in the grammar. + + Returns: + A jump forward helper which may be used in `jump_forward_str_state`. + None if the jump forward is not possible. + """ + raise NotImplementedError + + @abstractmethod + def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]: + """ + Jump forward for the grammar. + + Returns: + A tuple of the jump forward string and the next state of the grammar + (which can be used in `jump_and_retokenize` if needed). + """ + raise NotImplementedError + + @abstractmethod + def jump_and_retokenize( + self, old_output_ids: List[int], new_output_ids: List[int], next_state: int + ) -> None: + """ + Jump forward occurs, and update the grammar state if needed. + """ + raise NotImplementedError + + @abstractmethod + def allocate_vocab_mask( + self, vocab_size: int, batch_size: int, device + ) -> torch.Tensor: + raise NotImplementedError + + @abstractmethod + def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None: + raise NotImplementedError + + @staticmethod + @abstractmethod + def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor: + raise NotImplementedError + + @staticmethod + @abstractmethod + def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None: + raise NotImplementedError + + @abstractmethod + def copy(self) -> "BaseGrammarObject": + raise NotImplementedError + + +@dataclass +class CacheEntry: + value: Optional[BaseGrammarObject] + event: Event + + +class BaseGrammarBackend(ABC): + def __init__(self): + self.executor = ThreadPoolExecutor() + self.cache: Dict[Tuple[str, str], CacheEntry] = {} + self.cache_lock = Lock() + + def _not_supported(self, key_type: str, key_string: str) -> None: + logger.warning(f"Skip unsupported {key_type}: {key_type}={key_string}") + + def dispatch_fallback( + self, key_type: str, key_string: str + ) -> Optional[BaseGrammarObject]: + """ + This function should not be reached in any case. + """ + raise ValueError(f"Invalid key_type: {key_type}={key_string}") + + @abstractmethod + def dispatch_json(self, key_string: str) -> Optional[BaseGrammarObject]: + return self._not_supported("json", key_string) + + @abstractmethod + def dispatch_regex(self, key_string: str) -> Optional[BaseGrammarObject]: + return self._not_supported("regex", key_string) + + @abstractmethod + def dispatch_ebnf(self, key_string: str) -> Optional[BaseGrammarObject]: + return self._not_supported("ebnf", key_string) + + @abstractmethod + def dispatch_structural_tag(self, key_string: str) -> Optional[BaseGrammarObject]: + return self._not_supported("structural_tag", key_string) + + def _init_value_dispatch(self, key: Tuple[str, str]) -> Optional[BaseGrammarObject]: + key_type, key_string = key + if key_type == "json": + return self.dispatch_json(key_string) + elif key_type == "regex": + return self.dispatch_regex(key_string) + elif key_type == "ebnf": + return self.dispatch_ebnf(key_string) + elif key_type == "structural_tag": + return self.dispatch_structural_tag(key_string) + else: + return self.dispatch_fallback(key_type, key_string) + + def _init_value(self, key: Tuple[str, str]) -> Optional[BaseGrammarObject]: + with self.cache_lock: + if key in self.cache: + cache_hit = True + entry = self.cache[key] + else: + cache_hit = False + entry = CacheEntry(None, Event()) + self.cache[key] = entry + + if cache_hit: + entry.event.wait() + else: + entry.value = self._init_value_dispatch(key) + entry.event.set() + return entry.value.copy() if entry.value else None + + def get_cached_value(self, key: Tuple[str, str]) -> Optional[BaseGrammarObject]: + with self.cache_lock: + entry = self.cache.get(key) + if not entry or not entry.event.is_set(): + return None + val = self.cache[key].value + return val.copy() if val else None + + def get_future_value(self, key: Tuple[str, str]) -> Future: + return self.executor.submit(self._init_value, key) + + def reset(self): + with self.cache_lock: + self.cache.clear() + + +def create_grammar_backend(server_args: ServerArgs, tokenizer, vocab_size): + if server_args.grammar_backend == "outlines": + from sglang.srt.constrained.outlines_backend import OutlinesGrammarBackend + + grammar_backend = OutlinesGrammarBackend( + tokenizer, + whitespace_pattern=server_args.constrained_json_whitespace_pattern, + ) + elif server_args.grammar_backend == "xgrammar": + from sglang.srt.constrained.xgrammar_backend import XGrammarGrammarBackend + + grammar_backend = XGrammarGrammarBackend(tokenizer, vocab_size=vocab_size) + elif server_args.grammar_backend == "llguidance": + from sglang.srt.constrained.llguidance_backend import GuidanceBackend + + grammar_backend = GuidanceBackend( + tokenizer=tokenizer, + whitespace_pattern=server_args.constrained_json_whitespace_pattern, + ) + else: + raise ValueError(f"Invalid grammar backend: {server_args.grammar_backend}") + + return grammar_backend diff --git a/sglang/python/sglang/srt/constrained/llguidance_backend.py b/sglang/python/sglang/srt/constrained/llguidance_backend.py new file mode 100644 index 0000000..24893a4 --- /dev/null +++ b/sglang/python/sglang/srt/constrained/llguidance_backend.py @@ -0,0 +1,151 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Constrained decoding with llguidance backend.""" + +import json +import os +from typing import List, Optional, Tuple + +import llguidance +import llguidance.hf +import llguidance.torch +import torch +from llguidance.gbnf_to_lark import any_to_lark + +from sglang.srt.constrained.base_grammar_backend import ( + BaseGrammarBackend, + BaseGrammarObject, +) + + +class GuidanceGrammar(BaseGrammarObject): + def __init__( + self, llguidance_tokenizer: llguidance.LLTokenizer, serialized_grammar: str + ): + self.llguidance_tokenizer = llguidance_tokenizer + self.serialized_grammar = serialized_grammar + + # TODO: add support for fast-forward tokens in the future + self.ll_interpreter = llguidance.LLInterpreter( + self.llguidance_tokenizer, + self.serialized_grammar, + enable_backtrack=False, + enable_ff_tokens=False, + log_level=int(os.environ.get("LLGUIDANCE_LOG_LEVEL", "1")), + ) + self.pending_ff_tokens: list[int] = [] + self.finished = False + self.bitmask = None + + def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]: + if len(self.pending_ff_tokens) > 0: + s = self.llguidance_tokenizer.decode_str(self.pending_ff_tokens) + ff_tokens = self.pending_ff_tokens + self.pending_ff_tokens = [] + return (ff_tokens, s) + + return None + + def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]: + return "", -1 + + def jump_and_retokenize( + self, old_output_ids: List[int], new_output_ids: List[int], next_state: int + ): + pass + + def accept_token(self, token: int): + backtrack, ff_tokens = self.ll_interpreter.commit_token(token) + if len(ff_tokens) > 0 and backtrack == 0: + # first token is last generated token + ff_tokens = ff_tokens[1:] + self.pending_ff_tokens.extend(ff_tokens) + + def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None: + if len(self.pending_ff_tokens) > 0: + # if we have pending fast-forward tokens, + # just return them immediately + ff_token = self.pending_ff_tokens.pop(0) + vocab_mask[idx, :] = 0 + vocab_mask[idx, ff_token // 32] = 1 << (ff_token % 32) + return + + if self.ll_interpreter.has_pending_stop(): + self.finished = True + + llguidance.torch.fill_next_token_bitmask(self.ll_interpreter, vocab_mask, idx) + + def allocate_vocab_mask( + self, vocab_size: int, batch_size: int, device + ) -> torch.Tensor: + if self.bitmask is None or self.bitmask.shape[0] < batch_size: + # only create bitmask when batch gets larger + self.bitmask = llguidance.torch.allocate_token_bitmask( + batch_size, self.llguidance_tokenizer.vocab_size + ) + bitmask = self.bitmask + else: + bitmask = self.bitmask[:batch_size] + + return bitmask + + @staticmethod + def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor: + return vocab_mask.to(device, non_blocking=True) + + @staticmethod + def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None: + llguidance.torch.apply_token_bitmask_inplace(logits, vocab_mask) + + def copy(self): + return GuidanceGrammar( + llguidance_tokenizer=self.llguidance_tokenizer, + serialized_grammar=self.serialized_grammar, + ) + + +class GuidanceBackend(BaseGrammarBackend): + def __init__(self, tokenizer, whitespace_pattern: Optional[str] = None): + super().__init__() + + self.tokenizer = tokenizer + self.whitespace_flexible = ( + True if whitespace_pattern == "whitespace_flexible" else False + ) + self.llguidance_tokenizer = llguidance.hf.from_tokenizer(self.tokenizer, None) + + def _from_serialized(self, serialized_grammar) -> GuidanceGrammar: + return GuidanceGrammar( + llguidance_tokenizer=self.llguidance_tokenizer, + serialized_grammar=serialized_grammar, + ) + + def dispatch_json(self, key_string: str) -> GuidanceGrammar: + json_schema = key_string + compiler = llguidance.JsonCompiler(whitespace_flexible=self.whitespace_flexible) + serialized_grammar = compiler.compile(json_schema) + return self._from_serialized(serialized_grammar) + + def dispatch_regex(self, key_string: str) -> GuidanceGrammar: + compiler = llguidance.RegexCompiler() + serialized_grammar = compiler.compile(regex=key_string) + return self._from_serialized(serialized_grammar) + + def dispatch_ebnf(self, key_string: str) -> GuidanceGrammar: + compiler = llguidance.LarkCompiler() + serialized_grammar = compiler.compile(any_to_lark(key_string)) + return self._from_serialized(serialized_grammar) + + def dispatch_structural_tag(self, key_string: str): + return super().dispatch_structural_tag(key_string) diff --git a/sglang/python/sglang/srt/constrained/outlines_backend.py b/sglang/python/sglang/srt/constrained/outlines_backend.py new file mode 100644 index 0000000..7a05bb2 --- /dev/null +++ b/sglang/python/sglang/srt/constrained/outlines_backend.py @@ -0,0 +1,188 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Constrained decoding with outlines backend.""" + +import json +import logging +from typing import Dict, List, Optional, Tuple, Union + +import interegular +import torch +from outlines.fsm.guide import RegexGuide +from outlines.models.transformers import TransformerTokenizer +from pydantic import BaseModel + +from sglang.srt.constrained.base_grammar_backend import ( + BaseGrammarBackend, + BaseGrammarObject, +) +from sglang.srt.constrained.outlines_jump_forward import OutlinesJumpForwardMap + +try: + from outlines.fsm.json_schema import build_regex_from_schema +except ImportError: + from outlines_core.fsm.json_schema import build_regex_from_schema + + +logger = logging.getLogger(__name__) + + +class OutlinesGrammar(BaseGrammarObject): + def __init__( + self, + guide: RegexGuide, + jump_forward_map: Union[OutlinesJumpForwardMap, None], + ) -> None: + self.guide = guide + self.jump_forward_map = jump_forward_map + self.state = 0 + self.finished = False + + def accept_token(self, token: int): + self.state = self.guide.get_next_state(self.state, token) + + def try_jump_forward(self, tokenizer) -> Optional[Tuple]: + if not self.jump_forward_map: + return None + + jump_forward_bytes = self.jump_forward_map.jump_forward_byte(self.state) + if jump_forward_bytes is None or len(jump_forward_bytes) <= 1: + return None + + # preprocess the jump forward string + suffix_bytes = [] + continuation_range = range(0x80, 0xC0) + cur_state = self.state + while ( + len(jump_forward_bytes) and jump_forward_bytes[0][0] in continuation_range + ): + # continuation bytes + byte_edge = jump_forward_bytes.pop(0) + suffix_bytes.append(byte_edge[0]) + cur_state = byte_edge[1] + + suffix_tokens = [f"<0x{hex(b)[2:].upper()}>" for b in suffix_bytes] + suffix_ids = tokenizer.convert_tokens_to_ids(suffix_tokens) + return suffix_ids, cur_state + + def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]: + _, cur_state = helper + return self.jump_forward_map.jump_forward_symbol(cur_state) + + def jump_and_retokenize( + self, old_output_ids: List[int], new_output_ids: List[int], next_state: int + ): + self.state = next_state + + def allocate_vocab_mask( + self, vocab_size: int, batch_size: int, device + ) -> torch.Tensor: + return torch.zeros(batch_size, vocab_size, dtype=torch.bool, device=device) + + @staticmethod + def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor: + return vocab_mask + + def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None: + tokens = torch.tensor( + self.guide.get_next_instruction(self.state).tokens, dtype=torch.int64 + ).to(vocab_mask.device, non_blocking=True) + vocab_mask = vocab_mask[idx] + vocab_mask.fill_(1) + vocab_mask.scatter_(0, tokens, torch.zeros_like(tokens, dtype=torch.bool)) + + @staticmethod + def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor): + logits.masked_fill_(vocab_mask, float("-inf")) + + def copy(self): + return OutlinesGrammar(self.guide, self.jump_forward_map) + + +class OutlinesGrammarBackend(BaseGrammarBackend): + def __init__( + self, + tokenizer, + whitespace_pattern: bool, + ): + super().__init__() + + try: + self.outlines_tokenizer = TransformerTokenizer(tokenizer) + except AttributeError: + # FIXME: tmp fix for chatglm2 & chatglm3 (pad_token_id=0) + origin_pad_token_id = tokenizer.pad_token_id + + def fset(self, value): + self._value = value + + type(tokenizer).pad_token_id = property( + fget=type(tokenizer).pad_token_id.fget, fset=fset + ) + self.outlines_tokenizer = TransformerTokenizer(tokenizer) + self.outlines_tokenizer.tokenizer.pad_token_id = origin_pad_token_id + self.outlines_tokenizer.pad_token_id = origin_pad_token_id + self.outlines_tokenizer.pad_token = ( + self.outlines_tokenizer.tokenizer.pad_token + ) + self.outlines_tokenizer.vocabulary = ( + self.outlines_tokenizer.tokenizer.get_vocab() + ) + self.whitespace_pattern = whitespace_pattern + + def _compile_regex(self, regex: str) -> Optional[OutlinesGrammar]: + try: + if hasattr(RegexGuide, "from_regex"): + # outlines >= 0.1.1 + guide = RegexGuide.from_regex(regex, self.outlines_tokenizer) + else: + # outlines <= 0.0.46 + guide = RegexGuide(regex, self.outlines_tokenizer) + except interegular.patterns.InvalidSyntax as e: + logger.warning(f"skip invalid regex schema: {regex=}, {e=}") + return None + + jump_forward_map = None + return OutlinesGrammar(guide, jump_forward_map) + + def dispatch_ebnf(self, key_string: str): + return super().dispatch_ebnf(key_string) + + def dispatch_structural_tag(self, key_string: str): + return super().dispatch_structural_tag(key_string) + + def dispatch_json(self, key_string: str): + try: + regex = build_regex_from_object( + key_string, + whitespace_pattern=self.whitespace_pattern, + ) + except (NotImplementedError, json.decoder.JSONDecodeError) as e: + logger.warning(f"Skip invalid json_schema: json_schema={key_string}, {e=}") + return self._compile_regex(regex) + + def dispatch_regex(self, key_string: str): + return self._compile_regex(key_string) + + +def build_regex_from_object( + object: Union[str, BaseModel, Dict], whitespace_pattern: Optional[str] = None +): + if isinstance(object, type(BaseModel)): + schema = json.dumps(object.model_json_schema()) + elif isinstance(object, Dict): + schema = json.dumps(object) + else: + schema = object + return build_regex_from_schema(schema, whitespace_pattern) diff --git a/sglang/python/sglang/srt/constrained/outlines_jump_forward.py b/sglang/python/sglang/srt/constrained/outlines_jump_forward.py new file mode 100644 index 0000000..5e32c40 --- /dev/null +++ b/sglang/python/sglang/srt/constrained/outlines_jump_forward.py @@ -0,0 +1,187 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +""" +Faster constrained decoding with jump forward decoding / compressed finite state machine. +Reference: https://lmsys.org/blog/2024-02-05-compressed-fsm/ +""" + +import dataclasses +import logging +from collections import defaultdict + +import interegular +from interegular import InvalidSyntax +from outlines.caching import cache as disk_cache + +try: + # outlines >= 0.1.0 + from outlines_core.fsm.outlines_core_rs import FSMInfo + from outlines_core.fsm.regex import make_byte_level_fsm, make_deterministic_fsm +except ImportError: + # outlines <= 0.0.46 + from outlines.fsm.regex import FSMInfo, make_byte_level_fsm, make_deterministic_fsm + +IP_REGEX = r"((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)" + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class JumpEdge: + symbol: str = None + symbol_next_state: int = None + byte: int = None + byte_next_state: int = None + + +@disk_cache() +def init_state_to_jump_forward(regex_string): + try: + regex_pattern = interegular.parse_pattern(regex_string) + except InvalidSyntax as e: + logger.warning(f"skip invalid regex: {regex_string}, {e=}") + return + + byte_fsm = make_byte_level_fsm(regex_pattern.to_fsm().reduce(), keep_utf8=True) + regex_fsm, _ = make_deterministic_fsm(byte_fsm) + + fsm_info: FSMInfo = regex_fsm.fsm_info + + symbol_to_id = fsm_info.alphabet_symbol_mapping + id_to_symbol = {} + for symbol, id_ in symbol_to_id.items(): + id_to_symbol.setdefault(id_, []).append(symbol) + + transitions = fsm_info.transitions + + outgoings_ct = defaultdict(int) + # NOTE(lsyin): Final states can lead to terminate, so they have one outgoing edge naturally + for s in fsm_info.finals: + outgoings_ct[s] = 1 + + state_to_jump_forward = {} + for (state, id_), next_state in transitions.items(): + if id_ == fsm_info.alphabet_anything_value: + # Arbitrarily symbol cannot be recognized as jump forward + continue + + symbols = id_to_symbol[id_] + for c in symbols: + if len(c) > 1: + # Skip byte level transitions like c = "5E" + continue + + outgoings_ct[state] += 1 + if outgoings_ct[state] > 1: + if state in state_to_jump_forward: + del state_to_jump_forward[state] + break + + state_to_jump_forward[state] = JumpEdge( + symbol=c, + symbol_next_state=next_state, + ) + + # Process the byte level jump forward + outgoings_ct = defaultdict(int) + for s in fsm_info.finals: + outgoings_ct[s] = 1 + + for (state, id_), next_state in transitions.items(): + if id_ == fsm_info.alphabet_anything_value: + continue + symbols = id_to_symbol[id_] + for c in symbols: + byte_ = None + if len(c) == 1 and ord(c) < 0x80: + # ASCII character + byte_ = ord(c) + elif len(c) > 1: + # FIXME: This logic is due to the leading \x00 + # https://github.com/outlines-dev/outlines/pull/930 + byte_ = int(symbols[0][1:], 16) + + if byte_ is not None: + outgoings_ct[state] += 1 + if outgoings_ct[state] > 1: + if state in state_to_jump_forward: + del state_to_jump_forward[state] + break + e = state_to_jump_forward.get(state, JumpEdge()) + e.byte = byte_ + e.byte_next_state = next_state + state_to_jump_forward[state] = e + + return state_to_jump_forward + + +class OutlinesJumpForwardMap: + def __init__(self, regex_string): + self.state_to_jump_forward = init_state_to_jump_forward(regex_string) + + def jump_forward_symbol(self, state): + jump_forward_str = "" + next_state = state + while state in self.state_to_jump_forward: + e = self.state_to_jump_forward[state] + if e.symbol is None: + break + jump_forward_str += e.symbol + next_state = e.symbol_next_state + state = next_state + + return jump_forward_str, next_state + + def jump_forward_byte(self, state): + if state not in self.state_to_jump_forward: + return None + + jump_forward_bytes = [] + next_state = None + while state in self.state_to_jump_forward: + e = self.state_to_jump_forward[state] + assert e.byte is not None and e.byte_next_state is not None + jump_forward_bytes.append((e.byte, e.byte_next_state)) + next_state = e.byte_next_state + state = next_state + + return jump_forward_bytes + + def is_jump_forward_symbol_state(self, state): + return ( + state in self.state_to_jump_forward + and self.state_to_jump_forward[state].symbol is not None + ) + + +def test_main(regex_string): + jump_forward_map = OutlinesJumpForwardMap(regex_string) + for state, e in jump_forward_map.state_to_jump_forward.items(): + if e.symbol is not None: + jump_forward_str, next_state = jump_forward_map.jump_forward_symbol(state) + print(f"{state} -> {next_state}", jump_forward_str) + bytes_ = jump_forward_map.jump_forward_byte(state) + print(f"{state} -> {bytes_[-1][1]}", [hex(b) for b, _ in bytes_]) + + +if __name__ == "__main__": + import outlines + + outlines.caching.clear_cache() + test_main(r"The google's DNS sever address is " + IP_REGEX) + test_main(r"霍格沃茨特快列车|霍比特人比尔博") + # 霍格: \xe9\x9c\x8d \xe6\xa0\xbc ... + # 霍比: \xe9\x9c\x8d \xe6\xaf\x94 ... + + test_main(r"[-+]?[0-9]+[ ]*") diff --git a/sglang/python/sglang/srt/constrained/xgrammar_backend.py b/sglang/python/sglang/srt/constrained/xgrammar_backend.py new file mode 100644 index 0000000..4df3ae2 --- /dev/null +++ b/sglang/python/sglang/srt/constrained/xgrammar_backend.py @@ -0,0 +1,183 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Constrained decoding with xgrammar backend.""" + +import json +import logging +from typing import List, Optional, Tuple, Union + +import torch +from xgrammar import ( + CompiledGrammar, + GrammarCompiler, + GrammarMatcher, + StructuralTagItem, + TokenizerInfo, + allocate_token_bitmask, + apply_token_bitmask_inplace, +) + +from sglang.srt.constrained.base_grammar_backend import ( + BaseGrammarBackend, + BaseGrammarObject, +) + +logger = logging.getLogger(__name__) + + +MAX_ROLLBACK_TOKENS = 200 + + +class XGrammarGrammar(BaseGrammarObject): + + def __init__( + self, + matcher: GrammarMatcher, + vocab_size: int, + ctx: CompiledGrammar, + override_stop_tokens: Optional[Union[List[int], int]], + ) -> None: + self.matcher = matcher + self.vocab_size = vocab_size + self.ctx = ctx + self.override_stop_tokens = override_stop_tokens + self.finished = False + + def accept_token(self, token: int): + assert self.matcher.accept_token(token) + + def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]: + s = self.matcher.find_jump_forward_string() + if s: + return [], s + return None + + def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]: + _, data = helper + return data, -1 + + def jump_and_retokenize( + self, old_output_ids: List[int], new_output_ids: List[int], next_state: int + ): + k = 0 + for i, old_id in enumerate(old_output_ids): + if old_id == new_output_ids[i]: + k = i + 1 + else: + break + + # rollback to the last token that is the same + if k < len(old_output_ids): + self.matcher.rollback(len(old_output_ids) - k) + + for i in range(k, len(new_output_ids)): + assert self.matcher.accept_token(new_output_ids[i]) + + def allocate_vocab_mask( + self, vocab_size: int, batch_size: int, device + ) -> torch.Tensor: + return allocate_token_bitmask(batch_size, vocab_size) + + def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None: + self.matcher.fill_next_token_bitmask(vocab_mask, idx) + + @staticmethod + def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor: + return vocab_mask.to(device, non_blocking=True) + + @staticmethod + def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None: + apply_token_bitmask_inplace(logits, vocab_mask) + + def copy(self): + matcher = GrammarMatcher( + self.ctx, + max_rollback_tokens=MAX_ROLLBACK_TOKENS, + override_stop_tokens=self.override_stop_tokens, + ) + return XGrammarGrammar( + matcher, self.vocab_size, self.ctx, self.override_stop_tokens + ) + + +class XGrammarGrammarBackend(BaseGrammarBackend): + def __init__( + self, + tokenizer, + vocab_size: int, + ): + super().__init__() + + tokenizer_info = TokenizerInfo.from_huggingface( + tokenizer, vocab_size=vocab_size + ) + override_stop_tokens = None + + self.grammar_compiler = GrammarCompiler(tokenizer_info=tokenizer_info) + self.vocab_size = vocab_size + self.override_stop_tokens = override_stop_tokens + + def _from_context(self, ctx: CompiledGrammar) -> XGrammarGrammar: + matcher = GrammarMatcher(ctx, max_rollback_tokens=MAX_ROLLBACK_TOKENS) + return XGrammarGrammar(matcher, self.vocab_size, ctx, self.override_stop_tokens) + + def dispatch_json(self, key_string: str) -> Optional[XGrammarGrammar]: + try: + if key_string == "$$ANY$$": + ctx = self.grammar_compiler.compile_builtin_json_grammar() + else: + ctx = self.grammar_compiler.compile_json_schema(schema=key_string) + except RuntimeError as e: + logging.warning(f"Skip invalid json_schema: json_schema={key_string}, {e=}") + return None + return self._from_context(ctx) + + def dispatch_ebnf(self, key_string: str) -> Optional[XGrammarGrammar]: + try: + ctx = self.grammar_compiler.compile_grammar(key_string) + except RuntimeError as e: + logging.warning(f"Skip invalid ebnf: ebnf={key_string}, {e=}") + return None + return self._from_context(ctx) + + def dispatch_regex(self, key_string: str) -> Optional[XGrammarGrammar]: + try: + ctx = self.grammar_compiler.compile_regex(key_string) + except RuntimeError as e: + logging.warning(f"Skip invalid regex: regex={key_string}, {e=}") + return None + return self._from_context(ctx) + + def dispatch_structural_tag(self, key_string: str) -> Optional[XGrammarGrammar]: + try: + structural_tag = json.loads(key_string) + tags = [ + StructuralTagItem( + begin=structure["begin"], + schema=json.dumps(structure["schema"]), + end=structure["end"], + ) + for structure in structural_tag["structures"] + ] + ctx = self.grammar_compiler.compile_structural_tag( + tags, structural_tag["triggers"] + ) + except RuntimeError as e: + logging.warning(f"Skip invalid regex: regex={key_string}, {e=}") + return None + return self._from_context(ctx) + + def reset(self): + if self.grammar_compiler: + self.grammar_compiler.clear_cache() diff --git a/sglang/python/sglang/srt/conversation.py b/sglang/python/sglang/srt/conversation.py new file mode 100644 index 0000000..70152a6 --- /dev/null +++ b/sglang/python/sglang/srt/conversation.py @@ -0,0 +1,736 @@ +# Copyright 2023-2024 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Conversation chat templates.""" + +# Adapted from +# https://github.com/lm-sys/FastChat/blob/main/fastchat/conversation.py +import dataclasses +from enum import IntEnum, auto +from typing import Dict, List, Optional, Tuple, Union + +from sglang.srt.openai_api.protocol import ChatCompletionRequest + + +class SeparatorStyle(IntEnum): + """Separator styles.""" + + ADD_COLON_SINGLE = auto() + ADD_COLON_TWO = auto() + ADD_COLON_SPACE_SINGLE = auto() + NO_COLON_SINGLE = auto() + NO_COLON_TWO = auto() + ADD_NEW_LINE_SINGLE = auto() + LLAMA2 = auto() + LLAMA3 = auto() + CHATGLM = auto() + CHATML = auto() + CHATINTERN = auto() + DOLLY = auto() + RWKV = auto() + PHOENIX = auto() + ROBIN = auto() + FALCON_CHAT = auto() + CHATGLM3 = auto() + DEEPSEEK_CHAT = auto() + METAMATH = auto() + DeepSeekVL2 = auto() + QWEN2_VL_EMBED = auto() + GEMMA3 = auto() + + +@dataclasses.dataclass +class Conversation: + """A class that manages prompt templates and keeps all conversation history.""" + + # The name of this template + name: str + # The template of the system prompt + system_template: str = "{system_message}" + # The system message + system_message: str = "" + # The names of two roles + roles: Tuple[str] = ("USER", "ASSISTANT") + # All messages. Each item is (role, message). + messages: List[List[str]] = () + # The number of few shot examples + offset: int = 0 + # The separator style and configurations + sep_style: SeparatorStyle = SeparatorStyle.ADD_COLON_SINGLE + sep: str = "\n" + sep2: str = None + # Stop criteria (the default one is EOS token) + stop_str: Union[str, List[str]] = None + # The string that represents an image token in the prompt + image_token: str = "" + audio_token: str = "