Add files using upload-large-folder tool

bin/activate

@@ -0,0 +1,76 @@
+# This file must be used with "source bin/activate" *from bash*
+# You cannot run it directly
+
+deactivate () {
+    # reset old environment variables
+    if [ -n "${_OLD_VIRTUAL_PATH:-}" ] ; then
+        PATH="${_OLD_VIRTUAL_PATH:-}"
+        export PATH
+        unset _OLD_VIRTUAL_PATH
+    fi
+    if [ -n "${_OLD_VIRTUAL_PYTHONHOME:-}" ] ; then
+        PYTHONHOME="${_OLD_VIRTUAL_PYTHONHOME:-}"
+        export PYTHONHOME
+        unset _OLD_VIRTUAL_PYTHONHOME
+    fi
+
+    # Call hash to forget past locations. Without forgetting
+    # past locations the $PATH changes we made may not be respected.
+    # See "man bash" for more details. hash is usually a builtin of your shell
+    hash -r 2> /dev/null
+
+    if [ -n "${_OLD_VIRTUAL_PS1:-}" ] ; then
+        PS1="${_OLD_VIRTUAL_PS1:-}"
+        export PS1
+        unset _OLD_VIRTUAL_PS1
+    fi
+
+    unset VIRTUAL_ENV
+    unset VIRTUAL_ENV_PROMPT
+    if [ ! "${1:-}" = "nondestructive" ] ; then
+    # Self destruct!
+        unset -f deactivate
+    fi
+}
+
+# unset irrelevant variables
+deactivate nondestructive
+
+# on Windows, a path can contain colons and backslashes and has to be converted:
+case "$(uname)" in
+    CYGWIN*|MSYS*|MINGW*)
+        # transform D:\path\to\venv to /d/path/to/venv on MSYS and MINGW
+        # and to /cygdrive/d/path/to/venv on Cygwin
+        VIRTUAL_ENV=$(cygpath /storage/openpsi/users/zzy/.threadweaver)
+        export VIRTUAL_ENV
+        ;;
+    *)
+        # use the path as-is
+        export VIRTUAL_ENV=/storage/openpsi/users/zzy/.threadweaver
+        ;;
+esac
+
+_OLD_VIRTUAL_PATH="$PATH"
+PATH="$VIRTUAL_ENV/"bin":$PATH"
+export PATH
+
+VIRTUAL_ENV_PROMPT='(.threadweaver) '
+export VIRTUAL_ENV_PROMPT
+
+# unset PYTHONHOME if set
+# this will fail if PYTHONHOME is set to the empty string (which is bad anyway)
+# could use `if (set -u; : $PYTHONHOME) ;` in bash
+if [ -n "${PYTHONHOME:-}" ] ; then
+    _OLD_VIRTUAL_PYTHONHOME="${PYTHONHOME:-}"
+    unset PYTHONHOME
+fi
+
+if [ -z "${VIRTUAL_ENV_DISABLE_PROMPT:-}" ] ; then
+    _OLD_VIRTUAL_PS1="${PS1:-}"
+    PS1="("'(.threadweaver) '") ${PS1:-}"
+    export PS1
+fi
+
+# Call hash to forget past commands. Without forgetting
+# past commands the $PATH changes we made may not be respected
+hash -r 2> /dev/null

bin/deepspeed

@@ -0,0 +1,6 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+
+from deepspeed.launcher.runner import main
+
+if __name__ == '__main__':
+    main()

bin/ds_report

@@ -0,0 +1,6 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+
+from deepspeed.env_report import cli_main
+
+if __name__ == '__main__':
+    cli_main()

bin/fastapi

@@ -0,0 +1,8 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from fastapi.cli import main
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(main())

bin/huggingface-cli

@@ -0,0 +1,8 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from huggingface_hub.commands.huggingface_cli import main
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(main())

bin/ipython3

@@ -0,0 +1,8 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from IPython import start_ipython
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(start_ipython())

bin/isympy

@@ -0,0 +1,8 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from isympy import main
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(main())

bin/markdown_py

@@ -0,0 +1,8 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from markdown.__main__ import run
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(run())

bin/pip3

@@ -0,0 +1,8 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from pip._internal.cli.main import main
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(main())

bin/pip3.12

@@ -0,0 +1,8 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from pip._internal.cli.main import main
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(main())

bin/proton

@@ -0,0 +1,8 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from triton.profiler.proton import main
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(main())

bin/pyftmerge

@@ -0,0 +1,8 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from fontTools.merge import main
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(main())

bin/ray

@@ -0,0 +1,8 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from ray.scripts.scripts import main
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(main())

bin/serve

@@ -0,0 +1,8 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from ray.serve.scripts import cli
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(cli())

bin/wandb

@@ -0,0 +1,8 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from wandb.cli.cli import cli
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(cli())

bin/wb

@@ -0,0 +1,8 @@
+#!/storage/openpsi/users/zzy/.threadweaver/bin/python3
+# -*- coding: utf-8 -*-
+import re
+import sys
+from wandb.cli.cli import cli
+if __name__ == '__main__':
+    sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', '', sys.argv[0])
+    sys.exit(cli())

lib/python3.12/site-packages/decorator.py

@@ -0,0 +1,459 @@
+# #########################     LICENSE     ############################ #
+
+# Copyright (c) 2005-2025, Michele Simionato
+# All rights reserved.
+
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are
+# met:
+
+#   Redistributions of source code must retain the above copyright
+#   notice, this list of conditions and the following disclaimer.
+#   Redistributions in bytecode form must reproduce the above copyright
+#   notice, this list of conditions and the following disclaimer in
+#   the documentation and/or other materials provided with the
+#   distribution.
+
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
+# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
+# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
+# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
+# HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
+# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
+# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
+# OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR
+# TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE
+# USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
+# DAMAGE.
+
+"""
+Decorator module, see
+https://github.com/micheles/decorator/blob/master/docs/documentation.md
+for the documentation.
+"""
+import re
+import sys
+import inspect
+import operator
+import itertools
+import functools
+from contextlib import _GeneratorContextManager
+from inspect import getfullargspec, iscoroutinefunction, isgeneratorfunction
+
+__version__ = '5.2.1'
+
+DEF = re.compile(r'\s*def\s*([_\w][_\w\d]*)\s*\(')
+POS = inspect.Parameter.POSITIONAL_OR_KEYWORD
+EMPTY = inspect.Parameter.empty
+
+
+# this is not used anymore in the core, but kept for backward compatibility
+class FunctionMaker(object):
+    """
+    An object with the ability to create functions with a given signature.
+    It has attributes name, doc, module, signature, defaults, dict and
+    methods update and make.
+    """
+
+    # Atomic get-and-increment provided by the GIL
+    _compile_count = itertools.count()
+
+    # make pylint happy
+    args = varargs = varkw = defaults = kwonlyargs = kwonlydefaults = ()
+
+    def __init__(self, func=None, name=None, signature=None,
+                 defaults=None, doc=None, module=None, funcdict=None):
+        self.shortsignature = signature
+        if func:
+            # func can be a class or a callable, but not an instance method
+            self.name = func.__name__
+            if self.name == '<lambda>':  # small hack for lambda functions
+                self.name = '_lambda_'
+            self.doc = func.__doc__
+            self.module = func.__module__
+            if inspect.isroutine(func) or isinstance(func, functools.partial):
+                argspec = getfullargspec(func)
+                self.annotations = getattr(func, '__annotations__', {})
+                for a in ('args', 'varargs', 'varkw', 'defaults', 'kwonlyargs',
+                          'kwonlydefaults'):
+                    setattr(self, a, getattr(argspec, a))
+                for i, arg in enumerate(self.args):
+                    setattr(self, 'arg%d' % i, arg)
+                allargs = list(self.args)
+                allshortargs = list(self.args)
+                if self.varargs:
+                    allargs.append('*' + self.varargs)
+                    allshortargs.append('*' + self.varargs)
+                elif self.kwonlyargs:
+                    allargs.append('*')  # single star syntax
+                for a in self.kwonlyargs:
+                    allargs.append('%s=None' % a)
+                    allshortargs.append('%s=%s' % (a, a))
+                if self.varkw:
+                    allargs.append('**' + self.varkw)
+                    allshortargs.append('**' + self.varkw)
+                self.signature = ', '.join(allargs)
+                self.shortsignature = ', '.join(allshortargs)
+                self.dict = func.__dict__.copy()
+        # func=None happens when decorating a caller
+        if name:
+            self.name = name
+        if signature is not None:
+            self.signature = signature
+        if defaults:
+            self.defaults = defaults
+        if doc:
+            self.doc = doc
+        if module:
+            self.module = module
+        if funcdict:
+            self.dict = funcdict
+        # check existence required attributes
+        assert hasattr(self, 'name')
+        if not hasattr(self, 'signature'):
+            raise TypeError('You are decorating a non function: %s' % func)
+
+    def update(self, func, **kw):
+        """
+        Update the signature of func with the data in self
+        """
+        func.__name__ = self.name
+        func.__doc__ = getattr(self, 'doc', None)
+        func.__dict__ = getattr(self, 'dict', {})
+        func.__defaults__ = self.defaults
+        func.__kwdefaults__ = self.kwonlydefaults or None
+        func.__annotations__ = getattr(self, 'annotations', None)
+        try:
+            frame = sys._getframe(3)
+        except AttributeError:  # for IronPython and similar implementations
+            callermodule = '?'
+        else:
+            callermodule = frame.f_globals.get('__name__', '?')
+        func.__module__ = getattr(self, 'module', callermodule)
+        func.__dict__.update(kw)
+
+    def make(self, src_templ, evaldict=None, addsource=False, **attrs):
+        """
+        Make a new function from a given template and update the signature
+        """
+        src = src_templ % vars(self)  # expand name and signature
+        evaldict = evaldict or {}
+        mo = DEF.search(src)
+        if mo is None:
+            raise SyntaxError('not a valid function template\n%s' % src)
+        name = mo.group(1)  # extract the function name
+        names = set([name] + [arg.strip(' *') for arg in
+                              self.shortsignature.split(',')])
+        for n in names:
+            if n in ('_func_', '_call_'):
+                raise NameError('%s is overridden in\n%s' % (n, src))
+
+        if not src.endswith('\n'):  # add a newline for old Pythons
+            src += '\n'
+
+        # Ensure each generated function has a unique filename for profilers
+        # (such as cProfile) that depend on the tuple of (<filename>,
+        # <definition line>, <function name>) being unique.
+        filename = '<decorator-gen-%d>' % next(self._compile_count)
+        try:
+            code = compile(src, filename, 'single')
+            exec(code, evaldict)
+        except Exception:
+            print('Error in generated code:', file=sys.stderr)
+            print(src, file=sys.stderr)
+            raise
+        func = evaldict[name]
+        if addsource:
+            attrs['__source__'] = src
+        self.update(func, **attrs)
+        return func
+
+    @classmethod
+    def create(cls, obj, body, evaldict, defaults=None,
+               doc=None, module=None, addsource=True, **attrs):
+        """
+        Create a function from the strings name, signature and body.
+        evaldict is the evaluation dictionary. If addsource is true an
+        attribute __source__ is added to the result. The attributes attrs
+        are added, if any.
+        """
+        if isinstance(obj, str):  # "name(signature)"
+            name, rest = obj.strip().split('(', 1)
+            signature = rest[:-1]  # strip a right parens
+            func = None
+        else:  # a function
+            name = None
+            signature = None
+            func = obj
+        self = cls(func, name, signature, defaults, doc, module)
+        ibody = '\n'.join('    ' + line for line in body.splitlines())
+        caller = evaldict.get('_call_')  # when called from `decorate`
+        if caller and iscoroutinefunction(caller):
+            body = ('async def %(name)s(%(signature)s):\n' + ibody).replace(
+                'return', 'return await')
+        else:
+            body = 'def %(name)s(%(signature)s):\n' + ibody
+        return self.make(body, evaldict, addsource, **attrs)
+
+
+def fix(args, kwargs, sig):
+    """
+    Fix args and kwargs to be consistent with the signature
+    """
+    ba = sig.bind(*args, **kwargs)
+    ba.apply_defaults()  # needed for test_dan_schult
+    return ba.args, ba.kwargs
+
+
+def decorate(func, caller, extras=(), kwsyntax=False):
+    """
+    Decorates a function/generator/coroutine using a caller.
+    If kwsyntax is True calling the decorated functions with keyword
+    syntax will pass the named arguments inside the ``kw`` dictionary,
+    even if such argument are positional, similarly to what functools.wraps
+    does. By default kwsyntax is False and the the arguments are untouched.
+    """
+    sig = inspect.signature(func)
+    if isinstance(func, functools.partial):
+        func = functools.update_wrapper(func, func.func)
+    if iscoroutinefunction(caller):
+        async def fun(*args, **kw):
+            if not kwsyntax:
+                args, kw = fix(args, kw, sig)
+            return await caller(func, *(extras + args), **kw)
+    elif isgeneratorfunction(caller):
+        def fun(*args, **kw):
+            if not kwsyntax:
+                args, kw = fix(args, kw, sig)
+            for res in caller(func, *(extras + args), **kw):
+                yield res
+    else:
+        def fun(*args, **kw):
+            if not kwsyntax:
+                args, kw = fix(args, kw, sig)
+            return caller(func, *(extras + args), **kw)
+
+    fun.__name__ = func.__name__
+    fun.__doc__ = func.__doc__
+    fun.__wrapped__ = func
+    fun.__signature__ = sig
+    fun.__qualname__ = func.__qualname__
+    # builtin functions like defaultdict.__setitem__ lack many attributes
+    try:
+        fun.__defaults__ = func.__defaults__
+    except AttributeError:
+        pass
+    try:
+        fun.__kwdefaults__ = func.__kwdefaults__
+    except AttributeError:
+        pass
+    try:
+        fun.__annotations__ = func.__annotations__
+    except AttributeError:
+        pass
+    try:
+        fun.__module__ = func.__module__
+    except AttributeError:
+        pass
+    try:
+        fun.__name__ = func.__name__
+    except AttributeError:  # happens with old versions of numpy.vectorize
+        func.__name__ == 'noname'
+    try:
+        fun.__dict__.update(func.__dict__)
+    except AttributeError:
+        pass
+    return fun
+
+
+def decoratorx(caller):
+    """
+    A version of "decorator" implemented via "exec" and not via the
+    Signature object. Use this if you are want to preserve the `.__code__`
+    object properties (https://github.com/micheles/decorator/issues/129).
+    """
+    def dec(func):
+        return FunctionMaker.create(
+            func,
+            "return _call_(_func_, %(shortsignature)s)",
+            dict(_call_=caller, _func_=func),
+            __wrapped__=func, __qualname__=func.__qualname__)
+    return dec
+
+
+def decorator(caller, _func=None, kwsyntax=False):
+    """
+    decorator(caller) converts a caller function into a decorator
+    """
+    if _func is not None:  # return a decorated function
+        # this is obsolete behavior; you should use decorate instead
+        return decorate(_func, caller, (), kwsyntax)
+    # else return a decorator function
+    sig = inspect.signature(caller)
+    dec_params = [p for p in sig.parameters.values() if p.kind is POS]
+
+    def dec(func=None, *args, **kw):
+        na = len(args) + 1
+        extras = args + tuple(kw.get(p.name, p.default)
+                              for p in dec_params[na:]
+                              if p.default is not EMPTY)
+        if func is None:
+            return lambda func: decorate(func, caller, extras, kwsyntax)
+        else:
+            return decorate(func, caller, extras, kwsyntax)
+    dec.__signature__ = sig.replace(parameters=dec_params)
+    dec.__name__ = caller.__name__
+    dec.__doc__ = caller.__doc__
+    dec.__wrapped__ = caller
+    dec.__qualname__ = caller.__qualname__
+    dec.__kwdefaults__ = getattr(caller, '__kwdefaults__', None)
+    dec.__dict__.update(caller.__dict__)
+    return dec
+
+
+# ####################### contextmanager ####################### #
+
+
+class ContextManager(_GeneratorContextManager):
+    def __init__(self, g, *a, **k):
+        _GeneratorContextManager.__init__(self, g, a, k)
+
+    def __call__(self, func):
+        def caller(f, *a, **k):
+            with self.__class__(self.func, *self.args, **self.kwds):
+                return f(*a, **k)
+        return decorate(func, caller)
+
+
+_contextmanager = decorator(ContextManager)
+
+
+def contextmanager(func):
+    # Enable Pylint config: contextmanager-decorators=decorator.contextmanager
+    return _contextmanager(func)
+
+
+# ############################ dispatch_on ############################ #
+
+def append(a, vancestors):
+    """
+    Append ``a`` to the list of the virtual ancestors, unless it is already
+    included.
+    """
+    add = True
+    for j, va in enumerate(vancestors):
+        if issubclass(va, a):
+            add = False
+            break
+        if issubclass(a, va):
+            vancestors[j] = a
+            add = False
+    if add:
+        vancestors.append(a)
+
+
+# inspired from simplegeneric by P.J. Eby and functools.singledispatch
+def dispatch_on(*dispatch_args):
+    """
+    Factory of decorators turning a function into a generic function
+    dispatching on the given arguments.
+    """
+    assert dispatch_args, 'No dispatch args passed'
+    dispatch_str = '(%s,)' % ', '.join(dispatch_args)
+
+    def check(arguments, wrong=operator.ne, msg=''):
+        """Make sure one passes the expected number of arguments"""
+        if wrong(len(arguments), len(dispatch_args)):
+            raise TypeError('Expected %d arguments, got %d%s' %
+                            (len(dispatch_args), len(arguments), msg))
+
+    def gen_func_dec(func):
+        """Decorator turning a function into a generic function"""
+
+        # first check the dispatch arguments
+        argset = set(getfullargspec(func).args)
+        if not set(dispatch_args) <= argset:
+            raise NameError('Unknown dispatch arguments %s' % dispatch_str)
+
+        typemap = {}
+
+        def vancestors(*types):
+            """
+            Get a list of sets of virtual ancestors for the given types
+            """
+            check(types)
+            ras = [[] for _ in range(len(dispatch_args))]
+            for types_ in typemap:
+                for t, type_, ra in zip(types, types_, ras):
+                    if issubclass(t, type_) and type_ not in t.mro():
+                        append(type_, ra)
+            return [set(ra) for ra in ras]
+
+        def ancestors(*types):
+            """
+            Get a list of virtual MROs, one for each type
+            """
+            check(types)
+            lists = []
+            for t, vas in zip(types, vancestors(*types)):
+                n_vas = len(vas)
+                if n_vas > 1:
+                    raise RuntimeError(
+                        'Ambiguous dispatch for %s: %s' % (t, vas))
+                elif n_vas == 1:
+                    va, = vas
+                    mro = type('t', (t, va), {}).mro()[1:]
+                else:
+                    mro = t.mro()
+                lists.append(mro[:-1])  # discard t and object
+            return lists
+
+        def register(*types):
+            """
+            Decorator to register an implementation for the given types
+            """
+            check(types)
+
+            def dec(f):
+                check(getfullargspec(f).args, operator.lt, ' in ' + f.__name__)
+                typemap[types] = f
+                return f
+            return dec
+
+        def dispatch_info(*types):
+            """
+            An utility to introspect the dispatch algorithm
+            """
+            check(types)
+            lst = []
+            for ancs in itertools.product(*ancestors(*types)):
+                lst.append(tuple(a.__name__ for a in ancs))
+            return lst
+
+        def _dispatch(dispatch_args, *args, **kw):
+            types = tuple(type(arg) for arg in dispatch_args)
+            try:  # fast path
+                f = typemap[types]
+            except KeyError:
+                pass
+            else:
+                return f(*args, **kw)
+            combinations = itertools.product(*ancestors(*types))
+            next(combinations)  # the first one has been already tried
+            for types_ in combinations:
+                f = typemap.get(types_)
+                if f is not None:
+                    return f(*args, **kw)
+
+            # else call the default implementation
+            return func(*args, **kw)
+
+        return FunctionMaker.create(
+            func, 'return _f_(%s, %%(shortsignature)s)' % dispatch_str,
+            dict(_f_=_dispatch), register=register, default=func,
+            typemap=typemap, vancestors=vancestors, ancestors=ancestors,
+            dispatch_info=dispatch_info, __wrapped__=func)
+
+    gen_func_dec.__name__ = 'dispatch_on' + dispatch_str
+    return gen_func_dec

lib/python3.12/site-packages/ipython_pygments_lexers.py

