support QuickGELU (#3250)

2025-02-01 19:31:47 +08:00
parent 4eb4b401cc
commit 8db776f049
2 changed files with 10 additions and 1 deletions
--- a/python/sglang/srt/layers/activation.py
+++ b/python/sglang/srt/layers/activation.py
@@ -72,6 +72,15 @@ class GeluAndMul(CustomOp):
        return out


+class QuickGELU(CustomOp):
+    def forward_native(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.sigmoid(1.702 * x)
+
+    def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
+        # TODO(zhyncs): Implement the CUDA kernel for QuickGELU in sgl-kernel
+        return self.forward_native(x)
+
+
 class ScaledActivation(nn.Module):
    """An activation function with post-scale parameters.