[CUDA] Guard qmm_naive scale and bias loads at tile boundaries (#3509)

Author: Lyxot

mlx/backend/cuda/quantized/qmm/qmm_naive.cu

@@ -46,7 +46,6 @@ void qmm_naive_kernel(
   Tensor mC_mnl = make_tensor(make_gmem_ptr(C),        select<0,1,3>(shape_MNKL), dC); // (M,N,L)
 
   Tensor mS_nkl = make_tensor(make_gmem_ptr(S), S_layout); // (N,(group_size,K/group_size),L)
-  Tensor mZ_nkl = make_tensor(make_gmem_ptr(Z), S_layout); // (N,(group_size,K/group_size),L)
 
   // For gather, use index lookup for input batch slicing.
   uint32_t a_batch = lhs_indices ? lhs_indices[l_coord] : l_coord;
@@ -58,7 +57,6 @@ void qmm_naive_kernel(
   Tensor mC = mC_mnl(_,_,l_coord); // (M,N)
 
   Tensor mS = mS_nkl(_,_,b_batch); // (N,(group_size,K/group_size))
-  Tensor mZ = mZ_nkl(_,_,b_batch); // (N,(group_size,K/group_size))
 
   // Get the appropriate blocks for this thread block.
   auto cta_coord = make_coord(m_coord, n_coord, _); // (m,n,k)
@@ -67,7 +65,16 @@ void qmm_naive_kernel(
   Tensor gC = local_tile(mC, cta_tiler, cta_coord, Step<_1,_1, X>{}); // (BLK_M,BLK_N)
 
   Tensor gS = local_tile(mS, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
-  Tensor gZ = local_tile(mZ, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
+  auto gZ = [&]() {
+    if constexpr (quant_has_bias_v<Quant>) {
+      Tensor mZ_nkl = make_tensor(make_gmem_ptr(Z), S_layout); // (N,(group_size,K/group_size),L)
+      Tensor mZ = mZ_nkl(_,_,b_batch); // (N,(group_size,K/group_size))
+      return local_tile(mZ, cta_tiler, cta_coord, Step< X,_1,_1>{}); // (BLK_N,BLK_K,k)
+    } else {
+      // Dummy tensor; no-bias paths never offset or load gZ.
+      return gS;
+    }
+  }();
 
   // Compute tile residues for predication.
   int m_max_coord = size<0>(shape_MNKL) - size<0>(cta_tiler) * m_coord; // M - BLK_M * m_coord

mlx/backend/cuda/quantized/qmm/qmm_naive.cuh

@@ -114,7 +114,9 @@ CUTE_DEVICE void qmm_naive_mainloop(
     gB.data() = recast_ptr<Quant>(raw_pointer_cast(gB.data()) + gB.layout()(0, k_residue, 0) * cuda::std::min(8, sizeof_bits_v<Quant>) / 8);
   }
   gS = domain_offset(make_coord(0, k_residue, 0), gS);
-  gZ = domain_offset(make_coord(0, k_residue, 0), gZ);
+  if constexpr (quant_has_bias_v<Quant>) {
+    gZ = domain_offset(make_coord(0, k_residue, 0), gZ);
+  }
 
   // Define smem layouts.
   auto [sA_layout, sB_layout] = make_smem_layouts(cta_tiler);
@@ -177,8 +179,15 @@ CUTE_DEVICE void qmm_naive_mainloop(
   auto fetch_gmem = [&](int tile) {
     copy_if(copy_a, tApA, tAgA(_,_,_,tile), tArA);
     copy_if(copy_b, tBpB, tBgB(_,_,_,tile), tBrB);
-    copy(tBgS(_,_,_,tile), tBrS);
-    copy(tBgZ(_,_,_,tile), tBrZ);
+    CUTE_UNROLL
+    for (int n = 0; n < size<1>(tBrS); ++n) {
+      if (tBpB(n,0)) {
+        copy(tBgS(_,n,_,tile), tBrS(_,n,_));
+        if constexpr (quant_has_bias_v<Quant>) {
+          copy(tBgZ(_,n,_,tile), tBrZ(_,n,_));
+        }
+      }
+    }
   };
   // RMEM => SMEM.
   auto store_smem = [&]() {
@@ -219,8 +228,15 @@ CUTE_DEVICE void qmm_naive_mainloop(
     for (int k = 0; k < size<2>(tBrB); ++k) {
       if (get<1>(tBcB(0,0,k)) >= -k_residue) {
         copy_if(copy_b, tBpB(_,k), tBgB_k(_,_,k), tBrB(_,_,k));
-        copy(tBgS_k(_,_,k), tBrS(_,_,k));
-        copy(tBgZ_k(_,_,k), tBrZ(_,_,k));
+        CUTE_UNROLL
+        for (int n = 0; n < size<1>(tBrS); ++n) {
+          if (tBpB(n,k)) {
+            copy(tBgS_k(_,n,k), tBrS(_,n,k));
+            if constexpr (quant_has_bias_v<Quant>) {
+              copy(tBgZ_k(_,n,k), tBrZ(_,n,k));
+            }
+          }
+        }
       }
     }
   } else {