feat: moe kernel tuning

.github/actions/setup-uv-project/action.yml

@@ -12,4 +12,4 @@ runs:
         github-token: ${{ github.token }}
 
     - shell: bash
-      run: uv sync --extra dev
+      run: uv sync --extra dev --extra vllm

pyproject.toml

@@ -141,6 +141,10 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+vllm = [
+    "vllm>=0.11.0",
+    "ray",
+]
 stable-fast = [
     "xformers>=0.0.30",
     "stable-fast-pruna==1.0.8",

src/pruna/algorithms/base/tags.py

@@ -80,6 +80,24 @@ class AlgorithmTag(Enum):
         "Resamplers change the shape of image or video latents during generation to speed up inference.",
     )
 
+    @classmethod
+    def tags_compatible_with_moe_kernel(cls) -> list["AlgorithmTag"]:
+        """
+        Return tags that the MoE kernel tuner is compatible with (for ordering).
+
+        Used so that compatible_before / compatible_after can be derived from
+        the enum and stay in sync when new tags are added.
+        """
+        return [
+            cls.KERNEL,
+            cls.QUANTIZER,
+            cls.PRUNER,
+            cls.CACHER,
+            cls.FACTORIZER,
+            cls.BATCHER,
+            cls.COMPILER,
+        ]
+
     def __init__(self, name: str, description: str):
         """
         Initialize an algorithm tag with name and description.