@@ -0,0 +1,582 @@
+# -*- coding: utf-8 -*-
+"""
+Defines a variety of Pygments lexers for highlighting IPython code.
+
+This includes:
+
+    IPythonLexer, IPython3Lexer
+        Lexers for pure IPython (python + magic/shell commands)
+
+    IPythonPartialTracebackLexer, IPythonTracebackLexer
+        Supports 2.x and 3.x via keyword `python3`.  The partial traceback
+        lexer reads everything but the Python code appearing in a traceback.
+        The full lexer combines the partial lexer with an IPython lexer.
+
+    IPythonConsoleLexer
+        A lexer for IPython console sessions, with support for tracebacks.
+
+    IPyLexer
+        A friendly lexer which examines the first line of text and from it,
+        decides whether to use an IPython lexer or an IPython console lexer.
+        This is probably the only lexer that needs to be explicitly added
+        to Pygments.
+
+"""
+# -----------------------------------------------------------------------------
+# Copyright (c) 2013, the IPython Development Team.
+#
+# Distributed under the terms of the Modified BSD License.
+#
+# The full license is in the file COPYING.txt, distributed with this software.
+# -----------------------------------------------------------------------------
+
+__version__ = "1.1.1"
+
+# Standard library
+import re
+
+# Third party
+from pygments.lexers import (
+    BashLexer,
+    HtmlLexer,
+    JavascriptLexer,
+    RubyLexer,
+    PerlLexer,
+    Python2Lexer,
+    Python3Lexer,
+    TexLexer,
+)
+from pygments.lexer import (
+    Lexer,
+    DelegatingLexer,
+    RegexLexer,
+    do_insertions,
+    bygroups,
+    using,
+)
+from pygments.token import (
+    Generic,
+    Keyword,
+    Literal,
+    Name,
+    Operator,
+    Other,
+    Text,
+    Error,
+)
+
+
+line_re = re.compile(".*?\n")
+
+__all__ = [
+    "IPython3Lexer",
+    "IPythonLexer",
+    "IPythonPartialTracebackLexer",
+    "IPythonTracebackLexer",
+    "IPythonConsoleLexer",
+    "IPyLexer",
+]
+
+
+ipython_tokens = [
+    (
+        r"(?s)(\s*)(%%capture)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(Python3Lexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%debug)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(Python3Lexer)),
+    ),
+    (
+        r"(?is)(\s*)(%%html)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(HtmlLexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%javascript)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(JavascriptLexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%js)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(JavascriptLexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%latex)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(TexLexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%perl)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(PerlLexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%prun)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(Python3Lexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%pypy)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(Python3Lexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%python2)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(Python2Lexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%python3)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(Python3Lexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%python)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(Python3Lexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%ruby)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(RubyLexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%timeit)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(Python3Lexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%time)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(Python3Lexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%writefile)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(Python3Lexer)),
+    ),
+    (
+        r"(?s)(\s*)(%%file)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(Python3Lexer)),
+    ),
+    (r"(?s)(\s*)(%%)(\w+)(.*)", bygroups(Text, Operator, Keyword, Text)),
+    (
+        r"(?s)(^\s*)(%%!)([^\n]*\n)(.*)",
+        bygroups(Text, Operator, Text, using(BashLexer)),
+    ),
+    (r"(%%?)(\w+)(\?\??)$", bygroups(Operator, Keyword, Operator)),
+    (r"\b(\?\??)(\s*)$", bygroups(Operator, Text)),
+    (r"(%)(sx|sc|system)(.*)(\n)", bygroups(Operator, Keyword, using(BashLexer), Text)),
+    (r"(%)(\w+)(.*\n)", bygroups(Operator, Keyword, Text)),
+    (r"^(!!)(.+)(\n)", bygroups(Operator, using(BashLexer), Text)),
+    (r"(!)(?!=)(.+)(\n)", bygroups(Operator, using(BashLexer), Text)),
+    (r"^(\s*)(\?\??)(\s*%{0,2}[\w\.\*]*)", bygroups(Text, Operator, Text)),
+    (r"(\s*%{0,2}[\w\.\*]*)(\?\??)(\s*)$", bygroups(Text, Operator, Text)),
+]
+
+
+class IPython3Lexer(Python3Lexer):
+    """IPython code lexer (based on Python 3)"""
+
+    name = "IPython"
+    aliases = ["ipython", "ipython3"]
+
+    tokens = Python3Lexer.tokens.copy()
+    tokens["root"] = ipython_tokens + tokens["root"]
+
+
+IPythonLexer = IPython3Lexer
+
+
+class IPythonPartialTracebackLexer(RegexLexer):
+    """
+    Partial lexer for IPython tracebacks.
+
+    Handles all the non-python output.
+
+    """
+
+    name = "IPython Partial Traceback"
+
+    tokens = {
+        "root": [
+            # Tracebacks for syntax errors have a different style.
+            # For both types of tracebacks, we mark the first line with
+            # Generic.Traceback.  For syntax errors, we mark the filename
+            # as we mark the filenames for non-syntax tracebacks.
+            #
+            # These two regexps define how IPythonConsoleLexer finds a
+            # traceback.
+            #
+            ## Non-syntax traceback
+            (r"^(\^C)?(-+\n)", bygroups(Error, Generic.Traceback)),
+            ## Syntax traceback
+            (
+                r"^(  File)(.*)(, line )(\d+\n)",
+                bygroups(
+                    Generic.Traceback,
+                    Name.Namespace,
+                    Generic.Traceback,
+                    Literal.Number.Integer,
+                ),
+            ),
+            # (Exception Identifier)(Whitespace)(Traceback Message)
+            (
+                r"(?u)(^[^\d\W]\w*)(\s*)(Traceback.*?\n)",
+                bygroups(Name.Exception, Generic.Whitespace, Text),
+            ),
+            # (Module/Filename)(Text)(Callee)(Function Signature)
+            # Better options for callee and function signature?
+            (
+                r"(.*)( in )(.*)(\(.*\)\n)",
+                bygroups(Name.Namespace, Text, Name.Entity, Name.Tag),
+            ),
+            # Regular line: (Whitespace)(Line Number)(Python Code)
+            (
+                r"(\s*?)(\d+)(.*?\n)",
+                bygroups(Generic.Whitespace, Literal.Number.Integer, Other),
+            ),
+            # Emphasized line: (Arrow)(Line Number)(Python Code)
+            # Using Exception token so arrow color matches the Exception.
+            (
+                r"(-*>?\s?)(\d+)(.*?\n)",
+                bygroups(Name.Exception, Literal.Number.Integer, Other),
+            ),
+            # (Exception Identifier)(Message)
+            (r"(?u)(^[^\d\W]\w*)(:.*?\n)", bygroups(Name.Exception, Text)),
+            # Tag everything else as Other, will be handled later.
+            (r".*\n", Other),
+        ],
+    }
+
+
+class IPythonTracebackLexer(DelegatingLexer):
+    """
+    IPython traceback lexer.
+
+    For doctests, the tracebacks can be snipped as much as desired with the
+    exception to the lines that designate a traceback. For non-syntax error
+    tracebacks, this is the line of hyphens. For syntax error tracebacks,
+    this is the line which lists the File and line number.
+
+    """
+
+    # The lexer inherits from DelegatingLexer.  The "root" lexer is an
+    # appropriate IPython lexer, which depends on the value of the boolean
+    # `python3`.  First, we parse with the partial IPython traceback lexer.
+    # Then, any code marked with the "Other" token is delegated to the root
+    # lexer.
+    #
+    name = "IPython Traceback"
+    aliases = ["ipythontb", "ipython3tb"]
+
+    def __init__(self, **options):
+        """
+        A subclass of `DelegatingLexer` which delegates to the appropriate to either IPyLexer,
+        IPythonPartialTracebackLexer.
+        """
+        # note we need a __init__ doc, as otherwise it inherits the doc from the super class
+        # which will fail the documentation build as it references section of the pygments docs that
+        # do not exists when building IPython's docs.
+        DelegatingLexer.__init__(
+            self, IPython3Lexer, IPythonPartialTracebackLexer, **options
+        )
+
+
+class IPythonConsoleLexer(Lexer):
+    """
+    An IPython console lexer for IPython code-blocks and doctests, such as:
+
+    .. code-block:: rst
+
+        .. code-block:: ipythonconsole
+
+            In [1]: a = 'foo'
+
+            In [2]: a
+            Out[2]: 'foo'
+
+            In [3]: print(a)
+            foo
+
+
+    Support is also provided for IPython exceptions:
+
+    .. code-block:: rst
+
+        .. code-block:: ipythonconsole
+
+            In [1]: raise Exception
+            Traceback (most recent call last):
+            ...
+            Exception
+
+    """
+
+    name = "IPython console session"
+    aliases = ["ipythonconsole", "ipython3console"]
+    mimetypes = ["text/x-ipython-console"]
+
+    # The regexps used to determine what is input and what is output.
+    # The default prompts for IPython are:
+    #
+    #    in           = 'In [#]: '
+    #    continuation = '   .D.: '
+    #    template     = 'Out[#]: '
+    #
+    # Where '#' is the 'prompt number' or 'execution count' and 'D'
+    # D is a number of dots  matching the width of the execution count
+    #
+    in1_regex = r"In \[[0-9]+\]: "
+    in2_regex = r"   \.\.+\.: "
+    out_regex = r"Out\[[0-9]+\]: "
+
+    #: The regex to determine when a traceback starts.
+    ipytb_start = re.compile(r"^(\^C)?(-+\n)|^(  File)(.*)(, line )(\d+\n)")
+
+    def __init__(self, **options):
+        """Initialize the IPython console lexer.
+
+        Parameters
+        ----------
+        in1_regex : RegexObject
+            The compiled regular expression used to detect the start
+            of inputs. Although the IPython configuration setting may have a
+            trailing whitespace, do not include it in the regex. If `None`,
+            then the default input prompt is assumed.
+        in2_regex : RegexObject
+            The compiled regular expression used to detect the continuation
+            of inputs. Although the IPython configuration setting may have a
+            trailing whitespace, do not include it in the regex. If `None`,
+            then the default input prompt is assumed.
+        out_regex : RegexObject
+            The compiled regular expression used to detect outputs. If `None`,
+            then the default output prompt is assumed.
+
+        """
+        in1_regex = options.get("in1_regex", self.in1_regex)
+        in2_regex = options.get("in2_regex", self.in2_regex)
+        out_regex = options.get("out_regex", self.out_regex)
+
+        # So that we can work with input and output prompts which have been
+        # rstrip'd (possibly by editors) we also need rstrip'd variants. If
+        # we do not do this, then such prompts will be tagged as 'output'.
+        # The reason can't just use the rstrip'd variants instead is because
+        # we want any whitespace associated with the prompt to be inserted
+        # with the token. This allows formatted code to be modified so as hide
+        # the appearance of prompts, with the whitespace included. One example
+        # use of this is in copybutton.js from the standard lib Python docs.
+        in1_regex_rstrip = in1_regex.rstrip() + "\n"
+        in2_regex_rstrip = in2_regex.rstrip() + "\n"
+        out_regex_rstrip = out_regex.rstrip() + "\n"
+
+        # Compile and save them all.
+        attrs = [
+            "in1_regex",
+            "in2_regex",
+            "out_regex",
+            "in1_regex_rstrip",
+            "in2_regex_rstrip",
+            "out_regex_rstrip",
+        ]
+        for attr in attrs:
+            self.__setattr__(attr, re.compile(locals()[attr]))
+
+        Lexer.__init__(self, **options)
+
+        self.pylexer = IPython3Lexer(**options)
+        self.tblexer = IPythonTracebackLexer(**options)
+
+        self.reset()
+
+    def reset(self):
+        self.mode = "output"
+        self.index = 0
+        self.buffer = ""
+        self.insertions = []
+
+    def buffered_tokens(self):
+        """
+        Generator of unprocessed tokens after doing insertions and before
+        changing to a new state.
+
+        """
+        if self.mode == "output":
+            tokens = [(0, Generic.Output, self.buffer)]
+        elif self.mode == "input":
+            tokens = self.pylexer.get_tokens_unprocessed(self.buffer)
+        else:  # traceback
+            tokens = self.tblexer.get_tokens_unprocessed(self.buffer)
+
+        for i, t, v in do_insertions(self.insertions, tokens):
+            # All token indexes are relative to the buffer.
+            yield self.index + i, t, v
+
+        # Clear it all
+        self.index += len(self.buffer)
+        self.buffer = ""
+        self.insertions = []
+
+    def get_mci(self, line):
+        """
+        Parses the line and returns a 3-tuple: (mode, code, insertion).
+
+        `mode` is the next mode (or state) of the lexer, and is always equal
+        to 'input', 'output', or 'tb'.
+
+        `code` is a portion of the line that should be added to the buffer
+        corresponding to the next mode and eventually lexed by another lexer.
+        For example, `code` could be Python code if `mode` were 'input'.
+
+        `insertion` is a 3-tuple (index, token, text) representing an
+        unprocessed "token" that will be inserted into the stream of tokens
+        that are created from the buffer once we change modes. This is usually
+        the input or output prompt.
+
+        In general, the next mode depends on current mode and on the contents
+        of `line`.
+
+        """
+        # To reduce the number of regex match checks, we have multiple
+        # 'if' blocks instead of 'if-elif' blocks.
+
+        # Check for possible end of input
+        in2_match = self.in2_regex.match(line)
+        in2_match_rstrip = self.in2_regex_rstrip.match(line)
+        if (
+            in2_match and in2_match.group().rstrip() == line.rstrip()
+        ) or in2_match_rstrip:
+            end_input = True
+        else:
+            end_input = False
+        if end_input and self.mode != "tb":
+            # Only look for an end of input when not in tb mode.
+            # An ellipsis could appear within the traceback.
+            mode = "output"
+            code = ""
+            insertion = (0, Generic.Prompt, line)
+            return mode, code, insertion
+
+        # Check for output prompt
+        out_match = self.out_regex.match(line)
+        out_match_rstrip = self.out_regex_rstrip.match(line)
+        if out_match or out_match_rstrip:
+            mode = "output"
+            if out_match:
+                idx = out_match.end()
+            else:
+                idx = out_match_rstrip.end()
+            code = line[idx:]
+            # Use the 'heading' token for output.  We cannot use Generic.Error
+            # since it would conflict with exceptions.
+            insertion = (0, Generic.Heading, line[:idx])
+            return mode, code, insertion
+
+        # Check for input or continuation prompt (non stripped version)
+        in1_match = self.in1_regex.match(line)
+        if in1_match or (in2_match and self.mode != "tb"):
+            # New input or when not in tb, continued input.
+            # We do not check for continued input when in tb since it is
+            # allowable to replace a long stack with an ellipsis.
+            mode = "input"
+            if in1_match:
+                idx = in1_match.end()
+            else:  # in2_match
+                idx = in2_match.end()
+            code = line[idx:]
+            insertion = (0, Generic.Prompt, line[:idx])
+            return mode, code, insertion
+
+        # Check for input or continuation prompt (stripped version)
+        in1_match_rstrip = self.in1_regex_rstrip.match(line)
+        if in1_match_rstrip or (in2_match_rstrip and self.mode != "tb"):
+            # New input or when not in tb, continued input.
+            # We do not check for continued input when in tb since it is
+            # allowable to replace a long stack with an ellipsis.
+            mode = "input"
+            if in1_match_rstrip:
+                idx = in1_match_rstrip.end()
+            else:  # in2_match
+                idx = in2_match_rstrip.end()
+            code = line[idx:]
+            insertion = (0, Generic.Prompt, line[:idx])
+            return mode, code, insertion
+
+        # Check for traceback
+        if self.ipytb_start.match(line):
+            mode = "tb"
+            code = line
+            insertion = None
+            return mode, code, insertion
+
+        # All other stuff...
+        if self.mode in ("input", "output"):
+            # We assume all other text is output. Multiline input that
+            # does not use the continuation marker cannot be detected.
+            # For example, the 3 in the following is clearly output:
+            #
+            #    In [1]: print(3)
+            #    3
+            #
+            # But the following second line is part of the input:
+            #
+            #    In [2]: while True:
+            #        print(True)
+            #
+            # In both cases, the 2nd line will be 'output'.
+            #
+            mode = "output"
+        else:
+            mode = "tb"
+
+        code = line
+        insertion = None
+
+        return mode, code, insertion
+
+    def get_tokens_unprocessed(self, text):
+        self.reset()
+        for match in line_re.finditer(text):
+            line = match.group()
+            mode, code, insertion = self.get_mci(line)
+
+            if mode != self.mode:
+                # Yield buffered tokens before transitioning to new mode.
+                for token in self.buffered_tokens():
+                    yield token
+                self.mode = mode
+
+            if insertion:
+                self.insertions.append((len(self.buffer), [insertion]))
+            self.buffer += code
+
+        for token in self.buffered_tokens():
+            yield token
+
+
+class IPyLexer(Lexer):
+    r"""
+    Primary lexer for all IPython-like code.
+
+    This is a simple helper lexer.  If the first line of the text begins with
+    "In \[[0-9]+\]:", then the entire text is parsed with an IPython console
+    lexer. If not, then the entire text is parsed with an IPython lexer.
+
+    The goal is to reduce the number of lexers that are registered
+    with Pygments.
+
+    """
+
+    name = "IPy session"
+    aliases = ["ipy", "ipy3"]
+
+    def __init__(self, **options):
+        """
+        Create a new IPyLexer instance which dispatch to either an
+        IPythonCOnsoleLexer (if In prompts are present) or and IPythonLexer (if
+        In prompts are not present).
+        """
+        # init docstring is necessary for docs not to fail to build do to parent
+        # docs referenceing a section in pygments docs.
+        Lexer.__init__(self, **options)
+
+        self.IPythonLexer = IPythonLexer(**options)
+        self.IPythonConsoleLexer = IPythonConsoleLexer(**options)
+
+    def get_tokens_unprocessed(self, text):
+        # Search for the input prompt anywhere...this allows code blocks to
+        # begin with comments as well.
+        if re.match(r".*(In \[[0-9]+\]:)", text.strip(), re.DOTALL):
+            lex = self.IPythonConsoleLexer
+        else:
+            lex = self.IPythonLexer
+        for token in lex.get_tokens_unprocessed(text):
+            yield token

lib/python3.12/site-packages/isympy.py

@@ -0,0 +1,342 @@
+"""
+Python shell for SymPy.
+
+This is just a normal Python shell (IPython shell if you have the
+IPython package installed), that executes the following commands for
+the user:
+
+    >>> from __future__ import division
+    >>> from sympy import *
+    >>> x, y, z, t = symbols('x y z t')
+    >>> k, m, n = symbols('k m n', integer=True)
+    >>> f, g, h = symbols('f g h', cls=Function)
+    >>> init_printing()
+
+So starting 'isympy' is equivalent to starting Python (or IPython) and
+executing the above commands by hand.  It is intended for easy and quick
+experimentation with SymPy.  isympy is a good way to use SymPy as an
+interactive calculator. If you have IPython and Matplotlib installed, then
+interactive plotting is enabled by default.
+
+COMMAND LINE OPTIONS
+--------------------
+
+-c CONSOLE, --console=CONSOLE
+
+     Use the specified shell (Python or IPython) shell as the console
+     backend instead of the default one (IPython if present, Python
+     otherwise), e.g.:
+
+        $isympy -c python
+
+    CONSOLE must be one of 'ipython' or 'python'
+
+-p PRETTY, --pretty PRETTY
+
+    Setup pretty-printing in SymPy. When pretty-printing is enabled,
+    expressions can be printed with Unicode or ASCII. The default is
+    to use pretty-printing (with Unicode if the terminal supports it).
+    When this option is 'no', expressions will not be pretty-printed
+    and ASCII will be used:
+
+        $isympy -p no
+
+    PRETTY must be one of 'unicode', 'ascii', or 'no'
+
+-t TYPES, --types=TYPES
+
+    Setup the ground types for the polys.  By default, gmpy ground types
+    are used if gmpy2 or gmpy is installed, otherwise it falls back to python
+    ground types, which are a little bit slower.  You can manually
+    choose python ground types even if gmpy is installed (e.g., for
+    testing purposes):
+
+        $isympy -t python
+
+    TYPES must be one of 'gmpy', 'gmpy1' or 'python'
+
+    Note that the ground type gmpy1 is primarily intended for testing; it
+    forces the use of gmpy version 1 even if gmpy2 is available.
+
+    This is the same as setting the environment variable
+    SYMPY_GROUND_TYPES to the given ground type (e.g.,
+    SYMPY_GROUND_TYPES='gmpy')
+
+    The ground types can be determined interactively from the variable
+    sympy.polys.domains.GROUND_TYPES.
+
+-o ORDER, --order ORDER
+
+    Setup the ordering of terms for printing.  The default is lex, which
+    orders terms lexicographically (e.g., x**2 + x + 1). You can choose
+    other orderings, such as rev-lex, which will use reverse
+    lexicographic ordering (e.g., 1 + x + x**2):
+
+        $isympy -o rev-lex
+
+    ORDER must be one of 'lex', 'rev-lex', 'grlex', 'rev-grlex',
+    'grevlex', 'rev-grevlex', 'old', or 'none'.
+
+    Note that for very large expressions, ORDER='none' may speed up
+    printing considerably but the terms will have no canonical order.
+
+-q, --quiet
+
+    Print only Python's and SymPy's versions to stdout at startup.
+
+-d, --doctest
+
+    Use the same format that should be used for doctests.  This is
+    equivalent to -c python -p no.
+
+-C, --no-cache
+
+    Disable the caching mechanism.  Disabling the cache may slow certain
+    operations down considerably.  This is useful for testing the cache,
+    or for benchmarking, as the cache can result in deceptive timings.
+
+    This is equivalent to setting the environment variable
+    SYMPY_USE_CACHE to 'no'.
+
+-a, --auto-symbols (requires at least IPython 0.11)
+
+    Automatically create missing symbols.  Normally, typing a name of a
+    Symbol that has not been instantiated first would raise NameError,
+    but with this option enabled, any undefined name will be
+    automatically created as a Symbol.
+
+    Note that this is intended only for interactive, calculator style
+    usage. In a script that uses SymPy, Symbols should be instantiated
+    at the top, so that it's clear what they are.
+
+    This will not override any names that are already defined, which
+    includes the single character letters represented by the mnemonic
+    QCOSINE (see the "Gotchas and Pitfalls" document in the
+    documentation). You can delete existing names by executing "del
+    name".  If a name is defined, typing "'name' in dir()" will return True.
+
+    The Symbols that are created using this have default assumptions.
+    If you want to place assumptions on symbols, you should create them
+    using symbols() or var().
+
+    Finally, this only works in the top level namespace. So, for
+    example, if you define a function in isympy with an undefined
+    Symbol, it will not work.
+
+    See also the -i and -I options.
+
+-i, --int-to-Integer (requires at least IPython 0.11)
+
+    Automatically wrap int literals with Integer.  This makes it so that
+    things like 1/2 will come out as Rational(1, 2), rather than 0.5.  This
+    works by preprocessing the source and wrapping all int literals with
+    Integer.  Note that this will not change the behavior of int literals
+    assigned to variables, and it also won't change the behavior of functions
+    that return int literals.
+
+    If you want an int, you can wrap the literal in int(), e.g. int(3)/int(2)
+    gives 1.5 (with division imported from __future__).
+
+-I, --interactive (requires at least IPython 0.11)
+
+    This is equivalent to --auto-symbols --int-to-Integer.  Future options
+    designed for ease of interactive use may be added to this.
+
+-D, --debug
+
+    Enable debugging output.  This is the same as setting the
+    environment variable SYMPY_DEBUG to 'True'.  The debug status is set
+    in the variable SYMPY_DEBUG within isympy.
+
+-- IPython options
+
+    Additionally you can pass command line options directly to the IPython
+    interpreter (the standard Python shell is not supported).  However you
+    need to add the '--' separator between two types of options, e.g the
+    startup banner option and the colors option. You need to enter the
+    options as required by the version of IPython that you are using, too:
+
+    in IPython 0.11,
+
+        $isympy -q -- --colors=NoColor
+
+    or older versions of IPython,
+
+        $isympy -q -- -colors NoColor
+
+See also isympy --help.
+"""
+
+import os
+import sys
+
+# DO NOT IMPORT SYMPY HERE! Or the setting of the sympy environment variables
+# by the command line will break.
+
+def main() -> None:
+    from argparse import ArgumentParser, RawDescriptionHelpFormatter
+
+    VERSION = None
+    if '--version' in sys.argv:
+        # We cannot import sympy before this is run, because flags like -C and
+        # -t set environment variables that must be set before SymPy is
+        # imported. The only thing we need to import it for is to get the
+        # version, which only matters with the --version flag.
+        import sympy
+        VERSION = sympy.__version__
+
+    usage = 'isympy [options] -- [ipython options]'
+    parser = ArgumentParser(
+        usage=usage,
+        description=__doc__,
+        formatter_class=RawDescriptionHelpFormatter,
+    )
+
+    parser.add_argument('--version', action='version', version=VERSION)
+
+    parser.add_argument(
+        '-c', '--console',
+        dest='console',
+        action='store',
+        default=None,
+        choices=['ipython', 'python'],
+        metavar='CONSOLE',
+        help='select type of interactive session: ipython | python; defaults '
+        'to ipython if IPython is installed, otherwise python')
+
+    parser.add_argument(
+        '-p', '--pretty',
+        dest='pretty',
+        action='store',
+        default=None,
+        metavar='PRETTY',
+        choices=['unicode', 'ascii', 'no'],
+        help='setup pretty printing: unicode | ascii | no; defaults to '
+        'unicode printing if the terminal supports it, otherwise ascii')
+
+    parser.add_argument(
+        '-t', '--types',
+        dest='types',
+        action='store',
+        default=None,
+        metavar='TYPES',
+        choices=['gmpy', 'gmpy1', 'python'],
+        help='setup ground types: gmpy | gmpy1 | python; defaults to gmpy if gmpy2 '
+        'or gmpy is installed, otherwise python')
+
+    parser.add_argument(
+        '-o', '--order',
+        dest='order',
+        action='store',
+        default=None,
+        metavar='ORDER',
+        choices=['lex', 'grlex', 'grevlex', 'rev-lex', 'rev-grlex', 'rev-grevlex', 'old', 'none'],
+        help='setup ordering of terms: [rev-]lex | [rev-]grlex | [rev-]grevlex | old | none; defaults to lex')
+
+    parser.add_argument(
+        '-q', '--quiet',
+        dest='quiet',
+        action='store_true',
+        default=False,
+        help='print only version information at startup')
+
+    parser.add_argument(
+        '-d', '--doctest',
+        dest='doctest',
+        action='store_true',
+        default=False,
+        help='use the doctest format for output (you can just copy and paste it)')
+
+    parser.add_argument(
+        '-C', '--no-cache',
+        dest='cache',
+        action='store_false',
+        default=True,
+        help='disable caching mechanism')
+
+    parser.add_argument(
+        '-a', '--auto-symbols',
+        dest='auto_symbols',
+        action='store_true',
+        default=False,
+        help='automatically construct missing symbols')
+
+    parser.add_argument(
+        '-i', '--int-to-Integer',
+        dest='auto_int_to_Integer',
+        action='store_true',
+        default=False,
+        help="automatically wrap int literals with Integer")
+
+    parser.add_argument(
+        '-I', '--interactive',
+        dest='interactive',
+        action='store_true',
+        default=False,
+        help="equivalent to -a -i")
+
+    parser.add_argument(
+        '-D', '--debug',
+        dest='debug',
+        action='store_true',
+        default=False,
+        help='enable debugging output')
+
+    (options, ipy_args) = parser.parse_known_args()
+    if '--' in ipy_args:
+        ipy_args.remove('--')
+
+    if not options.cache:
+        os.environ['SYMPY_USE_CACHE'] = 'no'
+
+    if options.types:
+        os.environ['SYMPY_GROUND_TYPES'] = options.types
+
+    if options.debug:
+        os.environ['SYMPY_DEBUG'] = str(options.debug)
+
+    if options.doctest:
+        options.pretty = 'no'
+        options.console = 'python'
+
+    session = options.console
+
+    if session is not None:
+        ipython = session == 'ipython'
+    else:
+        try:
+            import IPython
+            ipython = True
+        except ImportError:
+            if not options.quiet:
+                from sympy.interactive.session import no_ipython
+                print(no_ipython)
+            ipython = False
+
+    args = {
+        'pretty_print': True,
+        'use_unicode':  None,
+        'use_latex':    None,
+        'order':        None,
+        'argv':         ipy_args,
+    }
+
+    if options.pretty == 'unicode':
+        args['use_unicode'] = True
+    elif options.pretty == 'ascii':
+        args['use_unicode'] = False
+    elif options.pretty == 'no':
+        args['pretty_print'] = False
+
+    if options.order is not None:
+        args['order'] = options.order
+
+    args['quiet'] = options.quiet
+    args['auto_symbols'] = options.auto_symbols or options.interactive
+    args['auto_int_to_Integer'] = options.auto_int_to_Integer or options.interactive
+
+    from sympy.interactive import init_session
+    init_session(ipython, **args)
+
+if __name__ == "__main__":
+    main()

lib/python3.12/site-packages/package_readme.md

