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 
47
48
49
50
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
) -> Optional[Node]
Source code in vllm/compilation/fx_utils.py 
38
39
40
41
42
43
def find_auto_fn_maybe(nodes: Iterable[fx.Node],
                       op: OpOverload) -> Optional[fx.Node]:
    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 
63
64
65
66
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) -> Optional[Node]
Source code in vllm/compilation/fx_utils.py 
55
56
57
58
59
def find_getitem_maybe(node: fx.Node, idx: int) -> Optional[fx.Node]:
    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, graph: Graph
) -> Iterator[Node]
Source code in vllm/compilation/fx_utils.py 
70
71
72
73
74
75
76
def find_op_nodes(op: OpOverload, graph: fx.Graph) -> Iterator[fx.Node]:
    if not op._schema.is_mutable:
        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 
31
32
33
34
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
) -> Optional[Node]
Source code in vllm/compilation/fx_utils.py 
22
23
24
25
26
27
def find_specified_fn_maybe(nodes: Iterable[fx.Node],
                            op: OpOverload) -> Optional[fx.Node]:
    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 
82
83
84
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 
17
18
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 
13
14
def is_func(node: fx.Node, target) -> bool:
    return node.op == "call_function" and node.target == target