diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 89b954a7d..8e3b273a1 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7085,9 +7085,10 @@ struct test_flash_attn_ext : public test_case { const ggml_type type_V; std::array permute; const bool kv_view; // create K/V as views of a larger buffer (like a KV cache) + const bool v_is_view_of_k; std::string vars() override { - return VARS_TO_STR15(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view); + return VARS_TO_STR16(hsk, hsv, nh, nr23, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_K, type_V, permute, kv_view, v_is_view_of_k); } double max_nmse_err() override { @@ -7104,9 +7105,9 @@ struct test_flash_attn_ext : public test_case { test_flash_attn_ext(int64_t hsk = 128, int64_t hsv = 128, int64_t nh = 32, std::array nr23 = {1, 1}, int64_t kv = 96, int64_t nb = 8, bool mask = true, bool sinks = false, float max_bias = 0.0f, float logit_softcap = 0.0f, ggml_prec prec = GGML_PREC_F32, ggml_type type_K = GGML_TYPE_F16, ggml_type type_V = GGML_TYPE_F16, std::array permute = {0, 1, 2, 3}, - bool kv_view = true) + bool kv_view = true, bool v_is_view_of_k = false) : hsk(hsk), hsv(hsv), nh(nh), nr23(nr23), kv(kv), nb(nb), mask(mask), sinks(sinks), max_bias(max_bias), logit_softcap(logit_softcap), prec(prec), - type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view) {} + type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view), v_is_view_of_k(v_is_view_of_k) {} ggml_tensor * build_graph(ggml_context * ctx) override { const int64_t hsk_padded = GGML_PAD(hsk, ggml_blck_size(type_K)); @@ -7138,14 +7139,14 @@ struct test_flash_attn_ext : public test_case { ggml_set_name(k, "k"); ggml_tensor * v = nullptr; - if (type_K == type_V && hsk_padded == 576 && hsv_padded == 512) { - // TODO: this branch should become a separate test case parameter instead of hardcoding this for these head shapes - - // in this branch, the V cache is sub-view of the K cache. this is used by some MLA-based models + if (v_is_view_of_k) { + // the V cache is a sub-view of the K cache. this is used by some MLA-based models // for more info: // - https://github.com/ggml-org/llama.cpp/pull/13435 // - https://github.com/ggml-org/llama.cpp/pull/18953#issuecomment-3774948392 // - https://github.com/ggml-org/llama.cpp/pull/18986 + GGML_ASSERT(type_K == type_V && hsv_padded <= hsk_padded); + v = ggml_view_4d(ctx, k, hsv_padded, kv, nh, nr23[1], k->nb[1], k->nb[2], k->nb[3], 0); } else { v = create_permuted(type_V, hsv_padded, kv, nh, nr23[1], kv_view); // the V tensor is usually a view of the V cache @@ -9906,12 +9907,14 @@ static std::vector> make_test_cases_eval() { if (hsk != 128 && prec == GGML_PREC_DEFAULT) continue; for (ggml_type type_KV : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0, GGML_TYPE_IQ4_NL}) { if (type_KV != GGML_TYPE_F16 && hsk != 64 && hsk != 72) continue; + // DeepSeek MLA: the V cache is a sub-view of the K cache + const bool v_is_view_of_k = hsk == 576; test_cases.emplace_back(new test_flash_attn_ext( - hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV)); + hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV, {0, 1, 2, 3}, true, v_is_view_of_k)); // run fewer test cases permuted if (mask == true && max_bias == 0.0f && logit_softcap == 0 && kv == 512) { test_cases.emplace_back(new test_flash_attn_ext( - hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV, {0, 2, 1, 3})); + hsk, hsv, nh, {nr2, nr3}, kv, nb, mask, sinks, max_bias, logit_softcap, prec, type_KV, type_KV, {0, 2, 1, 3}, true, v_is_view_of_k)); } } } @@ -9950,11 +9953,16 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 1025, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 2, 1, 3})); test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, {16, 1}, 16384, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); - // MLA shape (V is a view of K) with quantized KV - // (the test harness builds V as a view of K for this shape; see build_graph) - test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 113, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); - test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); - test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0)); + // MLA shape: the V cache is a sub-view of the K cache, with quantized KV + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 113, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {20, 1}, 1024, 64, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); + + // more V-is-sub-view-of-K cases: other head shapes, and full views with equal head sizes + test_cases.emplace_back(new test_flash_attn_ext(320, 256, 1, {32, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(192, 128, 4, {8, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(128, 128, 8, {4, 1}, 512, 8, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true)); + test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, true)); // large-KV F16 cases (Qwen3.6-27B geometry and a llama-class control): the upstream matrix // stops at kv=1024, blind to long-context FA bugs (e.g. the oneDNN SDPA ordering race on BMG).