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 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
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: Any | None = 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 
113
114
def __call__(self, graph: torch.fx.Graph):
    self.callable(graph)

__init__ 

__init__(
    callable: Callable[[Graph], None],
    uuid: Any | None = None,
)
Source code in vllm/compilation/inductor_pass.py 
109
110
111
def __init__(self, callable: Callable[[fx.Graph], None], uuid: Any | None = 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 
116
117
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 
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
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: 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, type)):
                src_str = inspect.getsource(src)
            else:
                # object instance
                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(self, shape: int | None):
        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 
90
91
92
93
94
95
96
97
@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: 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 
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
@staticmethod
def hash_source(*srcs: 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, type)):
            src_str = inspect.getsource(src)
        else:
            # object instance
            src_str = inspect.getsource(src.__class__)
        hasher.update(src_str.encode("utf-8"))
    return hasher.hexdigest()

is_applicable 

is_applicable(shape: int | None)
Source code in vllm/compilation/inductor_pass.py 
 99
100
def is_applicable(self, shape: int | None):
    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 
61
62
63
64
65
66
67
68
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 
30
31
32
class PassContext:
    def __init__(self, runtime_shape: int | None):
        self.runtime_shape = runtime_shape

runtime_shape instance-attribute 

runtime_shape = runtime_shape

__init__ 

__init__(runtime_shape: int | None)
Source code in vllm/compilation/inductor_pass.py 
31
32
def __init__(self, runtime_shape: int | None):
    self.runtime_shape = runtime_shape

enable_fake_mode 

enable_fake_mode(
    fn: Callable[..., Any],
) -> Callable[..., Any]

Applies a FakeTensorMode context. This is useful when you don't want to create or run things with real tensors.

Source code in vllm/compilation/inductor_pass.py 
120
121
122
123
124
125
126
127
128
129
130
131
132
133
def enable_fake_mode(fn: Callable[..., Any]) -> Callable[..., Any]:
    """
    Applies a FakeTensorMode context. This is useful when you don't want to
    create or run things with real tensors.
    """

    @functools.wraps(fn)
    def fn_new(*args, **kwargs) -> Any:
        with torch._guards.tracing(None), unset_fake_temporarily(), FakeTensorMode():
            result = fn(*args, **kwargs)

        return result

    return fn_new

get_pass_context 

get_pass_context() -> PassContext

Get the current pass context.

Source code in vllm/compilation/inductor_pass.py 
35
36
37
38
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: int | None)

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 
41
42
43
44
45
46
47
48
49
50
51
52
@contextmanager
def pass_context(runtime_shape: int | None):
    """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