Skip to content

vllm.compilation.inductor_pass

_pass_context module-attribute

_pass_context = None

CallableInductorPass

Bases: InductorPass

This class is a wrapper for a callable that automatically provides an implementation of the UUID.

Source code in vllm/compilation/inductor_pass.py
class CallableInductorPass(InductorPass):
    """
    This class is a wrapper for a callable that automatically provides an
    implementation of the UUID.
    """

    def __init__(self,
                 callable: Callable[[fx.Graph], None],
                 uuid: Optional[Any] = None):
        self.callable = callable
        self._uuid = self.hash_source(callable) if uuid is None else uuid

    def __call__(self, graph: torch.fx.Graph):
        self.callable(graph)

    def uuid(self) -> Any:
        return self._uuid

_uuid instance-attribute

_uuid = hash_source(callable) if uuid is None else uuid

callable instance-attribute

callable = callable

__call__

__call__(graph: Graph)
Source code in vllm/compilation/inductor_pass.py
def __call__(self, graph: torch.fx.Graph):
    self.callable(graph)

__init__

__init__(
    callable: Callable[[Graph], None],
    uuid: Optional[Any] = None,
)
Source code in vllm/compilation/inductor_pass.py
def __init__(self,
             callable: Callable[[fx.Graph], None],
             uuid: Optional[Any] = None):
    self.callable = callable
    self._uuid = self.hash_source(callable) if uuid is None else uuid

uuid

uuid() -> Any
Source code in vllm/compilation/inductor_pass.py
def uuid(self) -> Any:
    return self._uuid

InductorPass

Bases: Torch25CustomGraphPass

A custom graph pass that uses a hash of its source as the UUID. This is defined as a convenience and should work in most cases.

Source code in vllm/compilation/inductor_pass.py
class InductorPass(CustomGraphPass):
    """
    A custom graph pass that uses a hash of its source as the UUID.
    This is defined as a convenience and should work in most cases.
    """

    def uuid(self) -> Any:
        """
        Provide a unique identifier for the pass, used in Inductor code cache.
        This should depend on the pass implementation, so that changes to the
        pass result in recompilation.
        By default, the object source is hashed.
        """
        return InductorPass.hash_source(self)

    @staticmethod
    def hash_source(*srcs: Union[str, Any]):
        """
        Utility method to hash the sources of functions or objects.
        :param srcs: strings or objects to add to the hash.
        Objects and functions have their source inspected.
        :return:
        """
        hasher = hashlib.sha256()
        for src in srcs:
            if isinstance(src, str):
                src_str = src
            elif isinstance(src, types.FunctionType):
                src_str = inspect.getsource(src)
            else:
                src_str = inspect.getsource(src.__class__)
            hasher.update(src_str.encode("utf-8"))
        return hasher.hexdigest()

    @staticmethod
    def hash_dict(dict_: dict[Any, Any]):
        """
        Utility method to hash a dictionary, can alternatively be used for uuid.
        :return: A sha256 hash of the json rep of the dictionary.
        """
        encoded = json.dumps(dict_, sort_keys=True).encode("utf-8")
        return hashlib.sha256(encoded).hexdigest()

    def is_applicable_for_shape(self, shape: Optional[int]):
        return True

hash_dict staticmethod

hash_dict(dict_: dict[Any, Any])

Utility method to hash a dictionary, can alternatively be used for uuid. :return: A sha256 hash of the json rep of the dictionary.

Source code in vllm/compilation/inductor_pass.py
@staticmethod
def hash_dict(dict_: dict[Any, Any]):
    """
    Utility method to hash a dictionary, can alternatively be used for uuid.
    :return: A sha256 hash of the json rep of the dictionary.
    """
    encoded = json.dumps(dict_, sort_keys=True).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()

hash_source staticmethod

hash_source(*srcs: Union[str, Any])

Utility method to hash the sources of functions or objects. :param srcs: strings or objects to add to the hash. Objects and functions have their source inspected. :return:

Source code in vllm/compilation/inductor_pass.py
@staticmethod
def hash_source(*srcs: Union[str, Any]):
    """
    Utility method to hash the sources of functions or objects.
    :param srcs: strings or objects to add to the hash.
    Objects and functions have their source inspected.
    :return:
    """
    hasher = hashlib.sha256()
    for src in srcs:
        if isinstance(src, str):
            src_str = src
        elif isinstance(src, types.FunctionType):
            src_str = inspect.getsource(src)
        else:
            src_str = inspect.getsource(src.__class__)
        hasher.update(src_str.encode("utf-8"))
    return hasher.hexdigest()

is_applicable_for_shape

is_applicable_for_shape(shape: Optional[int])
Source code in vllm/compilation/inductor_pass.py
def is_applicable_for_shape(self, shape: Optional[int]):
    return True

uuid

uuid() -> Any

Provide a unique identifier for the pass, used in Inductor code cache. This should depend on the pass implementation, so that changes to the pass result in recompilation. By default, the object source is hashed.

Source code in vllm/compilation/inductor_pass.py
def uuid(self) -> Any:
    """
    Provide a unique identifier for the pass, used in Inductor code cache.
    This should depend on the pass implementation, so that changes to the
    pass result in recompilation.
    By default, the object source is hashed.
    """
    return InductorPass.hash_source(self)

PassContext

Source code in vllm/compilation/inductor_pass.py
class PassContext:

    def __init__(self, runtime_shape: Optional[int]):
        self.runtime_shape = runtime_shape

runtime_shape instance-attribute

runtime_shape = runtime_shape

__init__

__init__(runtime_shape: Optional[int])
Source code in vllm/compilation/inductor_pass.py
def __init__(self, runtime_shape: Optional[int]):
    self.runtime_shape = runtime_shape

get_pass_context

get_pass_context() -> PassContext

Get the current pass context.

Source code in vllm/compilation/inductor_pass.py
def get_pass_context() -> PassContext:
    """Get the current pass context."""
    assert _pass_context is not None
    return _pass_context

pass_context

pass_context(runtime_shape: Optional[int])

A context manager that stores the current pass context, usually it is a list of sizes to specialize.

Source code in vllm/compilation/inductor_pass.py
@contextmanager
def pass_context(runtime_shape: Optional[int]):
    """A context manager that stores the current pass context,
    usually it is a list of sizes to specialize.
    """
    global _pass_context
    prev_context = _pass_context
    _pass_context = PassContext(runtime_shape)
    try:
        yield
    finally:
        _pass_context = prev_context