diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 01e72041b..8030186fb 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -56,7 +56,7 @@ static void init_tensor_uniform(ggml_tensor * tensor, float min = -1.0f, float m std::vector data(nels); { // parallel initialization - static const size_t n_threads = N_THREADS; + static const size_t n_threads = std::max(1, std::min(nels/1024, std::min(4, N_THREADS/2))); auto init_thread = [&](size_t start, size_t end) { thread_local std::default_random_engine gen(std::random_device{}()); @@ -11384,10 +11384,17 @@ static bool op_names_filter_selects(const char * op_names_filter, const char * o // Covers padded rows, sinks, kvpad, multi-SIMDgroup reduction, quantized K/V, and MLA views. // The override is backend-global, so this runs after all parallel workers have joined. static bool run_fa_vec_slice(ggml_backend_t backend, ggml_backend_t backend_cpu, const char * op_names_filter) { + const char * LLAMA_TEST_FA_VEC_DISABLE = getenv("LLAMA_TEST_FA_VEC_DISABLE"); + if (LLAMA_TEST_FA_VEC_DISABLE) { + return true; + } + if (!op_names_filter_selects(op_names_filter, "FLASH_ATTN_EXT")) { return true; } + printf("Running FA vec slice tests (env LLAMA_TEST_FA_VEC_DISABLE=1 to skip)\n"); + auto * reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); auto set_ov = (set_fa_vec_override_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_metal_tuning_set_fa_vec_override");