Skip to content

vllm.ir

Modules:

Name Description
op
ops
util

enable_torch_wrap 

enable_torch_wrap(enable: bool = True)

Context manager to enable/disable torch custom op wrapping for vLLM IR ops. When torch wrapping is disabled, the torch custom op layer is skipped and IR ops dispatch directly to the implementation. Helpful for avoiding torch dispatch overhead in eager mode and avoiding the need for lowering for platforms not using Inductor.

Source code in vllm/ir/op.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
@contextlib.contextmanager
def enable_torch_wrap(enable: bool = True):
    """
    Context manager to enable/disable torch custom op wrapping for vLLM IR ops.
    When torch wrapping is disabled, the torch custom op layer is skipped
    and IR ops dispatch directly to the implementation.
    Helpful for avoiding torch dispatch overhead in eager mode
    and avoiding the need for lowering for platforms not using Inductor.
    """

    global _ENABLE_TORCH_WRAP
    old = _ENABLE_TORCH_WRAP
    try:
        _ENABLE_TORCH_WRAP = enable
        yield
    finally:
        _ENABLE_TORCH_WRAP = old

register_op 

register_op(f: Callable[..., Any]) -> IrOp

register_op(
    *,
    name: str | None = None,
    activations: list[str] | None = None,
    allow_inplace: Literal[False] = False,
) -> Callable[[Callable[..., Any]], IrOp]

register_op(
    *,
    name: str | None = None,
    activations: list[str] | None = None,
    allow_inplace: Literal[True],
) -> Callable[[Callable[..., Any]], IrOpInplace]

register_op(
    f: Callable | None = None,
    *,
    name: str | None = None,
    activations: list[str] | None = None,
    allow_inplace: bool = False,
) -> IrOp | Callable[[Callable], IrOp]

Register a new vLLM IR op.

:param f: the native implementation of the op :param name: the name of the op, defaults to the function name :param activations: list of activation params, defaults to params starting with 'x' :param allow_inplace: add a maybe_inplace overload that allows inplace impls :return: the IrOp object if f is provided, otherwise a decorator

Example usage: ```python @vllm.ir.register_op def my_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return x + y

@vllm.ir.register_op(name="custom_mul") def multiply(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: return x * y

Source code in vllm/ir/op.py 
 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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def register_op(
    f: Callable | None = None,
    *,
    name: str | None = None,
    activations: list[str] | None = None,
    allow_inplace: bool = False,
) -> "IrOp | Callable[[Callable], IrOp]":
    """
    Register a new vLLM IR op.

    :param f: the native implementation of the op
    :param name: the name of the op, defaults to the function name
    :param activations: list of activation params, defaults to params starting with 'x'
    :param allow_inplace: add a maybe_inplace overload that allows inplace impls
    :return: the IrOp object if f is provided, otherwise a decorator

    Example usage:
    ```python
    @vllm.ir.register_op
    def my_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return x + y


    @vllm.ir.register_op(name="custom_mul")
    def multiply(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        return x * y"""

    def decorator(_f: Callable):
        op_name: str = _f.__name__ if name is None else name
        assert op_name not in IrOp.registry
        if allow_inplace:
            op: IrOp = IrOpInplace(op_name, _f, activations)
        else:
            op = IrOp(op_name, _f, activations)
        IrOp.registry[op_name] = op
        return op

    if f is not None:
        return decorator(f)

    return decorator