fix

Author: Awni Hannun

.github/actions/build-linux/action.yml

@@ -9,6 +9,9 @@ inputs:
 runs:
   using: "composite"
   steps:
+    - name: Setup tmate session
+      uses: mxschmitt/action-tmate@v3
+
     - name: Install Python package
       id: python_build
       shell: sh

mlx/backend/cuda/CMakeLists.txt

@@ -69,11 +69,12 @@ add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/unary)
 # fp4 is not available on < 12.8
 if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 12.8.0)
   target_include_directories(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/)
-  target_sources(
-    mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/no_cublas_qqmm.cpp)
-else()
   target_sources(mlx
-                 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/cublas_qqmm.cpp)
+                 PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/no_qqmm_impl.cpp)
+else()
+  target_sources(
+    mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/quantized/qqmm_impl.cpp
+                ${CMAKE_CURRENT_SOURCE_DIR}/quantized/cublas_qqmm.cpp)
 endif()
 
 if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 12.9.0)

mlx/backend/cuda/quantized/cuda_fp4.h

@@ -85,10 +85,14 @@ struct __nv_fp4_e2m1 {
 struct __nv_fp4x4_e2m1 {
   __device__ operator float4() {
     float4 out;
-    out.x = float(*(__nv_fp4_e2m1*)(__high & 0xf));
-    out.y = float(*(__nv_fp4_e2m1*)((__high >> 4) & 0xf));
-    out.z = float(*(__nv_fp4_e2m1*)(__low & 0xf));
-    out.w = float(*(__nv_fp4_e2m1*)((__low >> 4) & 0xf));
+    auto bits = __high & 0xf;
+    out.x = float(*(__nv_fp4_e2m1*)(&bits));
+    bits = (__high >> 4) & 0xf;
+    out.y = float(*(__nv_fp4_e2m1*)(&bits));
+    bits = (__low) & 0xf;
+    out.z = float(*(__nv_fp4_e2m1*)(&bits));
+    bits = (__low >> 4) & 0xf;
+    out.w = float(*(__nv_fp4_e2m1*)(&bits));
     return out;
   }
   uint8_t __high{0};

mlx/backend/cuda/quantized/no_cublas_qqmm.cpp

@@ -0,0 +1,26 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/qqmm_impl.h"
+
+namespace mlx::core {
+void qqmm_impl(
+    cu::CommandEncoder&,
+    int,
+    int,
+    int,
+    bool,
+    int64_t,
+    bool,
+    int64_t,
+    array&,
+    const array&,
+    const array&,
+    const array&,
+    const array&,
+    Dtype,
+    QuantizationMode,
+    float) {
+  throw std::runtime_error(
+      "[QQMatmul::eval_gpu] QQMM is only supported with CUDA 12.8 or higher.");
+}
+} // namespace mlx::core

mlx/backend/cuda/quantized/no_qqmm_impl.cpp

@@ -1,8 +1,8 @@
 // Copyright © 2025 Apple Inc.
 
 #include "mlx/backend/cuda/device.h"
-#include "mlx/backend/cuda/quantized/cublas_qqmm.h"
 #include "mlx/backend/cuda/quantized/qmv.h"
+#include "mlx/backend/cuda/quantized/qqmm_impl.h"
 #include "mlx/backend/cuda/quantized/qqmm_utils.h"
 #include "mlx/backend/cuda/quantized/quantized.h"
 #include "mlx/backend/cuda/quantized/quantized_utils.h"
@@ -38,47 +38,6 @@ array pad_and_swizzle_scales(
   return scale_tiled;
 }
 
-void qqmm_impl(
-    cu::CommandEncoder& encoder,
-    int M,
-    int N,
-    int K,
-    bool a_transposed,
-    int64_t lda,
-    bool b_transposed,
-    int64_t ldb,
-    array& out,
-    const array& a,
-    const array& b,
-    const array& a_scale,
-    const array& b_scale,
-    Dtype out_dtype,
-    QuantizationMode mode,
-    float alpha = 1.0f) {
-  // Invoke CublasQQMM
-  std::string qmode = quantization_mode_to_string(mode);
-
-  // Currently only supports non-batched QQMM operations
-  // that covers all use cases for training, we will just collapse (batch,
-  // seq_len) into (tokens)
-  CublasQQMM qqmm(
-      encoder.device(),
-      a_transposed,
-      M,
-      K,
-      lda,
-      b_transposed,
-      K,
-      N,
-      ldb,
-      1, // batch_count
-      0, // a_batch_stride
-      0, // b_batch_stride
-      out_dtype,
-      qmode);
-
-  qqmm.run(encoder, out, a, b, a_scale, b_scale, alpha);
-}
 } // namespace
 
 void QQMatmul::eval_gpu(const std::vector<array>& inputs, array& out) {

mlx/backend/cuda/quantized/qqmm.cpp

@@ -0,0 +1,50 @@
+// Copyright © 2026 Apple Inc.
+
+#include "mlx/backend/cuda/quantized/qqmm_impl.h"
+#include "mlx/backend/cuda/quantized/cublas_qqmm.h"
+
+namespace mlx::core {
+
+void qqmm_impl(
+    cu::CommandEncoder& encoder,
+    int M,
+    int N,
+    int K,
+    bool a_transposed,
+    int64_t lda,
+    bool b_transposed,
+    int64_t ldb,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& a_scale,
+    const array& b_scale,
+    Dtype out_dtype,
+    QuantizationMode mode,
+    float alpha) {
+  // Invoke CublasQQMM
+  std::string qmode = quantization_mode_to_string(mode);
+
+  // Currently only supports non-batched QQMM operations
+  // that covers all use cases for training, we will just collapse (batch,
+  // seq_len) into (tokens)
+  CublasQQMM qqmm(
+      encoder.device(),
+      a_transposed,
+      M,
+      K,
+      lda,
+      b_transposed,
+      K,
+      N,
+      ldb,
+      1, // batch_count
+      0, // a_batch_stride
+      0, // b_batch_stride
+      out_dtype,
+      qmode);
+
+  qqmm.run(encoder, out, a, b, a_scale, b_scale, alpha);
+}
+
+} // namespace mlx::core

mlx/backend/cuda/quantized/qqmm_impl.cpp

@@ -0,0 +1,26 @@
+// Copyright © 2026 Apple Inc.
+#pragma once
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+void qqmm_impl(
+    cu::CommandEncoder& encoder,
+    int M,
+    int N,
+    int K,
+    bool a_transposed,
+    int64_t lda,
+    bool b_transposed,
+    int64_t ldb,
+    array& out,
+    const array& a,
+    const array& b,
+    const array& a_scale,
+    const array& b_scale,
+    Dtype out_dtype,
+    QuantizationMode mode,
+    float alpha = 1.0f);
+
+} // namespace mlx::core