diff --git a/pyproject.toml b/pyproject.toml
index 0bbc3e91c..ef9bfaa6f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -122,3 +122,23 @@ docs = "cd docs && ./setup.sh"
 black = 'black --exclude=".*\.ipynb" --extend-exclude=".venv|.pixi" ./'
 isort = 'isort --profile black --skip .venv --skip .pixi ./'
 flake8 = 'flake8 --extend-exclude=.venv/,.pixi/ ./'
+
+[tool.pixi.feature.numba-cuda]
+platforms = ["linux-64", "linux-aarch64"]
+
+[tool.pixi.feature.numba-cuda-mlir]
+platforms = ["linux-64", "linux-aarch64"]
+
+[tool.pixi.feature.numba-cuda.dependencies]
+cuda-nvcc = "*"
+numba-cuda = "*"
+numpy = "<2.5"
+
+[tool.pixi.feature.numba-cuda-mlir.dependencies]
+cuda-nvcc = "*"
+numba-cuda-mlir = "*"
+
+[tool.pixi.environments]
+numba-cuda = ["numba-cuda"]
+numba-cuda-mlir = ["numba-cuda-mlir"]
+
diff --git a/stumpy/core.py b/stumpy/core.py
index bb36fa7ee..060b4fdd8 100644
--- a/stumpy/core.py
+++ b/stumpy/core.py
@@ -9,7 +9,12 @@
 import warnings
 
 import numpy as np
-from numba import cuda, njit, prange
+from numba import njit, prange
+
+try:
+    from numba_cuda_mlir import cuda
+except ModuleNotFoundError:
+    from numba import cuda
 from scipy import linalg
 from scipy.ndimage import maximum_filter1d, minimum_filter1d
 from scipy.spatial.distance import cdist
diff --git a/stumpy/gpu_aamp.py b/stumpy/gpu_aamp.py
index 1aebee850..155d6ee85 100644
--- a/stumpy/gpu_aamp.py
+++ b/stumpy/gpu_aamp.py
@@ -6,7 +6,11 @@
 import os
 
 import numpy as np
-from numba import cuda
+
+try:
+    from numba_cuda_mlir import cuda
+except ModuleNotFoundError:
+    from numba import cuda
 
 from . import config, core
 from .mparray import mparray
diff --git a/stumpy/gpu_stump.py b/stumpy/gpu_stump.py
index c52176a17..9b651928e 100644
--- a/stumpy/gpu_stump.py
+++ b/stumpy/gpu_stump.py
@@ -6,7 +6,11 @@
 import os
 
 import numpy as np
-from numba import cuda
+
+try:
+    from numba_cuda_mlir import cuda
+except ModuleNotFoundError:
+    from numba import cuda
 
 from . import config, core
 from .gpu_aamp import gpu_aamp
diff --git a/tests/test_core.py b/tests/test_core.py
index 28589cfa3..6f779622b 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -9,7 +9,11 @@
 import numpy.testing as npt
 import pandas as pd
 import pytest
-from numba import cuda
+
+try:
+    from numba_cuda_mlir import cuda
+except ModuleNotFoundError:
+    from numba import cuda
 from scipy.spatial.distance import cdist
 
 from stumpy import config, core, rng
diff --git a/tests/test_gpu_aamp.py b/tests/test_gpu_aamp.py
index 7adfca221..ab6aa6974 100644
--- a/tests/test_gpu_aamp.py
+++ b/tests/test_gpu_aamp.py
@@ -3,7 +3,11 @@
 import numpy as np
 import numpy.testing as npt
 import pandas as pd
-from numba import cuda
+
+try:
+    from numba_cuda_mlir import cuda
+except ModuleNotFoundError:
+    from numba import cuda
 
 from stumpy import config, rng
 
diff --git a/tests/test_gpu_aamp_ostinato.py b/tests/test_gpu_aamp_ostinato.py
index d45d9982c..b5337cfe3 100644
--- a/tests/test_gpu_aamp_ostinato.py
+++ b/tests/test_gpu_aamp_ostinato.py
@@ -3,7 +3,11 @@
 import mersenne
 import numpy as np
 import numpy.testing as npt
-from numba import cuda
+
+try:
+    from numba_cuda_mlir import cuda
+except ModuleNotFoundError:
+    from numba import cuda
 
 try:
     from numba.errors import NumbaPerformanceWarning
