diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 28b97d4..efb900c 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -16,10 +16,12 @@ jobs:
       uses: actions/setup-python@v5
       with:
         python-version: '3.10'
+        cache: 'pip'
         
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
+        pip install pytest
         pip install -e .
         
     - name: Run tests
diff --git a/pyproject.toml b/pyproject.toml
index 8c0c383..5d2bdea 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -20,7 +20,6 @@ dependencies = [
     "decord==0.6.0",
     "tqdm==4.66.5",
     "pandas==2.2.2",
-    "grouped_gemm==0.1.6",
     "matplotlib==3.9.2",
     "fire==0.6.0",
 ]
@@ -35,6 +34,9 @@ dev = [
 vllm = [
     "vllm==0.6.2"
 ]
+grouped_gemm = [
+    "grouped_gemm==0.1.6"
+]
 
 [build-system]
 requires = ["setuptools>=45", "setuptools-scm[toml]>=6.2"]