From f32f1ceb4db68f9dc45708e0823a3c148980a36f Mon Sep 17 00:00:00 2001 From: moristeven477-ship-it <311459415+moristeven477-ship-it@users.noreply.github.com> Date: Tue, 15 Sep 2026 03:38:22 -0700 Subject: [PATCH] hip : skip fully masked KV tiles in WMMA Flash Attention Assisted-by: Codex --- ggml/src/ggml-cuda/fattn-mma-f16.cuh | 19 +++++++++++++ tests/test-backend-ops.cpp | 41 +++++++++++++++++++++++++--- 2 files changed, 56 insertions(+), 4 deletions(-) diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index a290655776aa..0aedfd754651 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -612,6 +612,25 @@ static __device__ __forceinline__ void flash_attn_ext_f16_iter( constexpr bool swz_V = V_is_K_view ? swz_K : ggml_cuda_fattn_smem_swizzle::enabled(nbatch_V2); const int k_VKQ_0 = kb0 * nbatch_fa; +#if defined(AMD_WMMA_AVAILABLE) + if constexpr (!use_sparse) { + const int j_probe = fastmodulo(jt*ncols1 + ncols1 - 1, ne01); + if (mask_h && __half2float(mask_h[int64_t(j_probe)*stride_mask + k_VKQ_0]) == -INFINITY) { + // A finite probe keeps dense tiles on the normal path. Otherwise check every mask element. + bool has_unmasked = false; + for (int ij = threadIdx.y*warp_size + threadIdx.x; ij < ncols1*nbatch_fa; ij += nwarps*warp_size) { + const int i = ij % nbatch_fa; + const int j = fastmodulo(jt*ncols1 + ij/nbatch_fa, ne01); + if ((!oob_check || i < k_VKQ_sup) && __half2float(mask_h[int64_t(j)*stride_mask + k_VKQ_0 + i]) != -INFINITY) { + has_unmasked = true; + } + } + if (!__syncthreads_or(has_unmasked)) { + return; + } + } + } +#endif // defined(AMD_WMMA_AVAILABLE) #if defined(TURING_MMA_AVAILABLE) T_C_KQ KQ_C[nbatch_fa/(np*(cols_per_warp == 8 ? T_C_KQ::I : T_C_KQ::J))]; #elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index f650e0123224..d6bf366cc223 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -7656,9 +7656,10 @@ struct test_flash_attn_ext : public test_case { const bool kv_view; // create K/V as views of a larger buffer (like a KV cache) const bool v_is_view_of_k; const int64_t n_kv_max; + const int mask_pattern; std::string vars() override { - return VARS_TO_STR17(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, n_kv_max); + return VARS_TO_STR17(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, n_kv_max) + "," + VAR_TO_STR(mask_pattern); } double max_nmse_err() override { @@ -7675,9 +7676,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 v_is_view_of_k = false, int64_t n_kv_max = 0) + bool kv_view = true, bool v_is_view_of_k = false, int64_t n_kv_max = 0, int mask_pattern = 0) : 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), v_is_view_of_k(v_is_view_of_k), n_kv_max(n_kv_max) {} + type_K(type_K), type_V(type_V), permute(permute), kv_view(kv_view), v_is_view_of_k(v_is_view_of_k), n_kv_max(n_kv_max), mask_pattern(mask_pattern) {} ggml_tensor * build_graph(ggml_context * ctx) override { const int64_t hsk_padded = GGML_PAD(hsk, ggml_blck_size(type_K)); @@ -7750,7 +7751,29 @@ struct test_flash_attn_ext : public test_case { // make the sink values more noticeable in order to trigger a test failure when the implementation is wrong init_tensor_uniform(t, -10.0f, 10.0f); } else if (strcmp(t->name, "m") == 0) { - if (n_kv_max > 0) { + if (mask_pattern > 0) { + std::vector data(ggml_nelements(t), -INFINITY); + for (int64_t row = 0; row < ggml_nrows(t); ++row) { + if (mask_pattern == 3 && row % t->ne[1] == 0) { + continue; // Fully masked query with an attention sink. + } + if (mask_pattern == 2) { + data[row*t->ne[0]] = 0.0f; + if (row % t->ne[1] == 0) { + data[row*t->ne[0] + t->ne[0]/2 + 17] = 0.75f; + } + } else { + for (int64_t i = 0; i < t->ne[0]; ++i) { + if (i < t->ne[0]/4 || i >= 3*t->ne[0]/4) { + data[row*t->ne[0] + i] = -0.25f; + } + } + } + } + std::vector data_f16(data.size()); + ggml_fp32_to_fp16_row(data.data(), data_f16.data(), data.size()); + ggml_backend_tensor_set(t, data_f16.data(), 0, data_f16.size()*sizeof(ggml_fp16_t)); + } else if (n_kv_max > 0) { init_tensor_kq_mask_sparse(t, n_kv_max); } else { init_tensor_kq_mask(t); @@ -10744,6 +10767,16 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {2, 1}, 1024, 32, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); test_cases.emplace_back(new test_flash_attn_ext(512, 512, 4, {2, 1}, 1024, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); + // Masked interior tiles, one visible element in an otherwise masked tile, and fully masked queries with sinks. + for (int64_t hs : {64, 128, 256}) { + for (int mask_pattern : {1, 2, 3}) { + test_cases.emplace_back(new test_flash_attn_ext(hs, hs, 2, {8, 2}, 512, 17, true, mask_pattern == 3, 0, 0, + GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 0, mask_pattern)); + } + } + test_cases.emplace_back(new test_flash_attn_ext(256, 256, 4, {6, 1}, 1024, 64, true, false, 0, 0, + GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, false, 0, 2)); + test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, { 10, 5, 4, 3})); test_cases.emplace_back(new test_cross_entropy_loss (GGML_TYPE_F32, {30000, 1, 1, 1})); test_cases.emplace_back(new test_cross_entropy_loss_back(GGML_TYPE_F32, { 10, 5, 4, 3}));