Skip to content

vllm.compilation.fx_utils

find_auto_fn 

find_auto_fn(nodes: Iterable[Node], op: OpOverload) -> Node
Source code in vllm/compilation/fx_utils.py 
44
45
46
47
def find_auto_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node:
    node = find_auto_fn_maybe(nodes, op)
    assert node is not None, f"Could not find {op} in nodes {nodes}"
    return node

find_auto_fn_maybe 

find_auto_fn_maybe(
    nodes: Iterable[Node], op: OpOverload
) -> Node | None
Source code in vllm/compilation/fx_utils.py 
36
37
38
39
40
def find_auto_fn_maybe(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node | None:
    for node in nodes:
        if is_func(node, auto_functionalized) and node.args[0] == op:  # noqa
            return node
    return None

find_getitem 

find_getitem(node: Node, idx: int) -> Node
Source code in vllm/compilation/fx_utils.py 
60
61
62
63
def find_getitem(node: fx.Node, idx: int) -> fx.Node:
    ret = find_getitem_maybe(node, idx)
    assert ret is not None, f"Could not find getitem {idx} in node {node}"
    return ret

find_getitem_maybe 

find_getitem_maybe(node: Node, idx: int) -> Node | None
Source code in vllm/compilation/fx_utils.py 
52
53
54
55
56
def find_getitem_maybe(node: fx.Node, idx: int) -> fx.Node | None:
    for user in node.users:
        if is_func(user, operator.getitem) and user.args[1] == idx:
            return user
    return None

find_op_nodes 

find_op_nodes(
    op: OpOverload | OpOverloadPacket, graph: Graph
) -> Iterator[Node]
Source code in vllm/compilation/fx_utils.py 
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
def find_op_nodes(
    op: OpOverload | OpOverloadPacket, graph: fx.Graph
) -> Iterator[fx.Node]:
    if isinstance(op, OpOverloadPacket):
        for overload in op.overloads():
            overload_op = getattr(op, overload)
            yield from find_op_nodes(overload_op, graph)
        return

    assert isinstance(op, OpOverload)

    yield from graph.find_nodes(op="call_function", target=op)

    for n in graph.find_nodes(op="call_function", target=auto_functionalized):
        if n.args[0] == op:
            yield n

find_specified_fn 

find_specified_fn(
    nodes: Iterable[Node], op: OpOverload
) -> Node
Source code in vllm/compilation/fx_utils.py 
29
30
31
32
def find_specified_fn(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node:
    node = find_specified_fn_maybe(nodes, op)
    assert node is not None, f"Could not find {op} in nodes {nodes}"
    return node

find_specified_fn_maybe 

find_specified_fn_maybe(
    nodes: Iterable[Node], op: OpOverload
) -> Node | None
Source code in vllm/compilation/fx_utils.py 
21
22
23
24
25
def find_specified_fn_maybe(nodes: Iterable[fx.Node], op: OpOverload) -> fx.Node | None:
    for node in nodes:
        if node.target == op:
            return node
    return None

get_only_user 

get_only_user(node: Node) -> Node
Source code in vllm/compilation/fx_utils.py 
89
90
91
def get_only_user(node: fx.Node) -> fx.Node:
    assert len(node.users) == 1
    return next(iter(node.users))

is_auto_func 

is_auto_func(node: Node, op: OpOverload) -> bool
Source code in vllm/compilation/fx_utils.py 
16
17
def is_auto_func(node: fx.Node, op: OpOverload) -> bool:
    return is_func(node, auto_functionalized) and node.args[0] == op

is_func 

is_func(node: Node, target) -> bool
Source code in vllm/compilation/fx_utils.py 
12
13
def is_func(node: fx.Node, target) -> bool:
    return node.op == "call_function" and node.target == target