[torch.compile] Refactor duplicate torch op registrations

tests/compile/piecewise/test_multiple_graphs.py

@@ -19,7 +19,7 @@
 from vllm.utils import direct_register_custom_op
 
 # create a library to hold the custom op
-silly_lib = Library("silly", "FRAGMENT")  # noqa
+silly_lib = Library("silly_multiple", "FRAGMENT")  # noqa
 
 BATCH_SIZE = 32
 MLP_SIZE = 128
@@ -90,7 +90,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.pre_attn(x)
         x = self.rms_norm_ref(x)
         attn_output = torch.empty_like(x)
-        torch.ops.silly.attention(x, x, x, attn_output)
+        torch.ops.silly_multiple.attention(x, x, x, attn_output)
         x = attn_output
         x = self.rms_norm_ref(x)
         x = self.post_attn(x)
@@ -188,7 +188,7 @@ def __init__(self,
         def forward(self, x: torch.Tensor) -> torch.Tensor:
             x = x + x
             attn_output = torch.empty_like(x)
-            torch.ops.silly.attention(x, x, x, attn_output)
+            torch.ops.silly_multiple.attention(x, x, x, attn_output)
             x = attn_output
             x = x * 3
             return x

tests/compile/piecewise/test_simple.py

@@ -20,7 +20,7 @@
 global_counter = 0
 
 # create a library to hold the custom op
-silly_lib = Library("silly", "FRAGMENT")  # noqa
+silly_lib = Library("silly_simple", "FRAGMENT")  # noqa
 
 
 def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
@@ -66,12 +66,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = x + 1
         x = x + 2
         out = torch.empty_like(x)
-        torch.ops.silly.attention(x, x, x, out)
+        torch.ops.silly_simple.attention(x, x, x, out)
         x = out
         x = x - 2
         x = x - 1
         out = torch.empty_like(x)
-        torch.ops.silly.attention(x, x, x, out)
+        torch.ops.silly_simple.attention(x, x, x, out)
         x = out
         x = x + 1
         return x

tests/compile/piecewise/test_toy_llama.py

@@ -24,7 +24,7 @@
 from vllm.utils import direct_register_custom_op
 
 # create a library to hold the custom op
-silly_lib = Library("silly", "FRAGMENT")  # noqa
+silly_lib = Library("silly_toy_llama", "FRAGMENT")  # noqa
 
 
 def silly_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
@@ -160,7 +160,7 @@ def forward(
         k = k + positions.unsqueeze(1)
 
         attn_output = torch.empty_like(q)
-        torch.ops.silly.attention(q, k, v, attn_output)
+        torch.ops.silly_toy_llama.attention(q, k, v, attn_output)
 
         output = self.output_projection(attn_output)
         return output