diff --git a/vllm_fl/dispatch/backends/vendor/metax/impl/activation.py b/vllm_fl/dispatch/backends/vendor/metax/impl/activation.py index 49cd735e..325b6645 100644 --- a/vllm_fl/dispatch/backends/vendor/metax/impl/activation.py +++ b/vllm_fl/dispatch/backends/vendor/metax/impl/activation.py @@ -15,17 +15,12 @@ # SPDX-License-Identifier: Apache-2.0 # 2026 - Modified by MetaX Integrated Circuits (Shanghai) Co., Ltd. All Rights Reserved. import torch -from vllm.model_executor.layers.activation import ( - SiluAndMul, - GeluAndMul, -) +import torch.nn.functional as F def silu_and_mul_maca(obj, x: torch.Tensor) -> torch.Tensor: """ - SiLU activation followed by element-wise multiplication using CUDA. - - Uses vLLM's optimized CUDA kernel when available. + SiLU activation followed by element-wise multiplication. Args: obj: The calling obj (for interface consistency) @@ -34,15 +29,13 @@ def silu_and_mul_maca(obj, x: torch.Tensor) -> torch.Tensor: Returns: Output tensor of shape [..., d] """ - act_fn = SiluAndMul() - return act_fn.forward_cuda(x) + d = x.shape[-1] // 2 + return F.silu(x[..., :d]) * x[..., d:] def gelu_and_mul_maca(obj, x: torch.Tensor) -> torch.Tensor: """ - GELU activation followed by element-wise multiplication using CUDA. - - Uses vLLM's optimized CUDA kernel when available. + GELU activation followed by element-wise multiplication. Args: obj: The calling obj (for interface consistency) @@ -51,5 +44,5 @@ def gelu_and_mul_maca(obj, x: torch.Tensor) -> torch.Tensor: Returns: Output tensor of shape [..., d] """ - act_fn = GeluAndMul() - return act_fn.forward_cuda(x) + d = x.shape[-1] // 2 + return F.gelu(x[..., :d], approximate="none") * x[..., d:]