From 1bccbc97109f43041e402928d965fe82fa4d3e7b Mon Sep 17 00:00:00 2001
From: Cheng <git@zcbenz.com>
Date: Sun, 2 Aug 2026 23:13:03 -0700
Subject: [PATCH] Fix build error caused by TMA macro guard

---
 mlx/backend/cuda/ptx.cuh                   | 12 +++++-------
 mlx/backend/cuda/quantized/fp_quantize.cuh |  8 ++++----
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/mlx/backend/cuda/ptx.cuh b/mlx/backend/cuda/ptx.cuh
index 6ec7caedd6..b7b87e671e 100644
--- a/mlx/backend/cuda/ptx.cuh
+++ b/mlx/backend/cuda/ptx.cuh
@@ -1,14 +1,12 @@
 #pragma once
 
-#include <cuda.h>
-#include <cuda_runtime.h>
+#include <cute/arch/config.hpp>
 
 namespace mlx::core {
 
 namespace ptx {
 
-#if (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) && \
-    defined(__CUDA_ARCH_SPECIFIC__)
+#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
 
 __device__ __forceinline__ void mbarrier_init(uint64_t* mbar, uint32_t count) {
   uint32_t mbar_ptr = __cvta_generic_to_shared(mbar);
@@ -121,7 +119,7 @@ __device__ __forceinline__ void fence_proxy_async_shared_cta() {
   asm volatile("fence.proxy.async.shared::cta;");
 }
 
-#endif // (CUDART_VERSION >= 12080) && (__CUDA_ARCH__ >= 1000) &&
-       // (__CUDA_ARCH_FAMILY_SPECIFIC__ >= 1000)
+#endif // defined(CUTE_ARCH_TMA_SM90_ENABLED)
+
 } // namespace ptx
-} // namespace mlx::core
\ No newline at end of file
+} // namespace mlx::core
diff --git a/mlx/backend/cuda/quantized/fp_quantize.cuh b/mlx/backend/cuda/quantized/fp_quantize.cuh
index 769794129c..bea4e21979 100644
--- a/mlx/backend/cuda/quantized/fp_quantize.cuh
+++ b/mlx/backend/cuda/quantized/fp_quantize.cuh
@@ -55,7 +55,7 @@ __device__ __forceinline__ void copy_2d_to_shared(
     uint32_t num_bytes,
     uint64_t* barrier,
     const bool is_master_thread) {
-#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
   if (is_master_thread) {
     // Arrive and tell how many bytes are expected
     ptx::mbarrier_arrive_expect_tx(barrier, num_bytes);
@@ -66,7 +66,7 @@ __device__ __forceinline__ void copy_2d_to_shared(
     // Other threads just arrive
     ptx::mbarrier_arrive(barrier);
   }
-#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
+#endif // defined(CUTE_ARCH_TMA_SM90_ENABLED)
 }
 
 namespace cg = cooperative_groups;
@@ -228,7 +228,7 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK)
         uint8_t* __restrict__ scales,
         const size_t rows,
         const size_t cols) {
-#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000)
+#if defined(CUTE_ARCH_TMA_SM90_ENABLED)
   using Tx2 = Vector2_t<T>;
   using Tx4 = Vector4_t<T>;
 
@@ -408,7 +408,7 @@ __global__ void __launch_bounds__(THREADS_PER_BLOCK)
       ptx::mbarrier_invalidate(&mbar[iter]);
     }
   }
-#endif // __CUDA_ARCH__ >= 1000
+#endif // defined(CUTE_ARCH_TMA_SM90_ENABLED)
 }
 
 // TODO: add kernel with tma instructions