248 lines
8.9 KiB
Python
248 lines
8.9 KiB
Python
# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
|
|
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
# Adapted from vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/model_executor/layers/utils.py
|
|
"""Utility methods for model layers."""
|
|
import inspect
|
|
from typing import Any, Callable, List, Optional
|
|
|
|
import torch
|
|
from torch.library import Library
|
|
|
|
from sglang.multimodal_gen.runtime.platforms import current_platform
|
|
|
|
|
|
def get_token_bin_counts_and_mask(
|
|
tokens: torch.Tensor,
|
|
vocab_size: int,
|
|
num_seqs: int,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
# Compute the bin counts for the tokens.
|
|
# vocab_size + 1 for padding.
|
|
bin_counts = torch.zeros(
|
|
(num_seqs, vocab_size + 1), dtype=torch.long, device=tokens.device
|
|
)
|
|
bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens))
|
|
bin_counts = bin_counts[:, :vocab_size]
|
|
mask = bin_counts > 0
|
|
|
|
return bin_counts, mask
|
|
|
|
|
|
sglang_lib = Library("sglang", "FRAGMENT") # noqa
|
|
|
|
|
|
def direct_register_custom_op(
|
|
op_name: str,
|
|
op_func: Callable,
|
|
mutates_args: List[str],
|
|
fake_impl: Optional[Callable] = None,
|
|
target_lib: Optional[Library] = None,
|
|
):
|
|
"""
|
|
`torch.library.custom_op` can have significant overhead because it
|
|
needs to consider complicated dispatching logic. This function
|
|
directly registers a custom op and dispatches it to the CUDA backend.
|
|
See https://gist.github.com/youkaichao/ecbea9ec9fc79a45d2adce1784d7a9a5
|
|
for more details.
|
|
|
|
By default, the custom op is registered to the vLLM library. If you
|
|
want to register it to a different library, you can pass the library
|
|
object to the `target_lib` argument.
|
|
|
|
IMPORTANT: the lifetime of the operator is tied to the lifetime of the
|
|
library object. If you want to bind the operator to a different library,
|
|
make sure the library object is alive when the operator is used.
|
|
|
|
Note: This function will silently skip registration if the operator
|
|
with the same name is already registered to avoid RuntimeError in
|
|
multi-engine scenarios (e.g., VERL framework).
|
|
"""
|
|
import torch.library
|
|
|
|
my_lib = target_lib or sglang_lib
|
|
|
|
# Check if operator is already registered to avoid duplicate registration
|
|
# This is important for scenarios where multiple SGLang engines run in the same process
|
|
try:
|
|
# Try to access the operator to see if it's already registered
|
|
lib_name = my_lib.m.name if hasattr(my_lib.m, "name") else "sglang"
|
|
if hasattr(torch.ops, lib_name) and hasattr(
|
|
getattr(torch.ops, lib_name), op_name
|
|
):
|
|
# Operator already exists, skip registration
|
|
return
|
|
except (AttributeError, RuntimeError):
|
|
# Operator doesn't exist, proceed with registration
|
|
pass
|
|
|
|
if hasattr(torch.library, "infer_schema"):
|
|
schema_str = torch.library.infer_schema(op_func, mutates_args=mutates_args)
|
|
else:
|
|
# for pytorch 2.4
|
|
import torch._custom_op.impl
|
|
|
|
schema_str = torch._custom_op.impl.infer_schema(op_func, mutates_args)
|
|
|
|
try:
|
|
my_lib.define(op_name + schema_str)
|
|
my_lib.impl(
|
|
op_name, op_func, "CUDA" if not current_platform.is_npu() else "PrivateUse1"
|
|
)
|
|
if fake_impl is not None:
|
|
my_lib._register_fake(op_name, fake_impl)
|
|
except RuntimeError as error:
|
|
if "Tried to register an operator" in str(error) and "multiple times" in str(
|
|
error
|
|
):
|
|
# Silently ignore duplicate registration errors
|
|
# This can happen in multi-engine scenarios
|
|
pass
|
|
else:
|
|
# Re-raise other RuntimeErrors
|
|
raise error
|
|
except AttributeError as error:
|
|
# Always re-raise AttributeError as it indicates missing dependencies
|
|
raise error
|
|
|
|
|
|
class CustomOpWrapper:
|
|
def __init__(
|
|
self,
|
|
op_name: str,
|
|
op_func: Callable,
|
|
mutates_args: List[str],
|
|
**extra_kwargs,
|
|
):
|
|
self.op_name = op_name
|
|
self.op_func = op_func
|
|
self.mutates_args = mutates_args
|
|
self.extra_kwargs = extra_kwargs
|
|
self._impl: Optional[Callable] = None
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
return self.real_impl(*args, **kwargs)
|
|
|
|
@property
|
|
def real_impl(self) -> Callable:
|
|
if self._impl is None:
|
|
if not hasattr(torch.ops.sglang, self.op_name):
|
|
|
|
# NOTE(dark): if torch compile fail here, mark the decorator as eager
|
|
# lazy registration does not work with torch compile
|
|
direct_register_custom_op(
|
|
op_name=self.op_name,
|
|
op_func=self.op_func,
|
|
mutates_args=self.mutates_args,
|
|
fake_impl=self.fake_impl,
|
|
)
|
|
self._impl = getattr(torch.ops.sglang, self.op_name)
|
|
assert self._impl is not None
|
|
return self._impl
|
|
|
|
@property
|
|
def fake_impl(self) -> Callable:
|
|
if "fake_impl" in self.extra_kwargs:
|
|
return self.extra_kwargs["fake_impl"]
|
|
assert "out_shape" in self.extra_kwargs
|
|
signature = inspect.signature(self.op_func)
|
|
out_shape = self.extra_kwargs["out_shape"]
|
|
|
|
# check out_shape in signature
|
|
|
|
def fake_impl(*args, **kwargs):
|
|
if out_shape is None:
|
|
return None
|
|
bound = signature.bind(*args, **kwargs)
|
|
bound.apply_defaults()
|
|
try:
|
|
return torch.empty_like(
|
|
bound.args[out_shape]
|
|
if isinstance(out_shape, int)
|
|
else bound.arguments[out_shape]
|
|
)
|
|
except (IndexError, KeyError):
|
|
raise RuntimeError(
|
|
f"Cannot find output argument at position `{out_shape}` for "
|
|
f"custom operator `{self.op_name}` with signature `{signature}`."
|
|
)
|
|
|
|
return fake_impl
|
|
|
|
|
|
# Real implementation
|
|
def register_custom_op(
|
|
fn: Optional[Callable] = None,
|
|
*,
|
|
op_name: Optional[str] = None,
|
|
mutates_args: Optional[List[str]] = None,
|
|
eager: bool = True,
|
|
**extra_kwargs,
|
|
) -> Any:
|
|
"""
|
|
A decorator to register a custom operator.
|
|
|
|
Example usage:
|
|
```python
|
|
# inplace operator, out_shape is None by default
|
|
@register_custom_op(mutates_args=["x"])
|
|
def add_1_(x: torch.Tensor) -> None:
|
|
x.add_(1)
|
|
|
|
# operator with output, out_shape indicates the position of output
|
|
@register_custom_op(mutates_args=["x"], out_shape=0)
|
|
def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x.add_(y)
|
|
```
|
|
|
|
:param fn: The function to be registered as a custom operator.
|
|
If None, return a decorator.
|
|
:type fn: Callable
|
|
:param op_name: The name of the operator. If None, use the function name
|
|
:type op_name: Optional[str]
|
|
:param mutates_args: A list of argument names that are mutated in-place.
|
|
:type mutates_args: List[str]
|
|
:param out_shape: The position (int for positional, str for keyword) of the output-shape tensor.
|
|
It is used to generate a fake implementation for torch.compile compatibility.
|
|
If the operator is inplace and has no output, set to None.
|
|
:type out_shape: Optional[List[Union[int, str]]]
|
|
:param fake_impl: A fake implementation for the operator.
|
|
Only one of `out_shape` or `fake_impl` should be provided.
|
|
:type fake_impl: Optional[Callable]
|
|
:param eager: Whether to register the operator eagerly.
|
|
If False, the registration will be deferred until the first call.
|
|
If you met any issue with torch.compile, try to set eager=True.
|
|
Currently, to avoid misuse, we set eager=True by default.
|
|
:type eager: bool
|
|
:return: The registered JIT custom operator, or a decorator.
|
|
NOTE: the real register will occur at the first call of the function.
|
|
:rtype: Callable
|
|
"""
|
|
extra_kwarg_keys = set(extra_kwargs.keys())
|
|
expected_kwarg_keys = set({"out_shape", "fake_impl"})
|
|
assert (
|
|
expected_kwarg_keys >= extra_kwarg_keys
|
|
), f"Unexpected extra kwargs: {extra_kwarg_keys - expected_kwarg_keys}"
|
|
|
|
has_out_shape = "out_shape" in extra_kwargs
|
|
has_fake_impl = "fake_impl" in extra_kwargs
|
|
assert not (
|
|
has_out_shape and has_fake_impl
|
|
), "Only one of `out_shape` or `fake_impl` should be provided."
|
|
# Assume inplace if neither out_shape nor fake_impl is provided
|
|
if not (has_out_shape or has_fake_impl):
|
|
extra_kwargs["out_shape"] = None
|
|
|
|
def decorator(op_func: Callable) -> Callable:
|
|
wrapper = CustomOpWrapper(
|
|
op_name=op_name or op_func.__name__,
|
|
op_func=op_func,
|
|
mutates_args=mutates_args or [],
|
|
**extra_kwargs,
|
|
)
|
|
return wrapper.real_impl if eager else wrapper
|
|
|
|
if fn is not None:
|
|
return decorator(fn)
|
|
return decorator
|