wehub-resource-sync
767 lines
26 KiB
Python
767 lines
26 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
import importlib.util
|
|
import logging
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import pytest
|
|
import torch
|
|
from torch import fx
|
|
from torch.fx.experimental.proxy_tensor import make_fx
|
|
|
|
import vllm.ir.op
|
|
from vllm.ir.op import RESERVED_PROVIDERS, IrOp, IrOpImpl
|
|
|
|
|
|
class CustomError(Exception):
|
|
pass
|
|
|
|
|
|
@pytest.fixture
|
|
def custom_add_op(fake_vllm_ir):
|
|
"""Register ``_custom_add`` plus impl_a, impl_b, impl_even for this test."""
|
|
|
|
@vllm.ir.register_op(allow_inplace=True)
|
|
def _custom_add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x + y
|
|
|
|
@_custom_add.register_impl("impl_a")
|
|
def impl_a(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x + y + 10
|
|
|
|
@_custom_add.register_impl("impl_b", inplace=True)
|
|
def impl_b(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
"""Computes x+y+20"""
|
|
x.add_(y)
|
|
x.add_(20)
|
|
return x
|
|
|
|
@_custom_add.register_impl(
|
|
"impl_even", supports_args=lambda x, y: x.size(1) % 2 == 0
|
|
)
|
|
def impl_even(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x + y + 50
|
|
|
|
return _custom_add
|
|
|
|
|
|
def test_registration_overloads(fake_vllm_ir):
|
|
assert all(
|
|
n not in IrOp.registry for n in ["_custom_sub", "_custom_mul", "_custom_div"]
|
|
)
|
|
|
|
# Calling with decorator
|
|
@vllm.ir.register_op()
|
|
def _custom_sub(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x - y
|
|
|
|
assert _custom_sub.name == "_custom_sub"
|
|
assert _custom_sub is IrOp.registry["_custom_sub"]
|
|
|
|
# Custom name
|
|
@vllm.ir.register_op(name="_custom_mul")
|
|
def custom_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x * y
|
|
|
|
assert custom_mul.name == "_custom_mul"
|
|
assert custom_mul is IrOp.registry["_custom_mul"]
|
|
|
|
# Direct construction does not register directly
|
|
def _custom_div(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x / y
|
|
|
|
custom_div = IrOp("_custom_div", _custom_div)
|
|
assert custom_div.name == "_custom_div"
|
|
assert "_custom_div" not in IrOp.registry
|
|
|
|
# Duplicate op registration not allowed
|
|
with pytest.raises(AssertionError):
|
|
|
|
@vllm.ir.register_op
|
|
def _custom_mul(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x * y - 100
|
|
|
|
|
|
def test_no_kw_only_args(fake_vllm_ir):
|
|
# kw-only args not supported
|
|
with pytest.raises(ValueError, match="keyword-only arguments"):
|
|
|
|
@vllm.ir.register_op
|
|
def _custom_kwarg_op(
|
|
x: torch.Tensor, y: torch.Tensor, *, kwarg: int = 0
|
|
) -> torch.Tensor:
|
|
return x + y + kwarg
|
|
|
|
assert "_custom_kwarg_op" not in IrOp.registry
|
|
|
|
|
|
class TestIrOpCustomAdd:
|
|
# Registration invariants
|
|
def test_decorated_object(self, custom_add_op):
|
|
"""Make sure that referring directly to an op is correct"""
|
|
_custom_add = custom_add_op
|
|
assert isinstance(_custom_add, IrOp)
|
|
assert "_custom_add" in IrOp.registry
|
|
assert _custom_add is IrOp.registry["_custom_add"]
|
|
|
|
def test_torch_op_is_registered(self, custom_add_op):
|
|
_custom_add = custom_add_op
|
|
torch_ops = getattr(torch.ops, vllm.ir.op.vllm_ir_torch_lib.ns)
|
|
assert hasattr(torch_ops, "_custom_add")
|
|
assert callable(torch_ops._custom_add.default)
|
|
assert _custom_add.torch_op is torch_ops._custom_add.default
|
|
|
|
# Semantic correctness
|
|
def test_semantics_match_native(self, custom_add_op):
|
|
_custom_add = custom_add_op
|
|
x = torch.randn(4, 5)
|
|
y = torch.randn(4, 5)
|
|
|
|
# Calls native by default
|
|
out = _custom_add(x, y)
|
|
ref = x + y
|
|
|
|
torch.testing.assert_close(out, ref)
|
|
|
|
# -------------------------
|
|
# Implementation registration
|
|
# -------------------------
|
|
|
|
def test_register_impl_is_non_intrusive(self, custom_add_op):
|
|
_custom_add = custom_add_op
|
|
|
|
@_custom_add.register_impl("dummy_provider")
|
|
def dummy_impl(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x + y + 123
|
|
|
|
assert "dummy_provider" in _custom_add.impls
|
|
assert isinstance(_custom_add.impls["dummy_provider"], IrOpImpl)
|
|
|
|
x = torch.ones(2, 2)
|
|
y = torch.ones(2, 2)
|
|
|
|
# Native semantics must still hold
|
|
torch.testing.assert_close(_custom_add(x, y), x + y)
|
|
|
|
def test_schema_contains_tensor_signature(self, custom_add_op):
|
|
_custom_add = custom_add_op
|
|
schema = _custom_add._schema_str
|
|
|
|
assert "Tensor" in schema
|
|
assert "-> Tensor" in schema
|
|
|
|
# -------------------------
|
|
# FX visibility
|
|
# -------------------------
|
|
|
|
@pytest.mark.parametrize("enable_torch_wrap", [True, False])
|
|
@pytest.mark.parametrize("symbolic_trace", [True, False])
|
|
@pytest.mark.parametrize("overload", ["default", "maybe_inplace"])
|
|
def test_trace_sees_single_custom_op(
|
|
self,
|
|
custom_add_op,
|
|
symbolic_trace: bool,
|
|
enable_torch_wrap: bool,
|
|
overload: str,
|
|
):
|
|
_custom_add = custom_add_op
|
|
op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace
|
|
torch_op = (
|
|
_custom_add.torch_op
|
|
if overload == "default"
|
|
else _custom_add.maybe_inplace.torch_op
|
|
)
|
|
|
|
def fn(x, y):
|
|
return op_fn(x, y)
|
|
|
|
def find_fn(target: Any, gm: fx.GraphModule):
|
|
return gm.graph.find_nodes(op="call_function", target=target)
|
|
|
|
with pytest.raises(CustomError), vllm.ir.enable_torch_wrap(enable_torch_wrap):
|
|
if symbolic_trace:
|
|
gm = torch.fx.symbolic_trace(fn)
|
|
else:
|
|
gm = make_fx(fn)(torch.randn(2, 2), torch.randn(2, 2))
|
|
|
|
x1, y1 = torch.rand(5, 4), torch.rand(5, 4)
|
|
out_fx = gm(x1, y1)
|
|
out_eager = fn(x1, y1)
|
|
|
|
# raise error to check enable_torch_wrap context restored correctly
|
|
raise CustomError
|
|
|
|
# check behavior matches eager in all cases
|
|
torch.testing.assert_close(out_fx, out_eager)
|
|
|
|
# check that IR nodes only appear if enable_torch_wrap=True
|
|
ir_nodes = find_fn(torch_op, gm)
|
|
if enable_torch_wrap:
|
|
assert len(ir_nodes) == 1, gm.code
|
|
else:
|
|
assert len(ir_nodes) == 0, gm.code
|
|
|
|
# with torch wrapping enabled (default), IR nodes appear
|
|
if symbolic_trace:
|
|
gm = torch.fx.symbolic_trace(fn)
|
|
else:
|
|
gm = make_fx(fn)(torch.randn(2, 2), torch.randn(2, 2))
|
|
|
|
ir_nodes = find_fn(torch_op, gm)
|
|
assert len(ir_nodes) == 1, gm.code
|
|
|
|
|
|
class TestIrOpImplDispatch:
|
|
def test_register_impl(self, custom_add_op):
|
|
_custom_add = custom_add_op
|
|
assert "impl_a" in _custom_add.impls
|
|
impl = _custom_add.impls["impl_a"]
|
|
|
|
assert impl is _custom_add.impls["impl_a"]
|
|
assert impl.op is _custom_add
|
|
assert impl.provider == "impl_a"
|
|
assert callable(impl.impl_fn)
|
|
|
|
# Test duplicate registration rejected
|
|
with pytest.raises(AssertionError):
|
|
|
|
@_custom_add.register_impl("impl_a")
|
|
def impl_a_dup(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x + y + 30
|
|
|
|
# Check the original impl is still intact
|
|
assert _custom_add.impls["impl_a"] is impl
|
|
|
|
# Check support all args
|
|
assert _custom_add.impls["impl_a"].supports_all_args
|
|
assert _custom_add.impls["impl_b"].supports_all_args
|
|
assert not _custom_add.impls["impl_even"].supports_all_args
|
|
|
|
def test_reserved_provider_rejected(self, custom_add_op):
|
|
_custom_add = custom_add_op
|
|
for provider in RESERVED_PROVIDERS:
|
|
with pytest.raises(AssertionError):
|
|
|
|
@_custom_add.register_impl(provider)
|
|
def bad_impl(x, y):
|
|
return x + y
|
|
|
|
def test_set_priority_scoped(self, custom_add_op):
|
|
_custom_add = custom_add_op
|
|
assert _custom_add.get_priority() == []
|
|
|
|
with _custom_add.set_priority(["impl_even", "impl_b"]):
|
|
assert _custom_add.get_priority() == ["impl_even", "impl_b"]
|
|
|
|
# Check nesting
|
|
with _custom_add.set_priority(["impl_b"]):
|
|
assert _custom_add.get_priority() == ["impl_b"]
|
|
|
|
# Restored
|
|
assert _custom_add.get_priority() == ["impl_even", "impl_b"]
|
|
|
|
# Check that exception restores priority
|
|
with pytest.raises(CustomError), _custom_add.set_priority(["impl_a"]):
|
|
assert _custom_add.get_priority() == ["impl_a"]
|
|
raise CustomError
|
|
|
|
# Restored again
|
|
assert _custom_add.get_priority() == ["impl_even", "impl_b"]
|
|
|
|
# Restored to empty
|
|
assert _custom_add.get_priority() == []
|
|
|
|
@pytest.mark.parametrize(
|
|
"default,override",
|
|
[
|
|
(["impl_even", "impl_b"], ["impl_a"]),
|
|
(["impl_a"], ["impl_even", "impl_b"]),
|
|
],
|
|
)
|
|
def test_set_default_priority(
|
|
self, custom_add_op, default: list[str], override: list[str]
|
|
):
|
|
_custom_add = custom_add_op
|
|
assert _custom_add.get_priority() == []
|
|
|
|
_custom_add.set_default(default)
|
|
assert _custom_add.get_priority() == default
|
|
|
|
# Priority doesn't change after exiting the set_priority context.
|
|
with _custom_add.set_priority(override):
|
|
assert _custom_add.get_priority() == override
|
|
assert _custom_add.get_priority() == default
|
|
|
|
# Should override the previous default.
|
|
_custom_add.set_default(override)
|
|
assert _custom_add.get_priority() == override
|
|
|
|
@pytest.mark.parametrize("overload", ["default", "maybe_inplace"])
|
|
def test_dispatch_priority_order(self, custom_add_op, overload: str):
|
|
_custom_add = custom_add_op
|
|
op_fn = _custom_add if overload == "default" else _custom_add.maybe_inplace
|
|
torch_op = (
|
|
_custom_add.torch_op
|
|
if overload == "default"
|
|
else _custom_add.maybe_inplace.torch_op
|
|
)
|
|
|
|
x = torch.tensor(1, dtype=torch.int32)
|
|
y = torch.tensor(2, dtype=torch.int32)
|
|
|
|
with _custom_add.set_priority(["impl_b", "impl_a"]):
|
|
assert _custom_add.dispatch(x, y) is _custom_add.impls["impl_b"]
|
|
out1 = op_fn(x.clone(), y)
|
|
out2 = torch_op(x.clone(), y)
|
|
|
|
with _custom_add.set_priority(["impl_a"]):
|
|
assert _custom_add.dispatch(x, y) is _custom_add.impls["impl_a"]
|
|
out3 = op_fn(x.clone(), y)
|
|
out4 = torch_op(x.clone(), y)
|
|
|
|
# impl_b
|
|
assert out1.item() == 1 + 2 + 20
|
|
assert out2.item() == 1 + 2 + 20
|
|
# impl_a
|
|
assert out3.item() == 1 + 2 + 10
|
|
assert out4.item() == 1 + 2 + 10
|
|
|
|
def test_unsupported_impl_filtered(self, custom_add_op):
|
|
_custom_add = custom_add_op
|
|
|
|
@_custom_add.register_impl("impl_unsupported", supported=False)
|
|
def impl_unsupported(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x + y + 999
|
|
|
|
x = torch.tensor(1, dtype=torch.int32)
|
|
y = torch.tensor(2, dtype=torch.int32)
|
|
|
|
with _custom_add.set_priority(["impl_unsupported", "impl_a"]):
|
|
assert _custom_add.get_priority() == ["impl_a"]
|
|
out = _custom_add(x, y)
|
|
|
|
# impl_unsupported skipped → impl_a
|
|
assert out.item() == 1 + 2 + 10
|
|
|
|
def test_supports_args_runtime_dispatch_and_warning(
|
|
self, custom_add_op, caplog_vllm: pytest.LogCaptureFixture
|
|
):
|
|
_custom_add = custom_add_op
|
|
x1 = torch.ones((2, 2), dtype=torch.int32)
|
|
y1 = torch.full((2, 2), 2, dtype=torch.int32)
|
|
|
|
x2 = torch.ones((2, 3), dtype=torch.int32)
|
|
y2 = torch.full((2, 3), 2, dtype=torch.int32)
|
|
|
|
with (
|
|
caplog_vllm.at_level(logging.WARNING),
|
|
_custom_add.set_priority(["impl_even"]),
|
|
):
|
|
# Test the warning about native fallback is logged (before even dispatching)
|
|
assert len(caplog_vllm.records) == 1
|
|
message = caplog_vllm.records[0].message
|
|
assert "_custom_add" in message
|
|
assert "fallback to native" in message
|
|
assert "priority" in message
|
|
|
|
# Check dispatching
|
|
assert _custom_add.get_priority() == ["impl_even", "native"]
|
|
assert _custom_add.dispatch(x1, y1) is _custom_add.impls["impl_even"]
|
|
assert _custom_add.dispatch(x2, y2) is _custom_add.impls["native"]
|
|
|
|
out1 = _custom_add(x1, y1) # size(1) == 2 → impl_even
|
|
out2 = _custom_add(x2, y2) # size(1) == 3 → native fallback
|
|
|
|
# no other warnings
|
|
assert len(caplog_vllm.records) == 1
|
|
assert torch.all(out1 == 1 + 2 + 50)
|
|
assert torch.all(out2 == 1 + 2)
|
|
|
|
def test_default_priority(
|
|
self,
|
|
custom_add_op,
|
|
caplog_vllm: pytest.LogCaptureFixture,
|
|
disable_log_dedup,
|
|
):
|
|
_custom_add = custom_add_op
|
|
# Make sure logs are not deduplicated to properly test the warning
|
|
x = torch.tensor([3], dtype=torch.int32)
|
|
y = torch.tensor([4], dtype=torch.int32)
|
|
|
|
# No priority set → falls back to native
|
|
assert _custom_add.get_priority() == []
|
|
with caplog_vllm.at_level(logging.WARNING):
|
|
# Native by default
|
|
assert _custom_add.dispatch(x, y) is _custom_add.impls["native"]
|
|
out = _custom_add(x, y)
|
|
|
|
# Check dispatching to native by default
|
|
assert out.item() == 3 + 4
|
|
|
|
# Check warning
|
|
assert len(caplog_vllm.records) == 2
|
|
message = caplog_vllm.records[0].message.lower()
|
|
assert "_custom_add" in message
|
|
assert "priority not set" in message
|
|
|
|
|
|
@pytest.mark.parametrize("default", [True, False])
|
|
def test_set_default_torch_wrap(default: bool):
|
|
"""set_default_torch_wrap permanently flips the global flag."""
|
|
original = vllm.ir.op._ENABLE_TORCH_WRAP
|
|
try:
|
|
vllm.ir.set_default_torch_wrap(default)
|
|
assert vllm.ir.op._ENABLE_TORCH_WRAP is default
|
|
|
|
# Flag doesn't change after exiting the enable_torch_wrap context.
|
|
with vllm.ir.enable_torch_wrap(not default):
|
|
assert vllm.ir.op._ENABLE_TORCH_WRAP is (not default)
|
|
assert vllm.ir.op._ENABLE_TORCH_WRAP is default
|
|
|
|
# Should override the previous default.
|
|
vllm.ir.set_default_torch_wrap(not default)
|
|
assert vllm.ir.op._ENABLE_TORCH_WRAP is (not default)
|
|
finally:
|
|
vllm.ir.op._ENABLE_TORCH_WRAP = original
|
|
|
|
|
|
@pytest.fixture
|
|
def custom_mm_op(fake_vllm_ir):
|
|
"""Fixture that registers ``_custom_mm`` (isolated by ``fake_vllm_ir``)."""
|
|
|
|
@vllm.ir.register_op
|
|
def _custom_mm(
|
|
x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
|
|
) -> torch.Tensor:
|
|
tmp = x @ y
|
|
return tmp if bias is None else tmp + bias
|
|
|
|
return _custom_mm
|
|
|
|
|
|
def test_default_args(custom_mm_op):
|
|
_custom_mm = custom_mm_op
|
|
|
|
# Test that default args are properly applied when dispatching and calling
|
|
@_custom_mm.register_impl("impl_mm", supports_args=lambda x, y, bias=None: True)
|
|
def impl_mm(
|
|
x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
|
|
) -> torch.Tensor:
|
|
tmp = x @ y
|
|
return tmp + 50 if bias is None else tmp + bias + 100
|
|
|
|
x1 = torch.tensor([1, 2], dtype=torch.int32)
|
|
x2 = torch.tensor([3, 4], dtype=torch.int32)
|
|
|
|
# Test that supports_args receives the defaulted args
|
|
assert impl_mm.supports_args(x1, x2)
|
|
with _custom_mm.set_priority(["impl_mm", "native"]):
|
|
assert _custom_mm.dispatch(x1, x2) is impl_mm
|
|
|
|
|
|
def test_bad_impl_registrations(custom_mm_op):
|
|
_custom_mm = custom_mm_op
|
|
# Check bad schema
|
|
with pytest.raises(ValueError, match="does not match native schema"):
|
|
|
|
@_custom_mm.register_impl("impl_mm_bad_schema")
|
|
def impl_mm_bad_schema(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x @ y - 1
|
|
|
|
with pytest.raises(ValueError, match="does not match native schema"):
|
|
|
|
@_custom_mm.register_impl("impl_mm_bad_schema_2")
|
|
def impl_mm_bad_schema_2(
|
|
x: torch.Tensor, y: torch.Tensor, b: torch.Tensor | None = None
|
|
) -> torch.Tensor:
|
|
return x @ y + b - 2
|
|
|
|
with pytest.raises(ValueError, match="does not match native schema"):
|
|
|
|
@_custom_mm.register_impl("impl_mm_bad_schema_3")
|
|
def impl_mm_bad_schema_3(
|
|
x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor
|
|
) -> torch.Tensor:
|
|
return x @ y + bias - 5
|
|
|
|
# check supports_args with incorrect params
|
|
with pytest.raises(ValueError, match="supports_args must be a callable"):
|
|
|
|
@_custom_mm.register_impl("impl_mm_bad_supports_args", supports_args=True)
|
|
def impl_mm_bad_supports_args(
|
|
x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
|
|
) -> torch.Tensor:
|
|
return x @ y + 10
|
|
|
|
with pytest.raises(ValueError, match="number of parameters"):
|
|
|
|
@_custom_mm.register_impl(
|
|
"impl_mm_bad_supports_args_2", supports_args=lambda x, y: True
|
|
)
|
|
def impl_mm_bad_supports_args(
|
|
x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
|
|
) -> torch.Tensor:
|
|
return x @ y + 10
|
|
|
|
with pytest.raises(ValueError, match="keyword-only parameters"):
|
|
|
|
@_custom_mm.register_impl(
|
|
"impl_mm_bad_supports_args_3", supports_args=lambda x, y, *, b: True
|
|
)
|
|
def impl_mm_bad_supports_args_2(
|
|
x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
|
|
) -> torch.Tensor:
|
|
return x @ y + 20
|
|
|
|
with pytest.raises(ValueError, match="does not match native parameter"):
|
|
|
|
@_custom_mm.register_impl(
|
|
"impl_mm_bad_supports_args_4", supports_args=lambda x, y, b: True
|
|
)
|
|
def impl_mm_bad_supports_args_4(
|
|
x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
|
|
) -> torch.Tensor:
|
|
return x @ y + 30
|
|
|
|
with pytest.raises(ValueError, match="does not match native default"):
|
|
|
|
@_custom_mm.register_impl(
|
|
"impl_mm_bad_supports_args_5", supports_args=lambda x, y, bias=1: True
|
|
)
|
|
def impl_mm_bad_supports_args_5(
|
|
x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
|
|
) -> torch.Tensor:
|
|
return x @ y + 40
|
|
|
|
# With fixture, each test gets a fresh op with only "native" impl
|
|
assert set(_custom_mm.impls.keys()) == {"native"}
|
|
|
|
|
|
IMPL_OOT_SRC = """
|
|
import torch
|
|
|
|
@_custom_mm.register_impl("impl_mm_oot")
|
|
def impl_mm_oot(
|
|
x: torch.Tensor, y: torch.Tensor, bias: torch.Tensor | None = None
|
|
) -> torch.Tensor:
|
|
return x @ y - 99
|
|
"""
|
|
|
|
|
|
def load_custom_mm_module(file_path: Path, custom_mm_op):
|
|
spec = importlib.util.spec_from_file_location("_custom_mm_oot", file_path)
|
|
assert spec is not None
|
|
module = importlib.util.module_from_spec(spec)
|
|
|
|
# Inject the variable into the module's global namespace
|
|
# This allows the @_custom_mm.register_impl decorator to work
|
|
module._custom_mm = custom_mm_op # type: ignore[attr-defined]
|
|
|
|
# Execute the file; this triggers the decorator
|
|
assert spec.loader is not None
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def test_uuid_and_oot(custom_mm_op, tmp_path: Path):
|
|
_custom_mm = custom_mm_op
|
|
file_path = tmp_path / "_custom_mm_oot.py"
|
|
file_path.write_text(IMPL_OOT_SRC)
|
|
|
|
assert "impl_mm_oot" not in _custom_mm.impls
|
|
_ = load_custom_mm_module(file_path, _custom_mm)
|
|
assert "impl_mm_oot" in _custom_mm.impls
|
|
|
|
uuid = _custom_mm.impls["impl_mm_oot"].uuid()
|
|
del _custom_mm.impls["impl_mm_oot"]
|
|
|
|
# Replace file source
|
|
file_path.write_text(IMPL_OOT_SRC + " # added file source")
|
|
assert "impl_mm_oot" not in _custom_mm.impls
|
|
_ = load_custom_mm_module(file_path, _custom_mm)
|
|
assert "impl_mm_oot" in _custom_mm.impls
|
|
|
|
uuid1 = _custom_mm.impls["impl_mm_oot"].uuid()
|
|
assert uuid1 != uuid
|
|
del _custom_mm.impls["impl_mm_oot"]
|
|
|
|
# Back to original
|
|
file_path.write_text(IMPL_OOT_SRC)
|
|
assert "impl_mm_oot" not in _custom_mm.impls
|
|
_ = load_custom_mm_module(file_path, _custom_mm)
|
|
assert "impl_mm_oot" in _custom_mm.impls
|
|
|
|
uuid2 = _custom_mm.impls["impl_mm_oot"].uuid()
|
|
assert uuid2 == uuid
|
|
assert uuid2 != uuid1
|
|
del _custom_mm.impls["impl_mm_oot"]
|
|
|
|
|
|
def _test_native(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
|
|
return x + y
|
|
|
|
|
|
def _make_op_with_generator(name: str = "_ig_test"):
|
|
op = IrOp(name, _test_native)
|
|
|
|
@op.register_input_generator
|
|
def _gen(n: int = 4):
|
|
x = torch.randn(n, 3)
|
|
y = torch.randn(n, 3)
|
|
return x, y
|
|
|
|
return op
|
|
|
|
|
|
def _test_native_single(x: torch.Tensor) -> torch.Tensor:
|
|
return x
|
|
|
|
|
|
class TestInputGenerator:
|
|
def test_no_input_generator_by_default(self):
|
|
op = IrOp("_ig_test_no_gen", _test_native_single)
|
|
assert not op.has_input_generator
|
|
|
|
def test_register_input_generator(self):
|
|
op = _make_op_with_generator("_ig_test_reg")
|
|
assert op.has_input_generator
|
|
|
|
def test_generate_inputs_returns_tuple(self):
|
|
op = _make_op_with_generator("_ig_test_tuple")
|
|
result = op.generate_inputs(n=2)
|
|
assert isinstance(result, tuple)
|
|
assert len(result) == 2
|
|
assert result[0].shape == (2, 3)
|
|
assert result[1].shape == (2, 3)
|
|
|
|
def test_generate_inputs_default_kwargs(self):
|
|
op = _make_op_with_generator("_ig_test_default")
|
|
result = op.generate_inputs()
|
|
assert result[0].shape == (4, 3)
|
|
|
|
def test_generate_inputs_without_registration_raises(self):
|
|
op = IrOp("_ig_test_no_gen_raises", _test_native_single)
|
|
with pytest.raises(RuntimeError, match="No input generator"):
|
|
op.generate_inputs()
|
|
|
|
|
|
class TestTolerance:
|
|
def test_override_and_get_tolerance(self):
|
|
op = IrOp("_tol_test", _test_native)
|
|
|
|
tol = op.get_tolerance(torch.float32)
|
|
assert tol == {"atol": 1e-5, "rtol": 1.3e-6}
|
|
|
|
op.override_tolerance(torch.float32, atol=0.1, rtol=0.2)
|
|
assert op.get_tolerance(torch.float32) == {"atol": 0.1, "rtol": 0.2}
|
|
assert op.get_tolerance(torch.float16) == {"atol": 1e-3, "rtol": 1e-3}
|
|
|
|
def test_get_tolerance_raises_for_unknown_dtype(self):
|
|
op = IrOp("_tol_test_unknown", _test_native)
|
|
with pytest.raises(ValueError, match="No tolerance defined"):
|
|
op.get_tolerance(torch.complex64)
|
|
|
|
|
|
def test_naming_validation(fake_vllm_ir):
|
|
"""Test that op and provider names are validated ([a-z_][a-z_0-9]*)."""
|
|
|
|
# Valid op and provider names
|
|
@vllm.ir.register_op
|
|
def _valid_name_123(x: torch.Tensor) -> torch.Tensor:
|
|
return x
|
|
|
|
@_valid_name_123.register_impl("valid_provider_123")
|
|
def valid_impl(x: torch.Tensor) -> torch.Tensor:
|
|
return x + 1
|
|
|
|
# Invalid op names should fail
|
|
with pytest.raises(ValueError, match="name.*invalid"):
|
|
|
|
@vllm.ir.register_op
|
|
def InvalidName(x: torch.Tensor) -> torch.Tensor:
|
|
return x
|
|
|
|
with pytest.raises(ValueError, match="name.*invalid"):
|
|
|
|
@vllm.ir.register_op(name="123invalid")
|
|
def some_func(x: torch.Tensor) -> torch.Tensor:
|
|
return x
|
|
|
|
# Invalid provider names should fail
|
|
with pytest.raises(ValueError, match="name.*invalid"):
|
|
|
|
@_valid_name_123.register_impl("Invalid-Provider")
|
|
def invalid_impl(x: torch.Tensor) -> torch.Tensor:
|
|
return x + 1
|
|
|
|
|
|
def test_registration_stack_traces(fake_vllm_ir):
|
|
"""Test that stack traces are captured for ops and impls."""
|
|
|
|
@vllm.ir.register_op
|
|
def _test_stack(x: torch.Tensor) -> torch.Tensor:
|
|
return x
|
|
|
|
@_test_stack.register_impl("test_provider")
|
|
def test_impl(x: torch.Tensor) -> torch.Tensor:
|
|
return x + 1
|
|
|
|
# Verify op stack trace
|
|
assert hasattr(_test_stack, "_registration_stack")
|
|
assert len(_test_stack._registration_stack) > 0
|
|
op_stack_str = "".join(_test_stack._registration_stack)
|
|
assert "test_op.py" in op_stack_str
|
|
# Last frame should be the decorator in user code, not internal decorator logic
|
|
assert "@vllm.ir.register_op" in _test_stack._registration_stack[-1]
|
|
assert "return decorator(f)" not in op_stack_str
|
|
|
|
# Verify impl stack trace
|
|
impl = _test_stack.impls["test_provider"]
|
|
assert hasattr(impl, "_registration_stack")
|
|
assert len(impl._registration_stack) > 0
|
|
impl_stack_str = "".join(impl._registration_stack)
|
|
assert "test_op.py" in impl_stack_str
|
|
# Last frame should be the decorator in user code
|
|
assert '@_test_stack.register_impl("test_provider")' in impl._registration_stack[-1]
|
|
|
|
|
|
def test_op_repr_uses_docstring(fake_vllm_ir):
|
|
"""Test that __str__ uses the function's docstring and __repr__ is simple."""
|
|
|
|
@vllm.ir.register_op
|
|
def _test_repr_with_doc(x: torch.Tensor) -> torch.Tensor:
|
|
"""First line of docstring.
|
|
|
|
Additional details here.
|
|
"""
|
|
return x
|
|
|
|
@vllm.ir.register_op
|
|
def _test_repr_no_doc(x: torch.Tensor) -> torch.Tensor:
|
|
return x
|
|
|
|
# __str__ with docstring: uses first line only
|
|
str_with = str(_test_repr_with_doc)
|
|
assert "IrOp('_test_repr_with_doc')" in str_with
|
|
assert "First line of docstring." in str_with
|
|
assert "Additional details" not in str_with
|
|
|
|
# __str__ without docstring: simple format
|
|
assert str(_test_repr_no_doc) == "IrOp('_test_repr_no_doc')"
|
|
|
|
# __repr__ should be simple for both
|
|
assert repr(_test_repr_with_doc) == "IrOp('_test_repr_with_doc')"
|
|
assert repr(_test_repr_no_doc) == "IrOp('_test_repr_no_doc')"
|
|
|
|
|
|
def test_vllm_ir_fixture(fake_vllm_ir):
|
|
"""Test that the fake_vllm_ir fixture provides test isolation."""
|
|
|
|
@vllm.ir.register_op
|
|
def _test_fixture(x: torch.Tensor) -> torch.Tensor:
|
|
return x
|
|
|
|
assert "_test_fixture" in IrOp.registry
|
|
# Fixture will automatically clean up after test
|