[torch.compile] Compile `CustomOp.forward_native` for ReLU^2 to avoid raw torch ops inside opaque custom ops (#50244)
Signed-off-by: Roi Koren <roik@nvidia.com>
roikoren755 committed
Jul 29, 2026 at 21:27 UTC
93477454f9155d4f74352e3d725c83f91c6894bf
1 file changed
+2
-2
vllm/model_executor/layers/activation.py
+2
-2
@@ -658,8 +658,8 @@ class ReLUSquaredActivation(CustomOp):
658
659
# --8<-- [end:relu2]
660
661
- def __init__(self):
662
- super().__init__()
661
+ def __init__(self, *, compile_native: bool = True):
662
+ super().__init__(compile_native=compile_native)
663
if current_platform.is_cuda_alike():
664
self.op = torch.ops._C.relu_squared
665