metal : prevent int overflows [no ci]

ggerganov · ggerganov · commit 5d1a10d2751e · 2024-11-08T10:11:23.000+02:00
diff --git a/ggml/src/ggml-metal.metal b/ggml/src/ggml-metal.metal
@@ -3613,7 +3613,7 @@ kernel void kernel_flash_attn_ext_vec(
         const float S = ss[0];
 
         for (short i = tiisg; i < D16; i += NW) {
-            dst44[(iq3*ne2*ne1 + iq2 + (iq1)*ne1)*D16 + i] = (float4x4) sr4x4[i]/S;
+            dst44[((int64_t)iq3*ne2*ne1 + iq2 + (iq1)*ne1)*D16 + i] = (float4x4) sr4x4[i]/S;
         }
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -3613,7 +3613,7 @@ kernel void kernel_flash_attn_ext_vec(`
`3613`	`3613`	`const float S = ss[0];`
`3614`	`3614`
`3615`	`3615`	`for (short i = tiisg; i < D16; i += NW) {`
`3616`		`- dst44[(iq3ne2ne1 + iq2 + (iq1)ne1)D16 + i] = (float4x4) sr4x4[i]/S;`
	`3616`	`+ dst44[((int64_t)iq3ne2ne1 + iq2 + (iq1)ne1)D16 + i] = (float4x4) sr4x4[i]/S;`
`3617`	`3617`	`}`
`3618`	`3618`	`}`
`3619`	`3619`	`}`