hexagon: HMX flash-attention head_dim padding (support DK=DV=72) (#26539)

Allow HMX flash-attention to run with head_dim not a multiple of 64
(e.g. SigLIP head_dim=72), by operating on DK/DV rounded up to 64 with
zero-filled tail lanes.
This commit is contained in:
Todor Boinovski
2026-09-18 13:15:08 -07:00
committed by GitHub
parent ec92815050
commit 18a04f09c2
4 changed files with 256 additions and 38 deletions
+4
View File
@@ -10713,6 +10713,10 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
}
}
// asymmetric head_dim (hsk != hsv) with one or both sides not 64-aligned
test_cases.emplace_back(new test_flash_attn_ext(72, 64, 4, {1, 1}, 256, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
test_cases.emplace_back(new test_flash_attn_ext(64, 72, 4, {1, 1}, 256, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16));
// mixed quant and Q1_0 test cases
test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0));
test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_F16));