Fix looping limit in causal attention (#1999)

jagrit06 · web-flow · commit 6a40e1c1767f · 2025-03-24T12:28:00.000-07:00
diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention.h
@@ -237,6 +237,7 @@ template <
   if (do_causal) {
     int q_max = (tid.x + 1) * BQ + params->qL_off;
     kb_lim = (q_max + BK - 1) / BK;
+    kb_lim = min(params->NK, kb_lim);
   }
 
   // Loop over KV seq length
@@ -290,7 +291,7 @@ template <
     }
 
     // Mask out if causal
-    if (do_causal && kb >= (kb_lim - (BQ + BK - 1) / BK - int(!align_K))) {
+    if (do_causal && kb >= (kb_lim - ((BQ + BK - 1) / BK) - int(!align_K))) {
       using stile_t = decltype(Stile);
       using selem_t = typename stile_t::elem_type;
       constexpr auto neg_inf = -metal::numeric_limits<selem_t>::infinity();

Original file line number	Diff line number	Diff line change
`@@ -237,6 +237,7 @@ template <`
`237`	`237`	`if (do_causal) {`
`238`	`238`	`int q_max = (tid.x + 1) * BQ + params->qL_off;`
`239`	`239`	`kb_lim = (q_max + BK - 1) / BK;`
	`240`	`+ kb_lim = min(params->NK, kb_lim);`
`240`	`241`	`}`
`241`	`242`
`242`	`243`	`// Loop over KV seq length`
`@@ -290,7 +291,7 @@ template <`
`290`	`291`	`}`
`291`	`292`
`292`	`293`	`// Mask out if causal`
`293`		`- if (do_causal && kb >= (kb_lim - (BQ + BK - 1) / BK - int(!align_K))) {`
	`294`	`+ if (do_causal && kb >= (kb_lim - ((BQ + BK - 1) / BK) - int(!align_K))) {`
`294`	`295`	`using stile_t = decltype(Stile);`
`295`	`296`	`using selem_t = typename stile_t::elem_type;`
`296`	`297`	`constexpr auto neg_inf = -metal::numeric_limits<selem_t>::infinity();`