Skip to content

vllm.compilation.torch25_custom_graph_pass

Torch25CustomGraphPass

Bases: ABC

This class replaces CustomGraphPass from torch==2.6 when using torch<2.6. It conforms to the 2.6 interface but also supports pickling, as that's what the inductor code cache uses to determine the cache key before 2.6. (in 2.6 and above, uuid() is used.)

Subclasses can just "pretend" that uuid is used.

Source code in vllm/compilation/torch25_custom_graph_pass.py 
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
class Torch25CustomGraphPass(ABC):  # noqa (redefinition)
    """
    This class replaces CustomGraphPass from torch==2.6 when using torch<2.6.
    It conforms to the 2.6 interface but also supports pickling, as that's what
    the inductor code cache uses to determine the cache key before 2.6.
    (in 2.6 and above, uuid() is used.)

    Subclasses can just "pretend" that uuid is used.
    """

    @abstractmethod
    def __call__(self, graph: torch.fx.graph.Graph) -> None:
        """
        Implementation of the custom pass.
        """

    @abstractmethod
    def uuid(self) -> Optional[Any]:
        """
        Return an ID to uniquely identify your custom pass implementation.
        Return None to skip inductor code caching entirely.
        """

    def __getstate__(self):
        """
        Pickling is used instead of uuid() in torch<2.6. Just return uuid()
         to enable subclasses to only have to implement uuid.
        """
        return self.uuid()

    def __setstate__(self, state):
        raise ValueError("Cannot unpickle CustomGraphPass because pickling"
                         " is used for cache key uuid. Use torch>=2.6 with"
                         " native uuid support for custom passes.")

__call__ abstractmethod 

__call__(graph: Graph) -> None

Implementation of the custom pass.

Source code in vllm/compilation/torch25_custom_graph_pass.py 
19
20
21
22
23
@abstractmethod
def __call__(self, graph: torch.fx.graph.Graph) -> None:
    """
    Implementation of the custom pass.
    """

__getstate__ 

__getstate__()

Pickling is used instead of uuid() in torch<2.6. Just return uuid() to enable subclasses to only have to implement uuid.

Source code in vllm/compilation/torch25_custom_graph_pass.py 
32
33
34
35
36
37
def __getstate__(self):
    """
    Pickling is used instead of uuid() in torch<2.6. Just return uuid()
     to enable subclasses to only have to implement uuid.
    """
    return self.uuid()

__setstate__ 

__setstate__(state)
Source code in vllm/compilation/torch25_custom_graph_pass.py 
39
40
41
42
def __setstate__(self, state):
    raise ValueError("Cannot unpickle CustomGraphPass because pickling"
                     " is used for cache key uuid. Use torch>=2.6 with"
                     " native uuid support for custom passes.")

uuid abstractmethod 

uuid() -> Optional[Any]

Return an ID to uniquely identify your custom pass implementation. Return None to skip inductor code caching entirely.

Source code in vllm/compilation/torch25_custom_graph_pass.py 
25
26
27
28
29
30
@abstractmethod
def uuid(self) -> Optional[Any]:
    """
    Return an ID to uniquely identify your custom pass implementation.
    Return None to skip inductor code caching entirely.
    """