[backends] Add Quack (#289)

Author: xuzhao9

.gitmodules

@@ -19,3 +19,6 @@
 [submodule "submodules/aiter"]
 	path = submodules/aiter
 	url = https://github.com/ROCm/aiter.git
+[submodule "submodules/quack"]
+	path = submodules/quack
+	url = https://github.com/Dao-AILab/quack.git

install.py

@@ -198,6 +198,11 @@ def setup_hip(args: argparse.Namespace):
     if args.liger or args.all:
         logger.info("[tritonbench] installing liger-kernels...")
         install_liger()
+    if args.quack or args.all:
+        logger.info("[tritonbench] installing quack...")
+        from tools.quack.install import install_quack
+
+        install_quack()
     if args.xformers:
         logger.info("[tritonbench] installing xformers...")
         from tools.xformers.install import install_xformers

submodules/quack

@@ -0,0 +1,14 @@
+import os
+import subprocess
+
+from pathlib import Path
+
+
+REPO_PATH = Path(os.path.abspath(__file__)).parent.parent.parent
+CURRENT_DIR = Path(os.path.abspath(__file__)).parent
+QUACK_PATH = REPO_PATH.joinpath("submodules", "quack")
+
+
+def install_quack():
+    cmd = ["pip", "install", "-e", "."]
+    subprocess.check_call(cmd, cwd=QUACK_PATH)

tools/quack/install.py

@@ -23,6 +23,11 @@
 except ModuleNotFoundError:
     HAS_AITER = False
 
+try:
+    from .quack import QuackRMSNorm
+except ModuleNotFoundError:
+    QuackRMSNorm = None
+
 
 # Reference: https://github.com/linkedin/Liger-Kernel/
 # blob/main/benchmark/scripts/benchmark_rms_norm.py
@@ -72,6 +77,11 @@ def liger_rms(self, H, input) -> Callable:
         self.liger_rms_op = LigerRMSNorm(hidden_size=H, eps=self.eps).to(self.device)
         return lambda: self.liger_rms_op(input)
 
+    @register_benchmark(enabled=QuackRMSNorm)
+    def quack_rms(self, H, input) -> Callable:
+        self.quack_rms_op = QuackRMSNorm(hidden_size=H, eps=self.eps).to(self.device)
+        return lambda: self.quack_rms_op(input)
+
     @register_benchmark()
     def inductor_rms(self, H, input) -> Callable:
         if self.llama_rms_op is None:

tritonbench/operators/rms_norm/operator.py

@@ -0,0 +1,16 @@
+import torch
+
+from quack import rmsnorm as quack_rmsnorm
+
+
+class QuackRMSNorm(torch.nn.Module):
+    def __init__(self, hidden_size, eps=1e-6):
+        """
+        AITerRMSNorm
+        """
+        super().__init__()
+        self.weight = torch.nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states):
+        return quack_rmsnorm(hidden_states, self.weight, self.variance_epsilon)

tritonbench/operators/rms_norm/quack.py

@@ -15,6 +15,13 @@
     register_metric,
 )
 
+try:
+    from quack.softmax import softmax as quack_softmax
+
+    HAS_QUACK = True
+except ImportError:
+    HAS_QUACK = False
+
 
 class Operator(BenchmarkOperator):
     is_compute_bound = False
@@ -105,6 +112,11 @@ def _inner():
 
         return _inner
 
+    @register_benchmark(enabled=HAS_QUACK)
+    def quack(self, x):
+        inner = lambda: quack_softmax(x)
+        return inner
+
     def get_input_iter(self):
         M = 4096
         shapes = [(M, 128 * i) for i in range(2, 100)]