diff --git a/tests/test_gpu_aamp_stimp.py b/tests/test_gpu_aamp_stimp.py
index 2f4fe13d5..fdfac640c 100644
--- a/tests/test_gpu_aamp_stimp.py
+++ b/tests/test_gpu_aamp_stimp.py
@@ -2,7 +2,11 @@
 
 import numpy as np
 import numpy.testing as npt
-from numba import cuda
+
+try:
+    from numba_cuda_mlir import cuda
+except ModuleNotFoundError:
+    from numba import cuda
 
 if cuda.is_available():
     from stumpy.gpu_aamp_stimp import gpu_aamp_stimp
diff --git a/tests/test_gpu_aampdist.py b/tests/test_gpu_aampdist.py
index bd3b4838d..b92f8eace 100644
--- a/tests/test_gpu_aampdist.py
+++ b/tests/test_gpu_aampdist.py
@@ -2,7 +2,11 @@
 
 import numpy as np
 import numpy.testing as npt
-from numba import cuda
+
+try:
+    from numba_cuda_mlir import cuda
+except ModuleNotFoundError:
+    from numba import cuda
 
 if cuda.is_available():
     from stumpy.gpu_aampdist import gpu_aampdist
diff --git a/tests/test_gpu_mpdist.py b/tests/test_gpu_mpdist.py
index 21de32e90..b8acb1fcf 100644
--- a/tests/test_gpu_mpdist.py
+++ b/tests/test_gpu_mpdist.py
@@ -3,7 +3,11 @@
 
 import numpy as np
 import numpy.testing as npt
-from numba import cuda
+
+try:
+    from numba_cuda_mlir import cuda
+except ModuleNotFoundError:
+    from numba import cuda
 
 if cuda.is_available():
     from stumpy.gpu_mpdist import gpu_mpdist
diff --git a/tests/test_gpu_ostinato.py b/tests/test_gpu_ostinato.py
index 81d0dbf42..c98f92b0e 100644
--- a/tests/test_gpu_ostinato.py
+++ b/tests/test_gpu_ostinato.py
@@ -2,7 +2,11 @@
 
 import numpy as np
 import numpy.testing as npt
-from numba import cuda
+
+try:
+    from numba_cuda_mlir import cuda
+except ModuleNotFoundError:
+    from numba import cuda
 
 try:
     from numba.errors import NumbaPerformanceWarning
diff --git a/tests/test_gpu_stimp.py b/tests/test_gpu_stimp.py
index d93f77a27..45d1590ea 100644
--- a/tests/test_gpu_stimp.py
+++ b/tests/test_gpu_stimp.py
@@ -3,7 +3,11 @@
 
 import numpy as np
 import numpy.testing as npt
-from numba import cuda
+
+try:
+    from numba_cuda_mlir import cuda
+except ModuleNotFoundError:
+    from numba import cuda
 
 if cuda.is_available():
     from stumpy.gpu_stimp import gpu_stimp
diff --git a/tests/test_gpu_stump.py b/tests/test_gpu_stump.py
index eec2a0cc5..bbc7c41ca 100644
--- a/tests/test_gpu_stump.py
+++ b/tests/test_gpu_stump.py
@@ -4,7 +4,11 @@
 import numpy as np
 import numpy.testing as npt
 import pandas as pd
-from numba import cuda
+
+try:
+    from numba_cuda_mlir import cuda
+except ModuleNotFoundError:
+    from numba import cuda
 
 from stumpy import config, rng
 
diff --git a/tests/test_non_normalized_decorator.py b/tests/test_non_normalized_decorator.py
index 62687e814..5e7c8f664 100644
--- a/tests/test_non_normalized_decorator.py
+++ b/tests/test_non_normalized_decorator.py
@@ -3,7 +3,11 @@
 import numpy.testing as npt
 import tornado.ioloop
 from dask.distributed import Client, LocalCluster
-from numba import cuda
+
+try:
+    from numba_cuda_mlir import cuda
+except ModuleNotFoundError:
+    from numba import cuda
 
 from stumpy import core, rng
 from stumpy.aamp import aamp
diff --git a/tests/test_precision.py b/tests/test_precision.py
index 4edb3a93d..8c939cf55 100644
--- a/tests/test_precision.py
+++ b/tests/test_precision.py
@@ -6,7 +6,11 @@
 import numpy as np
 import numpy.testing as npt
 import pytest
-from numba import cuda
+
+try:
+    from numba_cuda_mlir import cuda
+except ModuleNotFoundError:
+    from numba import cuda
 
 from stumpy import config, core, rng, sdp