wehub-resource-sync
102 lines
4.2 KiB
Python
102 lines
4.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from collections import defaultdict
|
|
|
|
from torch import fx
|
|
from torch._inductor.pattern_matcher import (
|
|
PatternMatcherPass,
|
|
)
|
|
|
|
from vllm.config import VllmConfig
|
|
from vllm.logger import init_logger
|
|
|
|
from ..inductor_pass import get_pass_context
|
|
from ..vllm_inductor_pass import VllmInductorPass
|
|
from .lowering_pass import get_ir_op
|
|
from .utils import overload_or_default
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class VllmIRInplaceFunctionalizationPass(VllmInductorPass):
|
|
"""
|
|
This pass functionalizes maybe_inplace vLLM IR ops to the default overload.
|
|
The maybe_inplace overloads have the same signature as the default overload
|
|
so the pass simply replaces the called overload.
|
|
That makes the graph properly functional.
|
|
The pass also validates that activations passed to maybe_inplace have no later
|
|
uses in the graph: they are donated to the maybe_inplace op call,
|
|
and their contents are not defined afterward.
|
|
|
|
This pass operates pre-AOTAutograd,
|
|
so it must handle non-normalized and non-functional IR.
|
|
"""
|
|
|
|
def __init__(self, vllm_config: VllmConfig) -> None:
|
|
super().__init__(vllm_config)
|
|
self.patterns = PatternMatcherPass(self.pass_name)
|
|
self.functionalized_ops: dict[str, int] = defaultdict(lambda: 0)
|
|
|
|
@VllmInductorPass.time_and_log
|
|
def __call__(self, graph: fx.Graph) -> None:
|
|
# clear at the beginning instead of end, so that tests can inspect
|
|
self.functionalized_ops.clear()
|
|
assert graph.owning_module is not None
|
|
node_to_idx = {node: i for i, node in enumerate(graph.nodes)}
|
|
|
|
# Pass donated input via vLLM's pass context
|
|
pass_context = get_pass_context()
|
|
pass_context.donated_input_ids = set[int]()
|
|
|
|
for node in graph.nodes:
|
|
if (ir_op := get_ir_op(node)) is None:
|
|
continue
|
|
|
|
op_overload = overload_or_default(node.target)
|
|
overload_name = op_overload._overloadname
|
|
if overload_name != "maybe_inplace":
|
|
assert overload_name == "default", (
|
|
f"Found overload {overload_name} for op {ir_op.name}, "
|
|
f"expected maybe_inplace or default"
|
|
)
|
|
continue
|
|
|
|
# must have maybe_inplace overload and allow_inplace
|
|
assert ir_op.allow_inplace and hasattr(ir_op, "maybe_inplace")
|
|
|
|
# Check that activation inputs are not used after this op
|
|
for arg_idx in ir_op.activation_indices:
|
|
arg = node.args[arg_idx]
|
|
assert isinstance(arg, fx.Node), "Activation inputs must be fx.Node"
|
|
for user in arg.users:
|
|
if node_to_idx[user] > node_to_idx[node]:
|
|
raise ValueError(
|
|
f"Input {arg} to maybe_inplace node {node} "
|
|
f"is used again after the node. "
|
|
f"This is not allowed; activation inputs to maybe_inplace "
|
|
f"ops are donated to the op, meaning their memory may be "
|
|
f"recycled for outputs.\n\n"
|
|
f"To preserve the inputs, use the default overload or "
|
|
f"clone them manually beforehand."
|
|
)
|
|
|
|
if arg.op == "placeholder":
|
|
# Graph input that maybe_inplace might modify.
|
|
# Mark it so downstream passes know it's donated.
|
|
# TODO(luka) store in placeholder node meta once supported
|
|
pass_context.donated_input_ids.add(node_to_idx[arg])
|
|
|
|
# Same signature, just replace the overload that's called.
|
|
node.target = ir_op.torch_op
|
|
self.functionalized_ops[ir_op.name] += 1
|
|
|
|
count = sum(self.functionalized_ops.values())
|
|
ops = ",".join(self.functionalized_ops.keys())
|
|
logger.debug("Donated input IDs: %s", pass_context.donated_input_ids)
|
|
logger.debug(
|
|
"%s functionalized %d vLLM IR nodes for op(s) %s",
|
|
self.pass_name,
|
|
count,
|
|
ops,
|
|
)
|