metal : add FA kernels for HSK=96, HSV=64 (MiniCPM3) (#28599)

* metal : add FA kernels for HSK=96, HSV=64 (MiniCPM3)

MiniCPM3 sets attention.key_length to 96 and does not set
attention.value_length, which defaults to n_embd / n_head = 64. Metal had no
(96, 64) instantiation, so -fa auto aborted on the missing
kernel_flash_attn_ext_vec_f16_dk96_dv64.

Instantiate the tile kernel at (96, 64) for every K/V type that already has
(96, 96), and the vec kernel for the NE=4 configurations. Of the NE values the
vec dispatch considers, only NE=4 works here, because NL = 32/NE has to divide
both DK/4 = 24 and DV/4 = 16.

* tests : avoid redundant FA vec slice coverage
This commit is contained in:
Yanzhao Wang
2026-09-15 01:10:40 -07:00
committed by GitHub
parent 0ecb159c9e
commit 1e7bcf3da4
4 changed files with 37 additions and 1 deletions
+2 -1
View File
@@ -10602,7 +10602,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
for (int hsk : { 40, 64, 72, 80, 96, 128, 192, 256, 320, 512, 576 }) {
for (int hsv : { 40, 64, 72, 80, 96, 128, 192, 256, 512 }) {
if (hsk != 192 && hsk != 320 && hsk != 576 && hsk != hsv) continue;
if (hsk != 96 && hsk != 192 && hsk != 320 && hsk != 576 && hsk != hsv) continue;
if (hsk == 96 && (hsv != 64 && hsv != 96)) continue; // MiniCPM3
if (hsk == 192 && (hsv != 128 && hsv != 192)) continue;
if (hsk == 576 && hsv != 512) continue; // DeepSeek MLA
if (hsk == 320 && hsv != 256) continue; // Mistral4 MLA