A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 05:45 UTC
[Minor] avoid register new custom and just import silly_attn (#28578)
Signed-off-by: Boyuan Feng <boyuan@meta.com>
Boyuan Feng committed
Nov 14, 2025 at 01:32 UTC
fd75d3e8c0f522178e39845276fd57908760b4d0
1 file changed
+3
-9
tests/compile/test_config.py
+3
-9
@@ -15,6 +15,9 @@ from vllm.engine.arg_utils import EngineArgs
15
from vllm.platforms import current_platform
16
from vllm.utils.torch_utils import _is_torch_equal_or_newer
17
18
+# This import automatically registers `torch.ops.silly.attention`
19
+from . import silly_attention # noqa: F401
20
+
21
22
def test_version():
23
# Test the version comparison logic using the private function
@@ -257,15 +260,6 @@ def test_should_split():
260
splitting_ops = ["aten::add.Tensor"]
261
assert not should_split(node, splitting_ops)
262
260
- @torch.library.custom_op(
261
- "silly::attention",
262
- mutates_args=["out"],
263
- )
264
- def attention(
265
- q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, out: torch.Tensor
266
- ) -> None:
267
- out.copy_(q + k + v)
268
-
263
q, k, v, out = [torch.randn(1)] * 4
264
265
# supports custom ops as OpOverloadPacket