From 93477454f9155d4f74352e3d725c83f91c6894bf Mon Sep 17 00:00:00 2001 From: roikoren755 <26850796+roikoren755@users.noreply.github.com> Date: Wed, 29 Jul 2026 21:27:31 +0300 Subject: [PATCH] [torch.compile] Compile `CustomOp.forward_native` for ReLU^2 to avoid raw torch ops inside opaque custom ops (#50244) Signed-off-by: Roi Koren --- vllm/model_executor/layers/activation.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index c8a51f4c22..953a47e0fa 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -658,8 +658,8 @@ class ReLUSquaredActivation(CustomOp): # --8<-- [end:relu2] - def __init__(self): - super().__init__() + def __init__(self, *, compile_native: bool = True): + super().__init__(compile_native=compile_native) if current_platform.is_cuda_alike(): self.op = torch.ops._C.relu_squared