@@ -0,0 +1,97 @@
+<div align="center">
+  <img src="https://i.imgur.com/RUtiVzH.png" width="600" /><br><br>
+</div>
+
+# Weights and Biases [![PyPI](https://img.shields.io/pypi/v/wandb)](https://pypi.python.org/pypi/wandb) [![Conda (channel only)](https://img.shields.io/conda/vn/conda-forge/wandb)](https://anaconda.org/conda-forge/wandb) [![CircleCI](https://img.shields.io/circleci/build/github/wandb/wandb/main)](https://circleci.com/gh/wandb/wandb) [![Codecov](https://img.shields.io/codecov/c/gh/wandb/wandb)](https://codecov.io/gh/wandb/wandb)
+
+Use W&B to build better models faster. Track and visualize all the pieces of your machine learning pipeline, from datasets to production machine learning models. Get started with W&B today, [sign up for an account!](https://wandb.com?utm_source=github&utm_medium=code&utm_campaign=wandb&utm_content=readme)
+
+
+
+See the [W&B Developer Guide](https://docs.wandb.ai/?utm_source=github&utm_medium=code&utm_campaign=wandb&utm_content=documentation) and [API Reference Guide](https://docs.wandb.ai/ref?utm_source=github&utm_medium=code&utm_campaign=wandb&utm_content=documentation) for a full technical description of the W&B platform.
+
+&nbsp;
+
+# Quickstart
+
+Get started with W&B in four steps:
+
+1. First, sign up for a [W&B account](https://wandb.ai/login?utm_source=github&utm_medium=code&utm_campaign=wandb&utm_content=quickstart).
+
+2. Second, install the W&B SDK with [pip](https://pip.pypa.io/en/stable/). Navigate to your terminal and type the following command:
+
+```shell
+pip install wandb
+```
+
+3. Third, log into W&B:
+
+```python
+wandb.login()
+```
+
+4. Use the example code snippet below as a template to integrate W&B to your Python script:
+
+```python
+import wandb
+
+# Start a W&B Run with wandb.init
+run = wandb.init(project="my_first_project")
+
+# Save model inputs and hyperparameters in a wandb.config object
+config = run.config
+config.learning_rate = 0.01
+
+# Model training code here ...
+
+# Log metrics over time to visualize performance with wandb.log
+for i in range(10):
+    run.log({"loss": ...})
+
+# Mark the run as finished, and finish uploading all data
+run.finish()
+```
+
+For example, if the preceding code was stored in a script called train.py:
+
+```shell
+python train.py
+```
+
+You will see a URL in your terminal logs when your script starts and finishes. Data is staged locally in a directory named _wandb_ relative to your script. Navigate to the W&B App to view a dashboard of your first W&B Experiment. Use the W&B App to compare multiple experiments in a unified place, dive into the results of a single run, and much more!
+
+&nbsp;
+
+# Integrations
+
+Use your favorite framework with W&B. W&B integrations make it fast and easy to set up experiment tracking and data versioning inside existing projects. For more information on how to integrate W&B with the framework of your choice, see [W&B Integrations](https://docs.wandb.ai/guides/integrations) in the W&B Developer Guide.
+
+&nbsp;
+
+# Python Version Support
+
+We are committed to supporting our minimum required Python version for *at least* six months after its official end-of-life (EOL) date, as defined by the Python Software Foundation. You can find a list of Python EOL dates [here](https://devguide.python.org/versions/).
+
+When we discontinue support for a Python version, we will increment the library’s minor version number to reflect this change.
+
+&nbsp;
+
+# Contribution guidelines
+Weights & Biases ❤️ open source, and we welcome contributions from the community! See the [Contribution guide](https://github.com/wandb/wandb/blob/main/CONTRIBUTING.md) for more information on the development workflow and the internals of the wandb library. For wandb bugs and feature requests, visit [GitHub Issues](https://github.com/wandb/wandb/issues) or contact support@wandb.com.
+
+&nbsp;
+
+# Academic Researchers
+Reach out to W&B Support at support@wandb.com to get a [free academic license](https://www.wandb.com/academic) for you and your research group.
+
+&nbsp;
+
+# W&B Community
+
+Be a part of the growing W&B Community and interact with the W&B team in our [Discord](https://wandb.me/discord). Stay connected with the latest ML updates and tutorials with [W&B Fully Connected](https://wandb.ai/fully-connected).
+
+&nbsp;
+
+# License
+
+[MIT License](https://github.com/wandb/wandb/blob/main/LICENSE)

lib/python3.12/site-packages/pylab.py

@@ -0,0 +1,3 @@
+from matplotlib.pylab import *  # noqa: F401, F403
+import matplotlib.pylab
+__doc__ = matplotlib.pylab.__doc__

lib/python3.12/site-packages/sglang/__init__.py

@@ -0,0 +1,76 @@
+# 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.global_config import global_config
+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
+from sglang.version import __version__
+
+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")
+
+__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__",
+]

lib/python3.12/site-packages/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")

lib/python3.12/site-packages/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

lib/python3.12/site-packages/sglang/bench_one_batch.py

@@ -0,0 +1,568 @@
+"""
+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 ==========
+<s> The capital of France is Paris.
+The capital of the United States is Washington, D.C.
+
+
+========== Prompt 1 ==========
+<s> The capital of the United Kindom is London.
+The capital of the United Kingdom is London.
+The capital of the
+
+========== Prompt 2 ==========
+<s> 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.distributed.parallel_state import destroy_distributed_environment
+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.managers.scheduler import Scheduler
+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
+    log_decode_step: int = 0
+    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(
+            "--log-decode-step",
+            type=int,
+            default=BenchArgs.log_decode_step,
+            help="Log decode latency by step, default is set to zero to disable.",
+        )
+        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,
+        enable_multimodal=server_args.enable_multimodal,
+        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)
+        req.logprob_start_len = len(req.origin_input_ids) - 1
+        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)
+        req.logprob_start_len = len(req.origin_input_ids) - 1
+    return reqs
+
+
+def prepare_synthetic_inputs_for_latency_test(batch_size, input_len):
+    input_ids = np.random.randint(0, 10000, (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)
+        req.logprob_start_len = len(req.origin_input_ids) - 1
+        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()
+    _maybe_prepare_dp_attn_batch(batch, model_runner)
+    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()
+    _maybe_prepare_dp_attn_batch(batch, model_runner)
+    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 _maybe_prepare_dp_attn_batch(batch: ScheduleBatch, model_runner):
+    if model_runner.server_args.enable_dp_attention:
+        Scheduler.prepare_dp_attn_batch_raw(
+            batch,
+            dp_size=model_runner.server_args.dp_size,
+            attn_tp_size=1,
+            tp_cpu_group=model_runner.tp_group.cpu_group,
+            get_idle_batch=None,
+            disable_cuda_graph=model_runner.server_args.disable_cuda_graph,
+            spec_algorithm=SpeculativeAlgorithm.NONE,
+            speculative_num_draft_tokens=None,
+        )
+
+
+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,
+    log_decode_step,
+    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 or (log_decode_step > 0 and i % log_decode_step == 0):
+            rank_print(
+                f"Decode {i}. Batch size: {batch_size}, 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],
+        min(32, bench_args.output_len[0]),  # shorter decoding to speed up the warmup
+        server_args.device,
+        log_decode_step=0,
+        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.log_decode_step,
+            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")
+
+    if server_args.tp_size > 1:
+        destroy_distributed_environment()
+
+
+def main(server_args, bench_args):
+    server_args.cuda_graph_max_bs = max(bench_args.batch_size)
+
+    _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)

lib/python3.12/site-packages/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)

lib/python3.12/site-packages/sglang/bench_serving.py

@@ -0,0 +1,1623 @@
+# 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
+
+
+# don't want to import sglang package here
+def _get_bool_env_var(name: str, default: str = "false") -> bool:
+    value = os.getenv(name, default)
+    return value.lower() in ("true", "1")
+
+
+@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]["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.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.startswith("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,
+            random_sample=args.dataset_name == "random",
+        )
+    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,
+    random_sample: bool = True,
+) -> 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 random_sample:
+        # 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,
+    flush_cache: bool = False,
+    warmup_requests: int = 1,
+):
+    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(f"Starting warmup with {warmup_requests} sequences...")
+
+    # Use the first request for all warmup iterations
+    test_prompt, test_prompt_len, test_output_len = input_requests[0]
+    if lora_names is not None and len(lora_names) != 0:
+        lora_name = lora_names[0]
+    else:
+        lora_name = None
+
+    # Create the test input once
+    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,
+    )
+
+    # Run warmup requests
+    warmup_tasks = []
+    for _ in range(warmup_requests):
+        warmup_tasks.append(
+            asyncio.create_task(request_func(request_func_input=test_input))
+        )
+
+    warmup_outputs = await asyncio.gather(*warmup_tasks)
+
+    # Check if at least one warmup request succeeded
+    if warmup_requests > 0 and not any(output.success for output in warmup_outputs):
+        raise ValueError(
+            "Warmup failed - Please make sure benchmark arguments "
+            f"are correctly specified. Error: {warmup_outputs[0].error}"
+        )
+    else:
+        print(
+            f"Warmup completed with {args.warmup_requests} sequences. Starting main benchmark run..."
+        )
+
+    # Flush cache
+    if ("sglang" in backend and _get_bool_env_var("SGLANG_IS_IN_CI")) or flush_cache:
+        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 is not 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.startswith("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
+
+    # Set default value for warmup_requests if not present
+    if not hasattr(args, "warmup_requests"):
+        args.warmup_requests = 1
+
+    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)
+
+    # compatible with SimpleNamespace
+    if not hasattr(args, "flush_cache"):
+        args.flush_cache = False
+
+    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,
+            flush_cache=args.flush_cache,
+        )
+    )
+
+
+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", "random-ids", "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",
+    )
+    parser.add_argument(
+        "--flush-cache",
+        action="store_true",
+        help="Flush the cache before running the benchmark",
+    )
+    parser.add_argument(
+        "--warmup-requests",
+        type=int,
+        default=1,
+        help="Number of warmup requests to run before the benchmark",
+    )
+
+    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)

lib/python3.12/site-packages/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()

lib/python3.12/site-packages/sglang/compile_deep_gemm.py

@@ -0,0 +1,177 @@
+"""
+Compile DeepGEMM Kernels for a model with specify server arguments
+
+This script launches a server for capturing DeepGEMM calls and then compiles the kernels.
+It accepts server arguments (the same as launch_server.py).
+
+Usage:
+python3 -m sglang.compile_deep_gemm --model deepseek-ai/DeepSeek-V3 --tp 8 --trust-remote-code
+
+"""
+
+import argparse
+import dataclasses
+import multiprocessing
+import os
+import time
+
+import requests
+
+from sglang.srt.entrypoints.http_server import launch_server
+from sglang.srt.managers.io_struct import GenerateReqInput
+from sglang.srt.managers.tokenizer_manager import TokenizerManager
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import kill_process_tree
+from sglang.srt.warmup import warmup
+
+multiprocessing.set_start_method("spawn", force=True)
+
+# Reduce warning
+os.environ["SGL_IN_DEEPGEMM_PRECOMPILE_STAGE"] = "1"
+# Force enable deep gemm
+os.environ["SGL_ENABLE_JIT_DEEPGEMM"] = "1"
+# Force enable mha chunked kv for DeepSeek V3 to avoid missing kv_b_proj DeepGEMM case
+os.environ["SGL_CHUNKED_PREFIX_CACHE_THRESHOLD"] = "0"
+
+
+@dataclasses.dataclass
+class CompileArgs:
+    timeout: int = 3600
+
+    @staticmethod
+    def add_cli_args(parser: argparse.ArgumentParser):
+        parser.add_argument("--timeout", type=int, default=CompileArgs.timeout)
+
+    @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}
+        )
+
+
+@warmup("compile-deep-gemm")
+async def warm_up_compile(tokenizer_manager: TokenizerManager):
+    print("\nGenerate warm up request for compiling DeepGEMM...\n")
+    generate_req_input = GenerateReqInput(
+        input_ids=[0, 1, 2, 3],
+        sampling_params={
+            "temperature": 0.0,
+            "max_new_tokens": 8,
+            "ignore_eos": True,
+        },
+    )
+    await tokenizer_manager.generate_request(generate_req_input, None).__anext__()
+
+
+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_and_send_one_request(
+    server_args: ServerArgs, compile_args: CompileArgs
+):
+    proc = multiprocessing.Process(target=launch_server_internal, args=(server_args,))
+    proc.start()
+    base_url = f"http://{server_args.host}:{server_args.port}"
+    timeout = compile_args.timeout
+
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        try:
+            headers = {
+                "Content-Type": "application/json; charset=utf-8",
+            }
+            if server_args.node_rank == 0:
+                response = requests.get(f"{base_url}/v1/models", headers=headers)
+            else:
+                # This http api is created by launch_dummy_health_check_server for none-rank0 node.
+                response = requests.get(f"{base_url}/health", headers=headers)
+            if response.status_code == 200:
+                # Rank-0 node send a request to sync with other node and then return.
+                if server_args.node_rank == 0:
+                    response = requests.post(
+                        f"{base_url}/generate",
+                        json={
+                            "input_ids": [0, 1, 2, 3],
+                            "sampling_params": {
+                                "max_new_tokens": 8,
+                                "temperature": 0,
+                            },
+                        },
+                        timeout=600,
+                    )
+                    if response.status_code != 200:
+                        error = response.json()
+                        raise RuntimeError(f"Sync request failed: {error}")
+                # Other nodes should wait for the exit signal from Rank-0 node.
+                else:
+                    start_time_waiting = time.time()
+                    while proc.is_alive():
+                        if time.time() - start_time_waiting < timeout:
+                            time.sleep(10)
+                        else:
+                            raise TimeoutError("Waiting for main node timeout!")
+                return proc
+        except requests.RequestException:
+            pass
+        time.sleep(10)
+    raise TimeoutError(
+        "DeepGEMM Kernels compilation timeout."
+        "\n\nFeel free and please restart the command."
+    )
+
+
+def refine_server_args(server_args: ServerArgs, compile_args: CompileArgs):
+    # Disbale cuda graph and torch compile to save time
+    server_args.disable_cuda_graph = True
+    server_args.enable_torch_compile = False
+    print(f"Disable CUDA Graph and Torch Compile to save time...")
+
+    # Set watchdog timeout to compile_args.timeout because compilation will take a long time
+    server_args.watchdog_timeout = compile_args.timeout
+    server_args.warmups = "compile-deep-gemm"
+
+
+def run_compile(server_args: ServerArgs, compile_args: CompileArgs):
+    print(
+        "Begin DeepGEMM Kernels compilation...\n"
+        "It may take a long time and timeout maybe raised "
+        "while the compilation is still in progress.\n"
+        "Just feel free to restart the command "
+        "until the compilation is fully finished.\n"
+    )
+
+    proc = launch_server_process_and_send_one_request(server_args, compile_args)
+
+    print("\nDeepGEMM Kernels compilation finished successfully.")
+
+    # Sleep for safety
+    time.sleep(10)
+    if proc.is_alive():
+        # This is the rank0 node.
+        kill_process_tree(proc.pid)
+    else:
+        try:
+            kill_process_tree(proc.pid)
+        except Exception:
+            pass
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    ServerArgs.add_cli_args(parser)
+    CompileArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    server_args = ServerArgs.from_cli_args(args)
+    compile_args = CompileArgs.from_cli_args(args)
+
+    refine_server_args(server_args, compile_args)
+
+    run_compile(server_args, compile_args)

lib/python3.12/site-packages/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()

lib/python3.12/site-packages/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)

lib/python3.12/site-packages/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)

lib/python3.12/site-packages/sglang/srt/_custom_ops.py

@@ -0,0 +1,115 @@
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/_custom_ops.py
+import logging
+from typing import List, Tuple
+
+import torch
+
+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 not is_hip():
+    if use_vllm_custom_allreduce:
+        custom_op = torch.ops._C_custom_ar
+    else:
+        custom_op = sgl_kernel.allreduce
+
+    # custom allreduce
+    def init_custom_ar(
+        ipc_tensors: List[torch.Tensor],
+        rank_data: torch.Tensor,
+        rank: int,
+        full_nvlink: bool,
+    ) -> int:
+        return custom_op.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:
+        custom_op.all_reduce(fa, inp, out, reg_buffer, reg_buffer_sz_bytes)
+
+    def dispose(fa: int) -> None:
+        custom_op.dispose(fa)
+
+    def meta_size() -> int:
+        return custom_op.meta_size()
+
+    def register_buffer(fa: int, ipc_tensors: List[int]) -> None:
+        return custom_op.register_buffer(fa, ipc_tensors)
+
+    def get_graph_buffer_ipc_meta(fa: int) -> Tuple[List[int], List[int]]:
+        return custom_op.get_graph_buffer_ipc_meta(fa)
+
+    def register_graph_buffers(
+        fa: int, handles: List[List[int]], offsets: List[List[int]]
+    ) -> None:
+        custom_op.register_graph_buffers(fa, handles, offsets)
+
+else:
+    # 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)

lib/python3.12/site-packages/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()

lib/python3.12/site-packages/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 is not 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_prefix>",
+        fim_middle_token="<fim_middle>",
+        fim_end_token="<fim_suffix>",
+        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,
+    )
+)

lib/python3.12/site-packages/sglang/srt/conversation.py

@@ -0,0 +1,794 @@
+# 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()
+    LLAMA4 = 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 = "<image>"
+    audio_token: str = "<audio>"
+
+    image_data: Optional[List[str]] = None
+    modalities: Optional[List[str]] = None
+    stop_token_ids: Optional[int] = None
+
+    audio_data: Optional[List[str]] = None
+
+    def get_prompt(self) -> str:
+        """Get the prompt for generation."""
+        system_prompt = self.system_template.format(system_message=self.system_message)
+        if self.sep_style == SeparatorStyle.ADD_COLON_SINGLE:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_COLON_SPACE_SINGLE:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ": "  # must be end with a space
+            return ret
+        elif self.sep_style == SeparatorStyle.ADD_NEW_LINE_SINGLE:
+            ret = "" if system_prompt == "" else system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + "\n" + message + self.sep
+                else:
+                    ret += role + "\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.QWEN2_VL_EMBED:
+            ret = "" if system_prompt == "" else system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + "\n" + message + self.sep
+                else:
+                    ret += role + "\n"
+            ret += self.stop_str
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_SINGLE:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.NO_COLON_TWO:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + message + seps[i % 2]
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.RWKV:
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += (
+                        role
+                        + ": "
+                        + message.replace("\r\n", "\n").replace("\n\n", "\n")
+                    )
+                    ret += "\n\n"
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.LLAMA4:
+            # begin_of_text is added by default
+            if self.system_message:
+                ret = system_prompt
+            else:
+                ret = ""
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += f"<|header_start|>{role}<|header_end|>\n\n"
+                    ret += f"{message.strip()}<|eot|>"
+                else:
+                    ret += f"<|header_start|>{role}<|header_end|>\n\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.LLAMA3:
+            if self.system_message:
+                ret = system_prompt
+            else:
+                ret = ""
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += f"<|start_header_id|>{role}<|end_header_id|>\n\n"
+                    ret += f"{message.strip()}<|eot_id|>"
+                else:
+                    ret += f"<|start_header_id|>{role}<|end_header_id|>\n\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.LLAMA2:
+            seps = [self.sep, self.sep2]
+            if self.system_message:
+                ret = system_prompt
+            else:
+                ret = "[INST] "
+            for i, (role, message) in enumerate(self.messages):
+                tag = self.roles[i % 2]
+                if message:
+                    if i == 0:
+                        ret += message + " "
+                    else:
+                        ret += tag + " " + message + seps[i % 2]
+                else:
+                    ret += tag
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM:
+            # source: https://huggingface.co/THUDM/chatglm-6b/blob/1d240ba371910e9282298d4592532d7f0f3e9f3e/modeling_chatglm.py#L1302-L1308
+            # source2: https://huggingface.co/THUDM/chatglm2-6b/blob/e186c891cf64310ac66ef10a87e6635fa6c2a579/modeling_chatglm.py#L926
+            round_add_n = 1 if self.name == "chatglm2" else 0
+            if system_prompt:
+                ret = system_prompt + self.sep
+            else:
+                ret = ""
+
+            for i, (role, message) in enumerate(self.messages):
+                if i % 2 == 0:
+                    ret += f"[Round {i // 2 + round_add_n}]{self.sep}"
+
+                if message:
+                    ret += f"{role}：{message}{self.sep}"
+                else:
+                    ret += f"{role}："
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATML:
+            ret = "" if system_prompt == "" else system_prompt + self.sep + "\n"
+            for role, message in self.messages:
+                if message:
+                    ret += role + "\n" + message + self.sep + "\n"
+                else:
+                    ret += role + "\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATGLM3:
+            ret = ""
+            if self.system_message:
+                ret += system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + "\n" + message
+                else:
+                    ret += role
+            return ret
+        elif self.sep_style == SeparatorStyle.CHATINTERN:
+            # source: https://huggingface.co/internlm/internlm-chat-7b-8k/blob/bd546fa984b4b0b86958f56bf37f94aa75ab8831/modeling_internlm.py#L771
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if i % 2 == 0:
+                    ret += "<s>"
+                if message:
+                    ret += role + ":" + message + seps[i % 2] + "\n"
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.DOLLY:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ":\n" + message + seps[i % 2]
+                    if i % 2 == 1:
+                        ret += "\n\n"
+                else:
+                    ret += role + ":\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.PHOENIX:
+            ret = system_prompt
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + "<s>" + message + "</s>"
+                else:
+                    ret += role + ": " + "<s>"
+            return ret
+        elif self.sep_style == SeparatorStyle.ROBIN:
+            ret = system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ":\n" + message + self.sep
+                else:
+                    ret += role + ":\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.FALCON_CHAT:
+            ret = ""
+            if self.system_message:
+                ret += system_prompt + self.sep
+            for role, message in self.messages:
+                if message:
+                    ret += role + ": " + message + self.sep
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.METAMATH:
+            ret = "" if system_prompt == "" else system_prompt + self.sep
+            for i, (role, message) in enumerate(self.messages):
+                # For MetaMath, sep2 is used to prefix the message.
+                starting_sep = ":\n" if i % 2 == 0 else ": " + self.sep2
+                ending_sep = self.sep if i % 2 == 0 else ""
+                if message:
+                    ret += role + starting_sep + message + ending_sep
+                else:
+                    ret += role + starting_sep
+            return ret
+        elif self.sep_style == SeparatorStyle.DEEPSEEK_CHAT:
+            seps = [self.sep, self.sep2]
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.DeepSeekVL2:
+            seps = [self.sep, self.sep2]
+            if system_prompt == "" or system_prompt is None:
+                ret = ""
+            else:
+                ret = system_prompt + seps[0]
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+            return ret
+        elif self.sep_style == SeparatorStyle.GEMMA3:
+            ret = system_prompt
+            for i, (role, message) in enumerate(self.messages):
+                if message:
+                    if i == 0:
+                        ret += message + self.sep
+                    else:
+                        ret += role + message + self.sep
+                else:
+                    ret += role
+            return ret
+
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+    def set_system_message(self, system_message: str):
+        """Set the system message."""
+        self.system_message = system_message
+
+    def append_message(self, role: str, message: str):
+        """Append a new message."""
+        self.messages.append([role, message])
+
+    def append_image(self, image: str):
+        """Append a new message."""
+        self.image_data.append(image)
+
+    def append_audio(self, audio: str):
+        """Append a new message."""
+        self.audio_data.append(audio)
+
+    def update_last_message(self, message: str):
+        """Update the last output.
+
+        The last message is typically set to be None when constructing the prompt,
+        so we need to update it in-place after getting the response from a model.
+        """
+        self.messages[-1][1] = message
+
+    def to_gradio_chatbot(self):
+        """Convert the conversation to gradio chatbot format."""
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def to_openai_api_messages(self):
+        """Convert the conversation to OpenAI chat completion format."""
+        if self.system_message == "":
+            ret = []
+        else:
+            ret = [{"role": "system", "content": self.system_message}]
+
+        for i, (_, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                ret.append({"role": "user", "content": msg})
+            else:
+                if msg is not None:
+                    ret.append({"role": "assistant", "content": msg})
+        return ret
+
+    def copy(self):
+        return Conversation(
+            name=self.name,
+            system_template=self.system_template,
+            system_message=self.system_message,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            stop_str=self.stop_str,
+            image_token=self.image_token,
+            audio_token=self.audio_token,
+        )
+
+    def dict(self):
+        return {
+            "template_name": self.name,
+            "system_message": self.system_message,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+        }
+
+
+# A global registry for all conversation templates
+chat_templates: Dict[str, Conversation] = {}
+
+
+def register_conv_template(template: Conversation, override: bool = False):
+    """Register a new conversation template."""
+    if not override:
+        assert (
+            template.name not in chat_templates
+        ), f"{template.name} has been registered."
+
+    chat_templates[template.name] = template
+
+
+def chat_template_exists(template_name: str) -> bool:
+    return template_name in chat_templates
+
+
+def generate_embedding_convs(
+    texts: List[str], images: List[str], template_name: str
+) -> List[Conversation]:
+    conv_template = chat_templates[template_name].copy()
+    convs = []
+    for text, image in zip(texts, images):
+        conv = Conversation(
+            name=conv_template.name,
+            system_template=conv_template.system_template,
+            system_message=conv_template.system_message,
+            roles=conv_template.roles,
+            messages=list(conv_template.messages),  # prevent in-place modification
+            offset=conv_template.offset,
+            sep_style=SeparatorStyle(conv_template.sep_style),
+            sep=conv_template.sep,
+            sep2=conv_template.sep2,
+            stop_str=conv_template.stop_str,
+            image_data=[],
+            modalities=[],
+            image_token=conv_template.image_token,
+        )
+        real_content = ""
+
+        if image is not None:
+            image_token = (
+                conv.image_token + "\n"
+                if conv.name != "gme-qwen2-vl"
+                else conv.image_token
+            )
+            real_content += image_token
+        if text is not None:
+            real_content += text
+        conv.append_message(conv.roles[0], real_content)
+        # Add a blank message for the assistant.
+        conv.append_message(conv.roles[1], None)
+        convs.append(conv)
+
+    return convs
+
+
+# Models in which system adds modality tokens at prompt start automatically
+# when media inputs exceed modality tokens in prompt (e.g. 3 images but 2 <image> tokens)
+_MODELS_REQUIRING_MODALITY_SUPPLEMENT = {"deepseek-vl2"}
+
+
+# adapted from https://github.com/vllm-project/vllm/blob/5124f5bf51b83e6f344c1bc6652e8c4d81313b34/vllm/entrypoints/chat_utils.py#L856
+def _get_full_multimodal_text_prompt(
+    modality_token: str, modality_count: int, text_prompt: str
+) -> str:
+    """Combine multimodal prompts for a multimodal language model."""
+
+    # For any existing placeholder in the text prompt, we leave it as is
+    left: int = modality_count - text_prompt.count(modality_token)
+    if left < 0:
+        raise ValueError(
+            f"Found more '{modality_token}' placeholders in input prompt than "
+            "actual multimodal data items."
+        )
+
+    # NOTE: For now we always add missing modality_token at the front of
+    # the prompt. This may change to be customizable in the future.
+    return "\n".join([modality_token] * left + [text_prompt])
+
+
+def generate_chat_conv(
+    request: ChatCompletionRequest, template_name: str
+) -> Conversation:
+    conv = chat_templates[template_name].copy()
+    conv = Conversation(
+        name=conv.name,
+        system_template=conv.system_template,
+        system_message=conv.system_message,
+        roles=conv.roles,
+        messages=list(conv.messages),  # prevent in-place modification
+        offset=conv.offset,
+        sep_style=SeparatorStyle(conv.sep_style),
+        sep=conv.sep,
+        sep2=conv.sep2,
+        stop_str=conv.stop_str,
+        image_data=[],
+        audio_data=[],
+        modalities=[],
+        image_token=conv.image_token,
+        audio_token=conv.audio_token,
+    )
+
+    if isinstance(request.messages, str):
+        raise ValueError("The messages should be a list of dict.")
+    for message in request.messages:
+        msg_role = message.role
+        if msg_role == "system":
+            if isinstance(message.content, str):
+                conv.system_message = message.content
+            elif isinstance(message.content, list):
+                if (
+                    len(message.content) != 1
+                    or getattr(message.content[0], "type", None) != "text"
+                ):
+                    raise ValueError("The system message should be a single text.")
+                else:
+                    conv.system_message = getattr(message.content[0], "text", "")
+        elif msg_role == "user":
+            # Handle the various types of Chat Request content types here.
+            if isinstance(message.content, str):
+                conv.append_message(conv.roles[0], message.content)
+            else:
+                real_content = ""
+                # calculate number of image_url
+                num_image_url = 0
+                for content in message.content:
+                    if content.type == "image_url":
+                        num_image_url += 1
+                        conv.modalities.append(content.modalities)
+                if num_image_url > 1:
+                    image_token = conv.image_token
+                else:
+                    image_token = (
+                        conv.image_token + "\n"
+                        if conv.name != "qwen2-vl"
+                        else conv.image_token
+                    )
+                add_token_as_needed: bool = (
+                    conv.name in _MODELS_REQUIRING_MODALITY_SUPPLEMENT
+                )
+                if add_token_as_needed:
+                    image_token = ""
+
+                audio_token = conv.audio_token
+                for content in message.content:
+                    if content.type == "text":
+                        if num_image_url > 16:
+                            real_content += "\n"  # for video
+                        real_content += content.text
+                    elif content.type == "image_url":
+                        # NOTE: Only works for llava
+                        real_content += image_token
+                        conv.append_image(content.image_url.url)
+                    elif content.type == "audio_url":
+                        real_content += audio_token
+                        conv.append_audio(content.audio_url.url)
+                if add_token_as_needed:
+                    real_content = _get_full_multimodal_text_prompt(
+                        conv.image_token, num_image_url, real_content
+                    )
+                conv.append_message(conv.roles[0], real_content)
+        elif msg_role == "assistant":
+            parsed_content = ""
+            if isinstance(message.content, str):
+                parsed_content = message.content
+            elif isinstance(message.content, list):
+                if (
+                    len(message.content) != 1
+                    or getattr(message.content[0], "type", None) != "text"
+                ):
+                    raise ValueError(
+                        "The assistant's response should be a single text."
+                    )
+                else:
+                    parsed_content = getattr(message.content[0], "text", "")
+            conv.append_message(conv.roles[1], parsed_content)
+        else:
+            raise ValueError(f"Unknown role: {msg_role}")
+
+    # Add a blank message for the assistant.
+    conv.append_message(conv.roles[1], None)
+    return conv
+
+
+# llama2 template
+# reference: https://huggingface.co/blog/codellama#conversational-instructions
+# reference: https://github.com/facebookresearch/llama/blob/1a240688810f8036049e8da36b073f63d2ac552c/llama/generation.py#L212
+register_conv_template(
+    Conversation(
+        name="llama-2",
+        system_template="[INST] <<SYS>>\n{system_message}\n<</SYS>>\n\n",
+        roles=("[INST]", "[/INST]"),
+        sep_style=SeparatorStyle.LLAMA2,
+        sep=" ",
+        sep2=" </s><s>",
+        stop_str=["[INST]", "[/INST]", "<<SYS>>", "<</SYS>>"],
+    )
+)
+
+# reference: https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct/blob/main/chat_template.json
+register_conv_template(
+    Conversation(
+        name="llama-4",
+        system_template="<|header_start|>system<|header_end|>\n\n{system_message}<|eot|>",
+        roles=("user", "assistant"),
+        sep_style=SeparatorStyle.LLAMA4,
+        sep="",
+        stop_str=["<|end_of_text|>", "<|eot|>", "<|eom|>"],
+        image_token="<|image|>",
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="chatml",
+        system_template="<|im_start|>system\n{system_message}",
+        system_message="You are a helpful assistant.",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|im_end|>",
+        stop_str=["<|endoftext|>", "<|im_end|>"],
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="chatml-llava",
+        system_template="<|im_start|>system\n{system_message}",
+        system_message="You are a helpful assistant.",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep_style=SeparatorStyle.CHATML,
+        sep="<|im_end|>",
+        stop_str=["<|endoftext|>", "<|im_end|>"],
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="vicuna_v1.1",
+        system_message="A chat between a curious user and an artificial intelligence assistant. "
+        "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+        roles=("USER", "ASSISTANT"),
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        sep=" ",
+        sep2="</s>",
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="llama_3_vision",
+        system_message="You are a helpful language and vision assistant. You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language.",
+        system_template="<|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|>",
+        roles=("user", "assistant"),
+        sep_style=SeparatorStyle.LLAMA3,
+        sep="",
+        stop_str=["<|end_of_text|>", "<|eot_id|>"],
+        image_token="<|image|>",
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="llava_llama_3",
+        system_message="You are a helpful language and vision assistant. You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language.",
+        system_template="<|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|>",
+        roles=("user", "assistant"),
+        sep_style=SeparatorStyle.LLAMA3,
+        sep="",
+        stop_str=["<|end_of_text|>", "<|eot_id|>"],
+    )
+)
+# Reference: https://github.com/InternLM/lmdeploy/blob/387bf54b4f124e72aab30ae9755f562e435d3d01/lmdeploy/model.py#L425-L442
+register_conv_template(
+    Conversation(
+        name="internlm2-chat",
+        system_template="<|im_start|>system\n{system_message}",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep="\n",
+        stop_str=["<|im_end|>", "<|action_end|>"],
+    )
+)
+
+# Reference: https://huggingface.co/docs/transformers/main/model_doc/qwen2_vl#usage-example
+register_conv_template(
+    Conversation(
+        name="qwen2-vl",
+        system_message="You are a helpful assistant.",
+        system_template="<|im_start|>system\n{system_message}",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep="<|im_end|>\n",
+        sep_style=SeparatorStyle.ADD_NEW_LINE_SINGLE,
+        stop_str=["<|im_end|>"],
+        image_token="<|vision_start|><|image_pad|><|vision_end|>",
+    )
+)
+
+register_conv_template(
+    Conversation(
+        name="deepseek-vl2",
+        system_template="{system_message}",
+        # system_message="You are a helpful assistant. Please answer truthfully and write out your "
+        # "thinking step by step to be sure you get the right answer.",
+        system_message="",
+        roles=("<|User|>", "<|Assistant|>"),
+        messages=(),
+        offset=0,
+        sep_style=SeparatorStyle.DeepSeekVL2,
+        sep="\n\n",
+        sep2="<｜end▁of▁sentence｜>",
+        stop_str=["User:", "<｜end▁of▁sentence｜>"],
+    )
+)
+
+# Reference: https://huggingface.co/google/gemma-3-4b-it/blob/main/config.json
+register_conv_template(
+    Conversation(
+        name="gemma-it",
+        system_message="You are a helpful assistant.",
+        system_template="<start_of_turn>user{system_message}\n\n",
+        roles=("<start_of_turn>user\n", "<start_of_turn>model\n"),
+        sep="<end_of_turn>\n",
+        sep_style=SeparatorStyle.GEMMA3,
+        stop_str=["<end_of_turn>"],
+        image_token="<start_of_image>",
+    )
+)
+
+# Reference: https://huggingface.co/Alibaba-NLP/gme-Qwen2-VL-2B-Instruct#usage
+register_conv_template(
+    Conversation(
+        name="gme-qwen2-vl",
+        system_message="You are a helpful assistant.",
+        system_template="<|im_start|>system\n{system_message}",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep="<|im_end|>\n",
+        sep_style=SeparatorStyle.QWEN2_VL_EMBED,
+        stop_str="<|endoftext|>",
+        image_token="<|vision_start|><|image_pad|><|vision_end|>",
+    )
+)
+
+# Reference: https://huggingface.co/openbmb/MiniCPM-V-2_6#usage
+register_conv_template(
+    Conversation(
+        name="minicpmv",
+        system_message="You are a helpful assistant",
+        system_template="<|im_start|>system\n{system_message}.",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep="<|im_end|>\n",
+        sep_style=SeparatorStyle.ADD_NEW_LINE_SINGLE,
+        stop_str=("<|im_end|>", "<|endoftext|>"),
+        image_token="(<image>./</image>)",
+    )
+)
+
+# Reference: https://github.com/deepseek-ai/Janus?tab=readme-ov-file#janus-pro
+register_conv_template(
+    Conversation(
+        name="janus-pro",
+        system_message="You are a helpful language and vision assistant. You are able to understand the visual content that the user provides, and assist the user with a variety of tasks using natural language",
+        system_template="{system_message}.",
+        roles=("User", "Assistant"),
+        sep="\n\n",
+        sep2="<｜end▁of▁sentence｜>",
+        sep_style=SeparatorStyle.ADD_COLON_TWO,
+        stop_str=["<|User|>", "<｜end▁of▁sentence｜>"],
+        image_token="<image_placeholder>",
+    )
+)
+
+# Reference: https://huggingface.co/openbmb/MiniCPM-o-2_6#usage
+register_conv_template(
+    Conversation(
+        name="minicpmo",
+        system_message="You are Qwen, created by Alibaba Cloud. You are a helpful assistant.",
+        system_template="<|im_start|>system\n{system_message}",
+        roles=("<|im_start|>user", "<|im_start|>assistant"),
+        sep="<|im_end|>\n",
+        sep_style=SeparatorStyle.ADD_NEW_LINE_SINGLE,
+        stop_str=("<|im_end|>", "<|endoftext|>"),
+        image_token="(<image>./</image>)",
+        audio_token="(<audio>./</audio>)",
+    )
+)

lib/python3.12/site-packages/sglang/srt/custom_op.py

@@ -0,0 +1,44 @@
+from typing import Optional
+
+import torch
+from torch import nn
+
+from sglang.srt.utils import is_cuda, is_hip
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+
+
+class CustomOp(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self._forward_method = self.dispatch_forward()
+
+    def forward(self, *args, **kwargs):
+        return self._forward_method(*args, **kwargs)
+
+    def forward_native(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def forward_cuda(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def forward_hip(self, *args, **kwargs):
+        return self.forward_cuda(*args, **kwargs)
+
+    def forward_xpu(self, *args, **kwargs):
+        return self.forward_native(*args, **kwargs)
+
+    def forward_hpu(self, *args, **kwargs):
+        return self.forward_native(*args, **kwargs)
+
+    def forward_cpu(self, *args, **kwargs):
+        return self.forward_native(*args, **kwargs)
+
+    def dispatch_forward(self):
+        if _is_cuda:
+            return self.forward_cuda
+        elif _is_hip:
+            return self.forward_hip
+        else:
+            return self.forward_native

lib/python3.12/site-packages/sglang/srt/disaggregation/decode.py

@@ -0,0 +1,661 @@
+"""
+Life cycle of a request in the decode server
+
+1. PreallocQueue:
+    a. Initialize a receiver for each request
+    b. The request handshakes first, and pre-allocate kv once there is available kv.
+    c. Move the request to TransferQueue.
+
+2. TransferQueue:
+    a. Poll the receiver to check the transfer state
+    b. If the transfer has finished, move the request to waiting queue
+
+3. WaitingQueue:
+    a. Use the requests in the queue to construct a PrebuiltExtendBatch
+    b. Skip the prefill forward but only populate metadata
+
+4. RunningBatch:
+    a. Merge the resolved PrebuiltExtendBatch into running batch to run decoding
+"""
+
+from __future__ import annotations
+
+import logging
+from collections import deque
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
+import numpy as np
+import torch
+from torch.distributed import ProcessGroup
+
+from sglang.srt.disaggregation.base import BaseKVManager, BaseKVReceiver, KVArgs, KVPoll
+from sglang.srt.disaggregation.utils import (
+    DisaggregationMode,
+    FakeBootstrapHost,
+    KVClassType,
+    ReqToMetadataIdxAllocator,
+    TransferBackend,
+    get_kv_class,
+    kv_to_page_indices,
+    poll_and_all_reduce,
+)
+from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool, TokenToKVPoolAllocator
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
+
+logger = logging.getLogger(__name__)
+
+if TYPE_CHECKING:
+    from sglang.srt.configs.model_config import ModelConfig
+    from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+    from sglang.srt.managers.scheduler import Scheduler
+    from sglang.srt.server_args import ServerArgs
+
+
+@dataclass
+class DecodeRequest:
+    req: Req
+    kv_receiver: BaseKVReceiver
+    waiting_for_input: bool = False
+    metadata_buffer_index: int = -1
+
+
+class DecodePreallocQueue:
+    """
+    Store the requests that are preallocating.
+    """
+
+    def __init__(
+        self,
+        req_to_token_pool: ReqToTokenPool,
+        token_to_kv_pool_allocator: TokenToKVPoolAllocator,
+        req_to_metadata_buffer_idx_allocator: ReqToMetadataIdxAllocator,
+        metadata_buffers: List[torch.Tensor],
+        aux_dtype: torch.dtype,
+        scheduler: Scheduler,
+        transfer_queue: DecodeTransferQueue,
+        tree_cache: BasePrefixCache,
+        gloo_group: ProcessGroup,
+        tp_rank: int,
+        tp_size: int,
+        bootstrap_port: int,
+        transfer_backend: TransferBackend,
+    ):
+        self.req_to_token_pool = req_to_token_pool
+        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+        self.token_to_kv_pool = token_to_kv_pool_allocator.get_kvcache()
+        self.aux_dtype = aux_dtype
+        self.metadata_buffers = metadata_buffers
+        self.req_to_metadata_buffer_idx_allocator = req_to_metadata_buffer_idx_allocator
+        self.scheduler = scheduler
+        self.transfer_queue = transfer_queue
+        self.tree_cache = tree_cache  # this is always a chunk cache
+        self.gloo_group = gloo_group
+        self.tp_rank = tp_rank
+        self.tp_size = tp_size
+        self.bootstrap_port = bootstrap_port
+
+        self.num_reserved_decode_tokens = 512
+
+        # Queue for requests pending pre-allocation
+        self.queue: List[DecodeRequest] = []
+        self.transfer_backend = transfer_backend
+        self.kv_manager = self._init_kv_manager()
+
+    def _init_kv_manager(self) -> BaseKVManager:
+        kv_args = KVArgs()
+        kv_args.engine_rank = self.tp_rank
+        kv_data_ptrs, kv_data_lens, kv_item_lens = (
+            self.token_to_kv_pool.get_contiguous_buf_infos()
+        )
+
+        kv_args.kv_data_ptrs = kv_data_ptrs
+        kv_args.kv_data_lens = kv_data_lens
+        kv_args.kv_item_lens = kv_item_lens
+
+        kv_args.aux_data_ptrs = [
+            output_id_tensor.data_ptr() for output_id_tensor in self.metadata_buffers
+        ]
+        kv_args.aux_data_lens = [
+            metadata_buffer.nbytes for metadata_buffer in self.metadata_buffers
+        ]
+        kv_args.aux_item_lens = [
+            metadata_buffer[0].nbytes for metadata_buffer in self.metadata_buffers
+        ]
+        kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device
+        kv_args.gpu_id = self.scheduler.gpu_id
+        kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER)
+        kv_manager = kv_manager_class(
+            kv_args, DisaggregationMode.DECODE, self.scheduler.server_args
+        )
+        return kv_manager
+
+    def add(self, req: Req) -> None:
+        """Add a request to the pending queue."""
+        if req.bootstrap_host == FakeBootstrapHost:
+            # Fake transfer for warmup reqs
+            kv_receiver_class = get_kv_class(TransferBackend.FAKE, KVClassType.RECEIVER)
+        else:
+            kv_receiver_class = get_kv_class(
+                self.transfer_backend, KVClassType.RECEIVER
+            )
+        kv_receiver = kv_receiver_class(
+            mgr=self.kv_manager,
+            bootstrap_addr=f"{req.bootstrap_host}:{req.bootstrap_port}",
+            bootstrap_room=req.bootstrap_room,
+        )
+        self.queue.append(DecodeRequest(req=req, kv_receiver=kv_receiver))
+
+    def extend(self, reqs: List[Req]) -> None:
+        """Add a request to the pending queue."""
+        for req in reqs:
+            self.add(req)
+
+    def _update_handshake_waiters(self) -> None:
+        if not self.queue:
+            return
+
+        if all(decode_req.waiting_for_input for decode_req in self.queue):
+            return
+
+        polls = poll_and_all_reduce(
+            [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group
+        )
+
+        for i, (decode_req, poll) in enumerate(zip(self.queue, polls)):
+            if poll == KVPoll.Bootstrapping:
+                pass
+            elif poll == KVPoll.WaitingForInput:
+                decode_req.waiting_for_input = True
+            elif poll == KVPoll.Failed:
+                raise Exception("Handshake failed")
+
+    def pop_preallocated(self) -> List[DecodeRequest]:
+        """Pop the preallocated requests from the pending queue (FIFO)."""
+        self._update_handshake_waiters()
+
+        preallocated_reqs = []
+        indices_to_remove = set()
+        allocatable_tokens = self._allocatable_tokens()
+
+        for i, decode_req in enumerate(self.queue):
+            if not decode_req.waiting_for_input:
+                continue
+
+            if self.req_to_token_pool.available_size() <= 0:
+                break
+
+            if self.req_to_metadata_buffer_idx_allocator.available_size() <= 0:
+                break
+
+            required_tokens_for_request = (
+                len(decode_req.req.origin_input_ids) + self.num_reserved_decode_tokens
+            )
+
+            if required_tokens_for_request > allocatable_tokens:
+                break
+
+            allocatable_tokens -= required_tokens_for_request
+            self._pre_alloc(decode_req.req)
+
+            kv_indices = (
+                self.req_to_token_pool.req_to_token[decode_req.req.req_pool_idx][
+                    : len(decode_req.req.origin_input_ids)
+                ]
+                .cpu()
+                .numpy()
+                .astype(np.int64)
+            )
+
+            decode_req.metadata_buffer_index = (
+                self.req_to_metadata_buffer_idx_allocator.alloc()
+            )
+            assert decode_req.metadata_buffer_index is not None
+            page_indices = kv_to_page_indices(
+                kv_indices, self.token_to_kv_pool_allocator.page_size
+            )
+            decode_req.kv_receiver.init(page_indices, decode_req.metadata_buffer_index)
+            preallocated_reqs.append(decode_req)
+            indices_to_remove.add(i)
+
+        self.queue = [
+            entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
+        ]
+
+        return preallocated_reqs
+
+    def _allocatable_tokens(self) -> int:
+        allocatable_tokens = (
+            self.token_to_kv_pool_allocator.available_size()
+            - self.num_reserved_decode_tokens
+            * (
+                len(self.scheduler.running_batch.reqs)
+                + len(self.transfer_queue.queue)
+                + len(self.scheduler.waiting_queue)
+            )
+        )
+
+        # Note: if the last fake extend just finishes, and we enter `pop_preallocated` immediately in the next iteration
+        #       the extend batch is not in any queue, so we need to explicitly add the tokens slots here
+        if (
+            self.scheduler.last_batch
+            and self.scheduler.last_batch.forward_mode.is_extend()
+        ):
+            allocatable_tokens -= self.num_reserved_decode_tokens * len(
+                self.scheduler.last_batch.reqs
+            )
+
+        return allocatable_tokens
+
+    def _pre_alloc(self, req: Req) -> torch.Tensor:
+        """Pre-allocate the memory for req_to_token and token_kv_pool"""
+        req_pool_indices = self.req_to_token_pool.alloc(1)
+
+        assert req_pool_indices is not None
+
+        req.req_pool_idx = req_pool_indices[0]
+        if self.token_to_kv_pool_allocator.page_size == 1:
+            kv_loc = self.token_to_kv_pool_allocator.alloc(
+                len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
+            )
+        else:
+            num_tokens = len(req.origin_input_ids) + max(len(req.output_ids) - 1, 0)
+            kv_loc = self.token_to_kv_pool_allocator.alloc_extend(
+                prefix_lens=torch.tensor(
+                    [0],
+                    dtype=torch.int64,
+                    device=self.token_to_kv_pool_allocator.device,
+                ),
+                seq_lens=torch.tensor(
+                    [num_tokens],
+                    dtype=torch.int64,
+                    device=self.token_to_kv_pool_allocator.device,
+                ),
+                last_loc=torch.tensor(
+                    [-1],
+                    dtype=torch.int64,
+                    device=self.token_to_kv_pool_allocator.device,
+                ),
+                extend_num_tokens=num_tokens,
+            )
+        assert kv_loc is not None
+
+        self.req_to_token_pool.write((req.req_pool_idx, slice(0, len(kv_loc))), kv_loc)
+
+        # populate metadata
+        req.fill_ids = req.origin_input_ids + req.output_ids
+        req.extend_input_len = len(req.origin_input_ids)
+
+        return kv_loc
+
+
+class DecodeTransferQueue:
+    """
+    Store the requests that is polling kv
+    """
+
+    def __init__(
+        self,
+        gloo_group: ProcessGroup,
+        req_to_metadata_buffer_idx_allocator: ReqToMetadataIdxAllocator,
+        metadata_buffers: torch.Tensor,
+    ):
+        self.queue: List[DecodeRequest] = []
+        self.gloo_group = gloo_group
+        self.req_to_metadata_buffer_idx_allocator = req_to_metadata_buffer_idx_allocator
+        self.metadata_buffers = metadata_buffers
+
+    def add(self, req_conn: DecodeRequest) -> None:
+        self.queue.append(req_conn)
+
+    def extend(self, req_conns) -> None:
+        self.queue.extend(req_conns)
+
+    def pop_transferred(self) -> List[DecodeRequest]:
+        if not self.queue:
+            return []
+
+        polls = poll_and_all_reduce(
+            [decode_req.kv_receiver for decode_req in self.queue], self.gloo_group
+        )
+
+        transferred_reqs = []
+        indices_to_remove = set()
+        for i, (decode_req, poll) in enumerate(zip(self.queue, polls)):
+            if poll == KVPoll.Failed:
+                raise Exception("Transfer failed")
+            elif poll == KVPoll.Success:
+                # pop and push it to waiting queue
+                idx = decode_req.metadata_buffer_index
+                assert len(decode_req.req.output_ids) == 0
+                output_id_buffer = self.metadata_buffers[0]
+                # the last dimension is padded by the same values.
+                output_id = output_id_buffer[idx][0].item()
+                assert len(decode_req.req.output_ids) == 0
+                assert decode_req.req.transferred_output_id is None
+                decode_req.req.transferred_output_id = output_id
+                transferred_reqs.append(decode_req)
+                indices_to_remove.add(i)
+            elif poll in [
+                KVPoll.Bootstrapping,
+                KVPoll.WaitingForInput,
+                KVPoll.Transferring,
+            ]:
+                pass
+            else:
+                raise ValueError(f"Unexpected poll case: {poll}")
+
+        for i in indices_to_remove:
+            idx = self.queue[i].metadata_buffer_index
+            assert idx != -1
+            self.req_to_metadata_buffer_idx_allocator.free(idx)
+
+        self.queue = [
+            entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
+        ]
+
+        return transferred_reqs
+
+
+class ScheduleBatchDisaggregationDecodeMixin:
+
+    def prepare_for_prebuilt_extend(self: ScheduleBatch):
+        """
+        Prepare a prebuilt extend by populate metadata
+        Adapted from .prepare_for_extend().
+        """
+
+        self.forward_mode = ForwardMode.EXTEND
+        reqs = self.reqs
+        input_ids = [r.fill_ids[len(r.prefix_indices) :] for r in reqs]
+        extend_num_tokens = sum(len(ids) for ids in input_ids)
+        seq_lens = []
+        pre_lens = []
+        req_pool_indices = []
+
+        # Pre-calculate total size
+        total_size = sum(req.extend_input_len for req in reqs)
+        out_cache_loc = torch.empty(total_size, dtype=torch.int64, device=self.device)
+
+        # Fill the tensor in one pass
+        offset = 0
+        for i, req in enumerate(reqs):
+            req_pool_indices.append(req.req_pool_idx)
+
+            chunk = self.req_to_token_pool.req_to_token[req.req_pool_idx][
+                : req.extend_input_len
+            ]
+            assert (
+                offset + req.extend_input_len <= total_size
+            ), f"Exceeds total size: offset={offset}, req.extend_input_len={req.extend_input_len}, total_size={total_size}"
+            out_cache_loc[offset : offset + req.extend_input_len] = chunk
+            offset += req.extend_input_len
+
+            pre_len = len(req.prefix_indices)
+            seq_len = len(req.origin_input_ids) + max(0, len(req.output_ids) - 1)
+            seq_lens.append(seq_len)
+            if len(req.output_ids) == 0:
+                assert (
+                    seq_len - pre_len == req.extend_input_len
+                ), f"seq_len={seq_len}, pre_len={pre_len}, req.extend_input_len={req.extend_input_len}"
+
+            req.cached_tokens += pre_len - req.already_computed
+            req.already_computed = seq_len
+            req.is_retracted = False
+            pre_lens.append(pre_len)
+            req.extend_logprob_start_len = 0
+
+        extend_input_logprob_token_ids = None
+
+        # Set fields
+        self.input_ids = torch.tensor(
+            sum(input_ids, []), dtype=torch.int32, device=self.device
+        )
+        self.req_pool_indices = torch.tensor(
+            req_pool_indices, dtype=torch.int64, device=self.device
+        )
+        self.seq_lens = torch.tensor(seq_lens, dtype=torch.int64, device=self.device)
+        self.out_cache_loc = out_cache_loc
+        self.seq_lens_sum = sum(seq_lens)
+        self.extend_num_tokens = extend_num_tokens
+        self.prefix_lens = [len(r.prefix_indices) for r in reqs]
+        self.extend_lens = [r.extend_input_len for r in reqs]
+        self.extend_logprob_start_lens = [r.extend_logprob_start_len for r in reqs]
+        self.extend_input_logprob_token_ids = extend_input_logprob_token_ids
+
+        # Build sampling info
+        self.sampling_info = SamplingBatchInfo.from_schedule_batch(
+            self,
+            self.model_config.vocab_size,
+        )
+
+    def process_prebuilt_extend(
+        self: ScheduleBatch, server_args: ServerArgs, model_config: ModelConfig
+    ):
+        """Assign the buffered last input id to schedule batch"""
+        self.output_ids = []
+        for req in self.reqs:
+            if req.output_ids and len(req.output_ids) > 0:
+                # resumed retracted req
+                self.output_ids.append(req.output_ids[-1])
+            else:
+                assert req.transferred_output_id is not None
+                req.output_ids.append(req.transferred_output_id)
+                self.output_ids.append(req.transferred_output_id)
+            self.tree_cache.cache_unfinished_req(req)
+        self.output_ids = torch.tensor(self.output_ids, device=self.device)
+
+
+class SchedulerDisaggregationDecodeMixin:
+
+    def _prepare_idle_batch_and_run(self, batch, delay_process=False):
+        batch, _ = self.prepare_dp_attn_batch(batch)
+        result = None
+        if batch:
+            result = self.run_batch(batch)
+            if not delay_process:
+                self.process_batch_result(batch, result)
+        return batch, result
+
+    @torch.no_grad()
+    def event_loop_normal_disagg_decode(self: Scheduler):
+        """A normal scheduler loop for decode worker in disaggregation mode."""
+
+        while True:
+            recv_reqs = self.recv_requests()
+            self.process_input_requests(recv_reqs)
+            # polling and allocating kv cache
+            self.process_decode_queue()
+            batch = self.get_next_disagg_decode_batch_to_run()
+            self.cur_batch = batch
+
+            prepare_dp_attn_flag = (
+                self.server_args.enable_dp_attention
+                or self.server_args.enable_sp_layernorm
+            )
+
+            if batch:
+                # Generate fake extend output.
+                if batch.forward_mode.is_extend():
+                    # Note: Logprobs should be handled on the prefill engine.
+                    self.stream_output(batch.reqs, False)
+                    if prepare_dp_attn_flag:
+                        self._prepare_idle_batch_and_run(None)
+                else:
+                    if prepare_dp_attn_flag:
+                        self.prepare_dp_attn_batch(batch)
+                    result = self.run_batch(batch)
+                    self.process_batch_result(batch, result)
+            elif prepare_dp_attn_flag:
+                batch, _ = self._prepare_idle_batch_and_run(None)
+
+            if batch is None and (
+                len(self.disagg_decode_transfer_queue.queue)
+                + len(self.disagg_decode_prealloc_queue.queue)
+                == 0
+            ):
+                # When the server is idle, do self-check and re-init some states
+                self.check_memory()
+                self.new_token_ratio = self.init_new_token_ratio
+
+            self.last_batch = batch
+
+    @torch.no_grad()
+    def event_loop_overlap_disagg_decode(self: Scheduler):
+        result_queue = deque()
+        self.last_batch: Optional[ScheduleBatch] = None
+        self.last_batch_in_queue = False  # last batch is modifed in-place, so we need another variable to track if it's extend
+
+        while True:
+            recv_reqs = self.recv_requests()
+            self.process_input_requests(recv_reqs)
+            # polling and allocating kv cache
+            self.process_decode_queue()
+            batch = self.get_next_disagg_decode_batch_to_run()
+            self.cur_batch = batch
+            last_batch_in_queue = False
+
+            prepare_dp_attn_flag = (
+                self.server_args.enable_dp_attention
+                or self.server_args.enable_sp_layernorm
+            )
+
+            if batch:
+                # Generate fake extend output.
+                if batch.forward_mode.is_extend():
+                    # Note: Logprobs should be handled on the prefill engine.
+                    self.stream_output(batch.reqs, False)
+                    if prepare_dp_attn_flag:
+                        batch_, result = self._prepare_idle_batch_and_run(
+                            None, delay_process=True
+                        )
+                        if batch_:
+                            result_queue.append((batch_.copy(), result))
+                            last_batch_in_queue = True
+                else:
+                    if prepare_dp_attn_flag:
+                        self.prepare_dp_attn_batch(batch)
+                    result = self.run_batch(batch)
+                    result_queue.append((batch.copy(), result))
+                    last_batch_in_queue = True
+            elif prepare_dp_attn_flag:
+                batch, result = self._prepare_idle_batch_and_run(
+                    None, delay_process=True
+                )
+                if batch:
+                    result_queue.append((batch.copy(), result))
+                    last_batch_in_queue = True
+
+            # Process the results of the previous batch but skip if the last batch is extend
+            if self.last_batch and self.last_batch_in_queue:
+                tmp_batch, tmp_result = result_queue.popleft()
+                self.process_batch_result(tmp_batch, tmp_result)
+
+            if batch is None and (
+                len(self.disagg_decode_transfer_queue.queue)
+                + len(self.disagg_decode_prealloc_queue.queue)
+                == 0
+            ):
+                # When the server is idle, do self-check and re-init some states
+                self.check_memory()
+                self.new_token_ratio = self.init_new_token_ratio
+
+            self.last_batch = batch
+            self.last_batch_in_queue = last_batch_in_queue
+
+    def get_next_disagg_decode_batch_to_run(
+        self: Scheduler,
+    ) -> Optional[Tuple[ScheduleBatch, bool]]:
+        """Create fake completed prefill if possible and merge with running batch"""
+        # Merge the prefill batch into the running batch
+        last_batch = self.last_batch
+        if last_batch and last_batch.forward_mode.is_extend():
+            # chunked prefill doesn't happen in decode instance.
+            assert self.chunked_req is None
+            # Filter finished batches.
+            last_batch.filter_batch()
+            if not last_batch.is_empty():
+                if self.running_batch.is_empty():
+                    self.running_batch = last_batch
+                else:
+                    # merge running_batch with prefill batch
+                    self.running_batch.merge_batch(last_batch)
+
+        new_prebuilt_batch = self.get_new_prebuilt_batch()
+
+        ret: Optional[ScheduleBatch] = None
+        if new_prebuilt_batch:
+            ret = new_prebuilt_batch
+        else:
+            if self.running_batch.is_empty():
+                ret = None
+            else:
+                self.running_batch = self.update_running_batch(self.running_batch)
+                ret = self.running_batch if not self.running_batch.is_empty() else None
+
+        return ret
+
+    def get_new_prebuilt_batch(self: Scheduler) -> Optional[ScheduleBatch]:
+        """Create a schedulebatch for fake completed prefill"""
+        if len(self.waiting_queue) == 0:
+            return None
+
+        curr_batch_size = self.running_batch.batch_size()
+
+        batch_size = min(self.req_to_token_pool.size, self.max_running_requests)
+
+        num_not_used_batch = batch_size - curr_batch_size
+
+        # pop req from waiting queue
+        can_run_list: List[Req] = []
+        waiting_queue: List[Req] = []
+
+        for i in range(len(self.waiting_queue)):
+            req = self.waiting_queue[i]
+            # we can only add at least `num_not_used_batch` new batch to the running queue
+            if i < num_not_used_batch:
+                can_run_list.append(req)
+                req.init_next_round_input(self.tree_cache)
+            else:
+                waiting_queue.append(req)
+
+        self.waiting_queue = waiting_queue
+        if len(can_run_list) == 0:
+            return None
+        # local import to avoid circular import
+        from sglang.srt.managers.schedule_batch import ScheduleBatch
+
+        # construct a schedule batch with those requests and mark as decode
+        new_batch = ScheduleBatch.init_new(
+            can_run_list,
+            self.req_to_token_pool,
+            self.token_to_kv_pool_allocator,
+            self.tree_cache,
+            self.model_config,
+            self.enable_overlap,
+            self.spec_algorithm,
+            self.server_args.enable_custom_logit_processor,
+        )
+
+        # construct fake completed prefill
+        new_batch.prepare_for_prebuilt_extend()
+        new_batch.process_prebuilt_extend(self.server_args, self.model_config)
+
+        return new_batch
+
+    def process_decode_queue(self: Scheduler):
+        req_conns = self.disagg_decode_prealloc_queue.pop_preallocated()
+
+        def _num_pre_alloc(req):
+            return len(req.req.origin_input_ids) + max(len(req.req.output_ids) - 1, 0)
+
+        self.num_tokens_pre_allocated += sum(_num_pre_alloc(req) for req in req_conns)
+        self.disagg_decode_transfer_queue.extend(req_conns)
+        alloc_reqs = (
+            self.disagg_decode_transfer_queue.pop_transferred()
+        )  # the requests which kv has arrived
+        self.num_tokens_pre_allocated -= sum(_num_pre_alloc(req) for req in alloc_reqs)
+
+        self.waiting_queue.extend([req.req for req in alloc_reqs])

lib/python3.12/site-packages/sglang/srt/disaggregation/fake/__init__.py

@@ -0,0 +1 @@
+from .conn import FakeKVReceiver, FakeKVSender

lib/python3.12/site-packages/sglang/srt/disaggregation/fake/conn.py

@@ -0,0 +1,88 @@
+import logging
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import numpy.typing as npt
+
+from sglang.srt.disaggregation.base.conn import (
+    BaseKVManager,
+    BaseKVReceiver,
+    BaseKVSender,
+    KVArgs,
+    KVPoll,
+)
+
+logger = logging.getLogger(__name__)
+
+
+# For warmup reqs, we don't kv transfer, we use the fake sender and receiver
+class FakeKVSender(BaseKVSender):
+    def __init__(self, mgr: BaseKVManager, bootstrap_addr: str, bootstrap_room: int):
+        self.has_sent = False
+
+    def poll(self) -> KVPoll:
+        if self.has_sent is False:
+            # Assume handshake completed instantly
+            return KVPoll.WaitingForInput
+        else:
+            # Assume transfer completed instantly
+            logger.info("FakeKVSender poll success")
+            return KVPoll.Success
+
+    def init(
+        self,
+        kv_indices: list[int],
+        aux_index: Optional[int] = None,
+        dest_ranks: Optional[list[int]] = None,
+    ):
+        logger.info(
+            f"FakeKVSender init with kv_indices: {kv_indices}, aux_index: {aux_index}, dest_ranks: {dest_ranks}"
+        )
+        pass
+
+    def send(
+        self,
+        kv_indices: npt.NDArray[np.int64],
+        index_slice: slice,
+        is_last: bool,
+    ):
+        logger.info(
+            f"FakeKVSender send with kv_indices: {kv_indices}, index_slice: {index_slice}, is_last: {is_last}"
+        )
+        if is_last:
+            self.has_sent = True
+            logger.info(f"FakeKVSender send success")
+        else:
+            self.has_sent = False
+            logger.info(f"FakeKVSender send fake transfering")
+
+    def failure_exception(self):
+        raise Exception("Fake KVSender Exception")
+
+
+class FakeKVReceiver(BaseKVReceiver):
+    def __init__(
+        self,
+        mgr: BaseKVManager,
+        bootstrap_addr: str,
+        bootstrap_room: Optional[int] = None,
+    ):
+        self.has_init = False
+
+    def poll(self) -> KVPoll:
+        if self.has_init is False:
+            # Assume handshake completed instantly
+            return KVPoll.WaitingForInput
+        else:
+            # Assume transfer completed instantly
+            logger.info("FakeKVReceiver poll success")
+            return KVPoll.Success
+
+    def init(self, kv_indices: list[int], aux_index: Optional[int] = None):
+        self.has_init = True
+        logger.info(
+            f"FakeKVReceiver init with kv_indices: {kv_indices}, aux_index: {aux_index}"
+        )
+
+    def failure_exception(self):
+        raise Exception("Fake KVReceiver Exception")

lib/python3.12/site-packages/sglang/srt/disaggregation/mini_lb.py

@@ -0,0 +1,318 @@
+"""
+Minimal HTTP load balancer for prefill and decode servers for testing.
+"""
+
+import asyncio
+import random
+import urllib
+from itertools import chain
+from typing import List
+
+import aiohttp
+import orjson
+import uvicorn
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import ORJSONResponse, Response, StreamingResponse
+
+
+class PrefillConfig:
+    def __init__(self, url: str, bootstrap_port: int):
+        self.url = url
+        self.bootstrap_port = bootstrap_port
+
+
+class MiniLoadBalancer:
+    def __init__(self, prefill_configs: List[PrefillConfig], decode_servers: List[str]):
+        self.prefill_configs = prefill_configs
+        self.prefill_servers = [p.url for p in prefill_configs]
+        self.decode_servers = decode_servers
+
+    def select_pair(self):
+        prefill_config = random.choice(self.prefill_configs)
+        decode_server = random.choice(self.decode_servers)
+        return prefill_config.url, prefill_config.bootstrap_port, decode_server
+
+    async def generate(
+        self, modified_request, prefill_server, decode_server, endpoint
+    ) -> ORJSONResponse:
+        assert endpoint[0] != "/", f"Endpoint should not start with '/': {endpoint}"
+
+        async with aiohttp.ClientSession(
+            timeout=aiohttp.ClientTimeout(
+                total=3600
+            )  # Add timeout for request reliability
+        ) as session:
+            tasks = [
+                session.post(f"{prefill_server}/{endpoint}", json=modified_request),
+                session.post(f"{decode_server}/{endpoint}", json=modified_request),
+            ]
+            # Wait for both responses to complete. Prefill should end first.
+            prefill_response, decode_response = await asyncio.gather(*tasks)
+
+            return ORJSONResponse(
+                content=await decode_response.json(),
+                status_code=decode_response.status,
+            )
+
+    async def generate_stream(
+        self, modified_request, prefill_server, decode_server, endpoint="generate"
+    ):
+        assert endpoint[0] != "/", f"Endpoint should not start with '/': {endpoint}"
+
+        async def stream_results():
+            async with aiohttp.ClientSession(
+                timeout=aiohttp.ClientTimeout(
+                    total=3600
+                )  # Add timeout for request reliability
+            ) as session:
+                try:
+                    # Create the tasks for both prefill and decode requests
+                    tasks = [
+                        session.post(
+                            f"{prefill_server}/{endpoint}", json=modified_request
+                        ),
+                        session.post(
+                            f"{decode_server}/{endpoint}", json=modified_request
+                        ),
+                    ]
+                    # Wait for both responses to complete. Since this is streaming, they return immediately.
+                    prefill_response, decode_response = await asyncio.gather(*tasks)
+                    async for chunk in decode_response.content:
+                        yield chunk
+                except Exception as e:
+                    error_msg = {
+                        "error": {"message": f"Stream processing error: {str(e)}"}
+                    }
+                    yield b"data: " + orjson.dumps(
+                        error_msg, option=orjson.OPT_NON_STR_KEYS
+                    ) + b"\n\n"
+                finally:
+                    if prefill_response is not None:
+                        await prefill_response.release()
+
+        return StreamingResponse(
+            stream_results(),
+            media_type="text/event-stream",
+        )
+
+
+app = FastAPI()
+load_balancer = None
+
+
+@app.get("/health")
+async def health_check():
+    return Response(status_code=200)
+
+
+@app.get("/health_generate")
+async def health_check():
+    prefill_servers, decode_servers = (
+        load_balancer.prefill_servers,
+        load_balancer.decode_servers,
+    )
+    async with aiohttp.ClientSession() as session:
+        # Create the tasks
+        tasks = []
+        for server in chain(prefill_servers, decode_servers):
+            tasks.append(session.post(f"{server}/health_generate"))
+        for i, response in enumerate(asyncio.as_completed(tasks)):
+            await response
+    return Response(status_code=200)
+
+
+@app.post("/flush_cache")
+async def flush_cache():
+    prefill_servers, decode_servers = (
+        load_balancer.prefill_servers,
+        load_balancer.decode_servers,
+    )
+    async with aiohttp.ClientSession() as session:
+        # Create the tasks
+        tasks = []
+        for server in chain(prefill_servers, decode_servers):
+            tasks.append(session.post(f"{server}/flush_cache"))
+        for i, response in enumerate(asyncio.as_completed(tasks)):
+            await response
+    return Response(status_code=200)
+
+
+@app.get("/get_server_info")
+async def get_server_info():
+    prefill_servers, decode_servers = (
+        load_balancer.prefill_servers,
+        load_balancer.decode_servers,
+    )
+    prefill_infos = []
+    decode_infos = []
+    async with aiohttp.ClientSession() as session:
+        for server in chain(prefill_servers):
+            server_info = await session.get(f"{server}/get_server_info")
+            prefill_infos.append(await server_info.json())
+        for server in chain(decode_servers):
+            server_info = await session.get(f"{server}/get_server_info")
+            decode_infos.append(await server_info.json())
+
+    return {"prefill": prefill_infos, "decode": decode_infos}
+
+
+@app.get("/get_model_info")
+async def get_model_info():
+    # Dummy model information
+    model_info = {
+        "model_path": "/path/to/dummy/model",
+        "tokenizer_path": "/path/to/dummy/tokenizer",
+        "is_generation": True,
+        "preferred_sampling_params": {"temperature": 0.7, "max_new_tokens": 128},
+    }
+    return ORJSONResponse(content=model_info)
+
+
+@app.post("/generate")
+async def handle_generate_request(request_data: dict):
+    prefill_server, bootstrap_port, decode_server = load_balancer.select_pair()
+
+    # Parse and transform prefill_server for bootstrap data
+    parsed_url = urllib.parse.urlparse(prefill_server)
+    hostname = parsed_url.hostname
+    modified_request = request_data.copy()
+
+    batch_size = _get_request_batch_size(modified_request)
+    if batch_size is not None:
+        modified_request.update(
+            {
+                "bootstrap_host": [hostname] * batch_size,
+                "bootstrap_port": [bootstrap_port] * batch_size,
+                "bootstrap_room": [
+                    _generate_bootstrap_room() for _ in range(batch_size)
+                ],
+            }
+        )
+    else:
+        modified_request.update(
+            {
+                "bootstrap_host": hostname,
+                "bootstrap_port": bootstrap_port,
+                "bootstrap_room": _generate_bootstrap_room(),
+            }
+        )
+
+    if request_data.get("stream", False):
+        return await load_balancer.generate_stream(
+            modified_request, prefill_server, decode_server, "generate"
+        )
+    else:
+        return await load_balancer.generate(
+            modified_request, prefill_server, decode_server, "generate"
+        )
+
+
+@app.post("/v1/chat/completions")
+async def handle_completion_request(request_data: dict):
+    prefill_server, bootstrap_port, decode_server = load_balancer.select_pair()
+
+    # Parse and transform prefill_server for bootstrap data
+    parsed_url = urllib.parse.urlparse(prefill_server)
+    hostname = parsed_url.hostname
+    modified_request = request_data.copy()
+    modified_request.update(
+        {
+            "bootstrap_host": hostname,
+            "bootstrap_port": bootstrap_port,
+            "bootstrap_room": random.randint(0, 2**63 - 1),
+        }
+    )
+
+    if request_data.get("stream", False):
+        return await load_balancer.generate_stream(
+            modified_request,
+            prefill_server,
+            decode_server,
+            endpoint="v1/chat/completions",
+        )
+    else:
+        return await load_balancer.generate(
+            modified_request,
+            prefill_server,
+            decode_server,
+            endpoint="v1/chat/completions",
+        )
+
+
+def _generate_bootstrap_room():
+    return random.randint(0, 2**63 - 1)
+
+
+# We may utilize `GenerateReqInput`'s logic later
+def _get_request_batch_size(request):
+    if (text := request.get("text")) is not None:
+        return None if isinstance(text, str) else len(text)
+    if (input_ids := request.get("input_ids")) is not None:
+        return None if isinstance(input_ids[0], int) else len(input_ids)
+    return None
+
+
+@app.get("/v1/models")
+async def get_models():
+    prefill_server = load_balancer.prefill_servers[0]  # Get the first prefill server
+    async with aiohttp.ClientSession() as session:
+        try:
+            response = await session.get(f"{prefill_server}/v1/models")
+            if response.status != 200:
+                raise HTTPException(
+                    status_code=response.status,
+                    detail=f"Prefill server error: Status {response.status}",
+                )
+            return ORJSONResponse(content=await response.json())
+        except Exception as e:
+            raise HTTPException(status_code=500, detail=str(e))
+
+
+def run(prefill_configs, decode_addrs, host, port):
+    global load_balancer
+    load_balancer = MiniLoadBalancer(prefill_configs, decode_addrs)
+    uvicorn.run(app, host=host, port=port)
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Mini Load Balancer Server")
+    parser.add_argument(
+        "--prefill", required=True, help="Comma-separated URLs for prefill servers"
+    )
+    parser.add_argument(
+        "--prefill-bootstrap-ports",
+        help="Comma-separated bootstrap ports for prefill servers",
+        default="8998",
+    )
+    parser.add_argument(
+        "--decode", required=True, help="Comma-separated URLs for decode servers"
+    )
+    parser.add_argument(
+        "--host", default="0.0.0.0", help="Host to bind the server (default: 0.0.0.0)"
+    )
+    parser.add_argument(
+        "--port", type=int, default=8000, help="Port to bind the server (default: 8000)"
+    )
+    args = parser.parse_args()
+
+    prefill_urls = args.prefill.split(",")
+    bootstrap_ports = [int(p) for p in args.prefill_bootstrap_ports.split(",")]
+
+    if len(bootstrap_ports) == 1:
+        bootstrap_ports = bootstrap_ports * len(prefill_urls)
+    else:
+        if len(bootstrap_ports) != len(prefill_urls):
+            raise ValueError(
+                "Number of prefill URLs must match number of bootstrap ports"
+            )
+            exit(1)
+
+    prefill_configs = []
+    for url, port in zip(prefill_urls, bootstrap_ports):
+        prefill_configs.append(PrefillConfig(url, port))
+
+    decode_addrs = args.decode.split(",")
+
+    run(prefill_configs, decode_addrs, args.host, args.port)

lib/python3.12/site-packages/sglang/srt/disaggregation/mooncake/__init__.py

@@ -0,0 +1,6 @@
+from .conn import (
+    MooncakeKVBootstrapServer,
+    MooncakeKVManager,
+    MooncakeKVReceiver,
+    MooncakeKVSender,
+)

lib/python3.12/site-packages/sglang/srt/disaggregation/mooncake/conn.py

@@ -0,0 +1,752 @@
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import dataclasses
+import logging
+import os
+import queue
+import socket
+import struct
+import threading
+from functools import cache
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import numpy.typing as npt
+import requests
+import zmq
+from aiohttp import web
+
+from sglang.srt.disaggregation.base.conn import (
+    BaseKVBootstrapServer,
+    BaseKVManager,
+    BaseKVReceiver,
+    BaseKVSender,
+    KVArgs,
+    KVPoll,
+)
+from sglang.srt.disaggregation.mooncake.transfer_engine import MooncakeTransferEngine
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import get_free_port, get_ip, get_local_ip_by_remote
+
+logger = logging.getLogger(__name__)
+
+
+def group_concurrent_contiguous(
+    src_indices: npt.NDArray[np.int64], dst_indices: npt.NDArray[np.int64]
+) -> Tuple[List[npt.NDArray[np.int64]], List[npt.NDArray[np.int64]]]:
+    src_groups = []
+    dst_groups = []
+    current_src = [src_indices[0]]
+    current_dst = [dst_indices[0]]
+
+    for i in range(1, len(src_indices)):
+        src_contiguous = src_indices[i] == src_indices[i - 1] + 1
+        dst_contiguous = dst_indices[i] == dst_indices[i - 1] + 1
+        if src_contiguous and dst_contiguous:
+            current_src.append(src_indices[i])
+            current_dst.append(dst_indices[i])
+        else:
+            src_groups.append(current_src)
+            dst_groups.append(current_dst)
+            current_src = [src_indices[i]]
+            current_dst = [dst_indices[i]]
+
+    src_groups.append(current_src)
+    dst_groups.append(current_dst)
+
+    return src_groups, dst_groups
+
+
+@dataclasses.dataclass
+class TransferKVChunk:
+    room: int
+    prefill_kv_indices: npt.NDArray[np.int64]
+    index_slice: slice
+    is_last: bool
+    prefill_aux_index: Optional[int]
+
+
+@dataclasses.dataclass
+class TransferInfo:
+    room: int
+    endpoint: str
+    dst_port: int
+    mooncake_session_id: str
+    dst_kv_indices: npt.NDArray[np.int64]
+    dst_aux_index: int
+
+    @classmethod
+    def from_zmq(cls, msg: List[bytes]):
+        return cls(
+            room=int(msg[0].decode("ascii")),
+            endpoint=msg[1].decode("ascii"),
+            dst_port=int(msg[2].decode("ascii")),
+            mooncake_session_id=msg[3].decode("ascii"),
+            dst_kv_indices=np.frombuffer(msg[4], dtype=np.int64),
+            dst_aux_index=int(msg[5].decode("ascii")),
+        )
+
+
+@dataclasses.dataclass
+class KVArgsRegisterInfo:
+    room: str
+    endpoint: str
+    dst_port: int
+    mooncake_session_id: str
+    dst_kv_ptrs: list[int]
+    dst_aux_ptrs: list[int]
+
+    @classmethod
+    def from_zmq(cls, msg: List[bytes]):
+        return cls(
+            room=str(msg[0].decode("ascii")),
+            endpoint=msg[1].decode("ascii"),
+            dst_port=int(msg[2].decode("ascii")),
+            mooncake_session_id=msg[3].decode("ascii"),
+            dst_kv_ptrs=list(struct.unpack(f"{len(msg[4])//8}Q", msg[4])),
+            dst_aux_ptrs=list(struct.unpack(f"{len(msg[5])//8}Q", msg[5])),
+        )
+
+
+class MooncakeKVManager(BaseKVManager):
+    def __init__(
+        self,
+        args: KVArgs,
+        disaggregation_mode: DisaggregationMode,
+        server_args: ServerArgs,
+    ):
+        self.kv_args = args
+        self.engine = MooncakeTransferEngine(
+            hostname=get_local_ip_by_remote(),
+            gpu_id=self.kv_args.gpu_id,
+            ib_device=self.kv_args.ib_device,
+        )
+        self.disaggregation_mode = disaggregation_mode
+        # for p/d multi node infer
+        self.bootstrap_port = server_args.disaggregation_bootstrap_port
+        self.dist_init_addr = server_args.dist_init_addr
+        self.tp_size = server_args.tp_size
+        self.dp_size = server_args.dp_size
+        self.enable_dp_attention = server_args.enable_dp_attention
+        if not server_args.enable_dp_attention and server_args.dp_size != 1:
+            raise ValueError(
+                "If dp_attention is not enabled, dp size must be 1 in disaggregation mode."
+            )
+        self.request_status: Dict[int, KVPoll] = {}
+        self.rank_port = None
+        self.server_socket = zmq.Context().socket(zmq.PULL)
+        self.register_buffer_to_engine()
+        if self.disaggregation_mode == DisaggregationMode.PREFILL:
+            self.transfer_queue = queue.Queue()
+            self.transfer_infos: Dict[int, TransferInfo] = {}
+            self.decode_kv_args_table: Dict[str, KVArgsRegisterInfo] = {}
+            self.start_prefill_thread()
+            self._register_to_bootstrap()
+
+            # Determine the number of threads to use for kv sender
+            cpu_count = os.cpu_count()
+            self.executor = concurrent.futures.ThreadPoolExecutor(
+                min(cpu_count // 4, 16)
+            )
+        elif self.disaggregation_mode == DisaggregationMode.DECODE:
+            self.start_decode_thread()
+            self.connection_pool: Dict[str, Dict[str, Union[str, int]]] = {}
+            self.prefill_dp_size_table: Dict[str, int] = {}
+        else:
+            raise ValueError(
+                f"Unsupported DisaggregationMode: {self.disaggregation_mode}"
+            )
+
+    def register_buffer_to_engine(self):
+        for kv_data_ptr, kv_data_len in zip(
+            self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens
+        ):
+            self.engine.register(kv_data_ptr, kv_data_len)
+
+        for aux_data_ptr, aux_data_len in zip(
+            self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens
+        ):
+            self.engine.register(aux_data_ptr, aux_data_len)
+
+    @cache
+    def _connect(self, endpoint: str):
+        socket = zmq.Context().socket(zmq.PUSH)
+        socket.connect(endpoint)
+        return socket
+
+    def send_kvcache(
+        self,
+        mooncake_session_id: str,
+        prefill_kv_indices: npt.NDArray[np.int64],
+        dst_kv_ptrs: list[int],
+        dst_kv_indices: npt.NDArray[np.int64],
+    ):
+        # Group by indices
+        prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
+            prefill_kv_indices, dst_kv_indices
+        )
+
+        num_layers = len(self.kv_args.kv_data_ptrs)
+        layers_params = [
+            (
+                self.kv_args.kv_data_ptrs[layer_id],
+                dst_kv_ptrs[layer_id],
+                self.kv_args.kv_item_lens[layer_id],
+            )
+            for layer_id in range(num_layers)
+        ]
+
+        # Worker function for processing a single layer
+        def process_layer(src_ptr: int, dst_ptr: int, item_len: int) -> int:
+            for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks):
+                src_addr = src_ptr + int(prefill_index[0]) * item_len
+                dst_addr = dst_ptr + int(decode_index[0]) * item_len
+                length = item_len * len(prefill_index)
+
+                status = self.engine.transfer_sync(
+                    mooncake_session_id, src_addr, dst_addr, length
+                )
+                if status != 0:
+                    return status
+            return 0
+
+        futures = [
+            self.executor.submit(
+                process_layer,
+                src_ptr,
+                dst_ptr,
+                item_len,
+            )
+            for (src_ptr, dst_ptr, item_len) in layers_params
+        ]
+
+        for future in concurrent.futures.as_completed(futures):
+            status = future.result()
+            if status != 0:
+                # Immediate shutdown on first error (existing tasks will finish)
+                executor.shutdown(wait=False)
+                for f in futures:
+                    f.cancel()
+                return status
+
+        return 0
+
+    def send_aux(
+        self,
+        mooncake_session_id: str,
+        prefill_aux_index: int,
+        dst_aux_ptrs: list[int],
+        dst_aux_index: int,
+    ):
+        aux_item_len = self.kv_args.aux_item_lens[0]
+        prefill_aux_addr = (
+            self.kv_args.aux_data_ptrs[0] + prefill_aux_index * aux_item_len
+        )
+        decode_aux_addr = dst_aux_ptrs[0] + dst_aux_index * aux_item_len
+        # TODO: mooncake transfer engine can do async transfer. Do async later
+        # Not sure about the amount of aux data, maybe transfer it by zmq is more effective
+        status = self.engine.transfer_sync(
+            mooncake_session_id, prefill_aux_addr, decode_aux_addr, aux_item_len
+        )
+        return status
+
+    def sync_status_to_decode_endpoint(self, remote: str, dst_port: int, room: int):
+        if ":" in remote:
+            remote = remote.split(":")[0]
+        self._connect("tcp://" + remote + ":" + str(dst_port)).send_multipart(
+            [
+                str(room).encode("ascii"),
+                str(self.request_status[room]).encode("ascii"),
+            ]
+        )
+
+    def start_prefill_thread(self):
+        self.rank_port = get_free_port()
+        self.server_socket.bind(f"tcp://{get_local_ip_by_remote()}:{self.rank_port}")
+
+        def bootstrap_thread():
+            """This thread recvs pre-alloc notification from the decode engine"""
+            # KVPoll.Bootstrapping -> KVPoll.WaitingForInput
+            while True:
+                waiting_req_bytes = self.server_socket.recv_multipart()
+                room = waiting_req_bytes[0].decode("ascii")
+                if room == "None":
+                    mooncake_session_id = waiting_req_bytes[3].decode("ascii")
+                    self.decode_kv_args_table[mooncake_session_id] = (
+                        KVArgsRegisterInfo.from_zmq(waiting_req_bytes)
+                    )
+                    logger.debug(
+                        f"Register KVArgs from {mooncake_session_id} successfully"
+                    )
+                    continue
+                room = int(room)
+                self.transfer_infos[room] = TransferInfo.from_zmq(waiting_req_bytes)
+
+                # NOTE: after bootstrapping we can mark the req as waiting for input
+                self.request_status[room] = KVPoll.WaitingForInput
+
+        def transfer_thread():
+            # TODO: Shall we use KVPoll.Transferring state?
+            while True:
+                try:
+                    kv_chunk: TransferKVChunk = self.transfer_queue.get(timeout=0.01)
+                    req = self.transfer_infos[kv_chunk.room]
+                    chunked_dst_kv_indice = req.dst_kv_indices[kv_chunk.index_slice]
+                    assert len(chunked_dst_kv_indice) == len(
+                        kv_chunk.prefill_kv_indices
+                    ), f"len(chunked_dst_kv_indice) = {len(chunked_dst_kv_indice)}, len(kv_chunk.prefill_kv_indices) = {len(kv_chunk.prefill_kv_indices)}"
+
+                    ret = self.send_kvcache(
+                        req.mooncake_session_id,
+                        kv_chunk.prefill_kv_indices,
+                        self.decode_kv_args_table[req.mooncake_session_id].dst_kv_ptrs,
+                        chunked_dst_kv_indice,
+                    )
+                    if ret != 0:
+                        self.request_status[kv_chunk.room] = KVPoll.Failed
+                        self.sync_status_to_decode_endpoint(
+                            req.endpoint, req.dst_port, req.room
+                        )
+                        continue
+
+                    if kv_chunk.is_last:
+                        # Only the last chunk we need to send the aux data
+                        ret = self.send_aux(
+                            req.mooncake_session_id,
+                            kv_chunk.prefill_aux_index,
+                            self.decode_kv_args_table[
+                                req.mooncake_session_id
+                            ].dst_aux_ptrs,
+                            req.dst_aux_index,
+                        )
+                        self.request_status[req.room] = (
+                            KVPoll.Success if ret == 0 else KVPoll.Failed
+                        )
+                        self.sync_status_to_decode_endpoint(
+                            req.endpoint, req.dst_port, req.room
+                        )
+                        self.transfer_infos.pop(req.room)
+
+                except queue.Empty:
+                    continue
+
+        threading.Thread(target=bootstrap_thread).start()
+        threading.Thread(target=transfer_thread).start()
+
+    def start_decode_thread(self):
+        self.rank_port = get_free_port()
+        self.server_socket.bind(f"tcp://{get_local_ip_by_remote()}:{self.rank_port}")
+
+        def decode_thread():
+            while True:
+                (bootstrap_room, status) = self.server_socket.recv_multipart()
+                status = int(status.decode("ascii"))
+                bootstrap_room = int(bootstrap_room.decode("ascii"))
+                self.request_status[bootstrap_room] = status
+
+        threading.Thread(target=decode_thread).start()
+
+    def add_transfer_request(
+        self,
+        bootstrap_room: int,
+        kv_indices: npt.NDArray[np.int64],
+        index_slice: slice,
+        is_last: bool,
+        aux_index: Optional[int] = None,
+    ):
+        assert self.disaggregation_mode == DisaggregationMode.PREFILL
+        assert not is_last or (is_last and aux_index is not None)
+
+        self.transfer_queue.put(
+            TransferKVChunk(
+                room=bootstrap_room,
+                prefill_kv_indices=kv_indices,
+                index_slice=index_slice,
+                is_last=is_last,
+                prefill_aux_index=aux_index,
+            )
+        )
+        self.request_status[bootstrap_room] = KVPoll.WaitingForInput
+
+    def check_status(self, bootstrap_room: int):
+        # TOOD: do we really need the poll()?
+
+        return self.request_status[bootstrap_room]
+
+    def update_status(self, bootstrap_room: int, status: KVPoll):
+        if bootstrap_room not in self.request_status:
+            self.request_status[bootstrap_room] = status
+        else:
+            # NOTE: The prefill engine could recv bootstrapping first
+            self.request_status[bootstrap_room] = max(
+                self.request_status[bootstrap_room], status
+            )
+
+    def get_session_id(self):
+        return self.engine.get_session_id()
+
+    def _register_to_bootstrap(self):
+        """Register KVSender to bootstrap server via HTTP POST."""
+        if self.dist_init_addr:
+            ip_address = socket.gethostbyname(self.dist_init_addr.split(":")[0])
+        else:
+            ip_address = get_ip()
+
+        bootstrap_server_url = f"{ip_address}:{self.bootstrap_port}"
+        url = f"http://{bootstrap_server_url}/route"
+        payload = {
+            "role": "Prefill",
+            "tp_size": self.tp_size,
+            "dp_size": self.dp_size,
+            "rank_ip": get_local_ip_by_remote(),
+            "rank_port": self.rank_port,
+            "engine_rank": self.kv_args.engine_rank,
+        }
+
+        try:
+            response = requests.put(url, json=payload)
+            if response.status_code == 200:
+                logger.debug("Prefill successfully registered to bootstrap server.")
+            else:
+                logger.error(
+                    f"Prefill Failed to connect to bootstrap server: {response.status_code}, {response.text}"
+                )
+        except Exception as e:
+            logger.error(f"Prefill Failed to register to bootstrap server: {e}")
+
+
+class MooncakeKVSender(BaseKVSender):
+
+    def __init__(
+        self, mgr: MooncakeKVManager, bootstrap_addr: str, bootstrap_room: int
+    ):
+        self.kv_mgr = mgr
+        self.bootstrap_room = bootstrap_room
+        self.kv_mgr.update_status(bootstrap_room, KVPoll.Bootstrapping)
+        self.aux_index = None
+        self.bootstrap_server_url = bootstrap_addr
+        self.session_id = self.kv_mgr.get_session_id()
+
+    def init(self, num_kv_indices: int, aux_index: Optional[int] = None):
+        self.num_kv_indices = num_kv_indices
+        self.aux_index = aux_index
+
+    def send(
+        self,
+        kv_indices: npt.NDArray[np.int64],
+        index_slice: slice,
+        is_last: bool,
+    ):
+        if not is_last:
+            self.kv_mgr.add_transfer_request(
+                self.bootstrap_room, kv_indices, index_slice, False
+            )
+        else:
+            self.kv_mgr.add_transfer_request(
+                self.bootstrap_room,
+                kv_indices,
+                index_slice,
+                True,
+                aux_index=self.aux_index,
+            )
+
+    def poll(self) -> KVPoll:
+        return self.kv_mgr.check_status(self.bootstrap_room)
+
+    def failure_exception(self):
+        raise Exception("Fake KVSender Exception")
+
+
+class MooncakeKVReceiver(BaseKVReceiver):
+    _ctx = zmq.Context()
+    _socket_cache = {}
+    _socket_locks = {}
+    _global_lock = threading.Lock()
+
+    def __init__(
+        self,
+        mgr: MooncakeKVManager,
+        bootstrap_addr: str,
+        bootstrap_room: Optional[int] = None,
+    ):
+        self.bootstrap_room = bootstrap_room
+        self.bootstrap_addr = bootstrap_addr
+        self.kv_mgr = mgr
+        self.session_id = self.kv_mgr.get_session_id()
+        self.kv_mgr.update_status(bootstrap_room, KVPoll.Bootstrapping)
+
+        if not self.kv_mgr.enable_dp_attention:
+            # We assume dp_attention should be activated simultaneously for
+            # both prefill role and decode role. If the decode instance does
+            # not enable dp_attention, then dp_attention is not enabled on the
+            # prefill instance as well. Therefore, we should skip questioning
+            # the prefill dp size to reduce bootstrap overhead.
+            self.prefill_dp_size = 1
+        elif self.bootstrap_addr not in self.kv_mgr.prefill_dp_size_table:
+            self.prefill_dp_size, tp_size_per_dp_rank = (
+                self._get_prefill_dp_size_from_server()
+            )
+            # Currently, we don't allow prefill instance and decode instance to
+            # have different TP sizes per DP rank.
+            assert tp_size_per_dp_rank == self.kv_mgr.tp_size // self.kv_mgr.dp_size
+            if self.prefill_dp_size is None:
+                logger.error(
+                    f"Could not fetch prefill dp_size for bootstrap_addr: {self.bootstrap_addr}"
+                )
+            else:
+                self.kv_mgr.prefill_dp_size_table[self.bootstrap_addr] = (
+                    self.prefill_dp_size
+                )
+        else:
+            self.prefill_dp_size = self.kv_mgr.prefill_dp_size_table[
+                self.bootstrap_addr
+            ]
+
+        # NOTE: key distinguished by bootstrap_addr and engine_rank
+        self.target_dp_group = bootstrap_room % self.prefill_dp_size
+        bootstrap_key = f"{self.bootstrap_addr}_{self.kv_mgr.kv_args.engine_rank}"
+
+        if bootstrap_key not in self.kv_mgr.connection_pool:
+            self.bootstrap_info = self._get_bootstrap_info_from_server(
+                self.kv_mgr.kv_args.engine_rank,
+                self.target_dp_group,
+            )
+            if self.bootstrap_info is None:
+                logger.error(
+                    f"Could not fetch bootstrap info for engine rank: {self.kv_mgr.kv_args.engine_rank}"
+                )
+            else:
+                self.kv_mgr.connection_pool[bootstrap_key] = self.bootstrap_info
+                # Register kv_args only once to prefill KVManager according to the info fetched from the bootstrap server
+                self._register_kv_args()
+        else:
+            self.bootstrap_info = self.kv_mgr.connection_pool[bootstrap_key]
+
+        assert self.bootstrap_info is not None
+        self.kv_mgr.update_status(bootstrap_room, KVPoll.WaitingForInput)
+
+    def _get_bootstrap_info_from_server(self, engine_rank, target_dp_group):
+        """Fetch the bootstrap info from the bootstrap server."""
+        try:
+            url = f"http://{self.bootstrap_addr}/route?engine_rank={engine_rank}&target_dp_group={target_dp_group}"
+            response = requests.get(url)
+            if response.status_code == 200:
+                bootstrap_info = response.json()
+                return bootstrap_info
+            else:
+                logger.error(
+                    f"Failed to get prefill server info: {response.status_code}, {response.text}"
+                )
+                return None
+        except Exception as e:
+            logger.error(f"Error fetching prefill info from bootstrap: {e}")
+            return None
+
+    def _get_prefill_dp_size_from_server(self) -> int:
+        """Fetch the prefill parallel info from the bootstrap server."""
+        try:
+            url = f"http://{self.bootstrap_addr}/route?engine_rank={-1}&target_dp_group={-1}"
+            response = requests.get(url)
+            if response.status_code == 200:
+                prefill_parallel_info = response.json()
+                return int(prefill_parallel_info["prefill_dp_size"]), int(
+                    prefill_parallel_info["tp_size_per_dp_rank"]
+                )
+            else:
+                logger.error(
+                    f"Failed to get prefill parallel info: {response.status_code}, {response.text}"
+                )
+                return None
+        except Exception as e:
+            logger.error(f"Error fetching prefill parallel info from bootstrap: {e}")
+            return None
+
+    def _register_kv_args(self):
+        self.prefill_server_url = (
+            f"{self.bootstrap_info['rank_ip']}:{self.bootstrap_info['rank_port']}"
+        )
+
+        packed_kv_data_ptrs = b"".join(
+            struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.kv_data_ptrs
+        )
+        packed_aux_data_ptrs = b"".join(
+            struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.aux_data_ptrs
+        )
+        sock, lock = self._connect("tcp://" + self.prefill_server_url)
+        with lock:
+            sock.send_multipart(
+                [
+                    "None".encode("ascii"),
+                    get_local_ip_by_remote().encode("ascii"),
+                    str(self.kv_mgr.rank_port).encode("ascii"),
+                    self.session_id.encode("ascii"),
+                    packed_kv_data_ptrs,
+                    packed_aux_data_ptrs,
+                ]
+            )
+
+    @classmethod
+    def _connect(cls, endpoint: str):
+        with cls._global_lock:
+            if endpoint not in cls._socket_cache:
+                sock = cls._ctx.socket(zmq.PUSH)
+                sock.connect(endpoint)
+                cls._socket_cache[endpoint] = sock
+                cls._socket_locks[endpoint] = threading.Lock()
+            return cls._socket_cache[endpoint], cls._socket_locks[endpoint]
+
+    def init(self, kv_indices: npt.NDArray[np.int64], aux_index: Optional[int] = None):
+        self.prefill_server_url = (
+            f"{self.bootstrap_info['rank_ip']}:{self.bootstrap_info['rank_port']}"
+        )
+        logger.debug(
+            f"Fetched bootstrap info: {self.bootstrap_info} for engine rank: {self.kv_mgr.kv_args.engine_rank}"
+        )
+
+        sock, lock = self._connect("tcp://" + self.prefill_server_url)
+        with lock:
+            sock.send_multipart(
+                [
+                    str(self.bootstrap_room).encode("ascii"),
+                    get_local_ip_by_remote().encode("ascii"),
+                    str(self.kv_mgr.rank_port).encode("ascii"),
+                    self.session_id.encode("ascii"),
+                    kv_indices.tobytes(),
+                    str(aux_index).encode("ascii"),
+                ]
+            )
+
+    def poll(self) -> KVPoll:
+        return self.kv_mgr.check_status(self.bootstrap_room)
+
+    def failure_exception(self):
+        raise Exception("Fake KVReceiver Exception")
+
+
+class MooncakeKVBootstrapServer(BaseKVBootstrapServer):
+    def __init__(self, port: int):
+        self.port = port
+        self.app = web.Application()
+        self.store = dict()
+        self.lock = asyncio.Lock()
+        self._setup_routes()
+        self.dp_size = None
+        self.tp_size_per_dp_rank = None
+        self.prefill_port_table: Dict[int, Dict[int, Dict[str, Union[str, int]]]] = {}
+
+        # Start bootstrap server
+        self.thread = threading.Thread(target=self._run_server, daemon=True)
+        self.run()
+
+    def run(self):
+        self.thread.start()
+
+    def _setup_routes(self):
+        self.app.router.add_route("*", "/route", self._handle_route)
+
+    async def _handle_route(self, request: web.Request):
+        method = request.method
+        if method == "PUT":
+            return await self._handle_route_put(request)
+        elif method == "GET":
+            return await self._handle_route_get(request)
+        else:
+            return web.Response(
+                text="Method not allowed", status=405, content_type="application/json"
+            )
+
+    async def _handle_route_put(self, request: web.Request):
+        data = await request.json()
+        role = data["role"]
+        tp_size = data["tp_size"]
+        dp_size = data["dp_size"]
+        rank_ip = data["rank_ip"]
+        rank_port = int(data["rank_port"])
+        engine_rank = int(data["engine_rank"])
+
+        if self.dp_size is None:
+            self.dp_size = dp_size
+
+        tp_size_per_dp_rank = tp_size // dp_size
+        if self.tp_size_per_dp_rank == None:
+            self.tp_size_per_dp_rank = tp_size_per_dp_rank
+
+        # Add lock to make sure thread-safe
+        if role == "Prefill":
+            dp_group = engine_rank // tp_size_per_dp_rank
+            tp_rank_in_dp_group = engine_rank % tp_size_per_dp_rank
+
+            async with self.lock:
+                if dp_group not in self.prefill_port_table:
+                    self.prefill_port_table[dp_group] = {}
+
+            self.prefill_port_table[dp_group][tp_rank_in_dp_group] = {
+                "rank_ip": rank_ip,
+                "rank_port": rank_port,
+            }
+            logger.debug(
+                f"Register Prefill bootstrap: {engine_rank} with rank_ip: {rank_ip} and rank_port: {rank_port}"
+            )
+
+        return web.Response(text="OK", status=200)
+
+    async def _handle_route_get(self, request: web.Request):
+        engine_rank = request.query.get("engine_rank")
+        target_dp_group = request.query.get("target_dp_group")
+        if not engine_rank or not target_dp_group:
+            return web.Response(text="Missing inputs for bootstrap server.", status=400)
+
+        # Currently we use engine_rank == -1 and target_dp_group == -1 to sync dp size
+        if int(engine_rank) == -1 and int(target_dp_group) == -1:
+            prefill_parallel_info = {
+                "prefill_dp_size": self.dp_size,
+                "tp_size_per_dp_rank": self.tp_size_per_dp_rank,
+            }
+            return web.json_response(prefill_parallel_info, status=200)
+
+        # Find corresponding prefill info
+        tp_rank_in_dp_group = int(engine_rank) % self.tp_size_per_dp_rank
+
+        async with self.lock:
+            bootstrap_info = self.prefill_port_table[int(target_dp_group)][
+                tp_rank_in_dp_group
+            ]
+
+        if bootstrap_info is not None:
+            return web.json_response(bootstrap_info, status=200)
+        else:
+            return web.Response(text="Bootstrap info not Found", status=404)
+
+    def _run_server(self):
+        try:
+            # Event Loop
+            self._loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(self._loop)
+
+            self._runner = web.AppRunner(self.app)
+            self._loop.run_until_complete(self._runner.setup())
+
+            site = web.TCPSite(self._runner, port=self.port)
+            self._loop.run_until_complete(site.start())
+            self._loop.run_forever()
+        except Exception as e:
+            logger.error(f"Server error: {str(e)}")
+        finally:
+            # Cleanup
+            self._loop.run_until_complete(self._runner.cleanup())
+            self._loop.close()
+
+    def close(self):
+        """Shutdown"""
+        if self._loop is not None and self._loop.is_running():
+            self._loop.call_soon_threadsafe(self._loop.stop)
+            logger.info("Stopping server loop...")
+
+        if self.thread.is_alive():
+            self.thread.join(timeout=2)
+            logger.info("Server thread stopped")
+
+    def poll(self) -> KVPoll: ...

lib/python3.12/site-packages/sglang/srt/disaggregation/mooncake/transfer_engine.py

@@ -0,0 +1,77 @@
+import json
+import logging
+from dataclasses import dataclass
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+class MooncakeTransferEngine:
+
+    def __init__(self, hostname: str, gpu_id: int, ib_device: Optional[str] = None):
+        try:
+            from mooncake.engine import TransferEngine
+        except ImportError as e:
+            raise ImportError(
+                "Please install mooncake by following the instructions at "
+                "https://github.com/kvcache-ai/Mooncake/blob/main/doc/en/build.md "  # noqa: E501
+                "to run SGLang with MooncakeTransferEngine."
+            ) from e
+
+        self.engine = TransferEngine()
+        self.hostname = hostname
+        self.gpu_id = gpu_id
+        self.ib_device = ib_device
+
+        self.initialize(
+            hostname=self.hostname,
+            device_name=self.ib_device,
+        )
+        self.session_id = f"{self.hostname}:{self.engine.get_rpc_port()}"
+
+    def register(self, ptr, length):
+        ret_value = self.engine.register_memory(ptr, length)
+        if ret_value != 0:
+            logger.error("Mooncake memory registration failed.")
+            raise RuntimeError("Mooncake memory registration failed.")
+
+    def deregister(self, ptr):
+        ret_value = self.engine.unregister_memory(ptr)
+        if ret_value != 0:
+            logger.error("Mooncake memory deregistration failed.")
+            raise RuntimeError("Mooncake memory deregistration failed.")
+
+    def initialize(
+        self,
+        hostname: str,
+        device_name: Optional[str],
+    ) -> None:
+        """Initialize the mooncake instance."""
+        ret_value = self.engine.initialize(
+            hostname,
+            "P2PHANDSHAKE",
+            "rdma",
+            device_name if device_name is not None else "",
+        )
+        if ret_value != 0:
+            logger.error("Mooncake Transfer Engine initialization failed.")
+            raise RuntimeError("Mooncake Transfer Engine initialization failed.")
+
+    def transfer_sync(
+        self, session_id: str, buffer: int, peer_buffer_address: int, length: int
+    ) -> int:
+        """Synchronously transfer data to the specified address."""
+
+        ret = self.engine.transfer_sync_write(
+            session_id, buffer, peer_buffer_address, length
+        )
+        if ret < 0:
+            logger.error("Mooncake Transfer Engine Return Error.")
+            raise RuntimeError("Mooncake Transfer Engine Return Error.")
+        return ret
+
+    def get_localhost(self):
+        return self.hostname
+
+    def get_session_id(self):
+        return self.session_id

lib/python3.12/site-packages/sglang/srt/disaggregation/prefill.py

@@ -0,0 +1,417 @@
+"""
+Life cycle of a request in the prefill server
+
+1. Bootstrap Queue
+    a. Initialize a sender for each request
+    b. Use the queue to store requests whose bootstrap (handshake and preallocation) has not finished
+    c. Poll senders to check bootstrap state
+    d. Once bootstrap is complete, move request to Waiting Queue
+
+2. Waiting Queue
+    a. Use PrefillAdder to pop requests
+    b. Run forward
+    c. Add the request to Inflight Queue
+
+3. Inflight Queue
+    a. Poll (non-blocking) the sender of the request
+    b. Once the transfer has finished, return the request
+"""
+
+from __future__ import annotations
+
+import logging
+import threading
+from collections import deque
+from typing import TYPE_CHECKING, List, Optional
+
+import torch
+
+from sglang.srt.disaggregation.base import BaseKVManager, KVArgs, KVPoll
+from sglang.srt.disaggregation.utils import (
+    DisaggregationMode,
+    FakeBootstrapHost,
+    KVClassType,
+    ReqToMetadataIdxAllocator,
+    TransferBackend,
+    get_kv_class,
+    kv_to_page_indices,
+    kv_to_page_num,
+    poll_and_all_reduce,
+)
+from sglang.srt.managers.schedule_batch import FINISH_LENGTH, Req, ScheduleBatch
+
+if TYPE_CHECKING:
+    from torch.distributed import ProcessGroup
+
+    from sglang.srt.managers.scheduler import GenerationBatchResult, Scheduler
+    from sglang.srt.mem_cache.memory_pool import KVCache
+
+
+logger = logging.getLogger(__name__)
+
+
+class PrefillBootstrapQueue:
+    """
+    Store the requests in bootstrapping
+    """
+
+    def __init__(
+        self,
+        token_to_kv_pool: KVCache,
+        req_to_metadata_buffer_idx_allocator: ReqToMetadataIdxAllocator,
+        metadata_buffers: List[torch.Tensor],
+        aux_dtype: torch.dtype,
+        tp_rank: int,
+        tp_size: int,
+        bootstrap_port: int,
+        gloo_group: ProcessGroup,
+        transfer_backend: TransferBackend,
+        scheduler: Scheduler,
+    ):
+        self.token_to_kv_pool = token_to_kv_pool
+        self.aux_dtype = aux_dtype
+
+        self.metadata_buffers = metadata_buffers
+        self.req_to_metadata_buffer_idx_allocator = req_to_metadata_buffer_idx_allocator
+        self.tp_rank = tp_rank
+        self.tp_size = tp_size
+        self.transfer_backend = transfer_backend
+        self.scheduler = scheduler
+        self.kv_manager = self._init_kv_manager()
+        self.queue: List[Req] = []
+        self.gloo_group = gloo_group
+        self.bootstrap_port = bootstrap_port
+
+    def store_prefill_results(self, idx: int, token_id: int):
+        assert token_id >= 0, f"token_id: {token_id} is negative"
+        output_id_buffer = self.metadata_buffers[0]
+        output_id_buffer[idx] = token_id
+
+    def _init_kv_manager(self) -> BaseKVManager:
+        kv_args = KVArgs()
+        kv_args.engine_rank = self.tp_rank
+        kv_data_ptrs, kv_data_lens, kv_item_lens = (
+            self.token_to_kv_pool.get_contiguous_buf_infos()
+        )
+
+        kv_args.kv_data_ptrs = kv_data_ptrs
+        kv_args.kv_data_lens = kv_data_lens
+        kv_args.kv_item_lens = kv_item_lens
+
+        # Define req -> input ids buffer
+        kv_args.aux_data_ptrs = [
+            metadata_buffer.data_ptr() for metadata_buffer in self.metadata_buffers
+        ]
+        kv_args.aux_data_lens = [
+            metadata_buffer.nbytes for metadata_buffer in self.metadata_buffers
+        ]
+        kv_args.aux_item_lens = [
+            metadata_buffer[0].nbytes for metadata_buffer in self.metadata_buffers
+        ]
+        kv_args.ib_device = self.scheduler.server_args.disaggregation_ib_device
+        kv_args.gpu_id = self.scheduler.gpu_id
+        kv_manager_class = get_kv_class(self.transfer_backend, KVClassType.MANAGER)
+        kv_manager = kv_manager_class(
+            kv_args, DisaggregationMode.PREFILL, self.scheduler.server_args
+        )
+        return kv_manager
+
+    def add(self, req: Req) -> None:
+        if req.bootstrap_host == FakeBootstrapHost:
+            # Fake transfer for warmup reqs
+            kv_sender_class = get_kv_class(TransferBackend.FAKE, KVClassType.SENDER)
+        else:
+            kv_sender_class = get_kv_class(self.transfer_backend, KVClassType.SENDER)
+        req.disagg_kv_sender = kv_sender_class(
+            mgr=self.kv_manager,
+            bootstrap_addr=f"{req.bootstrap_host}:{self.bootstrap_port}",
+            bootstrap_room=req.bootstrap_room,
+        )
+        self._process_req(req)
+        self.queue.append(req)
+
+    def _process_req(self, req: Req) -> None:
+        """
+        Set max_new_tokens = 1, so PrefillAdder memory estimation is accurate
+        """
+        req.sampling_params.max_new_tokens = 1
+
+    def pop_bootstrapped(self) -> List[Req]:
+        """pop the reqs which has finished bootstrapping"""
+        bootstrapped_reqs = []
+        indices_to_remove = set()
+
+        if len(self.queue) == 0:
+            return []
+
+        polls = poll_and_all_reduce(
+            [req.disagg_kv_sender for req in self.queue], self.gloo_group
+        )
+
+        for i, (req, poll) in enumerate(zip(self.queue, polls)):
+            if poll == KVPoll.Bootstrapping:
+                continue
+            elif poll == KVPoll.Failed:
+                raise Exception("Bootstrap failed")
+
+            # KV.WaitingForInput
+            num_kv_indices = len(req.origin_input_ids)
+            if self.req_to_metadata_buffer_idx_allocator.available_size() == 0:
+                break
+
+            req.metadata_buffer_index = (
+                self.req_to_metadata_buffer_idx_allocator.alloc()
+            )
+            assert req.metadata_buffer_index is not None
+            num_pages = kv_to_page_num(num_kv_indices, self.token_to_kv_pool.page_size)
+            req.disagg_kv_sender.init(num_pages, req.metadata_buffer_index)
+
+            bootstrapped_reqs.append(req)
+            indices_to_remove.add(i)
+
+        self.queue = [
+            entry for i, entry in enumerate(self.queue) if i not in indices_to_remove
+        ]
+
+        return bootstrapped_reqs
+
+
+class SchedulerDisaggregationPrefillMixin:
+    """
+    Mixin for Scheduler to handle disaggregation prefill
+    """
+
+    @torch.no_grad()
+    def event_loop_normal_disagg_prefill(self: Scheduler):
+        """A normal scheduler loop for prefill worker in disaggregation mode."""
+
+        while True:
+            recv_reqs = self.recv_requests()
+            self.process_input_requests(recv_reqs)
+            self.waiting_queue.extend(
+                self.disagg_prefill_bootstrap_queue.pop_bootstrapped()
+            )
+            self.process_prefill_chunk()
+            batch = self.get_new_batch_prefill()
+
+            # Handle DP attention
+            if (
+                self.server_args.enable_dp_attention
+                or self.server_args.enable_sp_layernorm
+            ):
+                batch, _ = self.prepare_dp_attn_batch(batch)
+
+            self.cur_batch = batch
+
+            if batch:
+                result = self.run_batch(batch)
+                self.process_batch_result_disagg_prefill(batch, result)
+
+            if len(self.disagg_prefill_inflight_queue) > 0:
+                self.process_disagg_prefill_inflight_queue()
+
+            if batch is None and len(self.disagg_prefill_inflight_queue) == 0:
+                self.check_memory()
+                self.new_token_ratio = self.init_new_token_ratio
+
+            self.last_batch = batch
+            # HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it
+            # Otherwise, it hangs under high concurrency
+            self.running_batch.batch_is_full = False
+
+    @torch.no_grad()
+    def event_loop_overlap_disagg_prefill(self: Scheduler):
+        self.result_queue = deque()
+
+        while True:
+            recv_reqs = self.recv_requests()
+            self.process_input_requests(recv_reqs)
+            self.waiting_queue.extend(
+                self.disagg_prefill_bootstrap_queue.pop_bootstrapped()
+            )
+            self.process_prefill_chunk()
+            batch = self.get_new_batch_prefill()
+
+            # Handle DP attention
+            if (
+                self.server_args.enable_dp_attention
+                or self.server_args.enable_sp_layernorm
+            ):
+                batch, _ = self.prepare_dp_attn_batch(batch)
+
+            self.cur_batch = batch
+
+            if batch:
+                result = self.run_batch(batch)
+                self.result_queue.append((batch.copy(), result))
+
+            if self.last_batch:
+                tmp_batch, tmp_result = self.result_queue.popleft()
+                self.process_batch_result_disagg_prefill(tmp_batch, tmp_result)
+
+            if len(self.disagg_prefill_inflight_queue) > 0:
+                self.process_disagg_prefill_inflight_queue()
+
+            if batch is None and len(self.disagg_prefill_inflight_queue) == 0:
+                self.check_memory()
+                self.new_token_ratio = self.init_new_token_ratio
+
+            self.last_batch = batch
+            # HACK (byronhsu): reset the batch_is_full flag because we never enter update_running_batch which resets it
+            # Otherwise, it hangs under high concurrency
+            self.running_batch.batch_is_full = False
+
+    def process_batch_result_disagg_prefill(
+        self: Scheduler,
+        batch: ScheduleBatch,
+        result: GenerationBatchResult,
+        launch_done: Optional[threading.Event] = None,
+    ) -> None:
+        """
+        Transfer kv for prefill completed requests and add it into disagg_prefill_inflight_queue
+        Adapted from process_batch_result_prefill
+        """
+
+        (
+            logits_output,
+            next_token_ids,
+            extend_input_len_per_req,
+            extend_logprob_start_len_per_req,
+            bid,
+        ) = (
+            result.logits_output,
+            result.next_token_ids,
+            result.extend_input_len_per_req,
+            result.extend_logprob_start_len_per_req,
+            result.bid,
+        )
+
+        # Transfer kv for prefill completed requests and add it into disagg_prefill_infight_queue
+        if self.enable_overlap:
+            # wait
+            _, next_token_ids = self.tp_worker.resolve_last_batch_result(launch_done)
+        else:
+            next_token_ids = result.next_token_ids.tolist()
+
+        for req, next_token_id in zip(batch.reqs, next_token_ids, strict=True):
+            req: Req
+            if req.is_chunked <= 0:
+                # There is no output_ids for prefill
+                req.output_ids.append(next_token_id)
+                self.tree_cache.cache_unfinished_req(req)  # update the tree and lock
+                self.send_kv_chunk(req, token_id=next_token_id)
+                self.disagg_prefill_inflight_queue.append(req)
+            else:
+                # being chunked reqs' prefill is not finished
+                req.is_chunked -= 1
+
+                if self.enable_overlap:
+                    self.send_kv_chunk(req, end_idx=req.tmp_end_idx)
+
+    def process_disagg_prefill_inflight_queue(self: Scheduler) -> None:
+        """
+        Poll the requests in the middle of transfer. If done, return the request.
+        """
+        assert len(self.disagg_prefill_inflight_queue) > 0
+
+        done_reqs = []
+
+        polls = poll_and_all_reduce(
+            [req.disagg_kv_sender for req in self.disagg_prefill_inflight_queue],
+            self.attn_tp_cpu_group,
+        )
+
+        undone_reqs: List[Req] = []
+        # Check .poll() for the reqs in disagg_prefill_inflight_queue. If Success, respond to the client and remove it from the queue
+        for req, poll in zip(self.disagg_prefill_inflight_queue, polls):
+            if poll in [KVPoll.WaitingForInput, KVPoll.Transferring]:
+                undone_reqs.append(req)
+            elif poll == KVPoll.Success:  # transfer done
+                self.tree_cache.cache_finished_req(req)  # unlock the tree
+                req.finished_reason = FINISH_LENGTH(length=0)
+                # FIXME: clean up req's data in transfer engine
+                done_reqs.append(req)
+            elif poll == KVPoll.Failed:
+                raise Exception("Transferring failed")
+
+        for req in done_reqs:
+            self.disagg_prefill_bootstrap_queue.req_to_metadata_buffer_idx_allocator.free(
+                req.metadata_buffer_index
+            )
+
+        # Stream requests which have finished transfer
+        self.stream_output(done_reqs, False, None)
+
+        self.disagg_prefill_inflight_queue = undone_reqs
+
+    def process_prefill_chunk(self: Scheduler) -> None:
+        if self.last_batch and self.last_batch.forward_mode.is_extend():
+            if self.chunked_req:
+                # Move the chunked request out of the batch so that we can merge
+                # only finished requests to running_batch.
+                self.last_batch.filter_batch(chunked_req_to_exclude=self.chunked_req)
+                self.tree_cache.cache_unfinished_req(self.chunked_req)
+                if self.enable_overlap:
+                    # Delay KV transfer to process_batch_result_disagg_prefill when overlap is enabled to ensure results are resolved
+                    self.chunked_req.tmp_end_idx = min(
+                        len(self.chunked_req.fill_ids),
+                        len(self.chunked_req.origin_input_ids),
+                    )
+                else:
+                    self.send_kv_chunk(self.chunked_req)
+                # chunked request keeps its rid but will get a new req_pool_idx
+                self.req_to_token_pool.free(self.chunked_req.req_pool_idx)
+                self.running_batch.batch_is_full = False
+
+    def send_kv_chunk(
+        self: Scheduler,
+        req: Req,
+        token_id: Optional[int] = None,
+        end_idx: Optional[int] = None,
+    ) -> None:
+        """
+        Send a prefilled chunk to the decode server
+        """
+        page_size = self.token_to_kv_pool_allocator.page_size
+        start_idx = req.start_send_idx
+        # if end_idx is specified, use it as the end index of the kv chunk because in overlap schedule,
+        # the resolved length is not the same as fill_ids's length
+        end_idx = (
+            end_idx
+            if end_idx is not None
+            else min(len(req.fill_ids), len(req.origin_input_ids))
+        )
+        last_chunk = token_id is not None
+
+        if (not last_chunk) and (
+            end_idx % page_size != 0
+        ):  # todo: remove the second condition
+            # if not the last chunk and the last page is partial, delay the last partial page to the next send
+            end_idx = end_idx - end_idx % page_size
+
+        # Update next start_send_idx
+        req.start_send_idx = end_idx
+
+        kv_indices = (
+            self.req_to_token_pool.req_to_token[req.req_pool_idx, start_idx:end_idx]
+            .cpu()
+            .numpy()
+        )
+        if last_chunk is True:
+            self.disagg_prefill_bootstrap_queue.store_prefill_results(
+                req.metadata_buffer_index, token_id
+            )
+        page_indices = kv_to_page_indices(kv_indices, page_size)
+
+        page_start_idx = start_idx // page_size
+        page_end_idx = page_start_idx + len(page_indices)
+
+        if len(page_indices) == 0:
+            logger.info(
+                f"Skip sending kv chunk for request {req.rid=} {req.bootstrap_room=} because page_indices is empty"
+            )
+            return
+
+        req.disagg_kv_sender.send(
+            page_indices, slice(page_start_idx, page_end_idx), last_chunk
+        )

lib/python3.12/site-packages/sglang/srt/disaggregation/utils.py

@@ -0,0 +1,121 @@
+from __future__ import annotations
+
+from collections import deque
+from enum import Enum
+from typing import List
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+
+class DisaggregationMode(Enum):
+    NULL = "null"
+    PREFILL = "prefill"
+    DECODE = "decode"
+
+
+FakeBootstrapHost = "2.2.2.2"
+
+
+def poll_and_all_reduce(pollers, gloo_group):
+    polls = [int(poller.poll()) for poller in pollers]
+    tensor_to_reduce = torch.tensor(polls, dtype=torch.uint8, device="cpu")
+    dist.all_reduce(tensor_to_reduce, op=dist.ReduceOp.MIN, group=gloo_group)
+    return tensor_to_reduce.tolist()
+
+
+class ReqToMetadataIdxAllocator:
+    """A memory pool that maps a request to its first output token location."""
+
+    def __init__(
+        self,
+        size: int,
+    ):
+        self.size = size
+        self.free_slots = deque(list(range(size)))
+
+    def available_size(self):
+        return len(self.free_slots)
+
+    def alloc(self) -> List[int]:
+        if len(self.free_slots) == 0:
+            return None
+
+        return self.free_slots.popleft()
+
+    def free(self, free_index: int):
+        self.free_slots.append(free_index)
+
+
+class TransferBackend(Enum):
+    MOONCAKE = "mooncake"
+    NIXL = "nixl"
+    FAKE = "fake"
+
+
+class KVClassType(Enum):
+    MANAGER = "manager"
+    SENDER = "sender"
+    RECEIVER = "receiver"
+    BOOTSTRAP_SERVER = "bootstrap_server"
+
+
+def get_kv_class(transfer_backend: TransferBackend, class_type: KVClassType):
+    from sglang.srt.disaggregation.fake import FakeKVReceiver, FakeKVSender
+
+    if transfer_backend == TransferBackend.MOONCAKE:
+        from sglang.srt.disaggregation.mooncake import (
+            MooncakeKVBootstrapServer,
+            MooncakeKVManager,
+            MooncakeKVReceiver,
+            MooncakeKVSender,
+        )
+
+        class_mapping = {
+            KVClassType.MANAGER: MooncakeKVManager,
+            KVClassType.SENDER: MooncakeKVSender,
+            KVClassType.RECEIVER: (MooncakeKVReceiver),
+            KVClassType.BOOTSTRAP_SERVER: MooncakeKVBootstrapServer,
+        }
+        return class_mapping.get(class_type)
+    if transfer_backend == TransferBackend.NIXL:
+        from sglang.srt.disaggregation.nixl import (
+            NixlKVBootstrapServer,
+            NixlKVManager,
+            NixlKVReceiver,
+            NixlKVSender,
+        )
+
+        class_mapping = {
+            KVClassType.MANAGER: NixlKVManager,
+            KVClassType.SENDER: NixlKVSender,
+            KVClassType.RECEIVER: (NixlKVReceiver),
+            KVClassType.BOOTSTRAP_SERVER: NixlKVBootstrapServer,
+        }
+        return class_mapping.get(class_type)
+    if transfer_backend == TransferBackend.FAKE:
+        from sglang.srt.disaggregation.fake import FakeKVReceiver, FakeKVSender
+
+        class_mapping = {
+            KVClassType.SENDER: FakeKVSender,
+            KVClassType.RECEIVER: (FakeKVReceiver),
+        }
+        return class_mapping.get(class_type)
+
+    raise ValueError(f"Unsupported transfer backend: {transfer_backend}")
+
+
+def kv_to_page_indices(kv_indices: np.ndarray, page_size: int):
+    # 1. The page is guaruanteed to be full except the last page.
+    # 2. page index = kv_index // page_size
+    # The return vector is kv_indices[::page_size] // page_size
+    if page_size == 1:  # shortcut
+        return kv_indices
+
+    return kv_indices[::page_size] // page_size
+
+
+def kv_to_page_num(num_kv_indices: int, page_size: int):
+    # ceil(num_kv_indices / page_size)
+    return (num_kv_indices + page_size - 1) // page_size

lib/python3.12/site-packages/sglang/srt/function_call_parser.py

@@ -0,0 +1,761 @@
+import json
+import logging
+import re
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from json import JSONDecodeError, JSONDecoder
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Type
+
+import partial_json_parser
+from partial_json_parser.core.exceptions import MalformedJSON
+from partial_json_parser.core.options import Allow
+from pydantic import BaseModel
+
+from sglang.srt.openai_api.protocol import (
+    StructuralTagResponseFormat,
+    StructuresResponseFormat,
+    Tool,
+)
+
+logger = logging.getLogger(__name__)
+
+TOOLS_TAG_LIST = [
+    "<|plugin|>",
+    "<function=",
+    "<tool_call>",
+    "<|python_tag|>",
+    "[TOOL_CALLS]",
+    "<｜tool▁calls▁begin｜>",
+]
+
+
+class ToolCallItem(BaseModel):
+    """Simple encapsulation of the parsed ToolCall result for easier usage in streaming contexts."""
+
+    tool_index: int
+    name: Optional[str] = None
+    parameters: str  # JSON string
+
+
+def _find_common_prefix(s1: str, s2: str) -> str:
+    prefix = ""
+    min_length = min(len(s1), len(s2))
+    for i in range(0, min_length):
+        if s1[i] == s2[i]:
+            prefix += s1[i]
+        else:
+            break
+    return prefix
+
+
+def _partial_json_loads(input_str: str, flags: Allow) -> Tuple[Any, int]:
+    try:
+        return (partial_json_parser.loads(input_str, flags), len(input_str))
+    except JSONDecodeError as e:
+        if "Extra data" in e.msg:
+            dec = JSONDecoder()
+            return dec.raw_decode(input_str)
+        raise
+
+
+def _is_complete_json(input_str: str) -> bool:
+    try:
+        json.loads(input_str)
+        return True
+    except JSONDecodeError:
+        return False
+
+
+class StreamingParseResult:
+    """Result of streaming incremental parsing."""
+
+    def __init__(
+        self, normal_text: str = "", calls: Optional[List[ToolCallItem]] = None
+    ):
+        self.normal_text = normal_text
+        self.calls = calls or []
+
+
+@dataclass
+class StructureInfo:
+    begin: str
+    end: str
+    trigger: str
+
+
+_GetInfoFunc = Callable[[str], StructureInfo]
+"""
+helper alias of function
+ususally it is a function that takes a name string and returns a StructureInfo object,
+which can be used to construct a structural_tag object
+"""
+
+
+class BaseFormatDetector(ABC):
+    """Base class providing two sets of interfaces: one-time and streaming incremental."""
+
+    def __init__(self):
+        # initialize properties used for state when parsing tool calls in
+        self._buffer = ""
+        # streaming mode
+        self.prev_tool_call_arr: List[Dict] = []
+        self.current_tool_id: int = -1
+        self.current_tool_name_sent: bool = False
+        self.streamed_args_for_tool: List[str] = (
+            []
+        )  # map what has been streamed for each tool so far to a list
+        self.bot_token = ""
+        self.eot_token = ""
+
+    def parse_base_json(self, action: Any, tools: List[Tool]) -> List[ToolCallItem]:
+        tool_indices = {
+            tool.function.name: i for i, tool in enumerate(tools) if tool.function.name
+        }
+        if not isinstance(action, list):
+            action = [action]
+
+        results = []
+        for act in action:
+            name = act.get("name")
+            if name and name in tool_indices:
+                results.append(
+                    ToolCallItem(
+                        tool_index=tool_indices[name],
+                        name=name,
+                        parameters=json.dumps(
+                            act.get("parameters") or act.get("arguments", {}),
+                            ensure_ascii=False,
+                        ),
+                    )
+                )
+            else:
+                logger.warning(f"Model attempted to call undefined function: {name}")
+
+        return results
+
+    @abstractmethod
+    def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult:
+        """
+        Parses the text in one go. Returns success=True if the format matches, otherwise False.
+        Note that leftover_text here represents "content that this parser will not consume further".
+        """
+        action = json.loads(text)
+        return StreamingParseResult(calls=self.parse_base_json(action, tools))
+
+    def parse_streaming_increment(
+        self, new_text: str, tools: List[Tool]
+    ) -> StreamingParseResult:
+        """
+        Streaming incremental parsing with tool validation.
+        """
+        # Append new text to buffer
+        self._buffer += new_text
+        current_text = self._buffer
+        if not (self.bot_token in current_text or current_text.startswith("{")):
+            self._buffer = ""
+            if self.eot_token in new_text:
+                new_text = new_text.replace(self.eot_token, "")
+            return StreamingParseResult(normal_text=new_text)
+
+        # Build tool indices if not already built
+        if not hasattr(self, "_tool_indices"):
+            self._tool_indices = {
+                tool.function.name: i
+                for i, tool in enumerate(tools)
+                if tool.function and tool.function.name
+            }
+
+        flags = Allow.ALL if self.current_tool_name_sent else Allow.ALL & ~Allow.STR
+        try:
+            tool_call_arr = []
+            is_complete = []
+            try:
+                start_idx = (
+                    len(self.bot_token)
+                    if current_text.startswith(self.bot_token)
+                    else 0
+                )
+                while start_idx < len(current_text):
+                    (obj, end_idx) = _partial_json_loads(
+                        current_text[start_idx:], flags
+                    )
+                    is_complete.append(
+                        _is_complete_json(current_text[start_idx : start_idx + end_idx])
+                    )
+                    start_idx += end_idx + len("; ")
+
+                    # Validate tool name if present
+                    if "name" in obj and obj["name"] not in self._tool_indices:
+                        # Invalid tool name - reset state
+                        self._buffer = ""
+                        self.current_tool_id = -1
+                        self.current_tool_name_sent = False
+                        if self.streamed_args_for_tool:
+                            self.streamed_args_for_tool.pop()
+                        return StreamingParseResult()
+
+                    # Handle parameters/arguments consistency
+                    if "parameters" in obj:
+                        assert (
+                            "arguments" not in obj
+                        ), "model generated both parameters and arguments"
+                        obj["arguments"] = obj["parameters"]
+                    tool_call_arr.append(obj)
+
+            except MalformedJSON:
+                return StreamingParseResult()
+
+            if len(tool_call_arr) == 0:
+                return StreamingParseResult()
+
+            current_tool_call: Dict = (
+                tool_call_arr[self.current_tool_id] if len(tool_call_arr) > 0 else {}
+            )
+
+            # Handle new tool in array
+            if len(tool_call_arr) > 0 and len(tool_call_arr) > self.current_tool_id + 1:
+                if self.current_tool_id >= 0:
+                    cur_arguments = current_tool_call.get("arguments")
+                    if cur_arguments:
+                        cur_args_json = json.dumps(cur_arguments)
+                        sent = len(self.streamed_args_for_tool[self.current_tool_id])
+                        argument_diff = cur_args_json[sent:]
+
+                        res = StreamingParseResult(
+                            calls=[
+                                ToolCallItem(
+                                    tool_index=self.current_tool_id,
+                                    name="",
+                                    parameters=argument_diff,
+                                )
+                            ],
+                        )
+                        self.streamed_args_for_tool[
+                            self.current_tool_id
+                        ] += argument_diff
+                    else:
+                        res = StreamingParseResult()
+                else:
+                    res = StreamingParseResult()
+
+                self.current_tool_id = len(tool_call_arr) - 1
+                self.current_tool_name_sent = False
+                self.streamed_args_for_tool.append("")
+                return res
+
+            # Handle tool name
+            elif not self.current_tool_name_sent:
+                function_name = current_tool_call.get("name")
+                if function_name and function_name in self._tool_indices:
+                    res = StreamingParseResult(
+                        calls=[
+                            ToolCallItem(
+                                tool_index=self._tool_indices[function_name],
+                                name=function_name,
+                                parameters="",
+                            )
+                        ],
+                    )
+                    self.current_tool_name_sent = True
+                else:
+                    res = StreamingParseResult()
+
+            # Handle streaming arguments
+            else:
+                cur_arguments = current_tool_call.get("arguments")
+                res = StreamingParseResult()
+
+                if cur_arguments:
+                    sent = len(self.streamed_args_for_tool[self.current_tool_id])
+                    cur_args_json = json.dumps(cur_arguments)
+                    prev_arguments = self.prev_tool_call_arr[self.current_tool_id].get(
+                        "arguments"
+                    )
+
+                    argument_diff = None
+                    if is_complete[self.current_tool_id]:
+                        argument_diff = cur_args_json[sent:]
+                        self._buffer = ""
+                        self.prev_tool_call_arr[self.current_tool_id].clear()
+                        self.current_tool_name_sent = False
+                        self.streamed_args_for_tool[self.current_tool_id] = ""
+
+                    elif prev_arguments:
+                        prev_args_json = json.dumps(prev_arguments)
+                        if cur_args_json != prev_args_json:
+                            prefix = _find_common_prefix(prev_args_json, cur_args_json)
+                            argument_diff = prefix[sent:]
+
+                    if argument_diff is not None:
+                        res = StreamingParseResult(
+                            calls=[
+                                ToolCallItem(
+                                    tool_index=self.current_tool_id,
+                                    parameters=argument_diff,
+                                )
+                            ],
+                        )
+                        if not is_complete[self.current_tool_id]:
+                            self.streamed_args_for_tool[
+                                self.current_tool_id
+                            ] += argument_diff
+
+            self.prev_tool_call_arr = tool_call_arr
+            return res
+
+        except Exception as e:
+            logger.error(f"Error in parse_streaming_increment: {e}")
+            return StreamingParseResult()
+
+    @abstractmethod
+    def has_tool_call(self, text: str) -> bool:
+        raise NotImplementedError()
+
+    @abstractmethod
+    def structure_info(self) -> _GetInfoFunc:
+        raise NotImplementedError()
+
+
+class Qwen25Detector(BaseFormatDetector):
+    """
+    Detector for Qwen 2.5 models.
+    Assumes function call format:
+      <tool_call>{"name":"xxx", "arguments":{...}}</tool_call>
+    """
+
+    def __init__(self):
+        """
+        Initializes the detector with necessary state variables.
+        """
+        super().__init__()
+        self.bot_token = "<tool_call>"
+        self.eot_token = "</tool_call>"
+
+    def has_tool_call(self, text: str) -> bool:
+        """Check if the text contains a Qwen 2.5 format tool call."""
+        return self.bot_token in text
+
+    def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult:
+        """
+        One-time parsing: Detects and parses tool calls in the provided text.
+
+        :param text: The complete text to parse.
+        :param tools: List of available tools.
+        :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls.
+        """
+        idx = text.find(self.bot_token)
+        normal_text = text[:idx].strip() if idx != -1 else text
+        if self.bot_token not in text:
+            return StreamingParseResult(normal_text=normal_text, calls=[])
+        pattern = rf"{self.bot_token}(.*?){self.eot_token}"
+        match_result_list = re.findall(pattern, text, re.DOTALL)
+        calls = []
+        for match_result in match_result_list:
+            match_result = json.loads(match_result)
+            calls.extend(self.parse_base_json(match_result, tools))
+        return StreamingParseResult(normal_text=normal_text, calls=calls)
+
+    def structure_info(self) -> _GetInfoFunc:
+        return lambda name: StructureInfo(
+            begin='<tool_call>{"name":"' + name + '", "arguments":',
+            end="}</tool_call>",
+            trigger="<tool_call>",
+        )
+
+
+class MistralDetector(BaseFormatDetector):
+    """
+    Detector for Mistral models.
+    Assumes function call format:
+      <|action_start|><|plugin|>{"name":"xxx", "arguments":{...}}<|action_end|>
+    """
+
+    def __init__(self):
+        """
+        Initializes the detector with necessary state variables.
+        """
+        super().__init__()
+        self.bot_token = "[TOOL_CALLS] ["
+        self.tool_call_regex = re.compile(r"\[{.*}\]", re.DOTALL)
+
+    def has_tool_call(self, text: str) -> bool:
+        """Check if the text contains a Mistral format tool call."""
+        return self.bot_token in text
+
+    def _clean_text(self, text: str) -> str:
+        """
+        clean text to only leave ''[TOOL_CALLS] [{"name": xxx, "arguments": {xxx}}]'
+        for example,
+            text = '[TOOL_CALLS] [{"name": "get_current_weather", "arguments": {"location": "Boston, MA", "unit": "fahrenheit"}}]\n\nToday\'s weather in Boston is :{function call result} (in Fahrenheit)\n\nIf you prefer Celsius, please let me know.'
+            return '[TOOL_CALLS] [{"name": "get_current_weather", "arguments": {"location": "Boston, MA", "unit": "fahrenheit"}}]'
+        The key pattern is [TOOL_CALLS] [...]
+        """
+        find_results = re.findall(r"\[TOOL_CALLS\] \[.*?\]", text, re.DOTALL)
+        if len(find_results) > 0:
+            return find_results[0]
+        else:
+            return ""
+
+    def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult:
+        """
+        One-time parsing: Detects and parses tool calls in the provided text.
+
+        :param text: The complete text to parse.
+        :param tools: List of available tools.
+        :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls.
+        """
+        idx = text.find(self.bot_token)
+        normal_text = text[:idx].strip() if idx != -1 else text
+        text = self._clean_text(text)
+        tool_content = text.replace("[TOOL_CALLS]", "").strip()
+        raw_tool_calls = self.tool_call_regex.findall(tool_content)
+        calls = []
+        if len(raw_tool_calls) > 0:
+            raw_tool_call = raw_tool_calls[0]
+            function_call_arr = json.loads(raw_tool_call)
+            for match_result in function_call_arr:
+                calls.extend(self.parse_base_json(match_result, tools))
+        return StreamingParseResult(normal_text=normal_text, calls=calls)
+
+    def structure_info(self) -> _GetInfoFunc:
+        return lambda name: StructureInfo(
+            begin='[TOOL_CALLS] [{"name":"' + name + '", "arguments":',
+            end="}]",
+            trigger="[TOOL_CALLS]",
+        )
+
+
+class Llama32Detector(BaseFormatDetector):
+    """
+    Detector for Llama 3.2 models.
+    Assumes function call format:
+      <|python_tag|>{"name":"xxx", "arguments":{...}}
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.bot_token = "<|python_tag|>"
+
+    def has_tool_call(self, text: str) -> bool:
+        """Check if the text contains a Llama 3.2 format tool call."""
+        # depending on the prompt format the Llama model may or may not
+        # prefix the output with the <|python_tag|> token
+        return "<|python_tag|>" in text or text.startswith("{")
+
+    def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult:
+        """Parse function calls from text, handling multiple JSON objects."""
+        if "<|python_tag|>" not in text and not text.startswith("{"):
+            return StreamingParseResult(normal_text=text, calls=[])
+
+        if "<|python_tag|>" in text:
+            normal_text, action_text = text.split("<|python_tag|>")
+        else:
+            normal_text, action_text = "", text
+
+        # Split by semicolon and process each part
+        json_parts = [part.strip() for part in action_text.split(";") if part.strip()]
+        all_actions = []
+        for part in json_parts:
+            try:
+                # Parse each individual JSON object
+                action = json.loads(part)
+                all_actions.append(action)
+            except json.JSONDecodeError as e:
+                logger.warning(f"Failed to parse JSON part: {part}")
+                logger.warning(f"JSON parse error: {str(e)}")
+                continue
+        calls = []
+        # Only process if we found valid JSON objects
+        if all_actions:
+            calls = self.parse_base_json(all_actions, tools)
+        return StreamingParseResult(normal_text=normal_text, calls=calls)
+
+    def structure_info(self) -> _GetInfoFunc:
+        return lambda name: StructureInfo(
+            begin='<|python_tag|>{"name":"' + name + '", "arguments":',
+            end="}",
+            trigger="<|python_tag|>",
+        )
+
+
+class DeepSeekV3Detector(BaseFormatDetector):
+    """
+    Detector for DeepSeek models.
+    Assumes function call format:
+      '<｜tool▁calls▁begin｜><｜tool▁call▁begin｜>function<｜tool▁sep｜>get_current_weather\n```json\n{"location": "Tokyo"}\n```<｜tool▁call▁end｜>\n<｜tool▁call▁begin｜>function<｜tool▁sep｜>get_current_weather\n```json\n{"location": "Paris"}\n```<｜tool▁call▁end｜><｜tool▁calls▁end｜><｜end▁of▁sentence｜>
+    """
+
+    def __init__(self):
+        super().__init__()
+        self.bot_token = "<｜tool▁calls▁begin｜>"
+        self.eot_token = "<｜tool▁calls▁end｜>"
+        self.func_call_regex = r"<｜tool▁call▁begin｜>.*?<｜tool▁call▁end｜>"
+        self.func_detail_regex = r"<｜tool▁call▁begin｜>(.*)<｜tool▁sep｜>(.*)\n```json\n(.*)\n```<｜tool▁call▁end｜>"
+        self._last_arguments = ""
+
+    def has_tool_call(self, text: str) -> bool:
+        """Check if the text contains a deepseek format tool call."""
+        return self.bot_token in text
+
+    def detect_and_parse(self, text: str, tools: List[Tool]) -> StreamingParseResult:
+        """
+        One-time parsing: Detects and parses tool calls in the provided text.
+
+        :param text: The complete text to parse.
+        :param tools: List of available tools.
+        :return: ParseResult indicating success or failure, consumed text, leftover text, and parsed calls.
+        """
+        idx = text.find(self.bot_token)
+        normal_text = text[:idx].strip() if idx != -1 else text
+        if self.bot_token not in text:
+            return StreamingParseResult(normal_text=normal_text, calls=[])
+        match_result_list = re.findall(self.func_call_regex, text, re.DOTALL)
+        calls = []
+        try:
+            for match_result in match_result_list:
+                # Get function name
+                func_detail = re.search(self.func_detail_regex, match_result, re.DOTALL)
+                func_name = func_detail.group(2)
+                func_args = func_detail.group(3)
+                func_args = json.loads(func_args)
+                # construct match_result for parse_base_json
+                match_result = {"name": func_name, "parameters": func_args}
+                calls.extend(self.parse_base_json(match_result, tools))
+            return StreamingParseResult(normal_text=normal_text, calls=calls)
+        except Exception as e:
+            logger.error(f"Error in detect_and_parse: {e}")
+            # return the normal text if parsing fails
+            return StreamingParseResult(normal_text=text)
+
+    def structure_info(self) -> _GetInfoFunc:
+        return lambda name: StructureInfo(
+            begin=">" + name + "\n```json\n",
+            end="\n```<",
+            trigger=">" + name + "\n```json\n",
+        )
+
+    def parse_streaming_increment(
+        self, new_text: str, tools: List[Tool]
+    ) -> StreamingParseResult:
+        """
+        Streaming incremental parsing tool calls for DeepSeekV3 format.
+        """
+        self._buffer += new_text
+        current_text = self._buffer
+
+        if self.bot_token not in current_text:
+            self._buffer = ""
+            for e_token in [self.eot_token, "```", "<｜tool▁call▁end｜>"]:
+                if e_token in new_text:
+                    new_text = new_text.replace(e_token, "")
+            return StreamingParseResult(normal_text=new_text)
+
+        if not hasattr(self, "_tool_indices"):
+            self._tool_indices = {
+                tool.function.name: i
+                for i, tool in enumerate(tools)
+                if tool.function and tool.function.name
+            }
+
+        calls: list[ToolCallItem] = []
+        try:
+            partial_match = re.search(
+                pattern=r"<｜tool▁call▁begin｜>(.*)<｜tool▁sep｜>(.*)\n```json\n(.*)",
+                string=current_text,
+                flags=re.DOTALL,
+            )
+            if partial_match:
+                func_name = partial_match.group(2).strip()
+                func_args_raw = partial_match.group(3).strip()
+
+                if not self.current_tool_name_sent:
+                    calls.append(
+                        ToolCallItem(
+                            tool_index=self._tool_indices.get(func_name, 0),
+                            name=func_name,
+                            parameters="",
+                        )
+                    )
+                    self.current_tool_name_sent = True
+                else:
+                    argument_diff = (
+                        func_args_raw[len(self._last_arguments) :]
+                        if func_args_raw.startswith(self._last_arguments)
+                        else func_args_raw
+                    )
+
+                    if argument_diff:
+                        calls.append(
+                            ToolCallItem(
+                                tool_index=self._tool_indices.get(func_name, 0),
+                                name=None,
+                                parameters=argument_diff,
+                            )
+                        )
+                        self._last_arguments += argument_diff
+
+                    if _is_complete_json(func_args_raw):
+                        result = StreamingParseResult(normal_text="", calls=calls)
+                        self._buffer = ""
+                        self._last_arguments = ""
+                        self.current_tool_name_sent = False
+                        return result
+
+            return StreamingParseResult(normal_text="", calls=calls)
+
+        except Exception as e:
+            logger.error(f"Error in parse_streaming_increment: {e}")
+            return StreamingParseResult(normal_text=current_text)
+
+
+class MultiFormatParser:
+    def __init__(self, detectors: List[BaseFormatDetector]):
+        """
+        :param detectors: A series of available Detector instances passed in
+        """
+        self.detectors = detectors
+
+    def parse_once(
+        self, text: str, tools: List[Tool]
+    ) -> Tuple[str, list[ToolCallItem]]:
+        """
+        One-time parsing: Loop through detectors until there are no new matches or text is exhausted
+        Return: (final_text, all_calls)
+        - final_text: The remaining text after parsing that was not consumed by any Detector (can be treated as normal text)
+        - all_calls: All calls parsed by the Detectors
+        """
+        final_calls = []
+        final_normal_text = text
+        for detector in self.detectors:
+            parsed_result = detector.detect_and_parse(text, tools)
+            tool_call_list = parsed_result.calls
+            if len(tool_call_list) > 0:  # parsed successfully
+                final_calls = tool_call_list
+                final_normal_text = parsed_result.normal_text
+                break
+
+        # leftover_text is the normal text not consumed by any Detector
+        return final_normal_text, final_calls
+
+    def parse_streaming_increment(
+        self, new_text: str, tools: List[Tool]
+    ) -> Tuple[str, list[ToolCallItem]]:
+        """
+        Streaming incremental parsing: Feed new_text to each detector's parse_streaming_increment
+        and merge their produced normal_text/calls to return.
+        (The logic here can be "priority-based" or "parallel parsing" based on your needs)
+        """
+        final_normal_text = ""
+        final_calls = []
+
+        for detector in self.detectors:
+            sp_result = detector.parse_streaming_increment(new_text, tools)
+            # Merge normal_text and calls
+            # If one sp_result contains result call, this should be a successful parse
+            # If one sp_result only contains normal_text, this can either be a successful
+            # parse or it is not using the desired parsing tool.
+            if sp_result.normal_text:
+                final_normal_text = sp_result.normal_text
+            if sp_result.calls:
+                final_calls.extend(sp_result.calls)
+                final_normal_text = sp_result.normal_text
+                break
+
+        return final_normal_text, final_calls
+
+
+class FunctionCallParser:
+    """
+    In streaming scenarios, each time new_text is received, it calls multi_format_parser.parse_streaming_increment
+    and returns the resulting normal_text and calls to the upper layer (or SSE).
+    """
+
+    ToolCallParserEnum: Dict[str, Type[BaseFormatDetector]] = {
+        "llama3": Llama32Detector,
+        "qwen25": Qwen25Detector,
+        "mistral": MistralDetector,
+        "deepseekv3": DeepSeekV3Detector,
+    }
+
+    def __init__(self, tools: List[Tool], tool_call_parser: str):
+        detectors = []
+        if tool_call_parser:
+            detector_class = self.ToolCallParserEnum.get(tool_call_parser)
+            if detector_class:
+                detectors.append(detector_class())
+            else:
+                raise ValueError(f"Unsupported tool_call_parser: {tool_call_parser}")
+        else:
+            raise ValueError("Tool Call Parser Not Given!")
+
+        self.multi_format_parser = MultiFormatParser(detectors)
+        self.tools = tools
+
+    def has_tool_call(self, text: str) -> bool:
+        """
+        Check if the given text contains a tool call in the format supported by this parser.
+        This delegates to the detector's implementation.
+
+        :param text: The text to check for tool calls
+        :return: True if the text contains a tool call, False otherwise
+        """
+        # Check all detectors in the multi_format_parser
+        for detector in self.multi_format_parser.detectors:
+            if detector.has_tool_call(text):
+                return True
+        return False
+
+    def parse_non_stream(self, full_text: str) -> Tuple[str, list[ToolCallItem]]:
+        """
+        Non-streaming call: one-time parsing
+        """
+        full_normal_text, calls = self.multi_format_parser.parse_once(
+            full_text, self.tools
+        )
+        return full_normal_text, calls
+
+    def parse_stream_chunk(self, chunk_text: str) -> Tuple[str, list[ToolCallItem]]:
+        """
+        Streaming call: incremental parsing
+        """
+        normal_text, calls = self.multi_format_parser.parse_streaming_increment(
+            chunk_text, self.tools
+        )
+        return normal_text, calls
+
+    def structure_infos(self) -> List[_GetInfoFunc]:
+        """
+        Returns a list of structure_info functions for each detector
+        """
+        return [
+            detector.structure_info() for detector in self.multi_format_parser.detectors
+        ]
+
+    def get_structure_tag(self) -> StructuralTagResponseFormat:
+        tool_structures: List[StructuresResponseFormat] = list()
+        tool_trigger_set: Set[str] = set()
+
+        for wrapper in self.structure_infos():
+            for tool in self.tools:
+                function = tool.function
+                name = function.name
+                assert name is not None
+                info = wrapper(name)
+
+                # accept all if not strict, otherwise only accept the schema
+                schema = function.parameters if function.strict else {}
+
+                tool_structures.append(
+                    StructuresResponseFormat(
+                        begin=info.begin,
+                        schema=schema,  # type: ignore
+                        end=info.end,
+                    )
+                )
+                tool_trigger_set.add(info.trigger)
+
+        return StructuralTagResponseFormat(
+            type="structural_tag",
+            structures=tool_structures,
+            triggers=list(tool_trigger_set),
+        )