Fix edge checking in attention mask

jagrit06 · jagrit06 · commit 12f766efe96b · 2026-01-15T14:10:24.000-08:00
diff --git a/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h b/mlx/backend/metal/kernels/steel/attn/kernels/steel_attention_nax.h
@@ -362,7 +362,6 @@ template <
             const int M_str = int(mask_params->M_strides[2]);
             const int M_load_off = row_pos * M_str + col_pos;
             mfrag.load(mask + M_load_off, M_str, Int<1>{});
-
           }
 
           thread auto& fg = Ptile.subtile_at(iq, ik).frag_at(0, 0);

Original file line number	Diff line number	Diff line change
`@@ -362,7 +362,6 @@ template <`
`362`	`362`	`const int M_str = int(mask_params->M_strides[2]);`
`363`	`363`	`const int M_load_off = row_pos * M_str + col_pos;`
`364`	`364`	`mfrag.load(mask + M_load_off, M_str, Int<1>{});`
`365`		`-`
`366`	`365`	`}`
`367`	`366`
`368`	`367`	`thread auto& fg = Ptile.subtile_at(iq, ik).frag_at(0, 0);`