init v0.23.0

Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>

vllm_ascend/ops/activation.py

@@ -16,29 +16,65 @@
 #
 
 import torch
-from vllm.model_executor.layers.activation import QuickGELU, SiluAndMul
+import torch_npu
+from vllm.model_executor.layers.activation import (
+    QuickGELU,
+    SiluAndMul,
+    SiluAndMulWithClamp,
+    SwigluOAIAndMul,
+    SwigluStepAndMul,
+)
+
+from vllm_ascend.utils import get_weight_prefetch_method
 
 
 class AscendQuickGELU(QuickGELU):
-
     def forward_oot(self, x: torch.tensor) -> torch.Tensor:
-        import torch_npu
-
         out = torch_npu.npu_fast_gelu(x)
         return out
 
 
 class AscendSiluAndMul(SiluAndMul):
-
     def forward_oot(self, x: torch.Tensor) -> torch.Tensor:
-        import torch_npu
-
-        from vllm_ascend.utils import is_310p
-
-        torch.ops.vllm.maybe_prefetch_mlp_down_proj(x)
-        if is_310p():
-            out = torch_npu.npu_swiglu(x.to(torch.float32)).to(torch.float16)
-        else:
-            out = torch_npu.npu_swiglu(x)
-        torch.ops.vllm.maybe_wait_prefetch_done(out)
+        weight_prefetch_method = get_weight_prefetch_method()
+        weight_prefetch_method.maybe_prefetch_mlp_weight_preprocess(weight_prefetch_method.MLP_DOWN, x)
+        out = torch_npu.npu_swiglu(x)
+        weight_prefetch_method.maybe_prefetch_mlp_weight_postprocess(out)
         return out
+
+
+class AscendSiluAndMulWithClamp(SiluAndMulWithClamp):
+    def forward_oot(self, x: torch.Tensor) -> torch.Tensor:
+        weight_prefetch_method = get_weight_prefetch_method()
+        weight_prefetch_method.maybe_prefetch_mlp_weight_preprocess(weight_prefetch_method.MLP_DOWN, x)
+        d = x.shape[-1] // 2
+        gate = torch.clamp(x[..., :d], max=self.swiglu_limit)
+        up = torch.clamp(x[..., d:], min=-self.swiglu_limit, max=self.swiglu_limit)
+        x = torch.cat([gate, up], dim=-1)
+        out = torch_npu.npu_swiglu(x)
+        weight_prefetch_method.maybe_prefetch_mlp_weight_postprocess(out)
+        return out
+
+
+class AscendSwigluOAIAndMul:
+    def swiglu_oai_forward(x: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor:
+        class MinimalSwigluOAIAndMul:
+            def __init__(self):
+                self.alpha = alpha
+                self.limit = limit
+
+        layer = MinimalSwigluOAIAndMul()
+        return SwigluOAIAndMul.forward_native(layer, x)
+
+
+class AscendSwigluStepAndMul:
+    def swiglustep_forward(x: torch.Tensor, limit: float = 7.0) -> torch.Tensor:
+        if limit is None:
+            raise ValueError("SwigluStepAndMul requires limit to be set.")
+
+        class MinimalSwigluStepAndMul:
+            def __init__(self):
+                self.limit = limit
+
+        layer = MinimalSwigluStepAndMul()
+        return SwigluStepAndMul.forward_native(layer, x)