From 183667fe39323336213099501aa7459caf2806f9 Mon Sep 17 00:00:00 2001
From: Cheng <git@zcbenz.com>
Date: Sat, 23 May 2026 19:09:58 -0700
Subject: [PATCH] Cache JIT-compiled CUDA kernels

---
 .github/actions/setup-linux/action.yml | 9 +++++++++
 .github/actions/test-linux/action.yml  | 1 +
 2 files changed, 10 insertions(+)

diff --git a/.github/actions/setup-linux/action.yml b/.github/actions/setup-linux/action.yml
index f5b6268ee8..419abfb163 100644
--- a/.github/actions/setup-linux/action.yml
+++ b/.github/actions/setup-linux/action.yml
@@ -41,6 +41,15 @@ runs:
         # ccache-action bug: running "apt-get update" fails on large arm runner.
         update-package-index: false
 
+    - name: Cache JIT-compiled CUDA kernels
+      if: ${{ startsWith(inputs.toolkit, 'cuda') }}
+      uses: actions/cache@v5
+      with:
+        path: /tmp/mlx-ptx-cache
+        key: >-
+          ptx-${{ runner.os }}-${{ runner.arch }}-${{ inputs.toolkit }}-
+          ${{ hashFiles('mlx/backend/cuda/**') }}
+
     - uses: actions/setup-python@v6
       with:
         python-version: ${{ inputs.python-version }}
diff --git a/.github/actions/test-linux/action.yml b/.github/actions/test-linux/action.yml
index f552fe5de0..24cda6c103 100644
--- a/.github/actions/test-linux/action.yml
+++ b/.github/actions/test-linux/action.yml
@@ -44,6 +44,7 @@ runs:
       shell: bash
       env:
         DEVICE: gpu
+        MLX_PTX_CACHE_DIR: /tmp/mlx-ptx-cache
       run: |
         echo "::group::Python tests - GPU"
         python -m tests discover python/tests -v