From 370ed6d18c53a83f4849e29e49a3fb6dceb1bc2d Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Sun, 18 Oct 2026 02:10:10 -0300 Subject: [PATCH 1/2] tests: build the HRX tests against the ggml-hrx we ship (engine#422) f2099e9b put ggml/src/ggml-hrx back to our backend as of 522dab47 but left AMD's newer tests/hrx-backend-test.cpp, test-hrx-ops.cpp and test-hrx-loom-jit.cpp in place. They target AMD's ggml-hrx (activation publication, WorkloadSpecialization, symmetric-i4), which we do not ship, so none of the three compiled. Errors before, at 8d9ce2c1a with hrx-system d3c731e1 (clang -ferror-limit=0, the missing include removed so the rest of each file is checked): hrx-backend-test.cpp fatal include dispatch-symmetric-i4.h + 21 errors (workload_specialization/WorkloadSpecialization, CommandPlanMetadata::activation_publication_diagnostics, common_symmetric_i4_lowrow_mul_mat_eligible, one struct initializer) test-hrx-ops.cpp fatal include dispatch-activation-publication.h + 128 errors (CommonActivationPublication*, DispatchActivation*, add_activation_consumer, WorkloadSpecialization, find_generated_resource) test-hrx-loom-jit.cpp 26 errors (WorkloadSpecialization and changed argument counts of the JIT entry points) test-hrx-buffer.cpp and testing_suite.h compile unchanged. The three files are taken whole from 522dab47, the version matching our ggml-hrx. No commit of ours touched them after 522dab47, so nothing is re-applied. After this, hrx-backend-test, test-hrx-ops, test-hrx-loom-jit, test-hrx-buffer and test-backend-ops all build (-j2, 0 errors). Co-Authored-By: Claude Opus 5.5 --- tests/hrx-backend-test.cpp | 4474 ++++------------------------------- tests/test-hrx-loom-jit.cpp | 681 +----- tests/test-hrx-ops.cpp | 2472 +++---------------- 3 files changed, 859 insertions(+), 6768 deletions(-) diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index ecf8f131d542..6573d16702f0 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -6,11 +6,9 @@ #include "dispatch/command-program.h" #include "dispatch/dispatch-scheduler.h" #include "dispatch_registration/common/dispatch-gather-add.h" -#include "dispatch_registration/common/dispatch-symmetric-i4.h" #include "dispatch_registration/dispatch-registry.h" #include "ggml-alloc.h" #include "ggml-backend.h" -#include "ggml-cpu.h" #include "ggml-hrx.h" #include "ggml-impl.h" #include "ggml.h" @@ -25,7 +23,6 @@ #include "runtime/graph-program-cache.h" #include "runtime/graph-replay.h" #include "runtime/loom-kernel-jit.h" -#include "testing_suite.h" #include #include @@ -72,16 +69,10 @@ static bool contains_value_id(const std::vector & ids, ggml: } static bool command_program_verifies(const ggml::hrx::CommandProgram & program) { - const ggml::hrx::VerificationResult result = - ggml::hrx::verify_command_program(program, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - if (!result.valid()) { - for (const std::string & error : result.status.errors()) { - std::fprintf(stderr, "%s\n", error.c_str()); - } - } - return result.valid(); + return ggml::hrx::verify_command_program(program, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151").valid(); } + static bool async_jit_expected_from_environment() { const char * value = std::getenv("GGML_HRX_ASYNC_JIT"); return value == nullptr || @@ -191,36 +182,6 @@ static std::string kernel_name_for_id(uint64_t kernel_id) { return ggml::hrx::kernel_definition_name(*resolved.definition); } -template static void require_scheduled_command_program(ggml_cgraph * graph, Check && check) { - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - - const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( - imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - check(imported.graph, scheduler.plan(), commands); -} - -static std::vector scheduled_kernel_names(ggml_cgraph * graph) { - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - - std::vector result; - for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { - result.push_back(kernel_name_for_id(dispatch.kernel.kernel_id)); - } - return result; -} - static void require_compile_parameter(const ggml::hrx::Dispatch & dispatch, const char * name, const std::string & value) { @@ -282,10 +243,6 @@ static int64_t matmul_weight_format_config(ggml_type type) { return 50; case GGML_TYPE_Q5_1: return 51; - case GGML_TYPE_IQ1_S: - return 19; - case GGML_TYPE_IQ1_M: - return 29; case GGML_TYPE_IQ2_S: return 22; case GGML_TYPE_IQ3_S: @@ -710,37 +667,16 @@ static void run_dispatch_registry_checks() { REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL), "common.binary_f32")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_DIV), "common.binary_f32")); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_swiglu.tiled_pair_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_swiglu.tiled_pair_postops_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.packed_mul_mat_glu.tiled_pair_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "llm.attention_qkv_matmul_postprocess.tiled_vector_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat.skinny_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_swiglu.f32_f32_wmma", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat.skinny_f32_f32_decode", + "llm.attention_qkv_matmul_postprocess.f32_f32_wmma", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_add.skinny_f32_f32_decode", - ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_unary.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat.f32_f32_wmma", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_unary.skinny_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_postops.tiled_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat.f32_f32_decode", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_postops.skinny_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_postops.f32_f32_wmma", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT), "qwen.matmul.q6k_q8_1_x4")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_MUL_MAT), "llm.moe_router.projection_f32_four_row_wave32")); @@ -758,12 +694,9 @@ static void run_dispatch_registry_checks() { "qwen.attention_postprocess_f32_f16")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), "common.get_rows.f32")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), "common.get_rows.f32_next")); - REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), "common.get_rows_scale.f32")); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_GET_ROWS), "common.gather_add_f32")); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_ROPE), "common.rope_set_rows.f32", ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_ROPE), "common.rope_concat.f32", - ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(registry.registrations_for_root(GGML_OP_ROPE), "common.rope.f32", ggml::hrx::DispatchMatchKind::SingleOp)); REQUIRE(has_dispatch_registration(registry.registrations_for_root(GGML_OP_SET_ROWS), "common.set_rows")); @@ -811,37 +744,16 @@ static void run_dispatch_registry_checks() { REQUIRE(!has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_RMS_NORM), "qwen.rmsnorm_f32_quantize_q8_1_x4")); REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_swiglu.tiled_pair_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_swiglu.tiled_pair_postops_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.packed_mul_mat_glu.tiled_pair_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "llm.attention_qkv_matmul_postprocess.tiled_vector_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat.skinny_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat.skinny_f32_f32_decode", - ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_swiglu.f32_f32_wmma", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_add.skinny_f32_f32_decode", + "llm.attention_qkv_matmul_postprocess.f32_f32_wmma", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_unary.tiled_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); - REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_unary.skinny_f32_f32", ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat.f32_f32_wmma", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_postops.tiled_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat.f32_f32_decode", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(has_dispatch_registration_kind(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), - "common.mul_mat_postops.skinny_f32_f32", - ggml::hrx::DispatchMatchKind::Fused)); + "common.mul_mat_postops.f32_f32_wmma", ggml::hrx::DispatchMatchKind::Fused)); REQUIRE(!has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_MUL_MAT), "qwen.matmul.q6k_q8_1_x4")); REQUIRE(has_dispatch_registration(generic_registry.registrations_for_root(GGML_OP_FLASH_ATTN_EXT), @@ -1274,7 +1186,8 @@ static void run_graph_import_checks() { transient_flow.commands[1].bindings[0].origin = ggml::hrx::CommandBindingOrigin::Transient; transient_flow.commands[1].bindings[0].value = transient_value; transient_flow.commands[1].bindings[0].length = command.bindings[2].length; - verification = ggml::hrx::verify_command_program(transient_flow, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + verification = + ggml::hrx::verify_command_program(transient_flow, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(verification.valid()); ggml::hrx::CommandProgram read_before_write = copy_command_program_shape(transient_flow); @@ -1284,7 +1197,8 @@ static void run_graph_import_checks() { read_before_write.commands[1].ordinal = 1; read_before_write.commands[1].dependencies.clear(); read_before_write.commands[1].dependencies.push_back(0); - verification = ggml::hrx::verify_command_program(read_before_write, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + verification = + ggml::hrx::verify_command_program(read_before_write, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(!verification.valid()); REQUIRE(status_contains(verification.status, "reads transient value 100000 before write by command 1")); @@ -1528,216 +1442,6 @@ static void run_scale_f32_dispatch_checks() { ggml_free(ctx); } -static void require_scale_add_dispatch(ggml_context * ctx, - ggml_tensor * input, - ggml_tensor * residual, - int64_t expected_element_count, - int64_t expected_input_stride1, - int64_t expected_residual_stride1, - ggml::hrx::BinaryKind binary_kind = ggml::hrx::BinaryKind::Add, - bool scaled_lhs = true) { - ggml_tensor * scaled = ggml_scale_bias(ctx, input, 0.177800179f, 0.25f); - ggml_tensor * output = nullptr; - switch (binary_kind) { - case ggml::hrx::BinaryKind::Add: - output = ggml_add(ctx, scaled, residual); - break; - case ggml::hrx::BinaryKind::Sub: - output = scaled_lhs ? ggml_sub(ctx, scaled, residual) : ggml_sub(ctx, residual, scaled); - break; - case ggml::hrx::BinaryKind::Mul: - output = ggml_mul(ctx, scaled, residual); - break; - case ggml::hrx::BinaryKind::Div: - output = scaled_lhs ? ggml_div(ctx, scaled, residual) : ggml_div(ctx, residual, scaled); - break; - default: - REQUIRE(false); - } - REQUIRE(scaled != nullptr); - REQUIRE(output != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - - require_scheduled_command_program(graph, [&](const ggml::hrx::Graph &, const ggml::hrx::CommandPlan & plan, - const ggml::hrx::CommandProgram & commands) { - REQUIRE(plan.dispatches.size() == 1); - const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_scale_add_f32"); - REQUIRE(dispatch.kernel.integer_parameters.at("element_count") == expected_element_count); - require_compile_parameter(dispatch, "ggml.scale_add_f32.scale", expected_config_value(0.177800179f)); - require_compile_parameter(dispatch, "ggml.scale_add_f32.bias", expected_config_value(0.25f)); - require_compile_parameter(dispatch, "ggml.scale_add_f32.binary_op", - std::to_string(ggml::hrx::binary_kind_config_value(binary_kind))); - require_compile_parameter(dispatch, "ggml.scale_add_f32.scaled_lhs", scaled_lhs ? "1" : "0"); - require_compile_parameter(dispatch, "ggml.scale_add_f32.input_stride1", std::to_string(expected_input_stride1)); - require_compile_parameter(dispatch, "ggml.scale_add_f32.residual_stride1", - std::to_string(expected_residual_stride1)); - REQUIRE(dispatch.bindings.size() == 3); - REQUIRE(commands.commands.size() == 1); - REQUIRE(commands.commands.front().bindings[0].name == "input"); - REQUIRE(commands.commands.front().bindings[1].name == "residual"); - REQUIRE(commands.commands.front().bindings[2].name == "output"); - }); -} - -static void run_scale_add_f32_dispatch_checks() { - ggml_init_params params = {}; - params.mem_size = 2 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * packed_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2560, 64); - ggml_tensor * packed_residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2560, 64); - REQUIRE(packed_input != nullptr); - REQUIRE(packed_residual != nullptr); - require_scale_add_dispatch(ctx, packed_input, packed_residual, 163840, 2560, 2560); - require_scale_add_dispatch(ctx, packed_input, packed_residual, 163840, 2560, 2560, ggml::hrx::BinaryKind::Sub, - false); - require_scale_add_dispatch(ctx, packed_input, packed_residual, 163840, 2560, 2560, ggml::hrx::BinaryKind::Mul, - true); - require_scale_add_dispatch(ctx, packed_input, packed_residual, 163840, 2560, 2560, ggml::hrx::BinaryKind::Div, - true); - - ggml_tensor * input_storage = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 320, 3); - ggml_tensor * residual_storage = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 384, 3); - ggml_tensor * strided_input = ggml_view_2d(ctx, input_storage, 257, 3, input_storage->nb[1], 0); - ggml_tensor * strided_residual = ggml_view_2d(ctx, residual_storage, 257, 3, residual_storage->nb[1], 0); - REQUIRE(strided_input != nullptr); - REQUIRE(strided_residual != nullptr); - require_scale_add_dispatch(ctx, strided_input, strided_residual, 771, 320, 384); - - ggml_tensor * scaled = ggml_scale(ctx, packed_input, 0.5f); - ggml_tensor * alias_output = ggml_add_inplace(ctx, scaled, packed_residual); - REQUIRE(scaled != nullptr); - REQUIRE(alias_output != nullptr); - REQUIRE(!can_schedule_tensor(ctx, alias_output)); - - ggml_tensor * alias_scaled = ggml_scale_inplace(ctx, packed_input, 0.5f); - ggml_tensor * output_from_alias = ggml_add(ctx, alias_scaled, packed_residual); - REQUIRE(alias_scaled != nullptr); - REQUIRE(output_from_alias != nullptr); - ggml_cgraph * alias_graph = ggml_new_graph(ctx); - REQUIRE(alias_graph != nullptr); - ggml_build_forward_expand(alias_graph, output_from_alias); - ggml::hrx::GraphImportResult alias_imported = ggml::hrx::import_ggml_graph(*alias_graph); - REQUIRE(alias_imported.valid()); - ggml::hrx::DispatchScheduler alias_scheduler; - REQUIRE(alias_scheduler.schedule_graph(alias_imported.graph, test_dispatch_target())); - REQUIRE(alias_scheduler.plan().dispatches.size() == 2); - REQUIRE(kernel_name_for_id(alias_scheduler.plan().dispatches[0].kernel.kernel_id) == - "loom_libs:ggml_scale_bias_f32"); - REQUIRE(kernel_name_for_id(alias_scheduler.plan().dispatches[1].kernel.kernel_id) == "loom_libs:ggml_binary_f32"); - - ggml_free(ctx); -} - -static ggml_tensor * make_test_binary(ggml_context * ctx, - ggml_tensor * lhs, - ggml_tensor * rhs, - ggml::hrx::BinaryKind kind) { - switch (kind) { - case ggml::hrx::BinaryKind::Add: - return ggml_add(ctx, lhs, rhs); - case ggml::hrx::BinaryKind::Sub: - return ggml_sub(ctx, lhs, rhs); - case ggml::hrx::BinaryKind::Mul: - return ggml_mul(ctx, lhs, rhs); - case ggml::hrx::BinaryKind::Div: - return ggml_div(ctx, lhs, rhs); - default: - REQUIRE(false); - return nullptr; - } -} - -static const char * expected_plain_rmsnorm_binary_kernel(int64_t hidden_size, int64_t token_count) { - return token_count < 6 && hidden_size >= 2048 && hidden_size <= 5120 ? - "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" : - "loom_libs:ggml_rmsnorm_binary_f32"; -} - -static void run_rmsnorm_lowtoken_striped_dispatch_checks() { - const struct { - int64_t hidden_size; - int64_t token_count; - ggml::hrx::BinaryKind op; - const char * expected_kernel; - } cases[] = { - { 1024, 1, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32" }, - { 2048, 1, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" }, - { 3072, 2, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" }, - { 5120, 4, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" }, - { 5120, 6, ggml::hrx::BinaryKind::Mul, "loom_libs:ggml_rmsnorm_binary_f32" }, - { 2048, 1, ggml::hrx::BinaryKind::Add, "loom_libs:ggml_rmsnorm_binary_f32" }, - }; - - for (const auto & test : cases) { - ggml_init_params params = {}; - params.mem_size = static_cast(test.hidden_size * test.token_count * sizeof(float) * 8 + 1024 * 1024); - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden_size, test.token_count); - ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.hidden_size); - ggml_tensor * rms = ggml_rms_norm(ctx, input, 1.0e-6f); - ggml_tensor * output = make_test_binary(ctx, rms, weight, test.op); - REQUIRE(output != nullptr); - - const ggml::hrx::Dispatch dispatch = schedule_single_dispatch_for_tensor(ctx, output); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == test.expected_kernel); - REQUIRE(dispatch.bindings.size() == 3); - require_compile_parameter(dispatch, "ggml.rmsnorm_binary_f32.hidden_size", - std::to_string(test.hidden_size)); - require_compile_parameter(dispatch, "ggml.rmsnorm_binary_f32.op", - std::to_string(ggml::hrx::binary_kind_config_value(test.op))); - ggml_free(ctx); - } -} - -static void run_rmsnorm_two_binary_dispatch_checks() { - const struct { - ggml::hrx::BinaryKind normalized_op; - ggml::hrx::BinaryKind output_op; - bool binary_lhs; - } cases[] = { - { ggml::hrx::BinaryKind::Add, ggml::hrx::BinaryKind::Mul, true }, - { ggml::hrx::BinaryKind::Sub, ggml::hrx::BinaryKind::Sub, true }, - { ggml::hrx::BinaryKind::Mul, ggml::hrx::BinaryKind::Div, true }, - { ggml::hrx::BinaryKind::Div, ggml::hrx::BinaryKind::Sub, false }, - }; - - for (const auto & test : cases) { - ggml_init_params params = {}; - params.mem_size = 512 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 128, 4); - ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 128); - ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 128, 4); - ggml_tensor * rms = ggml_rms_norm(ctx, input, 1.0e-6f); - ggml_tensor * scaled = make_test_binary(ctx, rms, weight, test.normalized_op); - ggml_tensor * output = test.binary_lhs ? make_test_binary(ctx, scaled, residual, test.output_op) : - make_test_binary(ctx, residual, scaled, test.output_op); - REQUIRE(output != nullptr); - - const ggml::hrx::Dispatch dispatch = schedule_single_dispatch_for_tensor(ctx, output); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_mul_add_f32"); - REQUIRE(dispatch.bindings.size() == 4); - require_compile_parameter(dispatch, "ggml.rmsnorm_mul_add_f32.normalized_op", - std::to_string(ggml::hrx::binary_kind_config_value(test.normalized_op))); - require_compile_parameter(dispatch, "ggml.rmsnorm_mul_add_f32.output_op", - std::to_string(ggml::hrx::binary_kind_config_value(test.output_op))); - require_compile_parameter(dispatch, "ggml.rmsnorm_mul_add_f32.binary_lhs", test.binary_lhs ? "1" : "0"); - ggml_free(ctx); - } -} - static void run_cont_f32_dispatch_checks() { ggml_init_params params = {}; params.mem_size = 512 * 1024; @@ -1970,86 +1674,6 @@ static void run_binary_f32_broadcast_dispatch_checks() { } } -static void run_binary_q8_publication_checks() { - const struct { - int64_t hidden; - int64_t tokens; - bool broadcast; - bool multiply; - bool side_use; - int projections; - const char * expected_kernel; - } cases[] = { - { 2048, 2, false, true, false, 1, "loom_libs:ggml_binary_f32_publish_q8_1_x4" }, - { 2048, 4, false, false, true, 2, "loom_libs:ggml_binary_f32_publish_q8_1_x4" }, - { 256, 2, true, true, false, 1, "loom_libs:ggml_binary_bc_f32_publish_q8_1_x4" }, - { 2048, 6, false, true, false, 1, "loom_libs:ggml_binary_f32" }, - { 2048, 2, false, true, false, 0, "loom_libs:ggml_binary_f32" }, - }; - - for (const auto & test : cases) { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * lhs = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens); - ggml_tensor * rhs = test.broadcast ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.hidden) : - ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens); - ggml_tensor * binary = test.multiply ? ggml_mul(ctx, lhs, rhs) : ggml_add(ctx, lhs, rhs); - REQUIRE(lhs != nullptr); - REQUIRE(rhs != nullptr); - REQUIRE(binary != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - for (int i = 0; i < test.projections; ++i) { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, test.hidden, 512); - ggml_build_forward_expand(graph, ggml_mul_mat(ctx, weight, binary)); - } - if (test.projections == 0 || test.side_use) { - ggml_build_forward_expand(graph, ggml_scale(ctx, binary, 0.5f)); - } - - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); - const ggml::hrx::Value * binary_value = imported.graph.values().find_tensor(binary); - REQUIRE(binary_value != nullptr); - const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return dispatch.bindings.size() >= 3 && dispatch.bindings[2].value == binary_value->id; - }); - REQUIRE(producer != plan.dispatches.end()); - REQUIRE(kernel_name_for_id(producer->kernel.kernel_id) == test.expected_kernel); - - const bool publishes = std::string_view(test.expected_kernel).find("publish_q8_1_x4") != std::string_view::npos; - const size_t q8_bytes = qwen_q8_1_x4_size(test.tokens, test.hidden); - const auto * q8 = plan.metadata.find_alternate_value(binary_value->id, GGML_TYPE_Q8_1, q8_bytes); - REQUIRE((q8 != nullptr) == publishes); - if (publishes) { - REQUIRE(producer->bindings.size() == 4); - REQUIRE(producer->bindings.back().value == q8->alternate_value); - REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32"; - }) == 0); - REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return &dispatch != &*producer && !dispatch.bindings.empty() && - dispatch.bindings.front().value == q8->alternate_value; - }) == test.projections); - } - - const ggml::hrx::CommandProgram commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - ggml_free(ctx); - } -} - static void run_graph_snapshot_diagnostics_checks() { ggml_init_params params = {}; params.mem_size = 256 * 1024; @@ -2278,7 +1902,7 @@ static void run_graph_index_checks() { ggml_free(ctx); - params.mem_size = 2 * 1024 * 1024; + params.mem_size = 256 * 1024; params.no_alloc = true; ctx = ggml_init(params); REQUIRE(ctx != nullptr); @@ -2358,80 +1982,12 @@ static void run_graph_index_checks() { REQUIRE(command_program_verifies(commands)); } - { - constexpr int64_t hidden_size = 256; - constexpr int64_t token_count = 64; - constexpr int64_t input_stride = 288; - const size_t storage_elements = static_cast(token_count - 1) * input_stride + hidden_size; - ggml_tensor * storage = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, storage_elements); - ggml_tensor * view = ggml_view_2d(ctx, storage, hidden_size, token_count, input_stride * sizeof(float), 0); - ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - ggml_tensor * norm = ggml_rms_norm(ctx, view, 0.00001f); - ggml_tensor * scaled = ggml_mul(ctx, norm, weight); - REQUIRE(storage != nullptr); - REQUIRE(view != nullptr); - REQUIRE(weight != nullptr); - REQUIRE(norm != nullptr); - REQUIRE(scaled != nullptr); - - ggml_cgraph * strided_graph = ggml_new_graph(ctx); - REQUIRE(strided_graph != nullptr); - ggml_build_forward_expand(strided_graph, scaled); - ggml::hrx::GraphImportResult strided_imported = ggml::hrx::import_ggml_graph(*strided_graph); - REQUIRE(strided_imported.valid()); - REQUIRE(scheduler.schedule_graph(strided_imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 1); - const ggml::hrx::Dispatch & strided_dispatch = scheduler.plan().dispatches[0]; - REQUIRE(kernel_name_for_id(strided_dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_binary_strided_f32"); - require_compile_parameter(strided_dispatch, "ggml.rmsnorm_binary_f32.hidden_size", "256"); - require_compile_parameter(strided_dispatch, "ggml.rmsnorm_binary_f32.input_stride", "288"); - require_compile_parameter(strided_dispatch, "ggml.rmsnorm_binary_f32.op", "2"); - REQUIRE(strided_dispatch.bindings.size() == 3); - REQUIRE(strided_dispatch.bindings[0].length == - static_cast(token_count - 1) * view->nb[1] + hidden_size * sizeof(float)); - - ggml_tensor * full_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, token_count); - ggml_tensor * rejected_norm = ggml_rms_norm(ctx, view, 0.00001f); - ggml_tensor * rejected_scaled = ggml_mul(ctx, rejected_norm, full_weight); - REQUIRE(full_weight != nullptr); - REQUIRE(rejected_scaled != nullptr); - ggml_cgraph * rejected_graph = ggml_new_graph(ctx); - REQUIRE(rejected_graph != nullptr); - ggml_build_forward_expand(rejected_graph, rejected_scaled); - ggml::hrx::GraphImportResult rejected_imported = ggml::hrx::import_ggml_graph(*rejected_graph); - REQUIRE(rejected_imported.valid()); - REQUIRE(scheduler.schedule_graph(rejected_imported.graph, test_dispatch_target())); - for (const ggml::hrx::Dispatch & rejected_dispatch : scheduler.plan().dispatches) { - REQUIRE(kernel_name_for_id(rejected_dispatch.kernel.kernel_id) != - "loom_libs:ggml_rmsnorm_binary_strided_f32"); - } - - ggml_tensor * second_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - ggml_tensor * shared_norm = ggml_rms_norm(ctx, view, 0.00001f); - ggml_tensor * scaled_a = ggml_mul(ctx, shared_norm, weight); - ggml_tensor * scaled_b = ggml_add(ctx, shared_norm, second_weight); - ggml_tensor * joined = ggml_add(ctx, scaled_a, scaled_b); - REQUIRE(second_weight != nullptr); - REQUIRE(joined != nullptr); - ggml_cgraph * fanout_graph = ggml_new_graph(ctx); - REQUIRE(fanout_graph != nullptr); - ggml_build_forward_expand(fanout_graph, joined); - ggml::hrx::GraphImportResult fanout_imported = ggml::hrx::import_ggml_graph(*fanout_graph); - REQUIRE(fanout_imported.valid()); - REQUIRE(scheduler.schedule_graph(fanout_imported.graph, test_dispatch_target())); - for (const ggml::hrx::Dispatch & fanout_dispatch : scheduler.plan().dispatches) { - REQUIRE(kernel_name_for_id(fanout_dispatch.kernel.kernel_id) != - "loom_libs:ggml_rmsnorm_binary_strided_f32"); - } - } - ggml_free(ctx); } static void run_rope_set_rows_dispatch_checks() { ggml_init_params params = {}; - params.mem_size = 2 * 1024 * 1024; + params.mem_size = 1024 * 1024; params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); @@ -2442,65 +1998,6 @@ static void run_rope_set_rows_dispatch_checks() { constexpr int64_t cache_rows = 8; constexpr int64_t hidden_size = head_size * head_count; - const auto check_rope_concat = [&](int64_t concat_head_count, int64_t concat_token_count, int64_t n_dims) { - constexpr int64_t prefix_size = 64; - constexpr int64_t rope_size = 32; - constexpr int64_t source_size = prefix_size + rope_size; - const size_t source_elements = static_cast(source_size) * concat_head_count * concat_token_count; - ggml_tensor * storage = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, source_elements); - ggml_tensor * prefix = - ggml_view_3d(ctx, storage, prefix_size, concat_head_count, concat_token_count, source_size * sizeof(float), - source_size * concat_head_count * sizeof(float), 0); - ggml_tensor * rope_input = - ggml_view_3d(ctx, storage, rope_size, concat_head_count, concat_token_count, source_size * sizeof(float), - source_size * concat_head_count * sizeof(float), prefix_size * sizeof(float)); - ggml_tensor * positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, concat_token_count); - ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_dims / 2); - ggml_tensor * rope = ggml_rope_ext(ctx, rope_input, positions, freqs, n_dims, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, - 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); - ggml_tensor * output = ggml_concat(ctx, prefix, rope, 0); - REQUIRE(storage != nullptr); - REQUIRE(prefix != nullptr); - REQUIRE(rope_input != nullptr); - REQUIRE(positions != nullptr); - REQUIRE(freqs != nullptr); - REQUIRE(rope != nullptr); - REQUIRE(output != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 1); - const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rope_concat_f32"); - REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == concat_token_count); - require_compile_parameter(dispatch, "ggml.rope_concat_f32.prefix_size", std::to_string(prefix_size)); - require_compile_parameter(dispatch, "ggml.rope_concat_f32.rope_size", std::to_string(rope_size)); - require_compile_parameter(dispatch, "ggml.rope_concat_f32.n_dims", std::to_string(n_dims)); - require_compile_parameter(dispatch, "ggml.rope_concat_f32.head_count", std::to_string(concat_head_count)); - require_compile_parameter(dispatch, "ggml.rope_concat_f32.token_capacity", std::to_string(concat_token_count)); - require_compile_parameter(dispatch, "ggml.rope_concat_f32.input_stride1", std::to_string(source_size)); - require_compile_parameter(dispatch, "ggml.rope_concat_f32.input_stride2", - std::to_string(source_size * concat_head_count)); - require_compile_parameter(dispatch, "ggml.rope_concat_f32.mode", "2"); - REQUIRE(dispatch.bindings.size() == 5); - REQUIRE(dispatch.bindings[1].length == source_elements * sizeof(float)); - - const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( - imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - }; - - check_rope_concat(40, 64, 32); - check_rope_concat(3, 2, 16); - { ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); @@ -2705,50 +2202,13 @@ static void run_rope_set_rows_dispatch_checks() { } { - ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); - ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); - ggml_tensor * raw_scales = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, token_count); - ggml_tensor * scales = ggml_scale(ctx, raw_scales, 0.5f); - ggml_tensor * rope = ggml_rope(ctx, input, pos, head_size, GGML_ROPE_TYPE_NORMAL); - ggml_tensor * output = ggml_mul(ctx, rope, scales); - REQUIRE(input != nullptr); - REQUIRE(pos != nullptr); - REQUIRE(raw_scales != nullptr); - REQUIRE(scales != nullptr); - REQUIRE(rope != nullptr); - REQUIRE(output != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - const size_t rope_index = producer_index_for_tensor(imported.graph, rope); - const size_t scale_index = producer_index_for_tensor(imported.graph, scales); - std::vector covered_nodes(imported.graph.nodes().size(), false); - ggml::hrx::CommandPlan plan; - ggml::hrx::DispatchMatch match; - - REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, rope_index, match)); - REQUIRE(match.dispatches.size() == 1); - REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == "loom_libs:ggml_rope_f32"); - - covered_nodes[scale_index] = true; - match = {}; - REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, rope_index, match)); - REQUIRE(match.dispatches.size() == 1); - REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == "loom_libs:ggml_rope_token_scale_f32"); - } - - { - ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); - ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); - ggml_tensor * rope = ggml_rope(ctx, input, pos, head_size, GGML_ROPE_TYPE_NEOX); - ggml_tensor * output = ggml_scale(ctx, rope, 0.125f); - REQUIRE(input != nullptr); - REQUIRE(pos != nullptr); - REQUIRE(rope != nullptr); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hidden_size, cache_rows); + ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, token_count); + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); + ggml_tensor * output = ggml_set_rows(ctx, cache, rows, indices); + REQUIRE(cache != nullptr); + REQUIRE(rows != nullptr); + REQUIRE(indices != nullptr); REQUIRE(output != nullptr); ggml_cgraph * graph = ggml_new_graph(ctx); @@ -2757,114 +2217,22 @@ static void run_rope_set_rows_dispatch_checks() { ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 1); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_SET_ROWS); ggml::hrx::DispatchScheduler scheduler; REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); REQUIRE(scheduler.plan().valid()); REQUIRE(scheduler.plan().dispatches.size() == 1); + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rope_scale_f32"); - require_compile_parameter(dispatch, "ggml.rope_scale_f32.scale", "0.125"); - REQUIRE(dispatch.bindings.size() == 5); - - const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( - imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - } - - { - ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); - ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); - ggml_tensor * scales = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, token_count); - ggml_tensor * rope = ggml_rope(ctx, input, pos, head_size, GGML_ROPE_TYPE_NORMAL); - ggml_tensor * output = ggml_mul(ctx, rope, scales); - REQUIRE(input != nullptr); - REQUIRE(pos != nullptr); - REQUIRE(scales != nullptr); - REQUIRE(rope != nullptr); - REQUIRE(output != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 1); - const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rope_token_scale_f32"); - REQUIRE(dispatch.bindings.size() == 6); - - const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( - imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - } - - { - ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); - ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); - ggml_tensor * rope = ggml_rope(ctx, input, pos, head_size, GGML_ROPE_TYPE_NEOX); - ggml_tensor * output = ggml_scale(ctx, rope, 0.125f); - ggml_tensor * side = ggml_scale(ctx, rope, 0.25f); - REQUIRE(input != nullptr); - REQUIRE(pos != nullptr); - REQUIRE(rope != nullptr); - REQUIRE(output != nullptr); - REQUIRE(side != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - ggml_build_forward_expand(graph, side); - - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 3); - REQUIRE(kernel_name_for_id(scheduler.plan().dispatches.front().kernel.kernel_id) == "loom_libs:ggml_rope_f32"); - } - - { - ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, hidden_size, cache_rows); - ggml_tensor * rows = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, token_count); - ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); - ggml_tensor * output = ggml_set_rows(ctx, cache, rows, indices); - REQUIRE(cache != nullptr); - REQUIRE(rows != nullptr); - REQUIRE(indices != nullptr); - REQUIRE(output != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - REQUIRE(imported.graph.nodes().size() == 1); - REQUIRE(imported.graph.nodes()[0].op == GGML_OP_SET_ROWS); - - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 1); - - const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_set_rows"); - REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); - REQUIRE(dispatch.kernel.integer_parameters.at("cache_row_count") == cache_rows); - REQUIRE(dispatch.kernel.integer_parameters.at("hidden_size") == hidden_size); - require_compile_parameter(dispatch, "ggml.set_rows.input_format", "32"); - require_compile_parameter(dispatch, "ggml.set_rows.output_format", "16"); - REQUIRE(dispatch.bindings.size() == 3); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_set_rows"); + REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == token_count); + REQUIRE(dispatch.kernel.integer_parameters.at("cache_row_count") == cache_rows); + REQUIRE(dispatch.kernel.integer_parameters.at("hidden_size") == hidden_size); + require_compile_parameter(dispatch, "ggml.set_rows.input_format", "32"); + require_compile_parameter(dispatch, "ggml.set_rows.output_format", "16"); + REQUIRE(dispatch.bindings.size() == 3); const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); @@ -3211,8 +2579,6 @@ static void run_graph_traversal_checks() { } } -static bool matmul_graph_is_supported(ggml_context * ctx, ggml_tensor * output); - static void schedule_single_matmul_command(ggml_context * ctx, ggml_tensor * output, const char * expected_kernel_name, @@ -3232,20 +2598,9 @@ static void schedule_single_matmul_command(ggml_context * ctx, REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); REQUIRE(scheduler.plan().valid()); const ggml_type weight_type = output->src[0]->type; - REQUIRE(!scheduler.plan().dispatches.empty()); - const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.back(); - const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); - const bool vector_kernel = string_contains(kernel_name, "ggml_mul_mat_vector"); - const bool q6_vector_kernel = kernel_name == "loom_libs:ggml_mul_mat_vector_q6_f32_f32"; - const bool aligned_dynamic_kernel = kernel_name == "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_aligned"; - const bool codebook_kernel = string_contains(kernel_name, "_iq1_") || string_contains(kernel_name, "_iq2_s_") || - string_contains(kernel_name, "_iq3_s_") || - string_contains(kernel_name, "_iq3_xxs_"); - const std::string expected_kernel(expected_kernel_name); - const bool q8_capable_kernel = expected_kernel == "loom_libs:ggml_mul_mat_f32_f32_wmma" || - expected_kernel == "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32"; - const bool q8_input = !vector_kernel && expected_token_count <= 5 && expected_output_size % 64 == 0 && - (weight_type == GGML_TYPE_Q4_K || weight_type == GGML_TYPE_Q6_K) && q8_capable_kernel; + const bool q8_input = expected_token_count <= 5 && expected_output_size % 64 == 0 && + (weight_type == GGML_TYPE_Q4_K || weight_type == GGML_TYPE_Q6_K) && + std::string(expected_kernel_name) == "loom_libs:ggml_mul_mat_f32_f32_wmma"; const size_t command_count = q8_input ? 2 : 1; REQUIRE(scheduler.plan().dispatches.size() == command_count); if (q8_input) { @@ -3253,36 +2608,17 @@ static void schedule_single_matmul_command(ggml_context * ctx, "qwen3_moe:ggml_quantize_q8_1_x4_f32"); } - const bool accepts_vector_token1 = - expected_token_count == 1 && (expected_kernel == "loom_libs:ggml_mul_mat_f32_f32_wmma" || - expected_kernel == "loom_libs:ggml_mul_mat_f32_f32_decode_wave64" || - expected_kernel == "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma" || - expected_kernel == "qwen3_moe:qwen3_moe_router_projection_f32_four_row_wave32"); - if (vector_kernel && accepts_vector_token1) { - REQUIRE(string_contains(kernel_name, "ggml_mul_mat_vector")); - } else { - REQUIRE(kernel_name == expected_kernel_name); - } + const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.back(); + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == expected_kernel_name); REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); - if (vector_kernel) { - REQUIRE(dispatch.bindings.size() == (codebook_kernel ? 5 : q6_vector_kernel ? 3 : 4)); - if (!q6_vector_kernel) { - require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); - } - require_compile_parameter(dispatch, "ggml.matmul.vector.input_size", std::to_string(expected_input_size)); - require_compile_parameter(dispatch, "ggml.matmul.vector.output_size", std::to_string(expected_output_size)); - require_compile_parameter(dispatch, "ggml.matmul.vector.output_accumulation", "0"); - if (q6_vector_kernel) { - require_compile_parameter(dispatch, "ggml.matmul.vector.q6_storage", "0"); - } - } else if (string_contains(kernel_name, "ggml_mul_mat_q6_k_packed_token1_f16_wmma")) { - REQUIRE(dispatch.bindings.size() == 3); + REQUIRE(dispatch.bindings.size() == 3); + if (string_contains(kernel_name, "ggml_mul_mat_q6_k_packed_token1_f16_wmma")) { require_compile_parameter(dispatch, "ggml.mul_mat_q6_k_packed.input_size", std::to_string(expected_input_size)); require_compile_parameter(dispatch, "ggml.mul_mat_q6_k_packed.output_size", std::to_string(expected_output_size)); require_compile_parameter(dispatch, "ggml.mul_mat_q6_k_packed.output_accumulation", "0"); } else if (string_contains(kernel_name, "_decode")) { - REQUIRE(dispatch.bindings.size() == 3); REQUIRE(dispatch.kernel.integer_parameters.at("input_size") == expected_input_size); REQUIRE(dispatch.kernel.integer_parameters.at("output_size") == expected_output_size); require_compile_parameter(dispatch, "ggml.mul_mat_f32_f32_decode.token_capacity", @@ -3294,13 +2630,7 @@ static void schedule_single_matmul_command(ggml_context * ctx, require_compile_parameter(dispatch, "ggml.mul_mat_f32_f32_decode.weight_format", std::to_string(matmul_weight_format_config(weight->type))); } else if (string_contains(kernel_name, "ggml_mul_mat")) { - REQUIRE(dispatch.bindings.size() == (codebook_kernel ? 4 : 3)); - if (aligned_dynamic_kernel) { - REQUIRE(dispatch.kernel.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); - REQUIRE(dispatch.kernel.compile_parameters.count("ggml.workload.token_capacity") == 0); - } else { - require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); - } + require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); require_compile_parameter(dispatch, "ggml.mul_mat.input_size", std::to_string(expected_input_size)); require_compile_parameter(dispatch, "ggml.mul_mat.output_size", std::to_string(expected_output_size)); require_compile_parameter(dispatch, "ggml.mul_mat.output_accumulation", "0"); @@ -3312,7 +2642,6 @@ static void schedule_single_matmul_command(ggml_context * ctx, std::to_string(q8_input ? (weight_type == GGML_TYPE_Q4_K ? 44 : 46) : matmul_weight_format_config(weight->type))); } else if (string_contains(kernel_name, "dense_linear")) { - REQUIRE(dispatch.bindings.size() == 3); require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(expected_token_count)); require_compile_parameter(dispatch, "qwen3_moe.dense_quantized.input_size", std::to_string(expected_input_size)); @@ -3320,7 +2649,6 @@ static void schedule_single_matmul_command(ggml_context * ctx, std::to_string(expected_output_size)); require_compile_parameter(dispatch, "qwen3_moe.dense_quantized.output_accumulation", "0"); } else { - REQUIRE(dispatch.bindings.size() == 3); require_compile_parameter(dispatch, "qwen3_moe.workload.token_capacity", std::to_string(expected_token_count)); require_compile_parameter(dispatch, "qwen3_moe.model.hidden_size", std::to_string(expected_input_size)); require_compile_parameter(dispatch, "qwen3_moe.router.expert_count", std::to_string(expected_output_size)); @@ -3331,109 +2659,12 @@ static void schedule_single_matmul_command(ggml_context * ctx, REQUIRE(commands.valid()); REQUIRE(commands.commands.size() == command_count); REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands.back().bindings.size() == - (codebook_kernel ? (vector_kernel ? 5 : 4) : vector_kernel && !q6_vector_kernel ? 4 : 3)); + REQUIRE(commands.commands.back().bindings.size() == 3); REQUIRE(commands.commands.back().bindings[0].name == "input"); REQUIRE(commands.commands.back().bindings[1].name == "weight"); - const size_t output_binding = codebook_kernel ? 3 : 2; - if (codebook_kernel) { - REQUIRE(string_contains(commands.commands.back().bindings[2].name, "grid")); - } - REQUIRE(commands.commands.back().bindings[output_binding].name == "output"); - if (vector_kernel && !q6_vector_kernel) { - REQUIRE(commands.commands.back().bindings[output_binding + 1].name == "next_output"); - } + REQUIRE(commands.commands.back().bindings[2].name == "output"); } -static void run_iq1_codebook_matmul_dispatch_checks() { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - auto check_single = [&](ggml_type type, int64_t token_count, const char * expected_kernel) { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, type, 2048, 128); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, token_count); - ggml_tensor * output = ggml_mul_mat(ctx, weight, input); - REQUIRE(weight != nullptr); - REQUIRE(input != nullptr); - REQUIRE(output != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 1); - REQUIRE(kernel_name_for_id(scheduler.plan().dispatches[0].kernel.kernel_id) == expected_kernel); - REQUIRE(scheduler.plan().transients.size() == 1); - REQUIRE(scheduler.plan().transients[0].name == "common.iq1.grid.v1"); - REQUIRE(scheduler.plan().transients[0].size == 16384); - REQUIRE(scheduler.plan().constant_initializations.size() == 1); - REQUIRE(scheduler.plan().constant_initializations[0].name == "common.iq1.grid.v1"); - REQUIRE(scheduler.plan().constant_initializations[0].data.size() == 16384); - const std::array expected_grid_1542 = { 0, 1, 0, 0xff, 1, 0, 0xff, 1 }; - REQUIRE(std::equal(expected_grid_1542.begin(), expected_grid_1542.end(), - scheduler.plan().constant_initializations[0].data.begin() + 1542 * 8)); - - const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( - imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(commands.commands.size() == 1); - REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands[0].bindings.size() == (token_count == 1 ? 5 : 4)); - REQUIRE(commands.commands[0].bindings[0].name == "input"); - REQUIRE(commands.commands[0].bindings[1].name == "weight"); - REQUIRE(commands.commands[0].bindings[2].name == "iq1_grid"); - REQUIRE(commands.commands[0].bindings[3].name == "output"); - }; - - check_single(GGML_TYPE_IQ1_S, 1, "loom_libs:ggml_mul_mat_vector_iq1_s_f32_f32"); - check_single(GGML_TYPE_IQ1_S, 64, "loom_libs:ggml_mul_mat_tiled_input_f32_iq1_s_publish_f32"); - check_single(GGML_TYPE_IQ1_M, 1, "loom_libs:ggml_mul_mat_vector_iq1_m_f32_f32"); - check_single(GGML_TYPE_IQ1_M, 64, "loom_libs:ggml_mul_mat_tiled_input_f32_iq1_m_publish_f32"); - - for (const auto [type, disable_env] : - { std::pair{ GGML_TYPE_IQ1_S, "GGML_HRX_DISABLE_IQ1_S_CODEBOOK_MATMUL" }, - std::pair{ GGML_TYPE_IQ1_M, "GGML_HRX_DISABLE_IQ1_M_CODEBOOK_MATMUL" } }) { - EnvironmentVariableGuard guard(disable_env); - REQUIRE(setenv(disable_env, "1", 1) == 0); - ggml_tensor * weight = ggml_new_tensor_2d(ctx, type, 2048, 128); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); - ggml_tensor * output = ggml_mul_mat(ctx, weight, input); - REQUIRE(output != nullptr); - REQUIRE(!matmul_graph_is_supported(ctx, output)); - } - - ggml_tensor * iq1_s_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ1_S, 2048, 128); - ggml_tensor * iq1_m_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ1_M, 2048, 128); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); - ggml_tensor * iq1_s_output = ggml_mul_mat(ctx, iq1_s_weight, input); - ggml_tensor * iq1_m_output = ggml_mul_mat(ctx, iq1_m_weight, input); - REQUIRE(iq1_s_output != nullptr); - REQUIRE(iq1_m_output != nullptr); - ggml_cgraph * graph = ggml_new_graph_custom(ctx, GGML_DEFAULT_GRAPH_SIZE, false); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, iq1_s_output); - ggml_build_forward_expand(graph, iq1_m_output); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 2); - REQUIRE(scheduler.plan().transients.size() == 1); - REQUIRE(scheduler.plan().transients[0].name == "common.iq1.grid.v1"); - REQUIRE(scheduler.plan().constant_initializations.size() == 1); - REQUIRE(scheduler.plan().dispatches[0].bindings[2].value == scheduler.plan().transients[0].value); - REQUIRE(scheduler.plan().dispatches[1].bindings[2].value == scheduler.plan().transients[0].value); - - ggml_free(ctx); -} static void schedule_fused_matmul_unary_command(ggml_context * ctx, ggml_tensor * output, @@ -3518,8 +2749,8 @@ static void schedule_fused_matmul_swiglu_command( REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); REQUIRE(scheduler.plan().valid()); const bool low_token = expected_token_count <= 5 && expected_input_size % 256 == 0; - const bool q8_input = low_token && expected_gate_type == GGML_TYPE_Q4_K && expected_up_type == GGML_TYPE_Q4_K && - expected_output_size % 64 == 0; + const bool q8_input = low_token && expected_gate_type == GGML_TYPE_Q4_K && + expected_up_type == GGML_TYPE_Q4_K && expected_output_size % 64 == 0; const size_t command_count = q8_input ? 2 : 1; REQUIRE(scheduler.plan().dispatches.size() == command_count); if (q8_input) { @@ -3531,7 +2762,7 @@ static void schedule_fused_matmul_swiglu_command( const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); REQUIRE(kernel_name == (q8_input ? "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot" : low_token ? "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot" : - "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32")); + "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma")); REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); REQUIRE(dispatch.bindings.size() == 4); require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); @@ -3557,55 +2788,6 @@ static void schedule_fused_matmul_swiglu_command( REQUIRE(commands.commands.back().bindings[3].name == "output"); } -static void schedule_fused_matmul_swiglu_postops_command(ggml_context * ctx, - ggml_tensor * output, - const char * expected_kernel_name, - size_t expected_node_count, - const std::vector & expected_ops, - const std::vector expected_binding_names) { - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - REQUIRE(imported.graph.nodes().size() == expected_node_count); - REQUIRE(expected_ops.size() == expected_node_count); - for (size_t i = 0; i < expected_ops.size(); ++i) { - REQUIRE(imported.graph.nodes()[i].op == expected_ops[i]); - } - - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 1); - - const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == expected_kernel_name); - REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == 8); - REQUIRE(dispatch.bindings.size() == expected_binding_names.size()); - require_compile_parameter(dispatch, "ggml.workload.token_capacity", "8"); - require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.input_size", "640"); - require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.output_size", "256"); - require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.gate_weight_format", - std::to_string(matmul_weight_format_config(GGML_TYPE_IQ4_NL))); - require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.up_weight_format", - std::to_string(matmul_weight_format_config(GGML_TYPE_IQ4_NL))); - require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.op", - std::to_string(ggml::hrx::binary_kind_config_value(ggml::hrx::BinaryKind::SwiGLU))); - REQUIRE(dispatch.bindings.back().value == imported.graph.nodes().back().output); - - const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( - imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(commands.commands.size() == 1); - REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands.front().bindings.size() == expected_binding_names.size()); - for (size_t i = 0; i < expected_binding_names.size(); ++i) { - REQUIRE(commands.commands.front().bindings[i].name == expected_binding_names[i]); - } -} - static void schedule_packed_glu_command(ggml_context * ctx, ggml_tensor * output, ggml::hrx::BinaryKind expected_binary_op, @@ -3677,7 +2859,7 @@ static void schedule_packed_matmul_glu_command(ggml_context * ctx, const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); - REQUIRE(kernel_name == "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32"); + REQUIRE(kernel_name == "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma"); REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == expected_token_count); REQUIRE(dispatch.bindings.size() == 4); require_compile_parameter(dispatch, "ggml.workload.token_capacity", std::to_string(expected_token_count)); @@ -3690,8 +2872,8 @@ static void schedule_packed_matmul_glu_command(ggml_context * ctx, require_compile_parameter(dispatch, "ggml.mul_mat_swiglu.op", std::to_string(ggml::hrx::binary_kind_config_value(expected_binary_op))); - const size_t half_offset = - ggml_row_size(expected_weight_type, expected_input_size) * static_cast(expected_output_size); + const size_t half_offset = ggml_row_size(expected_weight_type, expected_input_size) * + static_cast(expected_output_size); REQUIRE(dispatch.bindings[1].offset == (expected_swapped ? half_offset : 0)); REQUIRE(dispatch.bindings[2].offset == (expected_swapped ? 0 : half_offset)); REQUIRE(dispatch.bindings[3].value == imported.graph.nodes()[1].output); @@ -3750,8 +2932,7 @@ static void schedule_fused_matmul_postops_command(ggml_context * REQUIRE(dispatch.kernel.compile_parameters.count("ggml.mul_mat_postops.rms_epsilon") == 0); if (expected_rmsnorm) { const ggml::hrx::Dispatch & rmsnorm_dispatch = scheduler.plan().dispatches[1]; - REQUIRE(kernel_name_for_id(rmsnorm_dispatch.kernel.kernel_id) == - expected_plain_rmsnorm_binary_kernel(expected_output_size, expected_token_count)); + REQUIRE(kernel_name_for_id(rmsnorm_dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_binary_f32"); } const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( @@ -3985,233 +3166,40 @@ static void schedule_qwen_terminal_q6k_q8_command(int64_t token_count) { ggml_free(ctx); } -static void run_qwen_decode_rmsnorm_publication_checks() { - const struct { - bool projection; - bool side_use; - bool q8; - } cases[] = { - { true, false, true }, - { true, true, true }, - { false, false, false }, - }; - for (const auto & test : cases) { - ggml_init_params params = {}; - params.mem_size = 8 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); +static void schedule_get_rows_q8_1_alternate_command(ggml_type embedding_weight_type) { + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); - ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2048); - ggml_tensor * normalized = ggml_mul(ctx, ggml_rms_norm(ctx, input, 0.000001f), norm_weight); - ggml_tensor * projection = nullptr; - if (test.projection) { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 512); - projection = ggml_mul_mat(ctx, weight, normalized); - } - ggml_tensor * side = test.side_use ? ggml_scale(ctx, normalized, 0.5f) : nullptr; + ggml_tensor * embedding_weight = ggml_new_tensor_2d(ctx, embedding_weight_type, 2048, 151936); + ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_tensor * vocab_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 151936); + REQUIRE(embedding_weight != nullptr); + REQUIRE(token_ids != nullptr); + REQUIRE(vocab_weight != nullptr); - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, projection != nullptr ? projection : normalized); - if (side != nullptr) { - ggml_build_forward_expand(graph, side); - } + ggml_tensor * embedding = ggml_get_rows(ctx, embedding_weight, token_ids); + ggml_tensor * logits = ggml_mul_mat(ctx, vocab_weight, embedding); + REQUIRE(embedding != nullptr); + REQUIRE(logits != nullptr); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, logits); - const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == - "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4"; - }); - REQUIRE((producer != plan.dispatches.end()) == test.q8); - const ggml::hrx::Value * output = imported.graph.values().find_tensor(normalized); - REQUIRE(output != nullptr); - const auto * q8 = plan.metadata.find_alternate_value( - output->id, GGML_TYPE_Q8_1, qwen_q8_1_x4_size(1, 2048)); - REQUIRE((q8 != nullptr) == test.q8); - if (test.q8) { - REQUIRE(producer->bindings.size() == 4); - REQUIRE(producer->bindings.back().value == q8->alternate_value); - REQUIRE(std::any_of(std::next(producer), plan.dispatches.end(), [&](const auto & dispatch) { - return !dispatch.bindings.empty() && dispatch.bindings.front().value == q8->alternate_value; - })); - } - const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( - imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - ggml_free(ctx); - } -} + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.nodes().size() == 2); + REQUIRE(imported.graph.nodes()[0].op == GGML_OP_GET_ROWS); + REQUIRE(imported.graph.nodes()[1].op == GGML_OP_MUL_MAT); -static void run_qwen_decode_publication_consumer_qualification_checks() { - enum class ConsumerCase { - Dense, - UnsupportedWeight, - WrongOperand, - IncompatibleGeometry, - Routed, - RoutedQ6, - RoutedShape, - RoutedRoute, - AliasedDense, - }; - const struct { - ConsumerCase consumer; - bool q8; - } cases[] = { - { ConsumerCase::Dense, true }, - { ConsumerCase::UnsupportedWeight, false }, - { ConsumerCase::WrongOperand, false }, - { ConsumerCase::IncompatibleGeometry, false }, - { ConsumerCase::Routed, true }, - { ConsumerCase::RoutedQ6, true }, - { ConsumerCase::RoutedShape, false }, - { ConsumerCase::RoutedRoute, false }, - { ConsumerCase::AliasedDense, true }, - }; - - for (const auto & test : cases) { - ggml_init_params params = {}; - params.mem_size = 32 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); - ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2048); - ggml_tensor * rms = ggml_rms_norm(ctx, input, 0.000001f); - ggml_tensor * normalized = ggml_mul(ctx, rms, norm_weight); - REQUIRE(rms != nullptr); - REQUIRE(normalized != nullptr); - - ggml_tensor * consumer_input = normalized; - if (test.consumer == ConsumerCase::AliasedDense) { - consumer_input = ggml_reshape_2d(ctx, normalized, 2048, 1); - REQUIRE(consumer_input != nullptr); - } - - ggml_tensor * terminal = nullptr; - if (test.consumer == ConsumerCase::Dense || test.consumer == ConsumerCase::AliasedDense) { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 512); - terminal = ggml_mul_mat(ctx, weight, consumer_input); - } else if (test.consumer == ConsumerCase::UnsupportedWeight) { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_BF16, 2048, 512); - terminal = ggml_mul_mat(ctx, weight, normalized); - } else if (test.consumer == ConsumerCase::WrongOperand) { - ggml_tensor * activation = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); - terminal = ggml_mul_mat(ctx, normalized, activation); - } else if (test.consumer == ConsumerCase::IncompatibleGeometry) { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 63); - terminal = ggml_mul_mat(ctx, weight, normalized); - } else { - const int64_t route_count = test.consumer == ConsumerCase::RoutedShape ? 4 : 8; - ggml_tensor * route_ids = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, route_count, 1); - ggml_tensor * gate_weight = ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, 2048, 768, 128); - ggml_tensor * up_weight = ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, 2048, 768, 128); - const ggml_type down_weight_type = test.consumer == ConsumerCase::RoutedQ6 ? GGML_TYPE_Q6_K : GGML_TYPE_Q4_K; - ggml_tensor * down_weight = ggml_new_tensor_3d(ctx, down_weight_type, 768, 2048, 128); - ggml_tensor * gate = ggml_mul_mat_id(ctx, gate_weight, normalized, route_ids); - ggml_tensor * up = ggml_mul_mat_id(ctx, up_weight, normalized, route_ids); - ggml_tensor * glu = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU); - ggml_tensor * down_route_ids = test.consumer == ConsumerCase::RoutedRoute ? - ggml_new_tensor_2d(ctx, GGML_TYPE_I32, route_count, 1) : - route_ids; - ggml_tensor * down = ggml_mul_mat_id(ctx, down_weight, glu, down_route_ids); - ggml_tensor * route_weights = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, route_count, 1); - terminal = ggml_mul(ctx, down, route_weights); - } - REQUIRE(terminal != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, terminal); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - - const ggml::hrx::Value * rms_value = imported.graph.values().find_tensor(rms); - const ggml::hrx::Value * normalized_value = imported.graph.values().find_tensor(normalized); - REQUIRE(rms_value != nullptr); - REQUIRE(normalized_value != nullptr); - const ggml::hrx::GraphNode * rms_node = imported.graph.index().producer(rms_value->id); - size_t rms_index = 0; - REQUIRE(rms_node != nullptr); - REQUIRE(imported.graph.index().node_index(rms_node, rms_index)); - - const ggml::hrx::CommandPlan empty_plan; - std::vector covered_nodes(imported.graph.nodes().size(), false); - ggml::hrx::DispatchMatch match; - REQUIRE(match_dispatch_at_index(imported.graph, empty_plan, covered_nodes, rms_index, match)); - const auto * q8 = match.metadata.find_alternate_value( - normalized_value->id, GGML_TYPE_Q8_1, qwen_q8_1_x4_size(1, 2048)); - REQUIRE((q8 != nullptr) == test.q8); - if (test.q8) { - REQUIRE(kernel_name_for_id(match.dispatches.back().kernel.kernel_id) == - "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4"); - } - ggml_free(ctx); - } -} - -static void schedule_qwen_terminal_q6k_vector_command() { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); - ggml_tensor * vocab_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 151936); - REQUIRE(input != nullptr); - REQUIRE(vocab_weight != nullptr); - ggml_tensor * logits = ggml_mul_mat(ctx, vocab_weight, input); - REQUIRE(logits != nullptr); - - schedule_single_matmul_command(ctx, logits, "loom_libs:ggml_mul_mat_vector_q6_f32_f32", 1, 2048, 151936); - ggml_free(ctx); -} - -static void schedule_get_rows_q8_1_alternate_command(ggml_type embedding_weight_type) { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * embedding_weight = ggml_new_tensor_2d(ctx, embedding_weight_type, 2048, 151936); - ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); - ggml_tensor * vocab_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 151936); - REQUIRE(embedding_weight != nullptr); - REQUIRE(token_ids != nullptr); - REQUIRE(vocab_weight != nullptr); - - ggml_tensor * embedding = ggml_get_rows(ctx, embedding_weight, token_ids); - ggml_tensor * logits = ggml_mul_mat(ctx, vocab_weight, embedding); - REQUIRE(embedding != nullptr); - REQUIRE(logits != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, logits); - - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - REQUIRE(imported.graph.nodes().size() == 2); - REQUIRE(imported.graph.nodes()[0].op == GGML_OP_GET_ROWS); - REQUIRE(imported.graph.nodes()[1].op == GGML_OP_MUL_MAT); - - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 2); - REQUIRE(scheduler.plan().transients.size() == 1); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + REQUIRE(scheduler.plan().dispatches.size() == 2); + REQUIRE(scheduler.plan().transients.size() == 1); const size_t q8_byte_count = qwen_q8_1_x4_size(1, 2048); @@ -4236,18 +3224,10 @@ static void schedule_get_rows_q8_1_alternate_command(ggml_type embedding_weight_ REQUIRE(alternate->alternate_value == q8_transient.value); const ggml::hrx::Dispatch & vocab_dispatch = scheduler.plan().dispatches[1]; - const std::string vocab_kernel_name = kernel_name_for_id(vocab_dispatch.kernel.kernel_id); + REQUIRE(kernel_name_for_id(vocab_dispatch.kernel.kernel_id) == "qwen3_moe:ggml_linear_q6k_q8_1_x4"); REQUIRE(vocab_dispatch.bindings.size() == 3); - if (vocab_kernel_name == "qwen3_moe:ggml_linear_q6k_q8_1_x4") { - REQUIRE(vocab_dispatch.bindings[0].value == alternate->alternate_value); - REQUIRE(vocab_dispatch.bindings[0].length == alternate->byte_count); - } else { - REQUIRE(vocab_kernel_name == "loom_libs:ggml_mul_mat_vector_q6_f32_f32"); - REQUIRE(vocab_dispatch.bindings[0].value == imported.graph.nodes()[0].output); - require_compile_parameter(vocab_dispatch, "ggml.matmul.vector.input_size", "2048"); - require_compile_parameter(vocab_dispatch, "ggml.matmul.vector.output_size", "151936"); - require_compile_parameter(vocab_dispatch, "ggml.matmul.vector.q6_storage", "0"); - } + REQUIRE(vocab_dispatch.bindings[0].value == alternate->alternate_value); + REQUIRE(vocab_dispatch.bindings[0].length == alternate->byte_count); const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); @@ -4258,8 +3238,7 @@ static void schedule_get_rows_q8_1_alternate_command(ggml_type embedding_weight_ REQUIRE(commands.commands[0].bindings[3].name == "next_output"); REQUIRE(commands.commands[0].bindings[3].origin == ggml::hrx::CommandBindingOrigin::Transient); REQUIRE(commands.commands[1].bindings.size() == 3); - REQUIRE(commands.commands[1].bindings[0].name == - (vocab_kernel_name == "qwen3_moe:ggml_linear_q6k_q8_1_x4" ? "q8_input" : "input")); + REQUIRE(commands.commands[1].bindings[0].name == "q8_input"); ggml_free(ctx); } @@ -4540,193 +3519,6 @@ static void run_qwen_token_embedding_dispatch_checks() { ggml_free(ctx); } -static void run_get_rows_rmsnorm_binary_dispatch_checks() { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * embedding_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 128256); - ggml_tensor * token_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 512); - ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2048); - REQUIRE(embedding_weight != nullptr); - REQUIRE(token_ids != nullptr); - REQUIRE(norm_weight != nullptr); - ggml_tensor * embedding = ggml_get_rows(ctx, embedding_weight, token_ids); - ggml_tensor * rms = ggml_rms_norm(ctx, embedding, 0.00001f); - ggml_tensor * output = ggml_mul(ctx, rms, norm_weight); - REQUIRE(embedding != nullptr); - REQUIRE(rms != nullptr); - REQUIRE(output != nullptr); - ggml_set_output(embedding); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - REQUIRE(imported.graph.nodes().size() == 3); - - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); - REQUIRE(plan.dispatches.size() == 1); - REQUIRE(plan.transients.size() == 2); - const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_get_rows_rmsnorm_binary_q8_1_x4_f16"); - REQUIRE(dispatch.bindings.size() == 7); - require_compile_parameter(dispatch, "ggml.get_rows_f32.weight_format", "6"); - require_compile_parameter(dispatch, "ggml.get_rows_rmsnorm.rms_epsilon", "9.99999975e-06"); - - const ggml::hrx::Value * output_value = imported.graph.values().find_tensor(output); - REQUIRE(output_value != nullptr); - const size_t q8_bytes = qwen_q8_1_x4_size(512, 2048); - const size_t f16_bytes = output_value->byte_count / 2; - const auto * q8 = plan.metadata.find_alternate_value(output_value->id, GGML_TYPE_Q8_1, q8_bytes); - const auto * f16 = plan.metadata.find_alternate_value(output_value->id, GGML_TYPE_F16, f16_bytes); - REQUIRE(q8 == nullptr); - REQUIRE(f16 == nullptr); - REQUIRE(dispatch.bindings[5].value == plan.transients[0].value); - REQUIRE(dispatch.bindings[6].value == plan.transients[1].value); - REQUIRE(plan.transients[0].size == q8_bytes); - REQUIRE(plan.transients[1].size == f16_bytes); - - const ggml::hrx::CommandProgram commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(commands.commands.size() == 1); - REQUIRE(command_program_verifies(commands)); - - ggml_free(ctx); -} - -static void run_get_rows_scale_dispatch_checks() { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2560, 64); - ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 5); - REQUIRE(weight != nullptr); - REQUIRE(ids != nullptr); - ggml_tensor * rows = ggml_get_rows(ctx, weight, ids); - ggml_tensor * output = ggml_scale(ctx, rows, 0.177800179f); - REQUIRE(rows != nullptr); - REQUIRE(output != nullptr); - ggml_set_output(rows); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - require_scheduled_command_program( - graph, [&](const ggml::hrx::Graph & imported_graph, const ggml::hrx::CommandPlan & plan, - const ggml::hrx::CommandProgram & commands) { - REQUIRE(imported_graph.nodes().size() == 2); - REQUIRE(plan.dispatches.size() == 1); - const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_get_rows_scale_f32"); - REQUIRE(dispatch.bindings.size() == 4); - require_compile_parameter(dispatch, "ggml.get_rows_f32.weight_format", "32"); - require_compile_parameter(dispatch, "ggml.get_rows_scale_f32.scale", expected_config_value(0.177800179f)); - require_compile_parameter(dispatch, "ggml.get_rows_scale_f32.bias", expected_config_value(0.0f)); - REQUIRE(commands.commands.size() == 1); - REQUIRE(commands.commands.front().bindings[0].name == "token_ids"); - REQUIRE(commands.commands.front().bindings[1].name == "weight"); - REQUIRE(commands.commands.front().bindings[2].name == "raw_output"); - REQUIRE(commands.commands.front().bindings[3].name == "output"); - }); - - ggml_tensor * q5_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_1, 640, 262144); - ggml_tensor * q5_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 64); - REQUIRE(q5_weight != nullptr); - REQUIRE(q5_ids != nullptr); - ggml_tensor * q5_rows = ggml_get_rows(ctx, q5_weight, q5_ids); - ggml_tensor * q5_output = ggml_scale(ctx, q5_rows, 25.2982216f); - REQUIRE(q5_rows != nullptr); - REQUIRE(q5_output != nullptr); - ggml_set_output(q5_rows); - - ggml_cgraph * q5_graph = ggml_new_graph(ctx); - REQUIRE(q5_graph != nullptr); - ggml_build_forward_expand(q5_graph, q5_output); - require_scheduled_command_program( - q5_graph, [&](const ggml::hrx::Graph & imported_graph, const ggml::hrx::CommandPlan & plan, - const ggml::hrx::CommandProgram & commands) { - REQUIRE(imported_graph.nodes().size() == 2); - REQUIRE(plan.dispatches.size() == 1); - const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_get_rows_scale_f32"); - REQUIRE(dispatch.bindings.size() == 4); - require_compile_parameter(dispatch, "ggml.get_rows_f32.weight_format", "51"); - require_compile_parameter(dispatch, "ggml.get_rows_scale_f32.scale", expected_config_value(25.2982216f)); - require_compile_parameter(dispatch, "ggml.get_rows_scale_f32.bias", expected_config_value(0.0f)); - REQUIRE(commands.commands.size() == 1); - }); - - ggml_tensor * q5_decode_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); - REQUIRE(q5_decode_ids != nullptr); - ggml_tensor * q5_decode_rows = ggml_get_rows(ctx, q5_weight, q5_decode_ids); - ggml_tensor * q5_decode_output = ggml_scale(ctx, q5_decode_rows, 25.2982216f); - REQUIRE(q5_decode_rows != nullptr); - REQUIRE(q5_decode_output != nullptr); - ggml_cgraph * q5_decode_graph = ggml_new_graph(ctx); - REQUIRE(q5_decode_graph != nullptr); - ggml_build_forward_expand(q5_decode_graph, q5_decode_output); - const std::vector q5_decode_kernels = scheduled_kernel_names(q5_decode_graph); - REQUIRE(std::find(q5_decode_kernels.begin(), q5_decode_kernels.end(), "loom_libs:ggml_get_rows_scale_f32") == - q5_decode_kernels.end()); - REQUIRE(std::find(q5_decode_kernels.begin(), q5_decode_kernels.end(), "loom_libs:ggml_get_rows_f32") != - q5_decode_kernels.end()); - REQUIRE(std::find(q5_decode_kernels.begin(), q5_decode_kernels.end(), "loom_libs:ggml_scale_f32") != - q5_decode_kernels.end()); - - ggml_tensor * q4_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_0, 640, 64); - ggml_tensor * q4_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 5); - REQUIRE(q4_weight != nullptr); - REQUIRE(q4_ids != nullptr); - ggml_tensor * q4_rows = ggml_get_rows(ctx, q4_weight, q4_ids); - ggml_tensor * q4_output = ggml_scale(ctx, q4_rows, 0.5f); - REQUIRE(q4_rows != nullptr); - REQUIRE(q4_output != nullptr); - ggml_cgraph * q4_graph = ggml_new_graph(ctx); - REQUIRE(q4_graph != nullptr); - ggml_build_forward_expand(q4_graph, q4_output); - const std::vector q4_kernels = scheduled_kernel_names(q4_graph); - REQUIRE(std::find(q4_kernels.begin(), q4_kernels.end(), "loom_libs:ggml_get_rows_scale_f32") == q4_kernels.end()); - REQUIRE(std::find(q4_kernels.begin(), q4_kernels.end(), "loom_libs:ggml_get_rows_f32") != q4_kernels.end()); - - ggml_tensor * side_use = ggml_add(ctx, rows, rows); - REQUIRE(side_use != nullptr); - ggml_cgraph * fallback_graph = ggml_new_graph(ctx); - REQUIRE(fallback_graph != nullptr); - ggml_build_forward_expand(fallback_graph, output); - ggml_build_forward_expand(fallback_graph, side_use); - const std::vector fallback_kernels = scheduled_kernel_names(fallback_graph); - REQUIRE(std::find(fallback_kernels.begin(), fallback_kernels.end(), "loom_libs:ggml_get_rows_scale_f32") == - fallback_kernels.end()); - REQUIRE(std::find(fallback_kernels.begin(), fallback_kernels.end(), "loom_libs:ggml_get_rows_f32") != - fallback_kernels.end()); - - ggml_tensor * inplace_rows = ggml_get_rows(ctx, weight, ids); - ggml_tensor * inplace_output = ggml_scale_inplace(ctx, inplace_rows, 0.5f); - REQUIRE(inplace_rows != nullptr); - REQUIRE(inplace_output != nullptr); - ggml_cgraph * inplace_graph = ggml_new_graph(ctx); - REQUIRE(inplace_graph != nullptr); - ggml_build_forward_expand(inplace_graph, inplace_output); - const std::vector inplace_kernels = scheduled_kernel_names(inplace_graph); - REQUIRE(std::find(inplace_kernels.begin(), inplace_kernels.end(), "loom_libs:ggml_get_rows_scale_f32") == - inplace_kernels.end()); - REQUIRE(std::find(inplace_kernels.begin(), inplace_kernels.end(), "loom_libs:ggml_get_rows_f32") != - inplace_kernels.end()); - - ggml_free(ctx); -} - static ggml::hrx::Graph build_manual_gather_add_graph(ggml_context * ctx, int64_t hidden_size, int64_t source_token_count, @@ -4776,193 +3568,6 @@ static ggml::hrx::Graph build_manual_gather_add_graph(ggml_context * ctx, return graph; } -struct GatherAddRmsNormGraph { - ggml_tensor * attention = nullptr; - ggml_tensor * residual = nullptr; - ggml_tensor * row_ids = nullptr; - ggml_tensor * weight = nullptr; - ggml_tensor * selected = nullptr; - ggml_tensor * raw_output = nullptr; - ggml_tensor * normalized_output = nullptr; - ggml_cgraph * graph = nullptr; -}; - -static GatherAddRmsNormGraph build_gather_add_rmsnorm_graph( - ggml_context * ctx, - int64_t hidden_size, - int64_t source_token_count, - int64_t output_token_count, - bool produced_row_ids = false, - bool produced_weight = false, - ggml::hrx::BinaryKind gather_op = ggml::hrx::BinaryKind::Add, - bool first_is_lhs = true) { - GatherAddRmsNormGraph result; - result.attention = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, source_token_count); - result.residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, source_token_count); - ggml_tensor * raw_row_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, output_token_count); - ggml_tensor * raw_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - result.row_ids = produced_row_ids ? ggml_dup(ctx, raw_row_ids) : raw_row_ids; - result.weight = produced_weight ? ggml_scale(ctx, raw_weight, 0.5f) : raw_weight; - ggml_tensor * selected_attention = ggml_get_rows(ctx, result.attention, result.row_ids); - ggml_tensor * selected_residual = ggml_get_rows(ctx, result.residual, result.row_ids); - result.selected = selected_attention; - switch (gather_op) { - case ggml::hrx::BinaryKind::Add: - result.raw_output = ggml_add(ctx, selected_attention, selected_residual); - break; - case ggml::hrx::BinaryKind::Sub: - result.raw_output = first_is_lhs ? ggml_sub(ctx, selected_attention, selected_residual) : - ggml_sub(ctx, selected_residual, selected_attention); - break; - case ggml::hrx::BinaryKind::Mul: - result.raw_output = ggml_mul(ctx, selected_attention, selected_residual); - break; - case ggml::hrx::BinaryKind::Div: - result.raw_output = first_is_lhs ? ggml_div(ctx, selected_attention, selected_residual) : - ggml_div(ctx, selected_residual, selected_attention); - break; - default: - REQUIRE(false); - } - ggml_tensor * rms = ggml_rms_norm(ctx, result.raw_output, 0.000001f); - result.normalized_output = ggml_mul(ctx, rms, result.weight); - REQUIRE(result.attention != nullptr); - REQUIRE(result.residual != nullptr); - REQUIRE(raw_row_ids != nullptr); - REQUIRE(raw_weight != nullptr); - REQUIRE(result.row_ids != nullptr); - REQUIRE(result.weight != nullptr); - REQUIRE(selected_attention != nullptr); - REQUIRE(selected_residual != nullptr); - REQUIRE(result.raw_output != nullptr); - REQUIRE(rms != nullptr); - REQUIRE(result.normalized_output != nullptr); - - ggml_set_output(result.raw_output); - result.graph = ggml_new_graph(ctx); - REQUIRE(result.graph != nullptr); - ggml_build_forward_expand(result.graph, result.normalized_output); - return result; -} - -static ggml::hrx::DispatchRegistry gather_add_test_registry() { - ggml::hrx::DispatchRegistryBuilder builder; - ggml::hrx::register_gather_add_dispatch(builder); - return builder.build(); -} - -static bool match_gather_add_at(const ggml::hrx::DispatchRegistry & registry, - const ggml::hrx::Graph & graph, - size_t root_index, - const std::vector & covered_nodes, - ggml::hrx::DispatchMatch & match) { - ggml::hrx::CommandPlan plan; - const ggml::hrx::DispatchMatchContext context = { - graph, &graph.nodes()[root_index], - root_index, covered_nodes, - plan, ggml::hrx::ValueId(static_cast(graph.values().size())), - }; - return registry.match(context, match); -} - -static void require_gather_add_rmsnorm_route(ggml_context * ctx, - ggml::hrx::BinaryKind gather_op = ggml::hrx::BinaryKind::Add, - bool first_is_lhs = true) { - GatherAddRmsNormGraph tensors = - build_gather_add_rmsnorm_graph(ctx, 128, 4, 5, false, false, gather_op, first_is_lhs); - require_scheduled_command_program(tensors.graph, [&](const ggml::hrx::Graph &, const ggml::hrx::CommandPlan & plan, - const ggml::hrx::CommandProgram & commands) { - REQUIRE(plan.dispatches.size() == 1); - const ggml::hrx::Dispatch & dispatch = plan.dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_gather_add_rmsnorm_binary_f32"); - REQUIRE(dispatch.bindings.size() == 6); - require_compile_parameter(dispatch, "ggml.gather_add_rmsnorm_binary_f32.hidden_size", "128"); - require_compile_parameter(dispatch, "ggml.gather_add_rmsnorm_binary_f32.rms_epsilon", "9.99999997e-07"); - require_compile_parameter(dispatch, "ggml.gather_add_rmsnorm_binary_f32.gather_op", - std::to_string(ggml::hrx::binary_kind_config_value(gather_op))); - REQUIRE(commands.commands.size() == 1); - REQUIRE(commands.commands.front().bindings.size() == 6); - REQUIRE(commands.commands.front().bindings[0].name == "attention"); - REQUIRE(commands.commands.front().bindings[1].name == "residual"); - REQUIRE(commands.commands.front().bindings[2].name == "output_ids"); - REQUIRE(commands.commands.front().bindings[3].name == "raw_output"); - REQUIRE(commands.commands.front().bindings[4].name == "weight"); - REQUIRE(commands.commands.front().bindings[5].name == "normalized_output"); - }); -} - -static void require_gather_add_rmsnorm_availability(ggml_context * ctx) { - GatherAddRmsNormGraph tensors = build_gather_add_rmsnorm_graph(ctx, 128, 4, 5, true, true); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*tensors.graph); - REQUIRE(imported.valid()); - const ggml::hrx::DispatchRegistry registry = gather_add_test_registry(); - const size_t root_index = producer_index_for_tensor(imported.graph, tensors.selected); - const size_t row_index = producer_index_for_tensor(imported.graph, tensors.row_ids); - const size_t weight_index = producer_index_for_tensor(imported.graph, tensors.weight); - std::vector covered_nodes(imported.graph.nodes().size(), false); - ggml::hrx::DispatchMatch match; - - REQUIRE(!match_gather_add_at(registry, imported.graph, root_index, covered_nodes, match)); - - covered_nodes[row_index] = true; - match = {}; - REQUIRE(match_gather_add_at(registry, imported.graph, root_index, covered_nodes, match)); - REQUIRE(match.dispatches.size() == 1); - REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == "hrx:ggml_gather_add_f32"); - - covered_nodes[weight_index] = true; - match = {}; - REQUIRE(match_gather_add_at(registry, imported.graph, root_index, covered_nodes, match)); - REQUIRE(match.dispatches.size() == 1); - REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == - "loom_libs:ggml_gather_add_rmsnorm_binary_f32"); -} - -static void require_gather_add_rmsnorm_row_ids_noalias(ggml_context * ctx) { - GatherAddRmsNormGraph tensors = build_gather_add_rmsnorm_graph(ctx, 128, 4, 5); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*tensors.graph); - REQUIRE(imported.valid()); - const ggml::hrx::Value * row_ids = imported.graph.values().find_tensor(tensors.row_ids); - const ggml::hrx::Value * attention = imported.graph.values().find_tensor(tensors.attention); - REQUIRE(row_ids != nullptr); - REQUIRE(attention != nullptr); - - ggml::hrx::Graph aliased_graph; - for (const ggml::hrx::ValueStorage & storage : imported.graph.values().storages()) { - REQUIRE(aliased_graph.values().add_snapshot_storage(storage).success()); - } - for (ggml::hrx::Value value : imported.graph.values().values()) { - if (value.id == row_ids->id) { - value.storage = attention->storage; - value.storage_root = attention->storage_root; - value.alias_source = attention->id; - value.storage_offset = attention->storage_offset; - value.storage_byte_count = attention->storage_byte_count; - } - REQUIRE(aliased_graph.values().add_snapshot_value(std::move(value)).success()); - } - for (const ggml::hrx::GraphNode & node : imported.graph.nodes()) { - ggml::hrx::GraphNode & copied = aliased_graph.add_node(node.op, node.output, node.inputs); - copied.params = node.params; - } - REQUIRE(aliased_graph.build_index().success()); - - const ggml::hrx::DispatchRegistry registry = gather_add_test_registry(); - size_t root_index = aliased_graph.nodes().size(); - for (size_t i = 0; i < aliased_graph.nodes().size(); ++i) { - if (aliased_graph.nodes()[i].op == GGML_OP_GET_ROWS) { - root_index = i; - break; - } - } - REQUIRE(root_index < aliased_graph.nodes().size()); - std::vector covered_nodes(aliased_graph.nodes().size(), false); - ggml::hrx::DispatchMatch match; - REQUIRE(match_gather_add_at(registry, aliased_graph, root_index, covered_nodes, match)); - REQUIRE(match.dispatches.size() == 1); - REQUIRE(kernel_name_for_id(match.dispatches.front().kernel.kernel_id) == "hrx:ggml_gather_add_f32"); -} - static bool manual_gather_add_graph_is_supported(ggml_context * ctx, int64_t hidden_size, int64_t source_token_count, @@ -5054,18 +3659,24 @@ static void require_gather_add_rejects_unavailable_sources(ggml_context * ctx) { ggml::hrx::DispatchMatch match; std::vector covered_nodes(graph.nodes().size(), false); ggml::hrx::DispatchMatchContext context = { - graph, &graph.nodes()[get_rows_index], - get_rows_index, covered_nodes, - plan, ggml::hrx::ValueId(static_cast(graph.values().size())), + graph, + &graph.nodes()[get_rows_index], + get_rows_index, + covered_nodes, + plan, + ggml::hrx::ValueId(static_cast(graph.values().size())), }; REQUIRE(!gather_add_registry.match(context, match)); covered_nodes[producer_index] = true; const ggml::hrx::DispatchMatchContext available_context = { - graph, &graph.nodes()[get_rows_index], - get_rows_index, covered_nodes, - plan, ggml::hrx::ValueId(static_cast(graph.values().size())), + graph, + &graph.nodes()[get_rows_index], + get_rows_index, + covered_nodes, + plan, + ggml::hrx::ValueId(static_cast(graph.values().size())), }; match = {}; REQUIRE(gather_add_registry.match(available_context, match)); @@ -5078,12 +3689,7 @@ static void schedule_gather_add_command(ggml::hrx::Graph & graph, int64_t expected_source_token_count, int64_t expected_output_token_count) { ggml::hrx::DispatchScheduler scheduler; - if (!scheduler.schedule_graph(graph, test_dispatch_target())) { - for (const std::string & error : scheduler.plan().status.errors()) { - std::fprintf(stderr, "%s\n", error.c_str()); - } - REQUIRE(false); - } + REQUIRE(scheduler.schedule_graph(graph, test_dispatch_target())); REQUIRE(scheduler.plan().valid()); REQUIRE(scheduler.plan().dispatches.size() == 1); @@ -5092,7 +3698,6 @@ static void schedule_gather_add_command(ggml::hrx::Graph & graph, REQUIRE(dispatch.kernel.integer_parameters.at("hidden_size") == expected_hidden_size); REQUIRE(dispatch.kernel.integer_parameters.at("source_token_count") == expected_source_token_count); REQUIRE(dispatch.kernel.integer_parameters.at("output_token_count") == expected_output_token_count); - require_compile_parameter(dispatch, "ggml.gather_add_f32.binary_op", "0"); REQUIRE(dispatch.bindings.size() == 4); const ggml::hrx::CommandProgram commands = @@ -5128,12 +3733,6 @@ static void run_gather_add_dispatch_checks() { REQUIRE(!manual_gather_add_graph_is_supported(ctx, 96, 13, 1, true)); REQUIRE(!partial_gather_add_graph_is_supported(ctx)); require_gather_add_rejects_unavailable_sources(ctx); - require_gather_add_rmsnorm_route(ctx); - require_gather_add_rmsnorm_route(ctx, ggml::hrx::BinaryKind::Sub, false); - require_gather_add_rmsnorm_route(ctx, ggml::hrx::BinaryKind::Mul); - require_gather_add_rmsnorm_route(ctx, ggml::hrx::BinaryKind::Div); - require_gather_add_rmsnorm_availability(ctx); - require_gather_add_rmsnorm_row_ids_noalias(ctx); ggml_free(ctx); } @@ -5189,15 +3788,10 @@ static void schedule_flash_attention_command(ggml_context * ctx, REQUIRE(scheduler.plan().transients.size() == 1); const auto & copy = scheduler.plan().dispatches.front(); REQUIRE(kernel_name_for_id(copy.kernel.kernel_id) == "loom_libs:ggml_copy_transpose_f16"); - REQUIRE(copy.kernel.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); - REQUIRE(copy.kernel.integer_parameters.at("row_count") == key_value_token_count); - REQUIRE(copy.kernel.integer_parameters.at("column_count") == key_value_head_count * value_head_size); REQUIRE(copy.bindings.size() == 2); REQUIRE(copy.bindings[1].value == dispatch.bindings[2].value); REQUIRE(copy.bindings[0].value != copy.bindings[1].value); - REQUIRE(copy.bindings[1].length == - static_cast(key_value_token_count * key_value_head_count * value_head_size) * - sizeof(ggml_fp16_t)); + REQUIRE(copy.bindings[1].length == static_cast(key_value_token_count * key_value_head_count * value_head_size) * sizeof(ggml_fp16_t)); require_compile_parameter(dispatch, "ggml.flash_attention.value_layout", "1"); } else { REQUIRE(dispatch.kernel.compile_parameters.count("ggml.flash_attention.value_layout") == 0); @@ -5236,15 +3830,16 @@ static void schedule_common_flash_attention_command(ggml_context * ctx, query_head_count, key_value_head_count, qk_head_size, value_head_size, scale); } -static void schedule_common_flash_attention_decode_fallback_command(ggml_context * ctx, - ggml_tensor * output, - int64_t query_token_count, - int64_t key_value_token_count, - int64_t query_head_count, - int64_t key_value_head_count, - int64_t qk_head_size = kQwenFlashHeadSize, - int64_t value_head_size = kQwenFlashHeadSize, - float scale = 1.0f / std::sqrt(128.0f)) { +static void schedule_common_flash_attention_decode_split_command(ggml_context * ctx, + ggml_tensor * output, + int64_t query_token_count, + int64_t key_value_token_count, + int64_t query_head_count, + int64_t key_value_head_count, + int64_t qk_head_size = kQwenFlashHeadSize, + int64_t value_head_size = kQwenFlashHeadSize, + float scale = 1.0f / std::sqrt(128.0f)) { + const int64_t key_value_capacity = ((key_value_token_count + 63) / 64) * 64; ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); ggml_build_forward_expand(graph, output); @@ -5258,22 +3853,21 @@ static void schedule_common_flash_attention_decode_fallback_command(ggml_context REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); REQUIRE(scheduler.plan().valid()); REQUIRE(scheduler.plan().dispatches.size() == static_cast(query_token_count)); - REQUIRE(scheduler.plan().transients.size() == 4); - REQUIRE(scheduler.plan().completion_counter_requests.size() == 1); - REQUIRE(scheduler.plan().completion_counter_requests.front().count == key_value_head_count); - REQUIRE(scheduler.plan().metadata.alternate_values().empty()); - const ggml::hrx::Dispatch & dispatch = scheduler.plan().dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == - "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); - REQUIRE(dispatch.kernel.integer_parameters.at("key_value_token_count") == key_value_token_count); - REQUIRE(dispatch.bindings.size() == 10); - require_compile_parameter(dispatch, "ggml.flash_attention.query_head_count", std::to_string(query_head_count)); - require_compile_parameter(dispatch, "ggml.flash_attention.key_value_head_count", - std::to_string(key_value_head_count)); - require_compile_parameter(dispatch, "ggml.flash_attention.qk_head_size", std::to_string(qk_head_size)); - require_compile_parameter(dispatch, "ggml.flash_attention.value_head_size", std::to_string(value_head_size)); - require_compile_parameter(dispatch, "ggml.flash_attention.attention_scale", expected_config_value(scale)); + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); + REQUIRE(kernel_name == "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); + REQUIRE(dispatch.kernel.integer_parameters.at("key_value_token_count") == key_value_token_count); + REQUIRE(dispatch.bindings.size() == 10); + require_compile_parameter(dispatch, "ggml.flash_attention.query_head_count", std::to_string(query_head_count)); + require_compile_parameter(dispatch, "ggml.flash_attention.key_value_head_count", + std::to_string(key_value_head_count)); + require_compile_parameter(dispatch, "ggml.flash_attention.qk_head_size", std::to_string(qk_head_size)); + require_compile_parameter(dispatch, "ggml.flash_attention.value_head_size", std::to_string(value_head_size)); + require_compile_parameter(dispatch, "ggml.flash_attention.attention_scale", expected_config_value(scale)); + require_compile_parameter(dispatch, "ggml.flash_attention.decode.key_value_token_capacity", + std::to_string(key_value_capacity)); + } const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); @@ -5285,105 +3879,14 @@ static void schedule_common_flash_attention_decode_fallback_command(ggml_context REQUIRE(commands.commands.front().bindings[1].name == "key"); REQUIRE(commands.commands.front().bindings[2].name == "value"); REQUIRE(commands.commands.front().bindings[3].name == "mask"); + REQUIRE(commands.commands.front().bindings[4].name == "partial_max"); + REQUIRE(commands.commands.front().bindings[5].name == "partial_sum"); + REQUIRE(commands.commands.front().bindings[6].name == "partial_output"); + REQUIRE(commands.commands.front().bindings[7].name == "completion_counter"); REQUIRE(commands.commands.front().bindings[8].name == "output"); REQUIRE(commands.commands.front().bindings[9].name == "next_q8_output"); } -static void schedule_common_flash_attention_decode_q8_command(ggml_context * ctx, ggml_type weight_type) { - constexpr int64_t query_token_count = 1; - constexpr int64_t key_value_token_count = 512; - constexpr int64_t query_head_count = 4; - constexpr int64_t key_value_head_count = 2; - constexpr int64_t hidden_size = query_head_count * kQwenFlashHeadSize; - constexpr int64_t output_size = 256; - const size_t q8_bytes = ggml_row_size(GGML_TYPE_Q8_1, hidden_size); - - ggml_tensor * output = build_qwen_flash_attention_graph( - ctx, query_token_count, key_value_token_count, query_head_count, key_value_head_count); - ggml_tensor * reshaped = ggml_reshape_2d(ctx, output, hidden_size, query_token_count); - ggml_tensor * weight = ggml_new_tensor_2d(ctx, weight_type, hidden_size, output_size); - ggml_tensor * projected = ggml_mul_mat(ctx, weight, reshaped); - REQUIRE(reshaped != nullptr); - REQUIRE(weight != nullptr); - REQUIRE(projected != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, projected); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); - REQUIRE(plan.dispatches.size() == 2); - REQUIRE(plan.transients.size() == 4); - REQUIRE(plan.completion_counter_requests.size() == 1); - - const ggml::hrx::Value * native_value = imported.graph.values().find_tensor(output); - const ggml::hrx::Value * reshape_value = imported.graph.values().find_tensor(reshaped); - REQUIRE(native_value != nullptr); - REQUIRE(reshape_value != nullptr); - const ggml::hrx::CommandPlanAlternateValue * q8 = - ggml::hrx::find_alternate_value(plan, reshape_value->id, GGML_TYPE_Q8_1, q8_bytes); - REQUIRE(q8 != nullptr); - REQUIRE(ggml::hrx::find_alternate_value(plan, native_value->id, GGML_TYPE_Q8_1, q8_bytes) == nullptr); - - const ggml::hrx::Dispatch & producer = plan.dispatches.front(); - const ggml::hrx::Dispatch & consumer = plan.dispatches.back(); - REQUIRE(kernel_name_for_id(producer.kernel.kernel_id) == - "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); - REQUIRE(producer.bindings.size() == 10); - REQUIRE(producer.bindings[9].value == q8->alternate_value); - REQUIRE(consumer.bindings[0].value == q8->alternate_value); - - const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( - imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(commands.commands.size() == 2); - REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands.front().bindings.size() == 10); - REQUIRE(commands.commands.front().bindings[8].name == "output"); - REQUIRE(commands.commands.front().bindings[9].name == "next_q8_output"); -} - -static void schedule_common_flash_attention_decode_unqualified_command(ggml_context * ctx) { - constexpr int64_t hidden_size = 4 * kQwenFlashHeadSize; - ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 1, 512, 4, 2); - ggml_tensor * reshaped = ggml_reshape_2d(ctx, output, hidden_size, 1); - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 256); - ggml_tensor * projected = ggml_mul_mat(ctx, weight, reshaped); - REQUIRE(reshaped != nullptr); - REQUIRE(weight != nullptr); - REQUIRE(projected != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, projected); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); - REQUIRE(plan.metadata.alternate_values().empty()); - REQUIRE(plan.dispatches.size() == 2); - REQUIRE(plan.transients.size() == 4); - REQUIRE(plan.completion_counter_requests.size() == 1); - REQUIRE(kernel_name_for_id(plan.dispatches.front().kernel.kernel_id) == - "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); - REQUIRE(plan.dispatches.front().bindings.size() == 10); - - const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( - imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands.size() == 2); - REQUIRE(commands.commands.front().bindings.size() == 10); -} - static void run_qwen_flash_attention_dispatch_checks() { static constexpr const char * kDisableQwenDispatchEnv = "GGML_HRX_DISABLE_QWEN_DISPATCH"; const char * original_env = std::getenv(kDisableQwenDispatchEnv); @@ -5397,16 +3900,15 @@ static void run_qwen_flash_attention_dispatch_checks() { ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - for (const auto & shape : { - std::array{ 256, 512, 64, 1 }, + for (const auto & shape : { std::array{ 256, 512, 64, 1 }, std::array{ 512, 512, 256, 1 }, std::array{ 512, 2048, 320, 1 }, std::array{ 255, 512, 256, 0 }, - std::array{ 512, 513, 256, 0 } - }) { + std::array{ 512, 513, 256, 0 } }) { const float scale = 1.0f / std::sqrt(static_cast(shape[2])); - ggml_tensor * output = build_qwen_flash_attention_graph(ctx, shape[0], shape[1], 4, 2, GGML_TYPE_F32, - GGML_TYPE_F16, true, false, shape[2], scale); + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, shape[0], shape[1], 4, 2, + GGML_TYPE_F32, GGML_TYPE_F16, true, false, + shape[2], scale); schedule_flash_attention_command(ctx, output, "loom_libs:ggml_flash_attention_f32_f16_wmma", "ggml.flash_attention.", shape[0], shape[1], 4, 2, shape[2], shape[2], scale, shape[3] != 0); @@ -5414,7 +3916,7 @@ static void run_qwen_flash_attention_dispatch_checks() { { ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2); - schedule_common_flash_attention_decode_fallback_command(ctx, output, 4, 8, 4, 2); + schedule_common_flash_attention_decode_split_command(ctx, output, 4, 8, 4, 2); } { ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 16, 16, 4, 2); @@ -5449,7 +3951,7 @@ static void run_qwen_flash_attention_dispatch_checks() { } { ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 1, 8, 4, 2); - schedule_common_flash_attention_decode_fallback_command(ctx, output, 1, 8, 4, 2); + schedule_common_flash_attention_decode_split_command(ctx, output, 1, 8, 4, 2); } { REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); @@ -5458,59 +3960,13 @@ static void run_qwen_flash_attention_dispatch_checks() { ggml_tensor * output = build_qwen_flash_attention_graph( ctx, 1, 512, 32, 4, GGML_TYPE_F32, GGML_TYPE_F16, true, false, qk_head_size, 1.0f / std::sqrt(static_cast(qk_head_size)), value_head_size); - schedule_common_flash_attention_decode_fallback_command(ctx, output, 1, 512, 32, 4, qk_head_size, - value_head_size, - 1.0f / std::sqrt(static_cast(qk_head_size))); + schedule_common_flash_attention_decode_split_command(ctx, output, 1, 512, 32, 4, qk_head_size, value_head_size, + 1.0f / std::sqrt(static_cast(qk_head_size))); REQUIRE(unsetenv(kDisableQwenDispatchEnv) == 0); } { - REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); - schedule_common_flash_attention_decode_q8_command(ctx, GGML_TYPE_Q4_K); - schedule_common_flash_attention_decode_q8_command(ctx, GGML_TYPE_Q6_K); - schedule_common_flash_attention_decode_unqualified_command(ctx); - REQUIRE(unsetenv(kDisableQwenDispatchEnv) == 0); - } - { - REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); - constexpr int64_t tokens = 512; - constexpr int64_t hidden = 512; - ggml_tensor * output = build_qwen_flash_attention_graph(ctx, tokens, tokens, 4, 2); - ggml_tensor * reshaped = ggml_reshape_2d(ctx, output, hidden, tokens); - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, hidden, hidden); - ggml_tensor * projected = ggml_mul_mat(ctx, weight, reshaped); - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, projected); - - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); - const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == - "loom_libs:ggml_flash_attention_f32_f16_wmma_publish_f16"; - }); - REQUIRE(producer != plan.dispatches.end()); - REQUIRE(producer->bindings.size() == 7); - const ggml::hrx::Value * output_value = imported.graph.values().find_tensor(output); - REQUIRE(output_value != nullptr); - const auto * alternate = - plan.metadata.find_alternate_value(output_value->id, GGML_TYPE_F16, output_value->byte_count / 2); - REQUIRE(alternate != nullptr); - REQUIRE(producer->bindings.back().value == alternate->alternate_value); - const auto consumer = std::next(producer); - REQUIRE(consumer != plan.dispatches.end()); - REQUIRE(consumer->bindings.front().value == alternate->alternate_value); - const ggml::hrx::CommandProgram commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - REQUIRE(unsetenv(kDisableQwenDispatchEnv) == 0); - } - { - ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F32); - REQUIRE(!graph_is_supported(ctx, output)); + ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F32); + REQUIRE(!graph_is_supported(ctx, output)); } { ggml_tensor * output = build_qwen_flash_attention_graph(ctx, 4, 8, 4, 2, GGML_TYPE_F32, GGML_TYPE_F16, false); @@ -5934,6 +4390,7 @@ static void run_qwen_attention_postprocess_dispatch_checks() { const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); + REQUIRE(command_program_verifies(commands)); REQUIRE(commands.initialization_commands.size() == 2); REQUIRE(commands.initialization_commands[0].bindings[0].name == "positions"); REQUIRE(commands.initialization_commands[0].bindings[1].name == "control"); @@ -5941,10 +4398,6 @@ static void run_qwen_attention_postprocess_dispatch_checks() { REQUIRE(commands.initialization_commands[1].bindings[4].name == "attention_mask"); REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, context_capture.bindings[1].value) != nullptr); - - ggml::hrx::CommandProgram initialization_commands = copy_command_program_shape(commands); - initialization_commands.commands.clear(); - REQUIRE(command_program_verifies(initialization_commands)); } ggml_free(ctx); @@ -5964,22 +4417,20 @@ static void run_packed_f16_producer_consumer_checks() { int64_t outputs; bool side_use; bool packed; - bool alternate; bool packed_copy = false; } cases[] = { - { GGML_TYPE_Q6_K, 512, 4096, 2048, false, true, false }, - { GGML_TYPE_Q6_K, 1024, 4096, 2048, false, true, false }, - { GGML_TYPE_Q6_K, 256, 4096, 2048, false, false, true }, - { GGML_TYPE_Q6_K, 512, 4096, 1024, false, false, true }, - { GGML_TYPE_Q4_K, 512, 4096, 4096, false, true, false }, - { GGML_TYPE_Q4_K, 512, 8192, 2048, false, false, true }, - { GGML_TYPE_Q4_K, 512,12288, 4096, false, true, false }, - { GGML_TYPE_Q4_K, 512,12288, 4096, true, false, true }, - { GGML_TYPE_Q4_K, 512,32768, 4096, false, false, false }, - { GGML_TYPE_Q4_K, 1024, 4096, 4096, false, false, true }, - { GGML_TYPE_Q6_K, 512, 4096, 2048, true, false, true, true }, + { GGML_TYPE_Q6_K, 512, 4096, 2048, false, true }, + { GGML_TYPE_Q6_K, 1024, 4096, 2048, false, true }, + { GGML_TYPE_Q6_K, 256, 4096, 2048, false, false }, + { GGML_TYPE_Q6_K, 512, 4096, 1024, false, false }, + { GGML_TYPE_Q4_K, 512, 4096, 4096, false, true }, + { GGML_TYPE_Q4_K, 512, 8192, 2048, false, false }, + { GGML_TYPE_Q4_K, 512,12288, 4096, false, true }, + { GGML_TYPE_Q4_K, 512,12288, 4096, true, false }, + { GGML_TYPE_Q4_K, 512,32768, 4096, false, false }, + { GGML_TYPE_Q4_K, 1024, 4096, 4096, false, false }, + { GGML_TYPE_Q6_K, 512, 4096, 2048, true, false, true }, }; - for (const auto & test : cases) { ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, test.tokens); ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 256, test.hidden); @@ -6002,14 +4453,12 @@ static void run_packed_f16_producer_consumer_checks() { REQUIRE(plan.valid()); const auto * value = imported.graph.values().find_tensor(glu); REQUIRE(value != nullptr); - const auto * generated = - plan.metadata.find_generated_resource(value->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + const auto * generated = plan.metadata.find_generated_resource(value->id, ggml::hrx::GeneratedResourceRole::F16K16Major); const auto * alternate = plan.metadata.find_alternate_value(value->id, GGML_TYPE_F16, value->byte_count / 2); REQUIRE((generated != nullptr) == (test.packed || test.packed_copy)); - REQUIRE((alternate != nullptr) == test.alternate); + REQUIRE((alternate != nullptr) != test.packed); const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == - "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32"; + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32"; }); REQUIRE(producer != plan.dispatches.end()); require_compile_parameter(*producer, "ggml.mul_mat_swiglu.f16_output_layout", test.packed ? "1" : "0"); @@ -6031,417 +4480,10 @@ static void run_packed_f16_producer_consumer_checks() { REQUIRE(std::next(copy) != plan.dispatches.end()); REQUIRE(std::next(copy)->bindings.front().value == generated->generated_value); } - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } - - ggml_free(ctx); -} - -static void run_generic_swiglu_k16_publish_checks() { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - constexpr int64_t input_size = 2048; - constexpr int64_t hidden_size = 8192; - constexpr int64_t output_size = 2048; - constexpr int64_t token_count = 512; - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); - ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, input_size, hidden_size); - ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, input_size, hidden_size); - ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); - ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); - ggml_tensor * glu = ggml_swiglu_split(ctx, gate, up); - ggml_tensor * projection_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, hidden_size, output_size); - ggml_tensor * output = ggml_mul_mat(ctx, projection_weight, glu); - REQUIRE(output != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); - - const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == - "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16"; - }); - REQUIRE(producer != plan.dispatches.end()); - REQUIRE(producer->bindings.size() == 5); - const ggml::hrx::Value * glu_value = imported.graph.values().find_tensor(glu); - REQUIRE(glu_value != nullptr); - const auto * packed = - plan.metadata.find_generated_resource(glu_value->id, ggml::hrx::GeneratedResourceRole::F16K16Major); - REQUIRE(packed != nullptr); - REQUIRE(packed->byte_count == glu_value->byte_count / 2); - REQUIRE(producer->bindings.back().value == packed->generated_value); - const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return &dispatch != &*producer && !dispatch.bindings.empty() && - dispatch.bindings.front().value == packed->generated_value; - }); - REQUIRE(consumer != plan.dispatches.end()); - REQUIRE(std::none_of(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_copy_f16_k16_major"; - })); - - const ggml::hrx::CommandProgram commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - ggml_free(ctx); -} - -static void run_generic_swiglu_k16_publication_edge_checks() { - { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - constexpr int64_t input_size = 2048; - constexpr int64_t hidden_size = 8192; - constexpr int64_t token_count = 512; - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); - ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, input_size, hidden_size); - ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, input_size, hidden_size); - ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); - ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); - ggml_tensor * glu = ggml_swiglu_split(ctx, gate, up); - REQUIRE(glu != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, glu); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); - - const ggml::hrx::Value * glu_value = imported.graph.values().find_tensor(glu); - REQUIRE(glu_value != nullptr); - REQUIRE(plan.metadata.find_generated_resource(glu_value->id, ggml::hrx::GeneratedResourceRole::F16K16Major) == - nullptr); - const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == - "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32"; - }); - REQUIRE(producer != plan.dispatches.end()); - REQUIRE(producer->bindings.size() == 4); - REQUIRE(producer->bindings.back().value == glu_value->id); - - const ggml::hrx::CommandProgram commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - ggml_free(ctx); - } - - { - ggml_init_params params = {}; - params.mem_size = 8 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - constexpr int64_t input_size = 2048; - constexpr int64_t hidden_size = 8192; - constexpr int64_t token_count = 512; - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); - ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, input_size, hidden_size); - ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, input_size, hidden_size); - ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); - ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); - ggml_tensor * glu = ggml_swiglu_split(ctx, gate, up); - ggml_tensor * reshaped = ggml_reshape_2d(ctx, glu, 4096, 1024); - ggml_tensor * projection_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 4096, 2048); - ggml_tensor * output = ggml_mul_mat(ctx, projection_weight, reshaped); - REQUIRE(output != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); - - const ggml::hrx::Value * glu_value = imported.graph.values().find_tensor(glu); - const ggml::hrx::Value * reshaped_value = imported.graph.values().find_tensor(reshaped); - REQUIRE(glu_value != nullptr); - REQUIRE(reshaped_value != nullptr); - REQUIRE(plan.metadata.find_generated_resource(glu_value->id, ggml::hrx::GeneratedResourceRole::F16K16Major) == - nullptr); - const auto * packed = - plan.metadata.find_generated_resource(reshaped_value->id, ggml::hrx::GeneratedResourceRole::F16K16Major); - REQUIRE(packed != nullptr); - - const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == - "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32"; - }); - REQUIRE(producer != plan.dispatches.end()); - REQUIRE(producer->bindings.size() == 4); - const auto copy = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_copy_f16_k16_major" && - dispatch.bindings.front().value == reshaped_value->id; - }); - REQUIRE(copy != plan.dispatches.end()); - REQUIRE(copy->bindings.back().value == packed->generated_value); - require_compile_parameter(*copy, "ggml.copy_f16_k16_major.input_is_f16", "0"); - const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return &dispatch != &*copy && !dispatch.bindings.empty() && - dispatch.bindings.front().value == packed->generated_value; - }); - REQUIRE(consumer != plan.dispatches.end()); - - const ggml::hrx::CommandProgram commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - ggml_free(ctx); - } -} - -static void run_tiled_matmul_alternate_publish_checks() { - struct Case { - ggml_type consumer_type; - const char * producer_kernel; - const char * consumer_kernel; - ggml_type alternate_type; - }; - - const Case cases[] = { - { GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_f16_alternate", - "loom_libs:ggml_mul_mat_q6_k_f16_wmma_prefill_wave32", GGML_TYPE_F16 }, - { GGML_TYPE_Q5_K, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", - "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", GGML_TYPE_Q8_1 }, - }; - - for (const Case & test : cases) { - ggml_init_params params = {}; - params.mem_size = 8 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 256); - ggml_tensor * producer_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 2048); - ggml_tensor * projection = ggml_mul_mat(ctx, producer_weight, input); - ggml_tensor * consumer_weight = ggml_new_tensor_2d(ctx, test.consumer_type, 2048, 1024); - ggml_tensor * output = ggml_mul_mat(ctx, consumer_weight, projection); - - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, output); - auto imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const auto & plan = scheduler.plan(); - REQUIRE(plan.valid()); - - const auto * projection_value = imported.graph.values().find_tensor(projection); - REQUIRE(projection_value != nullptr); - const size_t alternate_bytes = test.alternate_type == GGML_TYPE_F16 ? - projection_value->byte_count / 2 : - static_cast(256) * ggml_row_size(GGML_TYPE_Q8_1, 2048); - const auto * alternate = - plan.metadata.find_alternate_value(projection_value->id, test.alternate_type, alternate_bytes); - REQUIRE(alternate != nullptr); - - const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == test.producer_kernel; - }); - REQUIRE(producer != plan.dispatches.end()); - REQUIRE(producer->bindings[2].value == projection_value->id); - if (test.alternate_type == GGML_TYPE_F16) { - REQUIRE(producer->bindings.size() == 4); - REQUIRE(producer->bindings[3].value == alternate->alternate_value); - } else { - const auto quantize = - std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32" && - dispatch.bindings.front().value == projection_value->id; - }); - REQUIRE(quantize != plan.dispatches.end()); - REQUIRE(quantize->bindings.back().value == alternate->alternate_value); - } - - const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == test.consumer_kernel; - }); - REQUIRE(consumer != plan.dispatches.end()); - REQUIRE(consumer->bindings.front().value == alternate->alternate_value); - - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - ggml_free(ctx); - } - - { - ggml_init_params params = {}; - params.mem_size = 12 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 256); - ggml_tensor * producer_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 2048); - ggml_tensor * projection = ggml_mul_mat(ctx, producer_weight, input); - ggml_tensor * q5_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_K, 2048, 1024); - ggml_tensor * q6_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 1024); - ggml_tensor * q5_output = ggml_mul_mat(ctx, q5_weight, projection); - ggml_tensor * q6_output = ggml_mul_mat(ctx, q6_weight, projection); - - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, q5_output); - ggml_build_forward_expand(graph, q6_output); - auto imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const auto & plan = scheduler.plan(); - REQUIRE(plan.valid()); - - const auto * projection_value = imported.graph.values().find_tensor(projection); - REQUIRE(projection_value != nullptr); - const size_t q8_bytes = static_cast(256) * ggml_row_size(GGML_TYPE_Q8_1, 2048); - const auto * q8 = plan.metadata.find_alternate_value(projection_value->id, GGML_TYPE_Q8_1, q8_bytes); - REQUIRE(q8 != nullptr); - const auto publication_count = std::count_if( - plan.metadata.activation_publication_diagnostics().begin(), - plan.metadata.activation_publication_diagnostics().end(), [&](const auto & diagnostic) { - return diagnostic.source_value == projection_value->id; - }); - REQUIRE(publication_count == 1); - const auto publication = std::find_if( - plan.metadata.activation_publication_diagnostics().begin(), - plan.metadata.activation_publication_diagnostics().end(), [&](const auto & diagnostic) { - return diagnostic.source_value == projection_value->id; - }); - REQUIRE(publication != plan.metadata.activation_publication_diagnostics().end()); - REQUIRE(publication->requested_format == "q8-1-x4"); - - const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32" && - dispatch.bindings[2].value == projection_value->id; - }); - REQUIRE(producer != plan.dispatches.end()); - REQUIRE(producer->bindings.size() == 3); - - const auto q5_consumer = - std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == - "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256" && - dispatch.bindings.front().value == q8->alternate_value; - }); - REQUIRE(q5_consumer != plan.dispatches.end()); - const auto q6_consumer = - std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == - "loom_libs:ggml_mul_mat_q6_k_f16_wmma_prefill_wave32"; - }); - REQUIRE(q6_consumer != plan.dispatches.end()); - const auto f16_copy = - std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_copy_f32_f16" && - dispatch.bindings.front().value == projection_value->id; - }); - REQUIRE(f16_copy != plan.dispatches.end()); - REQUIRE(q6_consumer->bindings.front().value == f16_copy->bindings.back().value); - - const size_t q8_quantizer_count = std::count_if( - plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32" && - dispatch.bindings.front().value == projection_value->id; - }); - REQUIRE(q8_quantizer_count == 1); - - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - ggml_free(ctx); - } -} - -static void run_tiled_pair_matmul_postops_dispatch_checks() { - ggml_init_params params = {}; - params.mem_size = 16 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - auto build_glu = [&]() { - ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 256); - ggml_tensor * up_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 256); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 8); - REQUIRE(gate_weight != nullptr); - REQUIRE(up_weight != nullptr); - REQUIRE(input != nullptr); - ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); - ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); - REQUIRE(gate != nullptr); - REQUIRE(up != nullptr); - ggml_tensor * glu = ggml_glu_split(ctx, gate, up, GGML_GLU_OP_SWIGLU); - REQUIRE(glu != nullptr); - return glu; - }; - - { - ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); - REQUIRE(bias != nullptr); - ggml_tensor * output = ggml_add(ctx, build_glu(), bias); - REQUIRE(output != nullptr); - schedule_fused_matmul_swiglu_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", 4, - { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU, GGML_OP_ADD }, - { "input", "gate_weight", "up_weight", "bias", "residual_input", "residual_output" }); - } - - { - ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 8); - REQUIRE(residual != nullptr); - ggml_tensor * output = ggml_add(ctx, build_glu(), residual); - REQUIRE(output != nullptr); - schedule_fused_matmul_swiglu_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", 4, - { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU, GGML_OP_ADD }, - { "input", "gate_weight", "up_weight", "bias", "residual_input", "residual_output" }); - } - - { - ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 256); - ggml_tensor * residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 256, 8); - REQUIRE(bias != nullptr); - REQUIRE(residual != nullptr); - ggml_tensor * biased = ggml_add(ctx, build_glu(), bias); - REQUIRE(biased != nullptr); - ggml_tensor * output = ggml_add(ctx, biased, residual); - REQUIRE(output != nullptr); - schedule_fused_matmul_swiglu_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32", 5, - { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU, GGML_OP_ADD, GGML_OP_ADD }, - { "input", "gate_weight", "up_weight", "bias", "residual_input", "residual_output" }); - } - ggml_free(ctx); } @@ -6465,19 +4507,18 @@ static void run_swiglu_q8_output_checks() { { 5, 5120, 17408, GGML_TYPE_Q6_K, 2, false, true, true, false, false, true }, { 5, 4096, 16384, GGML_TYPE_Q4_K, 1, true, false, false, false, false, true }, { 5, 5120, 32768, GGML_TYPE_Q6_K, 1, false, false, false, false, false, true }, - { 1, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, true }, + { 1, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, false }, { 6, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, false }, - { 5, 2048, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, true }, - { 5, 6144, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, true }, - { 5, 5120, 8192, GGML_TYPE_Q4_K, 1, false, false, false, false, false, true }, + { 5, 2048, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, false }, + { 5, 6144, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, false, false }, + { 5, 5120, 8192, GGML_TYPE_Q4_K, 1, false, false, false, false, false, false }, { 5, 5120, 33024, GGML_TYPE_Q4_K, 1, false, false, false, false, false, false }, { 5, 5120, 17408, GGML_TYPE_F32, 1, false, true, false, false, false, false }, - { 5, 5120, 17408, GGML_TYPE_Q6_K, 1, true, false, false, false, false, true }, + { 5, 5120, 17408, GGML_TYPE_Q6_K, 1, true, false, false, false, false, false }, { 5, 5120, 17408, GGML_TYPE_Q4_K, 0, false, true, true, false, false, false }, { 5, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, true, false, false }, { 5, 5120, 17408, GGML_TYPE_Q4_K, 1, false, false, false, false, true, false }, }; - for (const auto & test : cases) { ggml_init_params params = {}; params.mem_size = 4 * 1024 * 1024; @@ -6486,8 +4527,7 @@ static void run_swiglu_q8_output_checks() { REQUIRE(ctx != nullptr); ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.inputs, test.tokens); ggml_tensor * gate_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, test.inputs, test.outputs); - ggml_tensor * up_weight = test.weight_alias ? - ggml_view_tensor(ctx, gate_weight) : + ggml_tensor * up_weight = test.weight_alias ? ggml_view_tensor(ctx, gate_weight) : ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, test.inputs, test.outputs); ggml_tensor * gate = ggml_mul_mat(ctx, gate_weight, input); ggml_tensor * up = ggml_mul_mat(ctx, up_weight, input); @@ -6511,8 +4551,7 @@ static void run_swiglu_q8_output_checks() { const bool scheduled = scheduler.schedule_graph(imported.graph, test_dispatch_target()); const auto & plan = scheduler.plan(); const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == - "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot_q8_output"; + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot_q8_output"; }); REQUIRE((producer != plan.dispatches.end()) == test.packed); if (test.outputs > 32768) { @@ -6544,15 +4583,13 @@ static void run_swiglu_q8_output_checks() { if (side != nullptr) { const auto * side_value = imported.graph.values().find_tensor(side); REQUIRE(side_value != nullptr); - const auto consumer = - std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { return !dispatch.bindings.empty() && dispatch.bindings.back().value == side_value->id; }); REQUIRE(consumer != plan.dispatches.end()); REQUIRE(consumer->bindings.front().value == output->id); } - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); ggml_free(ctx); @@ -6581,12 +4618,10 @@ static void run_k16_major_preparation_reuse_checks() { { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, false, true }, { GGML_TYPE_Q6_K, GGML_TYPE_Q6_K, false, true }, }; - for (const auto & test : cases) { ggml_tensor * first = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4096, 512); ggml_tensor * second = test.reshape ? ggml_reshape_2d(ctx, first, 2048, 1024) : - test.shared ? first : - ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4096, 512); + test.shared ? first : ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4096, 512); ggml_tensor * inputs[] = { first, second }; ggml_tensor * outputs[2]; ggml_cgraph * graph = ggml_new_graph(ctx); @@ -6602,8 +4637,7 @@ static void run_k16_major_preparation_reuse_checks() { REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); const auto & plan = scheduler.plan(); REQUIRE(plan.valid()); - const auto copy_count = - std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + const auto copy_count = std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_copy_f16_k16_major"; }); REQUIRE(copy_count == (test.shared ? 1 : 2)); @@ -6613,20 +4647,18 @@ static void run_k16_major_preparation_reuse_checks() { const auto * output = imported.graph.values().find_tensor(outputs[i]); REQUIRE(input != nullptr); REQUIRE(output != nullptr); - const auto * generated = - plan.metadata.find_generated_resource(input->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + const auto * generated = plan.metadata.find_generated_resource(input->id, ggml::hrx::GeneratedResourceRole::F16K16Major); REQUIRE(generated != nullptr); REQUIRE(generated->byte_count == input->byte_count / 2); packed[i] = generated->generated_value; - const auto consumer = - std::find_if(plan.dispatches.begin(), plan.dispatches.end(), - [&](const auto & dispatch) { return dispatch.bindings.back().value == output->id; }); + const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return dispatch.bindings.back().value == output->id; + }); REQUIRE(consumer != plan.dispatches.end()); REQUIRE(consumer->bindings.front().value == packed[i]); } REQUIRE((packed[0] == packed[1]) == test.shared); - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } @@ -6645,8 +4677,6 @@ static void run_rmsnorm_binary_k16_output_checks() { bool multiply; bool packed; } cases[] = { - { GGML_TYPE_Q4_K, 2048, 512, 1, false, false, false, true, true }, - { GGML_TYPE_Q4_K, 3072, 512, 1, false, false, false, true, true }, { GGML_TYPE_Q4_K, 4096, 512, 1, false, false, false, true, true }, { GGML_TYPE_Q4_K, 5120, 512, 2, false, false, false, true, true }, { GGML_TYPE_Q6_K, 8192, 512, 2, true, false, false, true, true }, @@ -6662,7 +4692,6 @@ static void run_rmsnorm_binary_k16_output_checks() { { GGML_TYPE_Q4_K, 5120, 512, 1, false, false, true, true, false }, { GGML_TYPE_Q4_K, 5120, 512, 1, false, false, false, false, false }, }; - for (const auto & test : cases) { ggml_init_params params = {}; params.mem_size = 4 * 1024 * 1024; @@ -6671,11 +4700,10 @@ static void run_rmsnorm_binary_k16_output_checks() { REQUIRE(ctx != nullptr); ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens); ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.hidden); - ggml_tensor * rms = - test.inplace ? ggml_rms_norm_inplace(ctx, input, 1.0e-6f) : ggml_rms_norm(ctx, input, 1.0e-6f); + ggml_tensor * rms = test.inplace ? ggml_rms_norm_inplace(ctx, input, 1.0e-6f) : + ggml_rms_norm(ctx, input, 1.0e-6f); ggml_tensor * norm = test.inplace ? ggml_mul_inplace(ctx, rms, weight) : - test.multiply ? ggml_mul(ctx, rms, weight) : - ggml_add(ctx, rms, weight); + test.multiply ? ggml_mul(ctx, rms, weight) : ggml_add(ctx, rms, weight); if (test.observed) { ggml_set_output(norm); } @@ -6701,8 +4729,7 @@ static void run_rmsnorm_binary_k16_output_checks() { const auto * output = imported.graph.values().find_tensor(norm); REQUIRE(output != nullptr); if (test.packed) { - const auto * generated = - plan.metadata.find_generated_resource(output->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + const auto * generated = plan.metadata.find_generated_resource(output->id, ggml::hrx::GeneratedResourceRole::F16K16Major); REQUIRE(generated != nullptr); REQUIRE(generated->byte_count == output->byte_count / 2); REQUIRE(producer->bindings.size() == 4); @@ -6716,182 +4743,17 @@ static void run_rmsnorm_binary_k16_output_checks() { return &dispatch != &*producer && !dispatch.bindings.empty() && dispatch.bindings.front().value == generated->generated_value; }) == test.projections); - } else if (test.hidden == 7168 && test.tokens == 512) { - const auto * f16 = - plan.metadata.find_alternate_value(output->id, GGML_TYPE_F16, output->byte_count / 2); - REQUIRE(f16 != nullptr); - REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_copy_f16_k16_major"; - }) == 1); - } else if (test.projections == 0) { - REQUIRE(plan.metadata.find_generated_resource( - output->id, ggml::hrx::GeneratedResourceRole::F16K16Major) == nullptr); - REQUIRE(plan.metadata.find_alternate_value(output->id) == nullptr); } if (side != nullptr) { const auto * side_value = imported.graph.values().find_tensor(side); REQUIRE(side_value != nullptr); - const auto consumer = - std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { return !dispatch.bindings.empty() && dispatch.bindings.back().value == side_value->id; }); REQUIRE(consumer != plan.dispatches.end()); REQUIRE(consumer->bindings.front().value == output->id); } - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - ggml_free(ctx); - } -} - -static void run_rmsnorm_binary_f16_output_checks() { - { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 7168, 512); - ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 7168); - ggml_tensor * rms = ggml_rms_norm(ctx, input, 1.0e-6f); - ggml_tensor * norm = ggml_mul(ctx, rms, norm_weight); - ggml_tensor * projection_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 7168, 2048); - ggml_tensor * output = ggml_mul_mat(ctx, projection_weight, norm); - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, output); - - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); - const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_binary_f32_f16"; - }); - REQUIRE(producer != plan.dispatches.end()); - REQUIRE(producer->bindings.size() == 4); - const ggml::hrx::Value * norm_value = imported.graph.values().find_tensor(norm); - REQUIRE(norm_value != nullptr); - const auto * f16 = - plan.metadata.find_alternate_value(norm_value->id, GGML_TYPE_F16, norm_value->byte_count / 2); - REQUIRE(f16 != nullptr); - REQUIRE(producer->bindings.back().value == f16->alternate_value); - const auto consumer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return &dispatch != &*producer && !dispatch.bindings.empty() && - dispatch.bindings.front().value == f16->alternate_value; - }); - REQUIRE(consumer != plan.dispatches.end()); - const ggml::hrx::CommandProgram commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - ggml_free(ctx); - } - - { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 512); - ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 2048); - ggml_tensor * rms = ggml_rms_norm(ctx, input, 1.0e-6f); - ggml_tensor * norm = ggml_mul(ctx, rms, norm_weight); - ggml_tensor * q8_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_K, 2048, 1024); - ggml_tensor * f16_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 2048); - ggml_tensor * q8_output = ggml_mul_mat(ctx, q8_weight, norm); - ggml_tensor * f16_output = ggml_mul_mat(ctx, f16_weight, norm); - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, q8_output); - ggml_build_forward_expand(graph, f16_output); - - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); - const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_binary_q8_1_x4_publish"; - }); - REQUIRE(producer != plan.dispatches.end()); - require_compile_parameter(*producer, "ggml.rmsnorm_binary_q8_1_x4.publish_f16", "1"); - REQUIRE(producer->bindings.size() == 5); - const ggml::hrx::Value * norm_value = imported.graph.values().find_tensor(norm); - REQUIRE(norm_value != nullptr); - const auto * q8 = - plan.metadata.find_alternate_value(norm_value->id, GGML_TYPE_Q8_1, qwen_q8_1_x4_size(512, 2048)); - const auto * f16 = - plan.metadata.find_alternate_value(norm_value->id, GGML_TYPE_F16, norm_value->byte_count / 2); - REQUIRE(q8 != nullptr); - REQUIRE(f16 != nullptr); - REQUIRE(producer->bindings[3].value == q8->alternate_value); - REQUIRE(producer->bindings[4].value == f16->alternate_value); - REQUIRE(std::any_of(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return &dispatch != &*producer && !dispatch.bindings.empty() && - dispatch.bindings.front().value == q8->alternate_value; - })); - REQUIRE(std::any_of(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return &dispatch != &*producer && !dispatch.bindings.empty() && - dispatch.bindings.front().value == f16->alternate_value; - })); - const ggml::hrx::CommandProgram commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - ggml_free(ctx); - } - - { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - constexpr int64_t hidden = 2048; - constexpr int64_t tokens = 2; - ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hidden, 1, tokens); - ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden); - ggml_tensor * norm = ggml_mul(ctx, ggml_rms_norm(ctx, input, 1.0e-5f), norm_weight); - ggml_tensor * reshaped = ggml_reshape_2d(ctx, norm, hidden, tokens); - ggml_tensor * projection = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, hidden, 6144); - ggml_tensor * output = ggml_silu(ctx, ggml_mul_mat(ctx, projection, reshaped)); - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, output); - - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); - const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_binary_q8_1_x4_publish"; - }); - REQUIRE(producer != plan.dispatches.end()); - require_compile_parameter(*producer, "ggml.rmsnorm_binary_q8_1_x4.publish_f16", "1"); - REQUIRE(producer->bindings.size() == 5); - const ggml::hrx::Value * norm_value = imported.graph.values().find_tensor(norm); - REQUIRE(norm_value != nullptr); - const size_t q8_bytes = qwen_q8_1_x4_size(tokens, hidden); - const auto * q8 = plan.metadata.find_alternate_value(norm_value->id, GGML_TYPE_Q8_1, q8_bytes); - REQUIRE(q8 != nullptr); - REQUIRE(producer->bindings[3].value == q8->alternate_value); - REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32"; - }) == 0); - REQUIRE(std::any_of(std::next(producer), plan.dispatches.end(), [&](const auto & dispatch) { - return !dispatch.bindings.empty() && dispatch.bindings.front().value == q8->alternate_value; - })); - const ggml::hrx::CommandProgram commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); ggml_free(ctx); @@ -6920,10 +4782,7 @@ static void run_rmsnorm_gate_packed_output_checks() { { GGML_TYPE_Q4_K, 128, 6144, 512, true, false, 1 }, { GGML_TYPE_Q4_K, 128, 4096, 256, false, false, 0 }, { GGML_TYPE_Q4_K, 256, 256, 512, false, true, 0 }, - { GGML_TYPE_Q5_K, 128, 6144, 19, false, false, 0 }, - { GGML_TYPE_IQ4_XS, 128, 6144, 19, false, false, 0 }, }; - for (const auto & test : cases) { const int64_t heads = test.channels / test.hidden; ggml_tensor * input = heads == 1 ? ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens) : @@ -6946,10 +4805,9 @@ static void run_rmsnorm_gate_packed_output_checks() { const auto & plan = scheduler.plan(); REQUIRE(plan.valid()); const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_gate_f32_publish"; + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_gate_f32_f16"; }); REQUIRE(producer != plan.dispatches.end()); - require_compile_parameter(*producer, "ggml.rmsnorm_gate_f32.publish_q8", "0"); require_compile_parameter(*producer, "ggml.rmsnorm_gate_f32.f16_output_row_width", std::to_string(test.packed ? test.channels : 0)); const auto * output = imported.graph.values().find_tensor(gated); @@ -6958,8 +4816,7 @@ static void run_rmsnorm_gate_packed_output_checks() { const auto * alternate = plan.metadata.find_alternate_value(output->id, GGML_TYPE_F16, output->byte_count / 2); REQUIRE((alternate == nullptr) == test.packed); if (test.packed) { - const auto * generated = - plan.metadata.find_generated_resource(matmul_input->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + const auto * generated = plan.metadata.find_generated_resource(matmul_input->id, ggml::hrx::GeneratedResourceRole::F16K16Major); REQUIRE(generated != nullptr); REQUIRE(producer->bindings.back().value == generated->generated_value); REQUIRE(std::next(producer) != plan.dispatches.end()); @@ -6968,44 +4825,7 @@ static void run_rmsnorm_gate_packed_output_checks() { REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_copy_f16_k16_major"; }) == test.copies); - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - } - - { - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 128, 2); - ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 128); - ggml_tensor * gate = ggml_dup_tensor(ctx, input); - ggml_tensor * norm = ggml_mul(ctx, ggml_rms_norm(ctx, input, 1.0e-6f), weight); - ggml_tensor * gated = ggml_mul(ctx, norm, ggml_silu(ctx, gate)); - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, ggml_scale(ctx, gated, 0.5f)); - - auto imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const auto & plan = scheduler.plan(); - REQUIRE(plan.valid()); - const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_gate_f32_publish"; - }); - REQUIRE(producer != plan.dispatches.end()); - require_compile_parameter(*producer, "ggml.rmsnorm_gate_f32.publish_q8", "0"); - REQUIRE(producer->bindings.size() == 5); - const auto * output = imported.graph.values().find_tensor(gated); - REQUIRE(output != nullptr); - REQUIRE(plan.metadata.find_alternate_value(output->id) == nullptr); - REQUIRE(plan.metadata.find_generated_resource( - output->id, ggml::hrx::GeneratedResourceRole::F16K16Major) == nullptr); - REQUIRE(std::any_of(plan.transients.begin(), plan.transients.end(), [&](const auto & transient) { - return transient.value == producer->bindings.back().value && - transient.size == output->byte_count / 2; - })); - const auto commands = ggml::hrx::build_command_program( - imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } @@ -7018,7 +4838,6 @@ static void run_rmsnorm_gate_q8_output_checks() { params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - const struct { ggml_type type; int64_t hidden; @@ -7045,7 +4864,6 @@ static void run_rmsnorm_gate_q8_output_checks() { { GGML_TYPE_Q4_K, 128, 6144, 5, 48, false, false, false }, { GGML_TYPE_Q6_K, 128, 6144, 5, 256, false, true, false }, }; - for (const auto & test : cases) { const int64_t heads = test.channels / test.hidden; ggml_tensor * input = heads == 1 ? ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens) : @@ -7069,13 +4887,10 @@ static void run_rmsnorm_gate_q8_output_checks() { const auto & plan = scheduler.plan(); REQUIRE(plan.valid()); const auto producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - const auto publish_q8 = dispatch.kernel.compile_parameters.find("ggml.rmsnorm_gate_f32.publish_q8"); - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_gate_f32_publish" && - publish_q8 != dispatch.kernel.compile_parameters.end() && publish_q8->second == "1"; + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_gate_f32_q8_1_x4"; }); REQUIRE((producer != plan.dispatches.end()) == test.q8); if (test.q8) { - require_compile_parameter(*producer, "ggml.rmsnorm_gate_f32.publish_q8", "1"); const auto * output = imported.graph.values().find_tensor(gated); REQUIRE(output != nullptr); const size_t bytes = static_cast(test.tokens) * ggml_row_size(GGML_TYPE_Q8_1, test.channels); @@ -7087,124 +4902,24 @@ static void run_rmsnorm_gate_q8_output_checks() { return kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32"; }) == 0); REQUIRE(std::any_of(std::next(producer), plan.dispatches.end(), [&](const auto & dispatch) { - return std::any_of(dispatch.bindings.begin(), dispatch.bindings.end(), - [&](const auto & binding) { return binding.value == alternate->alternate_value; }); + return std::any_of(dispatch.bindings.begin(), dispatch.bindings.end(), [&](const auto & binding) { + return binding.value == alternate->alternate_value; + }); })); } - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } ggml_free(ctx); } -static void run_symmetric_i4_consumer_qualification_checks() { - const struct { - bool direct_consumer; - bool reshape_consumer; - bool strided_consumer; - bool q8_publication; - } cases[] = { - { true, false, false, true }, - { false, true, false, true }, - { false, false, true, false }, - { false, false, false, false }, - }; - for (const auto & test : cases) { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - constexpr int64_t hidden = 2048; - constexpr int64_t tokens = 2; - const int64_t producer_hidden = test.strided_consumer ? hidden + 128 : hidden; - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, producer_hidden, tokens); - ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, producer_hidden); - ggml_tensor * normalized = ggml_mul(ctx, ggml_rms_norm(ctx, input, 1.0e-6f), weight); - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_tensor * consumer_input = normalized; - if (test.reshape_consumer) { - consumer_input = ggml_reshape_2d(ctx, normalized, hidden, tokens); - } else if (test.strided_consumer) { - consumer_input = ggml_view_2d( - ctx, normalized, hidden, tokens, static_cast(producer_hidden) * sizeof(float), 0); - } - if (test.direct_consumer || test.reshape_consumer || test.strided_consumer) { - ggml_tensor * i4_projection = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_K, hidden, hidden); - ggml_tensor * i4_projected = ggml_mul_mat(ctx, i4_projection, consumer_input); - ggml_build_forward_expand(graph, ggml_scale(ctx, i4_projected, 0.5f)); - if (!test.strided_consumer) { - ggml_tensor * q8_projection = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, hidden, 512); - ggml_tensor * q8_projected = ggml_mul_mat(ctx, q8_projection, consumer_input); - ggml_build_forward_expand(graph, ggml_scale(ctx, q8_projected, 0.5f)); - } - } else { - ggml_build_forward_expand(graph, ggml_scale(ctx, normalized, 0.5f)); - } - - auto imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - const ggml::hrx::Value * normalized_value = imported.graph.values().find_tensor(normalized); - const ggml::hrx::Value * qualified_input = imported.graph.values().find_tensor(consumer_input); - REQUIRE(normalized_value != nullptr); - REQUIRE(qualified_input != nullptr); - for (const ggml::hrx::GraphNode * consumer : imported.graph.index().consumers(qualified_input->id)) { - REQUIRE(consumer == nullptr || - !ggml::hrx::common_symmetric_i4_lowrow_mul_mat_eligible( - imported.graph, *consumer, *qualified_input)); - } - if (test.strided_consumer) { - ggml_free(ctx); - continue; - } - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const auto & plan = scheduler.plan(); - REQUIRE(plan.valid()); - const auto symmetric_producer = - std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == - "loom_libs:ggml_rmsnorm_binary_symmetric_i4_k32"; - }); - REQUIRE(symmetric_producer == plan.dispatches.end()); - const auto q8_producer = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_binary_q8_1_x4_publish"; - }); - REQUIRE((q8_producer != plan.dispatches.end()) == test.q8_publication); - if (test.q8_publication) { - require_compile_parameter(*q8_producer, "ggml.rmsnorm_binary_q8_1_x4.publish_f16", "1"); - const auto * q8 = plan.metadata.find_alternate_value( - normalized_value->id, GGML_TYPE_Q8_1, qwen_q8_1_x4_size(tokens, hidden)); - REQUIRE(q8 != nullptr); - REQUIRE(q8_producer->bindings.size() == 5); - REQUIRE(q8_producer->bindings[3].value == q8->alternate_value); - REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "qwen3_moe:ggml_quantize_q8_1_x4_f32"; - }) == 0); - REQUIRE(std::any_of(std::next(q8_producer), plan.dispatches.end(), [&](const auto & dispatch) { - return !dispatch.bindings.empty() && dispatch.bindings.front().value == q8->alternate_value; - })); - } else if (!test.strided_consumer) { - REQUIRE(plan.metadata.alternate_values().empty()); - } - const auto commands = ggml::hrx::build_command_program( - imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - ggml_free(ctx); - } -} - static void run_quantized_conv4_dispatch_checks() { ggml_init_params params = {}; params.mem_size = 8 * 1024 * 1024; params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - const struct { ggml_type type; int64_t inputs; @@ -7218,203 +4933,43 @@ static void run_quantized_conv4_dispatch_checks() { { GGML_TYPE_Q4_K, 5120, 10240, 512, false, false, true }, { GGML_TYPE_Q6_K, 5120, 10240, 512, false, false, true }, { GGML_TYPE_Q4_K, 4096, 8192, 512, false, false, true }, - { GGML_TYPE_Q6_K, 4096, 8192, 512, false, false, true }, - { GGML_TYPE_Q4_K, 5120, 10240, 512, true, false, false }, - { GGML_TYPE_Q4_K, 5120, 10240, 512, false, true, false }, - { GGML_TYPE_Q4_K, 5120, 10240, 256, false, false, false }, - { GGML_TYPE_Q4_K, 5120, 10240, 512, false, false, true, true }, - { GGML_TYPE_Q6_K, 5120, 10240, 512, false, false, true, true }, - { GGML_TYPE_Q4_K, 4096, 8192, 512, false, false, true, true }, - { GGML_TYPE_Q6_K, 4096, 8192, 512, false, false, true, true }, - { GGML_TYPE_Q4_K, 5120, 10240, 512, true, false, false, true }, - { GGML_TYPE_Q4_K, 5120, 10240, 512, false, true, false, true }, - }; - - for (const auto & test : cases) { - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.inputs, test.tokens); - ggml_tensor * weight = ggml_new_tensor_2d(ctx, test.type, test.inputs, test.channels); - ggml_tensor * raw = ggml_mul_mat(ctx, weight, input); - if (test.exposed) { - ggml_set_output(raw); - } - ggml_tensor * x = ggml_reshape_3d(ctx, raw, test.channels, test.tokens, 1); - ggml_tensor * state = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 3, test.channels, 1); - if (test.gathered_state) { - ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); - state = ggml_get_rows(ctx, ggml_reshape_2d(ctx, state, 3 * test.channels, 1), indices); - state = ggml_reshape_3d(ctx, state, 3, test.channels, 1); - } - ggml_tensor * filter = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, test.channels); - ggml_tensor * window = ggml_concat(ctx, state, ggml_transpose(ctx, x), 0); - ggml_tensor * tail = - ggml_view_3d(ctx, window, 3, test.channels, 1, window->nb[1], window->nb[2], test.tokens * sizeof(float)); - ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3 * test.channels, 1); - ggml_tensor * target = ggml_view_2d(ctx, cache, 3 * test.channels, 1, cache->nb[1], 0); - ggml_tensor * output = ggml_silu(ctx, ggml_ssm_conv(ctx, window, filter)); - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, ggml_cpy(ctx, tail, target)); - ggml_build_forward_expand(graph, output); - if (test.side_use) { - ggml_build_forward_expand(graph, ggml_scale(ctx, raw, 0.5f)); - } - auto imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const auto & plan = scheduler.plan(); - REQUIRE(plan.valid()); - const auto fused = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == - (test.gathered_state ? "loom_libs:ggml_mul_mat_quantized_f16_wmma_prefill_conv4_interior" : - "loom_libs:ggml_mul_mat_quantized_f16_wmma_prefill_conv4"); - }); - REQUIRE((fused != plan.dispatches.end()) == test.fused); - if (test.fused) { - REQUIRE(fused->bindings.size() == (test.gathered_state ? 5 : 6)); - require_compile_parameter(*fused, "ggml.mul_mat.weight_format", test.type == GGML_TYPE_Q4_K ? "4" : "6"); - if (test.gathered_state) { - const auto finish = - std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == - "loom_libs:llm_ssm_conv_dconv4_silu_prefill_finish_f32"; - }); - const auto gather = - std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id).find("get_rows") != std::string::npos; - }); - REQUIRE(finish != plan.dispatches.end()); - REQUIRE(gather != plan.dispatches.end()); - REQUIRE(fused < gather && gather < finish); - REQUIRE(finish->bindings[2].value == fused->bindings[4].value); - REQUIRE(finish->bindings[3].value == fused->bindings[3].value); - REQUIRE(finish->bindings[2].length == static_cast(6 * test.channels) * sizeof(float)); - } - REQUIRE(std::none_of(plan.transients.begin(), plan.transients.end(), [](const auto & transient) { - return transient.name == "llm.ssm_conv.window_snapshot"; - })); - } - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - } - ggml_free(ctx); -} - -static void run_generic_ssm_conv_binary_dispatch_checks() { - const struct { - ggml::hrx::BinaryKind kind; - bool conv_is_lhs; - } cases[] = { - { ggml::hrx::BinaryKind::Add, true }, - { ggml::hrx::BinaryKind::Sub, true }, - { ggml::hrx::BinaryKind::Sub, false }, - { ggml::hrx::BinaryKind::Mul, true }, - { ggml::hrx::BinaryKind::Div, true }, - { ggml::hrx::BinaryKind::Div, false }, - }; - - for (const auto & test : cases) { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - constexpr int64_t d_conv = 4; - constexpr int64_t d_inner = 64; - constexpr int64_t n_t = 8; - ggml_tensor * window = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_conv + n_t - 1, d_inner); - ggml_tensor * filter = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_conv, d_inner); - ggml_tensor * operand = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, d_inner, n_t); - ggml_tensor * conv = ggml_ssm_conv(ctx, window, filter); - ggml_tensor * output = nullptr; - switch (test.kind) { - case ggml::hrx::BinaryKind::Add: - output = ggml_add(ctx, conv, operand); - break; - case ggml::hrx::BinaryKind::Sub: - output = test.conv_is_lhs ? ggml_sub(ctx, conv, operand) : ggml_sub(ctx, operand, conv); - break; - case ggml::hrx::BinaryKind::Mul: - output = ggml_mul(ctx, conv, operand); - break; - case ggml::hrx::BinaryKind::Div: - output = test.conv_is_lhs ? ggml_div(ctx, conv, operand) : ggml_div(ctx, operand, conv); - break; - default: - REQUIRE(false); - } - - ggml_cgraph * graph = ggml_new_graph(ctx); - ggml_build_forward_expand(graph, output); - auto imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const auto & plan = scheduler.plan(); - REQUIRE(plan.valid()); - REQUIRE(plan.dispatches.size() == 1); - const auto & dispatch = plan.dispatches.front(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:llm_ssm_conv_binary_f32"); - REQUIRE(dispatch.kernel.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); - REQUIRE(dispatch.kernel.integer_parameters.at("n_t") == n_t); - REQUIRE(dispatch.kernel.integer_parameters.at("n_s") == 1); - REQUIRE(dispatch.kernel.compile_parameters.count("llm.ssm_conv.generic.n_t") == 0); - REQUIRE(dispatch.kernel.compile_parameters.count("llm.ssm_conv.generic.n_s") == 0); - REQUIRE(dispatch.bindings.size() == 4); - require_compile_parameter(dispatch, "llm.ssm_conv.generic.binary_op", - std::to_string(ggml::hrx::binary_kind_config_value(test.kind))); - require_compile_parameter(dispatch, "llm.ssm_conv.generic.binary_lhs", test.conv_is_lhs ? "1" : "0"); - REQUIRE(dispatch.bindings[2].value == imported.graph.values().find_tensor(operand)->id); - REQUIRE(dispatch.bindings[3].value == imported.graph.values().find_tensor(output)->id); - - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - ggml_free(ctx); - } -} - -static void run_lfm_dconv3_dispatch_checks() { - const struct { - int64_t tokens; - int64_t hidden; - int64_t state_rows; - int64_t filter_rows; - bool disabled; - bool fused; - } cases[] = { - {1, 2048, 2, 3, false, true}, - {64, 2048, 2, 3, false, true}, - {2, 2048, 2, 3, false, false}, - {64, 2016, 2, 3, false, false}, - {64, 2048, 3, 3, false, false}, - {64, 2048, 2, 4, false, false}, - {64, 2048, 2, 3, true, false}, + { GGML_TYPE_Q6_K, 4096, 8192, 512, false, false, true }, + { GGML_TYPE_Q4_K, 5120, 10240, 512, true, false, false }, + { GGML_TYPE_Q4_K, 5120, 10240, 512, false, true, false }, + { GGML_TYPE_Q4_K, 5120, 10240, 256, false, false, false }, + { GGML_TYPE_Q4_K, 5120, 10240, 512, false, false, true, true }, + { GGML_TYPE_Q6_K, 5120, 10240, 512, false, false, true, true }, + { GGML_TYPE_Q4_K, 4096, 8192, 512, false, false, true, true }, + { GGML_TYPE_Q6_K, 4096, 8192, 512, false, false, true, true }, + { GGML_TYPE_Q4_K, 5120, 10240, 512, true, false, false, true }, + { GGML_TYPE_Q4_K, 5120, 10240, 512, false, true, false, true }, }; - for (const auto & test : cases) { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - ggml_tensor * x = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, test.hidden, test.tokens, 1); - ggml_tensor * state = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, test.state_rows, test.hidden, 1); - ggml_tensor * filter = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.filter_rows, test.hidden); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.inputs, test.tokens); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, test.type, test.inputs, test.channels); + ggml_tensor * raw = ggml_mul_mat(ctx, weight, input); + if (test.exposed) { + ggml_set_output(raw); + } + ggml_tensor * x = ggml_reshape_3d(ctx, raw, test.channels, test.tokens, 1); + ggml_tensor * state = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 3, test.channels, 1); + if (test.gathered_state) { + ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + state = ggml_get_rows(ctx, ggml_reshape_2d(ctx, state, 3 * test.channels, 1), indices); + state = ggml_reshape_3d(ctx, state, 3, test.channels, 1); + } + ggml_tensor * filter = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 4, test.channels); ggml_tensor * window = ggml_concat(ctx, state, ggml_transpose(ctx, x), 0); - ggml_tensor * tail = ggml_view_3d(ctx, window, test.state_rows, test.hidden, 1, window->nb[1], window->nb[2], - test.tokens * sizeof(float)); - ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.state_rows * test.hidden, 1); - ggml_tensor * target = ggml_view_2d(ctx, cache, test.state_rows * test.hidden, 1, cache->nb[1], 0); - ggml_tensor * output = ggml_ssm_conv(ctx, window, filter); + ggml_tensor * tail = ggml_view_3d(ctx, window, 3, test.channels, 1, + window->nb[1], window->nb[2], test.tokens * sizeof(float)); + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3 * test.channels, 1); + ggml_tensor * target = ggml_view_2d(ctx, cache, 3 * test.channels, 1, cache->nb[1], 0); + ggml_tensor * output = ggml_silu(ctx, ggml_ssm_conv(ctx, window, filter)); ggml_cgraph * graph = ggml_new_graph(ctx); ggml_build_forward_expand(graph, ggml_cpy(ctx, tail, target)); ggml_build_forward_expand(graph, output); - if (test.disabled) { - REQUIRE(setenv("GGML_HRX_DISABLE_LFM_DCONV3_FUSION", "1", 1) == 0); + if (test.side_use) { + ggml_build_forward_expand(graph, ggml_scale(ctx, raw, 0.5f)); } auto imported = ggml::hrx::import_ggml_graph(*graph); REQUIRE(imported.valid()); @@ -7422,23 +4977,40 @@ static void run_lfm_dconv3_dispatch_checks() { REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); const auto & plan = scheduler.plan(); REQUIRE(plan.valid()); - const auto fused = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:llm_ssm_conv_lfm_dconv3_state_f32"; + const auto fused = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [&](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == + (test.gathered_state ? "loom_libs:ggml_mul_mat_quantized_f16_wmma_prefill_conv4_interior" : + "loom_libs:ggml_mul_mat_quantized_f16_wmma_prefill_conv4"); }); REQUIRE((fused != plan.dispatches.end()) == test.fused); if (test.fused) { - REQUIRE(fused->bindings.size() == 5); - require_compile_parameter(*fused, "llm.ssm_conv.lfm_dconv3.n_t", std::to_string(test.tokens)); + REQUIRE(fused->bindings.size() == (test.gathered_state ? 5 : 6)); + require_compile_parameter(*fused, "ggml.mul_mat.weight_format", test.type == GGML_TYPE_Q4_K ? "4" : "6"); + if (test.gathered_state) { + const auto finish = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:llm_ssm_conv_dconv4_silu_prefill_finish_f32"; + }); + const auto gather = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { + return kernel_name_for_id(dispatch.kernel.kernel_id).find("get_rows") != std::string::npos; + }); + REQUIRE(finish != plan.dispatches.end()); + REQUIRE(gather != plan.dispatches.end()); + REQUIRE(fused < gather && gather < finish); + REQUIRE(finish->bindings[2].value == fused->bindings[4].value); + REQUIRE(finish->bindings[3].value == fused->bindings[3].value); + REQUIRE(finish->bindings[2].length == static_cast(6 * test.channels) * sizeof(float)); + } + REQUIRE(std::none_of(plan.transients.begin(), plan.transients.end(), [](const auto & transient) { + return transient.name == "llm.ssm_conv.window_snapshot"; + })); } - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = ggml::hrx::build_command_program( + imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); - if (test.disabled) { - REQUIRE(unsetenv("GGML_HRX_DISABLE_LFM_DCONV3_FUSION") == 0); - } - ggml_free(ctx); } + ggml_free(ctx); } static void run_gdn_native_projection_pair_dispatch_checks() { @@ -7447,7 +5019,6 @@ static void run_gdn_native_projection_pair_dispatch_checks() { params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - const struct { int64_t tokens; int64_t inputs; @@ -7465,7 +5036,6 @@ static void run_gdn_native_projection_pair_dispatch_checks() { {2, 5120, 48, false, false, false, GGML_TYPE_Q4_K, false}, {5, 5120, 48, false, false, false, GGML_TYPE_Q4_K, false}, {16, 5120, 48, false, false, false, GGML_TYPE_Q4_K, false}, - {17, 5120, 48, false, false, false, GGML_TYPE_Q4_K, false}, {1, 3072, 48, false, false, false, GGML_TYPE_Q4_K, false}, {1, 4352, 48, false, false, false, GGML_TYPE_Q4_K, false}, {1, 5120, 16, false, false, false, GGML_TYPE_Q4_K, false}, @@ -7474,7 +5044,6 @@ static void run_gdn_native_projection_pair_dispatch_checks() { {1, 5120, 48, false, false, true, GGML_TYPE_Q4_K, false}, {1, 5120, 48, false, false, false, GGML_TYPE_Q6_K, false}, }; - for (const auto & test : cases) { const int64_t width = 128; const int64_t qheads = test.heads / 4; @@ -7492,15 +5061,14 @@ static void run_gdn_native_projection_pair_dispatch_checks() { ggml_tensor * beta_raw = ggml_mul_mat(ctx, second_weight, second_input); ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.heads); ggml_tensor * scale = ggml_dup_tensor(ctx, bias); - ggml_tensor * gate = ggml_mul( - ctx, ggml_softplus(ctx, ggml_add(ctx, ggml_reshape_3d(ctx, alpha, test.heads, test.tokens, 1), bias)), - scale); + ggml_tensor * gate = ggml_mul(ctx, ggml_softplus(ctx, ggml_add(ctx, + ggml_reshape_3d(ctx, alpha, test.heads, test.tokens, 1), bias)), scale); gate = ggml_reshape_4d(ctx, gate, 1, test.heads, test.tokens, 1); ggml_tensor * beta = ggml_sigmoid(ctx, ggml_reshape_4d(ctx, beta_raw, 1, test.heads, test.tokens, 1)); ggml_tensor * raw = ggml_scale(ctx, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden, test.tokens), 0.5f); const auto view = [&](int64_t heads, size_t offset) { - return ggml_view_4d(ctx, raw, width, heads, test.tokens, 1, width * sizeof(float), hidden * sizeof(float), - hidden * test.tokens * sizeof(float), offset); + return ggml_view_4d(ctx, raw, width, heads, test.tokens, 1, width * sizeof(float), + hidden * sizeof(float), hidden * test.tokens * sizeof(float), offset); }; ggml_tensor * q = ggml_l2_norm(ctx, view(qheads, 0), 1.e-6f); ggml_tensor * k = ggml_l2_norm(ctx, view(qheads, width * qheads * sizeof(float)), 1.e-6f); @@ -7512,10 +5080,10 @@ static void run_gdn_native_projection_pair_dispatch_checks() { ggml_tensor * gdn = ggml_gated_delta_net(ctx, q, k, v, gate, beta, state, 5); const int64_t count = std::min(test.tokens, 5); const size_t attention_bytes = width * test.heads * test.tokens * sizeof(float); - ggml_tensor * attention = ggml_view_4d(ctx, gdn, width, test.heads, test.tokens, 1, width * sizeof(float), - width * test.heads * sizeof(float), attention_bytes, 0); - ggml_tensor * snapshots = - ggml_view_3d(ctx, gdn, state_elements, 1, count, state_bytes, state_bytes, attention_bytes); + ggml_tensor * attention = ggml_view_4d(ctx, gdn, width, test.heads, test.tokens, 1, + width * sizeof(float), width * test.heads * sizeof(float), attention_bytes, 0); + ggml_tensor * snapshots = ggml_view_3d(ctx, gdn, state_elements, 1, count, + state_bytes, state_bytes, attention_bytes); ggml_tensor * target = ggml_view_3d(ctx, cache, state_elements, 1, count, state_bytes, 4 * state_bytes, 0); ggml_cgraph * graph = ggml_new_graph(ctx); ggml_build_forward_expand(graph, raw); @@ -7546,18 +5114,8 @@ static void run_gdn_native_projection_pair_dispatch_checks() { REQUIRE(fused->bindings[3].value == imported.graph.values().find_tensor(alpha)->id); REQUIRE(fused->bindings[4].value == imported.graph.values().find_tensor(beta_raw)->id); } - const auto epilogue = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == - "loom_libs:llm_gated_delta_net_projection_epilogue_f32"; - }); - REQUIRE((epilogue != plan.dispatches.end()) == (test.tokens > 16)); - if (epilogue != plan.dispatches.end()) { - REQUIRE(epilogue->kernel.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); - REQUIRE(epilogue->kernel.integer_parameters.at("element_count") == test.heads * test.tokens); - REQUIRE(epilogue->kernel.compile_parameters.count("llm.gated_delta_net.epilogue_element_count") == 0); - } - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = ggml::hrx::build_command_program(imported.graph, plan, + ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } @@ -7570,7 +5128,6 @@ static void run_gdn_selected_snapshot_dispatch_checks() { params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - const struct { int64_t tokens; int64_t sequences; @@ -7663,8 +5220,7 @@ static void run_gdn_selected_snapshot_dispatch_checks() { return kernel_name_for_id(dispatch.kernel.kernel_id)=="loom_libs:ggml_get_rows_f32"; })); } - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands=ggml::hrx::build_command_program(imported.graph,plan,ggml::hrx::get_qwen_kernel_corpus(),"gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } @@ -7693,7 +5249,6 @@ static void run_gdn_rmsnorm_gate_dispatch_checks() { { GGML_TYPE_Q4_K, 512, false, true, false, false }, { GGML_TYPE_Q4_K, 512, false, false, true, false }, }; - for (const auto & test : cases) { const int64_t width = 128; const int64_t heads = 2; @@ -7706,15 +5261,17 @@ static void run_gdn_rmsnorm_gate_dispatch_checks() { ggml_tensor * beta = ggml_dup_tensor(ctx, g); ggml_tensor * state = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, width, width, heads, 1); ggml_tensor * gdn = ggml_gated_delta_net(ctx, q, k, v, g, beta, state, 1); - ggml_tensor * attention = ggml_view_4d(ctx, gdn, width, heads, test.tokens, 1, width * sizeof(float), - channels * sizeof(float), channels * test.tokens * sizeof(float), 0); + ggml_tensor * attention = ggml_view_4d(ctx, gdn, width, heads, test.tokens, 1, + width * sizeof(float), channels * sizeof(float), + channels * test.tokens * sizeof(float), 0); if (test.attention_output) { ggml_set_output(attention); } if (test.gdn_output) { ggml_set_output(gdn); } - ggml_tensor * new_state = ggml_view_2d(ctx, gdn, state_elements, 1, state_elements * sizeof(float), + ggml_tensor * new_state = ggml_view_2d(ctx, gdn, state_elements, 1, + state_elements * sizeof(float), channels * test.tokens * sizeof(float)); ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, state_elements, 1); ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, width); @@ -7741,19 +5298,19 @@ static void run_gdn_rmsnorm_gate_dispatch_checks() { }); REQUIRE((fused != plan.dispatches.end()) == test.fused); REQUIRE(std::count_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_gate_f32_publish"; + return kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_rmsnorm_gate_f32_f16"; }) == (test.fused ? 0 : 1)); if (test.fused) { REQUIRE(fused->bindings.size() == 11); const auto * projected = imported.graph.values().find_tensor(input); REQUIRE(projected != nullptr); - const auto * packed = - plan.metadata.find_generated_resource(projected->id, ggml::hrx::GeneratedResourceRole::F16K16Major); + const auto * packed = plan.metadata.find_generated_resource( + projected->id, ggml::hrx::GeneratedResourceRole::F16K16Major); REQUIRE(packed != nullptr && fused->bindings.back().value == packed->generated_value); require_compile_parameter(*fused, "ggml.rmsnorm_gate_f32.rms_epsilon", "1.99999999e-06"); } - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = ggml::hrx::build_command_program( + imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); } @@ -7762,7 +5319,7 @@ static void run_gdn_rmsnorm_gate_dispatch_checks() { static void run_qwen_matmul_dispatch_checks() { ggml_init_params params = {}; - params.mem_size = 12 * 1024 * 1024; + params.mem_size = 8 * 1024 * 1024; params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); @@ -7774,8 +5331,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, - 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); } { @@ -7795,50 +5351,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, 2048, - 128); - } - - { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 3); - REQUIRE(weight != nullptr); - REQUIRE(input != nullptr); - ggml_tensor * output = ggml_mul_mat(ctx, weight, input); - REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 3, 2048, - 128); - } - - { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 5); - REQUIRE(weight != nullptr); - REQUIRE(input != nullptr); - ggml_tensor * output = ggml_mul_mat(ctx, weight, input); - REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, 2048, - 128); - } - - { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 6); - REQUIRE(weight != nullptr); - REQUIRE(input != nullptr); - ggml_tensor * output = ggml_mul_mat(ctx, weight, input); - REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 6, 2048, 128); - } - - { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_0, 640, 128); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 2); - REQUIRE(weight != nullptr); - REQUIRE(input != nullptr); - ggml_tensor * output = ggml_mul_mat(ctx, weight, input); - REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 2, 640, 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 2048, 128); } { @@ -7858,8 +5371,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, - 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); } { @@ -7879,8 +5391,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, - 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); } { @@ -7900,8 +5411,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, - 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); } { @@ -7921,8 +5431,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, - 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); } { @@ -7942,8 +5451,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, - 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); } { @@ -7963,8 +5471,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, - 256); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 256); } { @@ -7988,10 +5495,9 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(projection != nullptr); ggml_tensor * output = ggml_add(ctx, projection, bias); REQUIRE(output != nullptr); - schedule_fused_matmul_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, - 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, - { "input", "weight", "bias", "residual_input", "residual_output" }, false); + schedule_fused_matmul_postops_command(ctx, output, "loom_libs:ggml_mul_mat_bias_f32_f32_wmma", GGML_TYPE_F16, + 33, 2048, 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, + { "input", "weight", "bias", "output" }, false); } { @@ -8005,10 +5511,9 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(projection != nullptr); ggml_tensor * output = ggml_add(ctx, projection, bias); REQUIRE(output != nullptr); - schedule_fused_matmul_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_BF16, 33, 2048, - 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, - { "input", "weight", "bias", "residual_input", "residual_output" }, false); + schedule_fused_matmul_postops_command(ctx, output, "loom_libs:ggml_mul_mat_bias_f32_f32_wmma", GGML_TYPE_BF16, + 33, 2048, 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, + { "input", "weight", "bias", "output" }, false); } { @@ -8022,10 +5527,9 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(projection != nullptr); ggml_tensor * output = ggml_add(ctx, projection, residual); REQUIRE(output != nullptr); - schedule_fused_matmul_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, - 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, - { "input", "weight", "bias", "residual_input", "residual_output" }, false); + schedule_fused_matmul_postops_command(ctx, output, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", GGML_TYPE_F16, 33, + 2048, 256, 2, { GGML_OP_MUL_MAT, GGML_OP_ADD }, + { "input", "weight", "residual_input", "residual_output" }, false); } { @@ -8044,8 +5548,8 @@ static void run_qwen_matmul_dispatch_checks() { ggml_tensor * output = ggml_add(ctx, biased, residual); REQUIRE(output != nullptr); schedule_fused_matmul_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, - 256, 3, { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_ADD }, + ctx, output, "loom_libs:ggml_mul_mat_bias_add_f32_f32_wmma", GGML_TYPE_F16, 33, 2048, 256, 3, + { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_ADD }, { "input", "weight", "bias", "residual_input", "residual_output" }, false); } @@ -8066,10 +5570,10 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(rms != nullptr); ggml_tensor * output = ggml_mul(ctx, rms, norm_weight); REQUIRE(output != nullptr); - schedule_fused_matmul_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, - 256, 4, { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, - { "input", "weight", "bias", "residual_input", "residual_output" }, true); + schedule_fused_matmul_postops_command(ctx, output, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", GGML_TYPE_F16, 33, + 2048, 256, 4, + { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, + { "input", "weight", "residual_input", "residual_output" }, true); } { @@ -8094,8 +5598,8 @@ static void run_qwen_matmul_dispatch_checks() { ggml_tensor * output = ggml_mul(ctx, rms, norm_weight); REQUIRE(output != nullptr); schedule_fused_matmul_postops_command( - ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", GGML_TYPE_F16, 33, 2048, - 256, 5, { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, + ctx, output, "loom_libs:ggml_mul_mat_bias_add_f32_f32_wmma", GGML_TYPE_F16, 33, 2048, 256, 5, + { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }, { "input", "weight", "bias", "residual_input", "residual_output" }, true); } @@ -8194,8 +5698,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(packed != nullptr); ggml_tensor * output = ggml_glu(ctx, packed, GGML_GLU_OP_SWIGLU, false); REQUIRE(output != nullptr); - schedule_packed_matmul_glu_command(ctx, output, GGML_TYPE_Q4_K, 2, 256, 128, ggml::hrx::BinaryKind::SwiGLU, - false); + schedule_packed_matmul_glu_command(ctx, output, GGML_TYPE_Q4_K, 2, 256, 128, + ggml::hrx::BinaryKind::SwiGLU, false); } { @@ -8207,8 +5711,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(packed != nullptr); ggml_tensor * output = ggml_glu(ctx, packed, GGML_GLU_OP_SWIGLU, true); REQUIRE(output != nullptr); - schedule_packed_matmul_glu_command(ctx, output, GGML_TYPE_Q4_K, 2, 256, 128, ggml::hrx::BinaryKind::SwiGLU, - true); + schedule_packed_matmul_glu_command(ctx, output, GGML_TYPE_Q4_K, 2, 256, 128, + ggml::hrx::BinaryKind::SwiGLU, true); } { @@ -8326,7 +5830,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(matmul != nullptr); ggml_tensor * output = ggml_silu(ctx, matmul); REQUIRE(output != nullptr); - schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", + schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", ggml::hrx::UnaryKind::Silu, 4, 2048, 128); } @@ -8339,7 +5843,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(matmul != nullptr); ggml_tensor * output = ggml_relu(ctx, matmul); REQUIRE(output != nullptr); - schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", + schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", ggml::hrx::UnaryKind::Relu, 4, 2048, 128); } @@ -8352,23 +5856,10 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(matmul != nullptr); ggml_tensor * output = ggml_gelu(ctx, matmul); REQUIRE(output != nullptr); - schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", + schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", ggml::hrx::UnaryKind::Gelu, 4, 2048, 128); } - { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_0, 640, 128); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 2); - REQUIRE(weight != nullptr); - REQUIRE(input != nullptr); - ggml_tensor * matmul = ggml_mul_mat(ctx, weight, input); - REQUIRE(matmul != nullptr); - ggml_tensor * output = ggml_gelu(ctx, matmul); - REQUIRE(output != nullptr); - schedule_fused_matmul_unary_command(ctx, output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", - ggml::hrx::UnaryKind::Gelu, 2, 640, 128); - } - { ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); @@ -8410,8 +5901,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q4_input != nullptr); ggml_tensor * q4_output = ggml_mul_mat(ctx, q4_weight, q4_input); REQUIRE(q4_output != nullptr); - schedule_single_matmul_command(ctx, q4_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, - 128); + schedule_single_matmul_command(ctx, q4_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); ggml_tensor * q4_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 128); ggml_tensor * q4_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -8419,7 +5909,8 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q4_decode_input != nullptr); ggml_tensor * q4_decode_output = ggml_mul_mat(ctx, q4_decode_weight, q4_decode_input); REQUIRE(q4_decode_output != nullptr); - schedule_single_matmul_command(ctx, q4_decode_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 1, 2048, 128); + schedule_single_matmul_command(ctx, q4_decode_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 1, 2048, + 128); for (ggml_type legacy_type : { GGML_TYPE_Q1_0, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1 }) { @@ -8429,8 +5920,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(legacy_input != nullptr); ggml_tensor * legacy_output = ggml_mul_mat(ctx, legacy_weight, legacy_input); REQUIRE(legacy_output != nullptr); - schedule_single_matmul_command(ctx, legacy_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, - 2048, 128); + schedule_single_matmul_command(ctx, legacy_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); ggml_tensor * legacy_decode_weight = ggml_new_tensor_2d(ctx, legacy_type, 2048, 128); ggml_tensor * legacy_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -8448,8 +5938,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q3_input != nullptr); ggml_tensor * q3_output = ggml_mul_mat(ctx, q3_weight, q3_input); REQUIRE(q3_output != nullptr); - schedule_single_matmul_command(ctx, q3_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, - 128); + schedule_single_matmul_command(ctx, q3_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); ggml_tensor * q3_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q3_K, 2048, 128); ggml_tensor * q3_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -8466,8 +5955,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq2_s_input != nullptr); ggml_tensor * iq2_s_output = ggml_mul_mat(ctx, iq2_s_weight, iq2_s_input); REQUIRE(iq2_s_output != nullptr); - schedule_single_matmul_command(ctx, iq2_s_output, - "loom_libs:ggml_mul_mat_tiled_input_f32_iq2_s_publish_f32", 2, 2048, 640); + schedule_single_matmul_command(ctx, iq2_s_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 2048, 640); ggml_tensor * iq2_s_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ2_S, 2048, 640); ggml_tensor * iq2_s_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -8475,7 +5963,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq2_s_decode_input != nullptr); ggml_tensor * iq2_s_decode_output = ggml_mul_mat(ctx, iq2_s_decode_weight, iq2_s_decode_input); REQUIRE(iq2_s_decode_output != nullptr); - schedule_single_matmul_command(ctx, iq2_s_decode_output, "loom_libs:ggml_mul_mat_vector_iq2_s_f32_f32", 1, + schedule_single_matmul_command(ctx, iq2_s_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, 640); ggml_tensor * iq3_s_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ3_S, 2048, 128); @@ -8484,8 +5972,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq3_s_input != nullptr); ggml_tensor * iq3_s_output = ggml_mul_mat(ctx, iq3_s_weight, iq3_s_input); REQUIRE(iq3_s_output != nullptr); - schedule_single_matmul_command(ctx, iq3_s_output, - "loom_libs:ggml_mul_mat_tiled_input_f32_iq3_s_publish_f32", 4, 2048, 128); + schedule_single_matmul_command(ctx, iq3_s_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); ggml_tensor * iq3_s_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ3_S, 2048, 128); ggml_tensor * iq3_s_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -8493,7 +5980,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq3_s_decode_input != nullptr); ggml_tensor * iq3_s_decode_output = ggml_mul_mat(ctx, iq3_s_decode_weight, iq3_s_decode_input); REQUIRE(iq3_s_decode_output != nullptr); - schedule_single_matmul_command(ctx, iq3_s_decode_output, "loom_libs:ggml_mul_mat_vector_iq3_s_f32_f32", 1, + schedule_single_matmul_command(ctx, iq3_s_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, 128); ggml_tensor * iq4_nl_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 2048, 128); @@ -8502,8 +5989,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq4_nl_input != nullptr); ggml_tensor * iq4_nl_output = ggml_mul_mat(ctx, iq4_nl_weight, iq4_nl_input); REQUIRE(iq4_nl_output != nullptr); - schedule_single_matmul_command(ctx, iq4_nl_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, - 2048, 128); + schedule_single_matmul_command(ctx, iq4_nl_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); ggml_tensor * iq4_nl_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 2048, 128); ggml_tensor * iq4_nl_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -8520,8 +6006,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq4_nl_small_input != nullptr); ggml_tensor * iq4_nl_small_output = ggml_mul_mat(ctx, iq4_nl_small_weight, iq4_nl_small_input); REQUIRE(iq4_nl_small_output != nullptr); - schedule_single_matmul_command(ctx, iq4_nl_small_output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", - 2, 640, 1024); + schedule_single_matmul_command(ctx, iq4_nl_small_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 640, 1024); ggml_tensor * iq4_nl_small_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_NL, 640, 1024); ggml_tensor * iq4_nl_small_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 1); @@ -8539,8 +6024,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q5_0_small_input != nullptr); ggml_tensor * q5_0_small_output = ggml_mul_mat(ctx, q5_0_small_weight, q5_0_small_input); REQUIRE(q5_0_small_output != nullptr); - schedule_single_matmul_command(ctx, q5_0_small_output, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 2, - 640, 256); + schedule_single_matmul_command(ctx, q5_0_small_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 640, 256); ggml_tensor * q5_0_small_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q5_0, 640, 256); ggml_tensor * q5_0_small_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 640, 1); @@ -8557,8 +6041,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q6_input != nullptr); ggml_tensor * q6_output = ggml_mul_mat(ctx, q6_weight, q6_input); REQUIRE(q6_output != nullptr); - schedule_single_matmul_command(ctx, q6_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, 2048, - 128); + schedule_single_matmul_command(ctx, q6_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 2048, 128); ggml_tensor * q6_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 2048, 128); ggml_tensor * q6_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -8584,8 +6067,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq4_xs_input != nullptr); ggml_tensor * iq4_xs_output = ggml_mul_mat(ctx, iq4_xs_weight, iq4_xs_input); REQUIRE(iq4_xs_output != nullptr); - schedule_single_matmul_command(ctx, iq4_xs_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, - 2048, 128); + schedule_single_matmul_command(ctx, iq4_xs_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); ggml_tensor * iq4_xs_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ4_XS, 2048, 128); ggml_tensor * iq4_xs_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -8593,8 +6075,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(iq4_xs_decode_input != nullptr); ggml_tensor * iq4_xs_decode_output = ggml_mul_mat(ctx, iq4_xs_decode_weight, iq4_xs_decode_input); REQUIRE(iq4_xs_decode_output != nullptr); - schedule_single_matmul_command(ctx, iq4_xs_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, - 2048, 128); + schedule_single_matmul_command(ctx, iq4_xs_decode_output, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 2048, 128); ggml_tensor * q8_0_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, 2048, 128); ggml_tensor * q8_0_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 4); @@ -8602,8 +6083,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q8_0_input != nullptr); ggml_tensor * q8_0_output = ggml_mul_mat(ctx, q8_0_weight, q8_0_input); REQUIRE(q8_0_output != nullptr); - schedule_single_matmul_command(ctx, q8_0_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, - 128); + schedule_single_matmul_command(ctx, q8_0_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); ggml_tensor * q8_0_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, 2048, 128); ggml_tensor * q8_0_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -8638,8 +6118,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(q8_1_input != nullptr); ggml_tensor * q8_1_output = ggml_mul_mat(ctx, q8_1_weight, q8_1_input); REQUIRE(q8_1_output != nullptr); - schedule_single_matmul_command(ctx, q8_1_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, - 128); + schedule_single_matmul_command(ctx, q8_1_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); ggml_tensor * q8_1_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_1, 2048, 128); ggml_tensor * q8_1_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -8656,8 +6135,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(f16_input != nullptr); ggml_tensor * f16_output = ggml_mul_mat(ctx, f16_weight, f16_input); REQUIRE(f16_output != nullptr); - schedule_single_matmul_command(ctx, f16_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, - 128); + schedule_single_matmul_command(ctx, f16_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); ggml_tensor * f16_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 2048, 128); ggml_tensor * f16_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -8674,8 +6152,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(f32_input != nullptr); ggml_tensor * f32_output = ggml_mul_mat(ctx, f32_weight, f32_input); REQUIRE(f32_output != nullptr); - schedule_single_matmul_command(ctx, f32_output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, - 128); + schedule_single_matmul_command(ctx, f32_output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); ggml_tensor * f32_decode_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 128); ggml_tensor * f32_decode_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); @@ -8729,8 +6206,7 @@ static void run_qwen_matmul_dispatch_checks() { REQUIRE(input != nullptr); ggml_tensor * output = ggml_mul_mat(ctx, weight, input); REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, 2048, - 128); + schedule_single_matmul_command(ctx, output, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 2048, 128); } { @@ -8766,148 +6242,6 @@ static void run_qwen_matmul_dispatch_checks() { ggml_free(ctx); } -static void run_iq3_xxs_codebook_matmul_dispatch_checks() { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - auto check_plan = [&](int64_t token_count, const char * expected_kernel, size_t expected_binding_count, - const char * weight_format_parameter) { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ3_XXS, 2048, 128); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, token_count); - REQUIRE(weight != nullptr); - REQUIRE(input != nullptr); - ggml_tensor * output = ggml_mul_mat(ctx, weight, input); - REQUIRE(output != nullptr); - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.dispatches.size() == 1); - REQUIRE(kernel_name_for_id(plan.dispatches[0].kernel.kernel_id) == expected_kernel); - REQUIRE(plan.dispatches[0].bindings.size() == expected_binding_count); - REQUIRE(plan.dispatches[0].bindings[2].length == 1024); - REQUIRE(plan.transients.size() == 1); - REQUIRE(plan.constant_initializations.size() == 1); - REQUIRE(plan.constant_initializations[0].value == plan.dispatches[0].bindings[2].value); - REQUIRE(plan.constant_initializations[0].data.size() == 1024); - require_compile_parameter(plan.dispatches[0], weight_format_parameter, "18"); - }; - - check_plan(1, "loom_libs:ggml_mul_mat_vector_iq3_xxs_f32_f32", 5, "ggml.matmul.vector.weight_format"); - check_plan(64, "loom_libs:ggml_mul_mat_tiled_input_f32_iq3_xxs_publish_f32", 4, - "ggml.mul_mat.weight_format"); - - const char * disable_name = "GGML_HRX_DISABLE_IQ3_XXS_CODEBOOK_MATMUL"; - REQUIRE(setenv(disable_name, "1", 1) == 0); - ggml_tensor * disabled_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_IQ3_XXS, 2048, 128); - ggml_tensor * disabled_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); - REQUIRE(disabled_weight != nullptr); - REQUIRE(disabled_input != nullptr); - ggml_tensor * disabled_output = ggml_mul_mat(ctx, disabled_weight, disabled_input); - REQUIRE(disabled_output != nullptr); - ggml_cgraph * disabled_graph = ggml_new_graph(ctx); - REQUIRE(disabled_graph != nullptr); - ggml_build_forward_expand(disabled_graph, disabled_output); - ggml::hrx::GraphImportResult disabled_imported = ggml::hrx::import_ggml_graph(*disabled_graph); - REQUIRE(disabled_imported.valid()); - ggml::hrx::DispatchScheduler disabled_scheduler; - REQUIRE(!disabled_scheduler.schedule_graph(disabled_imported.graph, test_dispatch_target())); - REQUIRE(unsetenv(disable_name) == 0); - - ggml_free(ctx); -} - -static void run_iq2_codebook_resource_reuse_checks() { - for (const ggml_type type : { GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS }) { - ggml_init_params params = {}; - params.mem_size = 8 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * vector_weight = ggml_new_tensor_2d(ctx, type, 2048, 128); - ggml_tensor * vector_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 1); - ggml_tensor * tiled_weight = ggml_new_tensor_2d(ctx, type, 2048, 128); - ggml_tensor * tiled_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, 64); - REQUIRE(vector_weight != nullptr); - REQUIRE(vector_input != nullptr); - REQUIRE(tiled_weight != nullptr); - REQUIRE(tiled_input != nullptr); - ggml_tensor * vector_output = ggml_mul_mat(ctx, vector_weight, vector_input); - ggml_tensor * tiled_output = ggml_mul_mat(ctx, tiled_weight, tiled_input); - REQUIRE(vector_output != nullptr); - REQUIRE(tiled_output != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, vector_output); - ggml_build_forward_expand(graph, tiled_output); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - REQUIRE(scheduler.plan().valid()); - REQUIRE(scheduler.plan().dispatches.size() == 2); - REQUIRE(scheduler.plan().constant_initializations.size() == 1); - REQUIRE(scheduler.plan().transients.size() == 1); - - const char * expected_key = - type == GGML_TYPE_IQ2_XXS ? "common.iq2_xxs.grid.v1" : "common.iq2_xs.grid.v1"; - const ggml::hrx::CommandPlanConstantInitialization & initialization = - scheduler.plan().constant_initializations.front(); - REQUIRE(initialization.name == expected_key); - REQUIRE(scheduler.plan().transients.front().name == expected_key); - REQUIRE(scheduler.plan().transients.front().value == initialization.value); - REQUIRE(std::count_if(scheduler.plan().dispatches.begin(), scheduler.plan().dispatches.end(), - [&](const ggml::hrx::Dispatch & dispatch) { - return std::any_of(dispatch.bindings.begin(), dispatch.bindings.end(), - [&](const ggml::hrx::DispatchBinding & binding) { - return binding.value == initialization.value; - }); - }) == 2); - - const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program( - imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.constant_initializations.size() == 1); - ggml_free(ctx); - } -} - -static void run_iq4_matmul_dispatch_checks() { - ggml_init_params params = {}; - params.mem_size = 2 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - const auto check = [ctx](ggml_type weight_type, int64_t input_size, int64_t output_size, - int64_t token_count, const char * expected_kernel) { - ggml_tensor * weight = ggml_new_tensor_2d(ctx, weight_type, input_size, output_size); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); - REQUIRE(weight != nullptr); - REQUIRE(input != nullptr); - ggml_tensor * output = ggml_mul_mat(ctx, weight, input); - REQUIRE(output != nullptr); - schedule_single_matmul_command(ctx, output, expected_kernel, token_count, input_size, output_size); - }; - - check(GGML_TYPE_IQ4_NL, 640, 1024, 1, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64"); - check(GGML_TYPE_IQ4_NL, 640, 1024, 64, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_aligned"); - check(GGML_TYPE_IQ4_XS, 2048, 128, 1, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64"); - check(GGML_TYPE_IQ4_XS, 2048, 128, 64, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_aligned"); - - ggml_free(ctx); -} - static void run_quantized_value_projection_dispatch_checks() { ggml_init_params params = {}; params.mem_size = 16 * 1024 * 1024; @@ -8927,152 +6261,24 @@ static void run_quantized_value_projection_dispatch_checks() { REQUIRE(imported.valid()); ggml::hrx::DispatchScheduler scheduler; REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const auto & plan = scheduler.plan(); - REQUIRE(plan.valid()); - REQUIRE(plan.dispatches.size() == 1 || plan.dispatches.size() == 2); - const auto & dispatch = plan.dispatches.back(); - const std::string kernel_name = kernel_name_for_id(dispatch.kernel.kernel_id); - if (plan.dispatches.size() == 1) { - REQUIRE(kernel_name == "loom_libs:llm_attention_v_matmul_set_rows_vector_f32_f32"); - require_compile_parameter(dispatch, "llm.attention_qkv.weight_format", - std::to_string(matmul_weight_format_config(weight_type))); - } else { - REQUIRE(kernel_name_for_id(plan.dispatches.front().kernel.kernel_id) == - "qwen3_moe:ggml_quantize_q8_1_x4_f32"); - REQUIRE(kernel_name == "loom_libs:llm_attention_v_matmul_set_rows_legacy_f32_f32"); - require_compile_parameter(dispatch, "llm.attention_qkv.weight_format", - weight_type == GGML_TYPE_Q4_K ? "44" : "46"); - require_compile_parameter(dispatch, "ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_Q8_1)); - REQUIRE(!dispatch.bindings[1].layout.empty()); - } - require_compile_parameter(dispatch, "llm.attention_qkv.cache_output_format", - cache_type == GGML_TYPE_F16 ? "16" : "32"); - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - } - } - ggml_free(ctx); -} - -static void run_prefill_value_projection_cache_dispatch_checks() { - static constexpr const char * kDisableFusionEnv = "GGML_HRX_DISABLE_PREFILL_V_CACHE_FUSION"; - EnvironmentVariableGuard disable_fusion(kDisableFusionEnv); - disable_fusion.unset(); - - auto check_case = [&](ggml_type weight_type, int64_t input_size, int64_t output_size, int64_t token_count, - ggml_type cache_type, int64_t cache_row_count, bool publish_f16_alternate, - bool set_rows_available, bool disable, bool expect_fusion, bool observe_projection = false, - bool alias_weight = false) { - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * weight_storage = ggml_new_tensor_2d(ctx, weight_type, input_size, output_size); - REQUIRE(weight_storage != nullptr); - ggml_tensor * weight = - alias_weight ? ggml_view_2d(ctx, weight_storage, input_size, output_size, weight_storage->nb[1], 0) : - weight_storage; - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); - ggml_tensor * cache = ggml_new_tensor_2d(ctx, cache_type, output_size, cache_row_count); - ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); - REQUIRE(weight != nullptr); - REQUIRE(input != nullptr); - REQUIRE(cache != nullptr); - REQUIRE(indices != nullptr); - ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); - ggml_tensor * value = ggml_reshape_3d(ctx, projection, 128, output_size / 128, token_count); - ggml_tensor * rows = ggml_reshape_2d(ctx, value, output_size, token_count); - ggml_tensor * output = ggml_set_rows(ctx, cache, rows, indices); - REQUIRE(projection != nullptr); - REQUIRE(value != nullptr); - REQUIRE(rows != nullptr); - REQUIRE(output != nullptr); - if (observe_projection) { - ggml_set_output(projection); - } - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - - ggml::hrx::CommandPlan plan; - const ggml::hrx::Value * input_value = imported.graph.values().find_tensor(input); - REQUIRE(input_value != nullptr); - if (publish_f16_alternate) { - const size_t bytes = static_cast(input_size * token_count) * sizeof(ggml_fp16_t); - const ggml::hrx::ValueId alternate(static_cast(imported.graph.values().size())); - plan.transients.push_back({ alternate, "test.prefill_v_cache.f16", bytes, 256 }); - plan.constant_initializations.push_back( - { alternate, "test.prefill_v_cache.f16", 0, std::vector(bytes) }); - ggml::hrx::Status status; - REQUIRE(plan.metadata.append_alternate_value( - { input_value->id, alternate, GGML_TYPE_F16, bytes, "test.prefill_v_cache.f16" }, status)); - } - - std::vector covered_nodes(imported.graph.nodes().size(), false); - if (!set_rows_available) { - covered_nodes[producer_index_for_tensor(imported.graph, output)] = true; - } - if (disable) { - disable_fusion.set("1"); - } else { - disable_fusion.unset(); - } - - ggml::hrx::DispatchMatch match; - const size_t projection_index = producer_index_for_tensor(imported.graph, projection); - REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, projection_index, match)); - REQUIRE(!match.dispatches.empty()); - const ggml::hrx::Dispatch & dispatch = match.dispatches.back(); - const bool fused = - kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:llm_attention_v_matmul_set_rows_tiled_f32_f32"; - REQUIRE(fused == expect_fusion); - if (expect_fusion) { - REQUIRE(match.dispatches.size() == 1); - REQUIRE(match.covered_nodes.size() == 4); - REQUIRE(dispatch.bindings.size() == 4); - REQUIRE(dispatch.bindings[0].value == plan.transients.front().value); - REQUIRE(dispatch.bindings[0].length == plan.transients.front().size); - require_compile_parameter(dispatch, "ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_F16)); - require_compile_parameter(dispatch, "llm.attention_qkv.weight_format", "6"); - require_compile_parameter(dispatch, "llm.attention_qkv.input_size", "1536"); - require_compile_parameter(dispatch, "llm.attention_qkv.output_size", "256"); - require_compile_parameter(dispatch, "llm.attention_qkv.cache_row_count", "512"); - require_compile_parameter(dispatch, "llm.attention_qkv.cache_output_format", "16"); - - append_match_to_plan(plan, match, covered_nodes, &imported.graph); - const ggml::hrx::CommandProgram commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto & plan = scheduler.plan(); + REQUIRE(plan.valid()); + REQUIRE(plan.dispatches.size() == 2); + REQUIRE(kernel_name_for_id(plan.dispatches.front().kernel.kernel_id) == + "qwen3_moe:ggml_quantize_q8_1_x4_f32"); + const auto & dispatch = plan.dispatches.back(); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:llm_attention_v_matmul_set_rows_f32_f32_wmma"); + require_compile_parameter(dispatch, "llm.attention_qkv.weight_format", weight_type == GGML_TYPE_Q4_K ? "44" : "46"); + require_compile_parameter(dispatch, "ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_Q8_1)); + require_compile_parameter(dispatch, "llm.attention_qkv.cache_output_format", cache_type == GGML_TYPE_F16 ? "16" : "32"); + REQUIRE(!dispatch.bindings[1].layout.empty()); + const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(command_program_verifies(commands)); - REQUIRE(commands.commands.size() == 1); - REQUIRE(commands.commands.front().bindings.size() == 4); - REQUIRE(commands.commands.front().bindings[0].name == "input"); - REQUIRE(commands.commands.front().bindings[1].name == "weight"); - REQUIRE(commands.commands.front().bindings[2].name == "indices"); - REQUIRE(commands.commands.front().bindings[3].name == "cache"); } - ggml_free(ctx); - }; - - check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, false, true); - check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, true, false); - check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, false, true, false, false); - check_case(GGML_TYPE_Q4_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, false, false); - check_case(GGML_TYPE_Q6_K, 2048, 256, 64, GGML_TYPE_F16, 512, true, true, false, false); - check_case(GGML_TYPE_Q6_K, 1536, 512, 64, GGML_TYPE_F16, 512, true, true, false, false); - check_case(GGML_TYPE_Q6_K, 1536, 256, 32, GGML_TYPE_F16, 512, true, true, false, false); - check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F32, 512, true, true, false, false); - check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 1024, true, true, false, false); - check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, false, false, false); - check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, false, false, true); - check_case(GGML_TYPE_Q6_K, 1536, 256, 64, GGML_TYPE_F16, 512, true, true, false, false, false, true); + } + ggml_free(ctx); } static void run_llama_attention_matmul_dispatch_checks() { @@ -9108,9 +6314,8 @@ static void run_llama_attention_matmul_dispatch_checks() { REQUIRE(query != nullptr); REQUIRE(output != nullptr); schedule_fused_llama_attention_matmul_command( - ctx, output, "loom_libs:llm_attention_q_matmul_rope_tiled_f32_f32", weight_type, token_count, - input_size, output_size, head_size, head_count, 0, 0, 3, - { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_ROPE }, + ctx, output, "loom_libs:llm_attention_q_matmul_rope_f32_f32_wmma", weight_type, token_count, input_size, + output_size, head_size, head_count, 0, 0, 3, { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_ROPE }, { "input", "weight", "positions", "theta", "freq_factors", "output" }); } @@ -9139,7 +6344,7 @@ static void run_llama_attention_matmul_dispatch_checks() { REQUIRE(rows != nullptr); REQUIRE(output != nullptr); schedule_fused_llama_attention_matmul_command( - ctx, output, "loom_libs:llm_attention_k_matmul_rope_set_rows_tiled_f32_f32", weight_type, token_count, + ctx, output, "loom_libs:llm_attention_k_matmul_rope_set_rows_f32_f32_wmma", weight_type, token_count, input_size, output_size, head_size, head_count, cache_row_count, 16, 5, { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_ROPE, GGML_OP_RESHAPE, GGML_OP_SET_ROWS }, { "input", "weight", "positions", "indices", "theta", "freq_factors", "cache" }); @@ -9163,7 +6368,7 @@ static void run_llama_attention_matmul_dispatch_checks() { REQUIRE(rows != nullptr); REQUIRE(output != nullptr); schedule_fused_llama_attention_matmul_command( - ctx, output, "loom_libs:llm_attention_v_matmul_set_rows_tiled_f32_f32", weight_type, token_count, + ctx, output, "loom_libs:llm_attention_v_matmul_set_rows_f32_f32_wmma", weight_type, token_count, input_size, output_size, 0, 0, cache_row_count, 16, 4, { GGML_OP_MUL_MAT, GGML_OP_RESHAPE, GGML_OP_RESHAPE, GGML_OP_SET_ROWS }, { "input", "weight", "indices", "cache" }); @@ -9498,11 +6703,10 @@ static QwenRoutedGateUpTensors build_qwen_routed_gate_up_graph(ggml_context * ct static void append_qwen_weighted_reduce_tail(ggml_context * ctx, QwenRoutedGateUpTensors & tensors, - bool include_next_rmsnorm = false, - ggml_tensor * routed_input = nullptr) { + bool include_next_rmsnorm = false) { REQUIRE(tensors.output != nullptr); REQUIRE(tensors.route_weights != nullptr); - tensors.weighted = ggml_mul(ctx, routed_input != nullptr ? routed_input : tensors.output, tensors.route_weights); + tensors.weighted = ggml_mul(ctx, tensors.output, tensors.route_weights); REQUIRE(tensors.weighted != nullptr); tensors.route_views.clear(); for (int64_t route = 0; route < kQwenRouterRouteCount; ++route) { @@ -9569,7 +6773,6 @@ static bool match_dispatch_at_index(const ggml::hrx::Graph & graph, plan, ggml::hrx::ValueId(static_cast(graph.values().size() + plan.transients.size() + plan.completion_counter_requests.size())), - &test_dispatch_registry(), }; return test_dispatch_registry().match(context, match); } @@ -9716,25 +6919,6 @@ static std::string common_mul_mat_weight_format(ggml_type type) { } } -static const char * common_mul_mat_id_f32_kernel_name(int64_t token_count) { - return token_count <= 5 ? "loom_libs:ggml_mul_mat_id_skinny_input_f32_publish_f32" : - "loom_libs:ggml_mul_mat_id_tiled_input_f32_publish_f32"; -} - -static const char * common_mul_mat_id_postops_f32_kernel_name(int64_t token_count, bool has_rmsnorm) { - if (has_rmsnorm) { - return token_count <= 5 ? "loom_libs:ggml_mul_mat_id_skinny_input_f32_postops_next_rmsnorm_publish_f32" : - "loom_libs:ggml_mul_mat_id_tiled_input_f32_postops_next_rmsnorm_publish_f32"; - } - return token_count <= 5 ? "loom_libs:ggml_mul_mat_id_skinny_input_f32_postops_publish_f32" : - "loom_libs:ggml_mul_mat_id_tiled_input_f32_postops_publish_f32"; -} - -static const char * common_mul_mat_id_swiglu_f32_kernel_name(int64_t token_count) { - return token_count <= 5 ? "loom_libs:ggml_mul_mat_id_skinny_pair_input_f32_swiglu_publish_f32" : - "loom_libs:ggml_mul_mat_id_tiled_pair_input_f32_swiglu_publish_f32"; -} - static ggml::hrx::DispatchMatch require_common_mul_mat_id_for_graph(ggml::hrx::Graph & graph, ggml::hrx::CommandPlan & plan, std::vector & covered_nodes, @@ -9759,7 +6943,7 @@ static ggml::hrx::DispatchMatch require_common_mul_mat_id_for_graph(ggml::hrx::G REQUIRE(match_dispatch_at_index(graph, plan, covered_nodes, output_index, match)); REQUIRE(match.dispatches.size() == 1); const ggml::hrx::Dispatch & dispatch = match.dispatches[0]; - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == common_mul_mat_id_f32_kernel_name(output->ne[2])); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_id_f32_f32_wmma"); REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == output->ne[2]); REQUIRE(dispatch.bindings.size() == 5); REQUIRE(dispatch.bindings[0].value == input_value->id); @@ -9797,9 +6981,9 @@ static CommonMulMatIdSwiGLUTensors build_common_mul_mat_id_swiglu_graph(ggml_con ggml_type up_weight_type, ggml_glu_op glu_op = GGML_GLU_OP_SWIGLU, int64_t input_size = 256, - int64_t output_size = 128, - int64_t token_count = 4) { + int64_t output_size = 128) { CommonMulMatIdSwiGLUTensors tensors; + constexpr int64_t token_count = 4; constexpr int64_t route_count = 8; constexpr int64_t input_route_count = 1; constexpr int64_t expert_count = 16; @@ -9846,8 +7030,7 @@ static void require_common_mul_mat_id_swiglu_match(ggml_context * REQUIRE(match.metadata.moe_routing_bundles().size() == 1); const ggml::hrx::Dispatch & dispatch = match.dispatches.back(); - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == - common_mul_mat_id_swiglu_f32_kernel_name(tensors.output->ne[2])); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_id_swiglu_f32_f32_wmma"); REQUIRE(dispatch.kernel.integer_parameters.at("token_count") == tensors.output->ne[2]); REQUIRE(dispatch.bindings.size() == 6); require_compile_parameter(dispatch, "ggml.mul_mat_id_swiglu.input_size", @@ -9958,7 +7141,7 @@ static void run_common_mul_mat_id_swiglu_dispatch_checks() { ggml::hrx::DispatchMatch match; REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, match)); REQUIRE(kernel_name_for_id(match.dispatches.back().kernel.kernel_id) == - common_mul_mat_id_f32_kernel_name(tensors.output->ne[2])); + "loom_libs:ggml_mul_mat_id_f32_f32_wmma"); } ggml_free(ctx); @@ -9990,9 +7173,9 @@ static CommonMulMatIdPostOpsTensors build_common_mul_mat_id_postops_graph(ggml_c bool include_bias, bool include_residual, bool include_rmsnorm, - ggml_type weight_type = GGML_TYPE_Q4_K, - int64_t token_count = 4) { + ggml_type weight_type = GGML_TYPE_Q4_K) { CommonMulMatIdPostOpsTensors tensors; + constexpr int64_t token_count = 4; constexpr int64_t route_count = 8; constexpr int64_t input_size = 256; constexpr int64_t output_size = 128; @@ -10127,6 +7310,7 @@ static void require_common_mul_mat_id_postops_match(ggml_context * ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(commands.commands.size() == 3); + REQUIRE(command_program_verifies(commands)); } static void run_common_mul_mat_id_postops_dispatch_checks() { @@ -10138,39 +7322,34 @@ static void run_common_mul_mat_id_postops_dispatch_checks() { { const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, true, false, false); - require_common_mul_mat_id_postops_match( - ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], false), true, false, - false); + require_common_mul_mat_id_postops_match(ctx, tensors, "loom_libs:ggml_mul_mat_id_postops_f32_f32_wmma", true, + false, false); } { const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, false, true, false); - require_common_mul_mat_id_postops_match( - ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], false), false, true, - false); + require_common_mul_mat_id_postops_match(ctx, tensors, "loom_libs:ggml_mul_mat_id_postops_f32_f32_wmma", false, + true, false); } { const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, true, true, false); - require_common_mul_mat_id_postops_match( - ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], false), true, true, - false); + require_common_mul_mat_id_postops_match(ctx, tensors, "loom_libs:ggml_mul_mat_id_postops_f32_f32_wmma", true, + true, false); } { const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, true, true, false, GGML_TYPE_BF16); - require_common_mul_mat_id_postops_match( - ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], false), true, true, - false); + require_common_mul_mat_id_postops_match(ctx, tensors, "loom_libs:ggml_mul_mat_id_postops_f32_f32_wmma", true, + true, false); } { const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, false, true, true); require_common_mul_mat_id_postops_match( - ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], true), false, true, - true); + ctx, tensors, "loom_libs:ggml_mul_mat_id_postops_next_rmsnorm_f32_f32_wmma", false, true, true); } { const CommonMulMatIdPostOpsTensors tensors = build_common_mul_mat_id_postops_graph(ctx, true, true, true); require_common_mul_mat_id_postops_match( - ctx, tensors, common_mul_mat_id_postops_f32_kernel_name(tensors.projection->ne[2], true), true, true, true); + ctx, tensors, "loom_libs:ggml_mul_mat_id_postops_next_rmsnorm_f32_f32_wmma", true, true, true); } ggml_free(ctx); @@ -10276,6 +7455,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(commands.commands.size() == 4); + REQUIRE(command_program_verifies(commands)); REQUIRE(commands.commands[3].bindings[0].name == "input"); REQUIRE(commands.commands[3].bindings[2].name == "partition_table"); REQUIRE(commands.commands[3].bindings[3].name == "weight"); @@ -10298,13 +7478,14 @@ static void run_qwen_routed_gate_up_dispatch_checks() { ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(commands.commands.size() == 4); + REQUIRE(command_program_verifies(commands)); } { constexpr int64_t token_count = 4; - QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); - append_qwen_weighted_reduce_tail(ctx, tensors); + const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); + REQUIRE(imported.graph.nodes().size() == 14); const ggml::hrx::Value * route_ids_value = imported.graph.values().find_tensor(tensors.route_ids); const ggml::hrx::Value * route_weights_value = imported.graph.values().find_tensor(tensors.route_weights); @@ -10313,7 +7494,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(route_weights_value != nullptr); REQUIRE(glu_value != nullptr); REQUIRE(route_ids_value->kind == ggml::hrx::ValueKind::Transient); - REQUIRE(route_weights_value->kind == ggml::hrx::ValueKind::Transient); + REQUIRE(route_weights_value->kind == ggml::hrx::ValueKind::External); REQUIRE(glu_value->kind == ggml::hrx::ValueKind::Transient); std::vector covered_nodes(imported.graph.nodes().size(), false); @@ -10403,8 +7584,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { { constexpr int64_t token_count = 1; - QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); - append_qwen_weighted_reduce_tail(ctx, tensors); + const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); std::vector covered_nodes(imported.graph.nodes().size(), false); ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); @@ -10434,13 +7614,12 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(first_route_ids != nullptr); REQUIRE(first_route_weights != nullptr); - QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); - append_qwen_weighted_reduce_tail(ctx, tensors); + const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); ggml_build_forward_expand(graph, first_route_weights); ggml_build_forward_expand(graph, tensors.route_weights); - ggml_build_forward_expand(graph, tensors.residual); + ggml_build_forward_expand(graph, tensors.output); ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); REQUIRE(imported.valid()); @@ -10510,20 +7689,17 @@ static void run_qwen_routed_gate_up_dispatch_checks() { std::vector covered_nodes(imported.graph.nodes().size(), false); ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); append_qwen_routed_gate_up_for_graph(imported.graph, tensors, covered_nodes, plan); - REQUIRE(kernel_name_for_id(plan.dispatches.back().kernel.kernel_id) == - common_mul_mat_id_swiglu_f32_kernel_name(token_count)); - ggml::hrx::DispatchMatch down_match = require_common_mul_mat_id_for_graph( - imported.graph, plan, covered_nodes, tensors.output, tensors.glu, tensors.down_weight, tensors.route_ids); - append_match_to_plan(plan, down_match, covered_nodes, &imported.graph); + append_qwen_routed_down_for_graph(imported.graph, tensors, covered_nodes, plan, + "loom_libs:ggml_mul_mat_id_f16_f16_wmma"); REQUIRE(plan.dispatches.size() == 5); - REQUIRE(plan.transients.size() == 2); - REQUIRE(plan.metadata.alternate_values().empty()); + REQUIRE(plan.transients.size() == 4); const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); REQUIRE(commands.commands.size() == 5); REQUIRE(command_program_verifies(commands)); + REQUIRE(ggml::hrx::find_transient_allocation(commands.transients, plan.transients.back().value) != nullptr); } { @@ -10577,66 +7753,6 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(command_program_verifies(commands)); } - { - constexpr int64_t token_count = 4; - QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); - ggml_tensor * glu_reshape = ggml_reshape_4d(ctx, tensors.glu, tensors.glu->ne[0], tensors.glu->ne[1], - tensors.glu->ne[2], tensors.glu->ne[3]); - REQUIRE(glu_reshape != nullptr); - tensors.output = ggml_mul_mat_id(ctx, tensors.down_weight, glu_reshape, tensors.route_ids); - REQUIRE(tensors.output != nullptr); - ggml_tensor * routed_output_view = - ggml_view_4d(ctx, tensors.output, tensors.output->ne[0], tensors.output->ne[1], tensors.output->ne[2], - tensors.output->ne[3], tensors.output->nb[1], tensors.output->nb[2], tensors.output->nb[3], 0); - REQUIRE(routed_output_view != nullptr); - append_qwen_weighted_reduce_tail(ctx, tensors, false, routed_output_view); - - ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); - const ggml::hrx::Value * glu_reshape_value = imported.graph.values().find_tensor(glu_reshape); - const ggml::hrx::Value * routed_output_view_value = imported.graph.values().find_tensor(routed_output_view); - REQUIRE(glu_reshape_value != nullptr); - REQUIRE(routed_output_view_value != nullptr); - const ggml::hrx::GraphNode * glu_reshape_node = imported.graph.index().producer(glu_reshape_value->id); - const ggml::hrx::GraphNode * routed_output_view_node = - imported.graph.index().producer(routed_output_view_value->id); - REQUIRE(glu_reshape_node != nullptr); - REQUIRE(routed_output_view_node != nullptr); - REQUIRE(glu_reshape_node->op == GGML_OP_RESHAPE); - REQUIRE(routed_output_view_node->op == GGML_OP_VIEW); - std::vector covered_nodes(imported.graph.nodes().size(), false); - ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); - append_qwen_routed_gate_up_for_graph(imported.graph, tensors, covered_nodes, plan); - const ggml::hrx::Value * glu_value = imported.graph.values().find_tensor(tensors.glu); - REQUIRE(glu_value != nullptr); - const ggml::hrx::CommandPlanAlternateValue * gate_up_alternate = ggml::hrx::find_alternate_value( - imported.graph, plan, glu_value->id, GGML_TYPE_F16, qwen_routed_gate_up_f16_output_size(token_count)); - REQUIRE(gate_up_alternate != nullptr); - const ggml::hrx::ValueId gate_up_alternate_value = gate_up_alternate->alternate_value; - - append_qwen_routed_down_for_graph(imported.graph, tensors, covered_nodes, plan, - "loom_libs:ggml_mul_mat_id_f16_f16_wmma"); - REQUIRE(plan.dispatches.back().bindings[0].value == gate_up_alternate_value); - const ggml::hrx::Value * routed_output_value = imported.graph.values().find_tensor(tensors.output); - REQUIRE(routed_output_value != nullptr); - const ggml::hrx::CommandPlanAlternateValue * routed_down_alternate = ggml::hrx::find_alternate_value( - imported.graph, plan, routed_output_value->id, GGML_TYPE_F16, - qwen_routed_down_f16_output_size(token_count)); - REQUIRE(routed_down_alternate != nullptr); - const ggml::hrx::ValueId routed_down_alternate_value = routed_down_alternate->alternate_value; - - append_qwen_weighted_reduce_for_graph(imported.graph, tensors, covered_nodes, plan, - "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_f16_f32"); - REQUIRE(plan.dispatches.back().bindings[1].value == routed_down_alternate_value); - for (const ggml::hrx::Dispatch & dispatch : plan.dispatches) { - REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) != "loom_libs:ggml_copy_f32_f16"); - } - - const ggml::hrx::CommandProgram commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(commands.valid()); - REQUIRE(command_program_verifies(commands)); - } - { constexpr int64_t token_count = 4; QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); @@ -10711,9 +7827,8 @@ static void run_qwen_routed_gate_up_dispatch_checks() { { constexpr int64_t token_count = 4; - QwenRoutedGateUpTensors tensors = + const QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count, GGML_GLU_OP_SWIGLU, GGML_TYPE_Q4_K, true, GGML_TYPE_Q4_K); - append_qwen_weighted_reduce_tail(ctx, tensors); ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); std::vector covered_nodes(imported.graph.nodes().size(), false); ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); @@ -10745,7 +7860,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(kernel_name_for_id(gate_up_match.dispatches[1].kernel.kernel_id) == "loom_libs:ggml_moe_build_expert_partition_table"); REQUIRE(kernel_name_for_id(gate_up_match.dispatches[2].kernel.kernel_id) == - common_mul_mat_id_swiglu_f32_kernel_name(tensors.output->ne[2])); + "loom_libs:ggml_mul_mat_id_swiglu_f32_f32_wmma"); REQUIRE(gate_up_match.dispatches[2].bindings.size() == 6); REQUIRE(gate_up_match.dispatches[2].bindings[1].value == gate_up_match.transients[0].value); REQUIRE(gate_up_match.dispatches[2].bindings[2].value == gate_up_match.transients[1].value); @@ -10763,7 +7878,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); REQUIRE(gate_up_match.dispatches.size() == 1); REQUIRE(kernel_name_for_id(gate_up_match.dispatches[0].kernel.kernel_id) == - common_mul_mat_id_f32_kernel_name(tensors.gate->ne[2])); + "loom_libs:ggml_mul_mat_id_f32_f32_wmma"); } { @@ -10777,7 +7892,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); REQUIRE(gate_up_match.dispatches.size() == 1); REQUIRE(kernel_name_for_id(gate_up_match.dispatches[0].kernel.kernel_id) == - common_mul_mat_id_swiglu_f32_kernel_name(tensors.output->ne[2])); + "loom_libs:ggml_mul_mat_id_swiglu_f32_f32_wmma"); REQUIRE(gate_up_match.dispatches[0].bindings.size() == 6); } @@ -10792,7 +7907,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, gate_index, gate_up_match)); REQUIRE(gate_up_match.dispatches.size() == 1); REQUIRE(kernel_name_for_id(gate_up_match.dispatches[0].kernel.kernel_id) == - common_mul_mat_id_swiglu_f32_kernel_name(tensors.output->ne[2])); + "loom_libs:ggml_mul_mat_id_swiglu_f32_f32_wmma"); REQUIRE(gate_up_match.dispatches[0].bindings.size() == 6); } @@ -10806,7 +7921,7 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, down_index, down_match)); REQUIRE(down_match.dispatches.size() == 1); REQUIRE(kernel_name_for_id(down_match.dispatches[0].kernel.kernel_id) == - common_mul_mat_id_f32_kernel_name(tensors.output->ne[2])); + "loom_libs:ggml_mul_mat_id_f32_f32_wmma"); } { @@ -10821,95 +7936,12 @@ static void run_qwen_routed_gate_up_dispatch_checks() { REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, down_index, down_match)); REQUIRE(down_match.dispatches.size() == 1); REQUIRE(kernel_name_for_id(down_match.dispatches[0].kernel.kernel_id) == - common_mul_mat_id_f32_kernel_name(tensors.output->ne[2])); + "loom_libs:ggml_mul_mat_id_f32_f32_wmma"); } ggml_free(ctx); } -static void run_qwen_weighted_reduce_next_rmsnorm_q8_publication_checks() { - enum class ConsumerCase { - Direct, - Alias, - UnsupportedWeight, - WrongOperand, - IncompatibleGeometry, - }; - const struct { - ConsumerCase consumer; - bool publishes_q8; - } cases[] = { - { ConsumerCase::Direct, true }, - { ConsumerCase::Alias, true }, - { ConsumerCase::UnsupportedWeight, false }, - { ConsumerCase::WrongOperand, false }, - { ConsumerCase::IncompatibleGeometry, false }, - }; - - for (const auto & test : cases) { - ggml_init_params params = {}; - params.mem_size = 128 * 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - constexpr int64_t token_count = 2; - QwenRoutedGateUpTensors tensors = build_qwen_routed_gate_up_graph(ctx, token_count); - append_qwen_weighted_reduce_tail(ctx, tensors, true); - ggml_tensor * consumer_input = tensors.next_output; - if (test.consumer == ConsumerCase::Alias) { - consumer_input = ggml_reshape_2d(ctx, tensors.next_output, kQwenMoeHiddenSize, token_count); - REQUIRE(consumer_input != nullptr); - } - - if (test.consumer == ConsumerCase::WrongOperand) { - ggml_tensor * activation = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kQwenMoeHiddenSize, 64); - REQUIRE(activation != nullptr); - tensors.hidden_use = ggml_mul_mat(ctx, tensors.next_output, activation); - } else { - const ggml_type weight_type = - test.consumer == ConsumerCase::UnsupportedWeight ? GGML_TYPE_BF16 : GGML_TYPE_Q4_K; - const int64_t output_size = test.consumer == ConsumerCase::IncompatibleGeometry ? 63 : 512; - ggml_tensor * weight = ggml_new_tensor_2d(ctx, weight_type, kQwenMoeHiddenSize, output_size); - REQUIRE(weight != nullptr); - tensors.hidden_use = ggml_mul_mat(ctx, weight, consumer_input); - } - REQUIRE(tensors.hidden_use != nullptr); - - ggml::hrx::GraphImportResult imported = import_qwen_routed_gate_up_graph(ctx, tensors); - std::vector covered_nodes(imported.graph.nodes().size(), false); - ggml::hrx::CommandPlan plan = build_qwen_router_plan_for_graph(imported.graph, covered_nodes); - append_qwen_routed_gate_up_for_graph(imported.graph, tensors, covered_nodes, plan); - append_qwen_routed_down_for_graph(imported.graph, tensors, covered_nodes, plan, - "loom_libs:ggml_mul_mat_id_f16_f16_wmma"); - - const size_t weighted_index = producer_index_for_tensor(imported.graph, tensors.weighted); - ggml::hrx::DispatchMatch weighted_match; - REQUIRE(match_dispatch_at_index(imported.graph, plan, covered_nodes, weighted_index, weighted_match)); - const char * expected_kernel = - test.publishes_q8 ? - "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8" : - "qwen3_moe:qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32"; - REQUIRE(kernel_name_for_id(weighted_match.dispatches.back().kernel.kernel_id) == expected_kernel); - REQUIRE(weighted_match.dispatches.back().bindings.size() == (test.publishes_q8 ? 6 : 5)); - - const ggml::hrx::Value * next_output_value = imported.graph.values().find_tensor(tensors.next_output); - REQUIRE(next_output_value != nullptr); - const ggml::hrx::CommandPlanAlternateValue * q8 = weighted_match.metadata.find_alternate_value( - next_output_value->id, GGML_TYPE_Q8_1, qwen_q8_1_x4_size(token_count, kQwenMoeHiddenSize)); - REQUIRE((q8 != nullptr) == test.publishes_q8); - if (test.publishes_q8) { - REQUIRE(weighted_match.metadata.activation_publication_diagnostics().size() == 1); - const auto & diagnostic = weighted_match.metadata.activation_publication_diagnostics().front(); - REQUIRE(diagnostic.source_value == next_output_value->id); - REQUIRE(diagnostic.alternate_value == q8->alternate_value); - REQUIRE(diagnostic.publication_stage == "published"); - REQUIRE(diagnostic.requested_format == "q8-1-x4"); - } - ggml_free(ctx); - } -} - static void run_qwen_router_top8_dispatch_checks() { ggml_init_params params = {}; params.mem_size = 2 * 1024 * 1024; @@ -11347,8 +8379,8 @@ static void run_zero_output_scheduler_elision_checks() { REQUIRE(ids_value != nullptr); ggml::hrx::Command referenced; referenced.bindings.push_back({ "ids", ids_value->id }); - for (auto * list : - { &lookup.program->commands().initialization_commands, &lookup.program->commands().commands }) { + for (auto * list : { &lookup.program->commands().initialization_commands, + &lookup.program->commands().commands }) { list->push_back(referenced); const auto match = lookup.program->match_current_graph(*graph); REQUIRE(match.valid()); @@ -11359,8 +8391,8 @@ static void run_zero_output_scheduler_elision_checks() { const auto unreferenced = lookup.program->match_current_graph(*graph); REQUIRE(unreferenced.valid()); REQUIRE(unreferenced.external_bindings.size() == 1); - for (const auto origin : - { ggml::hrx::CommandBindingOrigin::Transient, ggml::hrx::CommandBindingOrigin::ProgramConstant }) { + for (const auto origin : { ggml::hrx::CommandBindingOrigin::Transient, + ggml::hrx::CommandBindingOrigin::ProgramConstant }) { ggml::hrx::Command unrelated; unrelated.bindings.push_back({ "negative", ggml::hrx::ValueId(-1) }); unrelated.bindings.push_back({ "large", ggml::hrx::ValueId(INT32_MAX) }); @@ -12140,51 +9172,6 @@ static void run_graph_program_cache_uid_mismatch_checks() { ggml_free(ctx); } -static void run_graph_match_empty_expert_views_checks() { - ggml_context * ctx = ggml_init({ 512 * 1024, nullptr, true }); - REQUIRE(ctx != nullptr); - - // Simulate MoE output layout: 2048 features, 8 selected expert slots, and no requested token rows. - // Scale supplies a producer node without needing a model or actual expert computation. - ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 2048, 8, 0); - ggml_tensor * experts = ggml_scale(ctx, input, 2.0f); - ggml_cgraph * graph = ggml_new_graph(ctx); - graph->uid = 3003; - ggml_build_forward_expand(graph, experts); - - // Build empty views to verify that the graph is accepted with each slot's - // offset preserved (0, 8192, ..., 57344 bytes). Previous check rejected - // empty views with positive offsets (which admittedly is a little odd but - // that's how GGML represents this). - for (int expert = 0; expert < 8; ++expert) { - ggml_tensor * view = ggml_view_2d(ctx, experts, 2048, 0, experts->nb[2], expert * experts->nb[1]); - ggml_build_forward_expand(graph, view); - } - - // Building and matching this graph must accept the empty views; an old bounds check rejected their positive offsets. - // gfx1151 selects a target configuration for this host test, not a GPU on which to execute it. - ggml::hrx::GraphProgramCache cache; - auto lookup = cache.get_or_build(*graph, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(lookup.valid()); - REQUIRE(lookup.program->match_current_graph(*graph).valid()); - // Empty operations require no numerical device work. - REQUIRE(lookup.program->commands().commands.empty()); - - // The same graph must also match when reusing the cached program. - lookup = cache.get_or_build(*graph, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(lookup.valid()); - REQUIRE(cache.stats().hits == 1); - - // Empty views still have to agree with the recorded alias offsets: changing 57344 to 57348 must fail. - // Restoring the offset must make the graph match again. - ggml_tensor * last_view = ggml_graph_node(graph, ggml_graph_n_nodes(graph) - 1); - last_view->view_offs += sizeof(float); - REQUIRE(!lookup.program->match_current_graph(*graph).valid()); - last_view->view_offs -= sizeof(float); - REQUIRE(lookup.program->match_current_graph(*graph).valid()); - ggml_free(ctx); -} - static void run_graph_match_hash_collision_checks() { ggml_context * ctx = ggml_init({ 2 * 1024 * 1024, nullptr, true }); REQUIRE(ctx != nullptr); @@ -12495,214 +9482,6 @@ static void write_transient_i32_value(ggml_backend_hrx_context * con hrx_synchronous_h2d(context->device->device, &value, arena.buffer, allocation.arena_offset, sizeof(value))); } -static std::vector read_transient_bytes(ggml_backend_hrx_context * context, - ggml::hrx::TransientArenaAllocationRef arena, - const ggml::hrx::TransientAllocation & allocation) { - REQUIRE(context != nullptr); - REQUIRE(arena.buffer != nullptr); - std::vector data(allocation.size); - require_hrx_status(hrx_synchronous_d2h(context->device->device, arena.buffer, allocation.arena_offset, data.data(), - allocation.size)); - return data; -} - -static void run_binary_q8_publication_numerics() { - const struct { - int64_t hidden; - int64_t tokens; - bool broadcast; - int op; - } cases[] = { - { 2048, 3, false, 0 }, - { 2176, 4, false, 2 }, - { 2048, 4, true, 0 }, - { 2176, 5, true, 2 }, - }; - - ggml_backend_t backend = ggml_backend_hrx_init(0); - REQUIRE(backend != nullptr); - auto * backend_context = static_cast(backend->context); - REQUIRE(backend_context != nullptr); - const ggml::hrx::KernelCorpus & corpus = ggml::hrx::get_qwen_kernel_corpus(); - const char * target = backend_context->device->architecture.c_str(); - - for (const auto & test : cases) { - ggml_init_params params = {}; - params.mem_size = 512 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - const int64_t element_count = test.hidden * test.tokens; - const int64_t rhs_count = test.broadcast ? test.hidden : element_count; - ggml_tensor * lhs = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens); - ggml_tensor * rhs = test.broadcast ? ggml_new_tensor_1d(ctx, GGML_TYPE_F32, test.hidden) : - ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.hidden, test.tokens); - REQUIRE(lhs != nullptr && rhs != nullptr); - ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); - REQUIRE(buffer != nullptr); - - std::vector lhs_data(element_count); - std::vector rhs_data(rhs_count); - std::vector expected(element_count); - for (int64_t i = 0; i < element_count; ++i) { - lhs_data[i] = static_cast((i * 17 + 3) % 251 - 125) / 31.0f; - } - for (int64_t i = 0; i < rhs_count; ++i) { - rhs_data[i] = static_cast((i * 29 + 11) % 239 - 119) / 37.0f; - } - for (int64_t token = 0; token < test.tokens; ++token) { - for (int64_t channel = 0; channel < test.hidden; ++channel) { - const int64_t i = token * test.hidden + channel; - const float rhs_v = rhs_data[test.broadcast ? channel : i]; - expected[i] = test.op == 0 ? lhs_data[i] + rhs_v : lhs_data[i] * rhs_v; - } - } - ggml_backend_tensor_set(lhs, lhs_data.data(), 0, lhs_data.size() * sizeof(float)); - ggml_backend_tensor_set(rhs, rhs_data.data(), 0, rhs_data.size() * sizeof(float)); - ggml_backend_synchronize(backend); - - ggml::hrx::Graph graph; - const ggml::hrx::ValueId lhs_value = - graph.values().get_or_add_tensor_value(lhs, ggml::hrx::ValueKind::External); - const ggml::hrx::ValueId rhs_value = - graph.values().get_or_add_tensor_value(rhs, ggml::hrx::ValueKind::External); - ggml::hrx::ValueBufferBinding lhs_binding; - ggml::hrx::ValueBufferBinding rhs_binding; - REQUIRE(ggml_backend_hrx_resolve_value_buffer(lhs, lhs_binding)); - REQUIRE(ggml_backend_hrx_resolve_value_buffer(rhs, rhs_binding)); - REQUIRE(graph.values().bind_buffer(lhs_value, lhs_binding)); - REQUIRE(graph.values().bind_buffer(rhs_value, rhs_binding)); - - const ggml::hrx::ValueId reference_f32(static_cast(graph.values().size())); - const ggml::hrx::ValueId reference_q8(reference_f32.value + 1); - const ggml::hrx::ValueId published_f32(reference_f32.value + 2); - const ggml::hrx::ValueId published_q8(reference_f32.value + 3); - const size_t f32_bytes = static_cast(element_count) * sizeof(float); - const size_t q8_bytes = qwen_q8_1_x4_size(test.tokens, test.hidden); - - ggml::hrx::CommandPlan plan; - plan.transients.push_back({ reference_f32, "test.binary.reference_f32", f32_bytes, 256 }); - plan.transients.push_back({ reference_q8, "test.binary.reference_q8", q8_bytes, 256 }); - plan.transients.push_back({ published_f32, "test.binary.published_f32", f32_bytes, 256 }); - plan.transients.push_back({ published_q8, "test.binary.published_q8", q8_bytes, 256 }); - - ggml::hrx::Dispatch legacy; - legacy.kernel = ggml::hrx::make_kernel_specialization(ggml::hrx::kernel_catalog_ref( - "loom_libs", test.broadcast ? "ggml_binary_bc_f32" : "ggml_binary_f32")); - legacy.kernel.integer_parameters.emplace("element_count", element_count); - legacy.kernel.compile_parameters.emplace( - test.broadcast ? "ggml.binary_bc_f32.op" : "ggml.binary_f32.op", std::to_string(test.op)); - if (test.broadcast) { - legacy.kernel.integer_parameters.emplace("ne0", test.hidden); - legacy.kernel.integer_parameters.emplace("ne1", test.tokens); - legacy.kernel.integer_parameters.emplace("ne2", 1); - legacy.kernel.integer_parameters.emplace("ne3", 1); - legacy.kernel.integer_parameters.emplace("src0_element_count", element_count); - legacy.kernel.integer_parameters.emplace("src1_element_count", rhs_count); - } else { - legacy.kernel.compile_parameters.emplace("ggml.binary_f32.ne0", std::to_string(test.hidden)); - legacy.kernel.compile_parameters.emplace("ggml.binary_f32.ne1", std::to_string(test.tokens)); - legacy.kernel.compile_parameters.emplace("ggml.binary_f32.ne2", "1"); - legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride1", std::to_string(test.hidden)); - legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride2", std::to_string(element_count)); - legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src0_stride3", std::to_string(element_count)); - legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride1", std::to_string(test.hidden)); - legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride2", std::to_string(element_count)); - legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src1_stride3", std::to_string(element_count)); - legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src0_span", std::to_string(element_count)); - legacy.kernel.compile_parameters.emplace("ggml.binary_f32.src1_span", std::to_string(element_count)); - } - if (test.broadcast) { - for (int dim = 0; dim < GGML_MAX_DIMS; ++dim) { - legacy.kernel.compile_parameters.emplace("ggml.binary_bc_f32.src0_broadcast_dim" + - std::to_string(dim), - "0"); - legacy.kernel.compile_parameters.emplace("ggml.binary_bc_f32.src1_broadcast_dim" + - std::to_string(dim), - dim == 1 ? "1" : "0"); - } - } - legacy.bindings.push_back({ lhs_value, 0, ggml_nbytes(lhs) }); - legacy.bindings.push_back({ rhs_value, 0, ggml_nbytes(rhs) }); - legacy.bindings.push_back({ reference_f32, 0, f32_bytes }); - plan.dispatches.push_back(std::move(legacy)); - - ggml::hrx::Dispatch quantize; - quantize.kernel = ggml::hrx::make_kernel_specialization( - ggml::hrx::kernel_catalog_ref("qwen3_moe", "ggml_quantize_q8_1_x4_f32")); - quantize.kernel.integer_parameters.emplace("token_count", test.tokens); - quantize.kernel.integer_parameters.emplace("input_size", test.hidden); - quantize.kernel.compile_parameters.emplace("ggml.quantize_q8_1_x4.group_capacity", - std::to_string(element_count / 128)); - quantize.bindings.push_back({ reference_f32, 0, f32_bytes }); - quantize.bindings.push_back({ reference_q8, 0, q8_bytes }); - plan.dispatches.push_back(std::move(quantize)); - - ggml::hrx::Dispatch publish; - publish.kernel = ggml::hrx::make_kernel_specialization(ggml::hrx::kernel_catalog_ref( - "loom_libs", test.broadcast ? "ggml_binary_bc_f32_publish_q8_1_x4" : - "ggml_binary_f32_publish_q8_1_x4")); - publish.kernel.integer_parameters.emplace("token_count", test.tokens); - if (test.broadcast) { - publish.kernel.integer_parameters.emplace("hidden_size", test.hidden); - publish.kernel.integer_parameters.emplace("src0_element_count", element_count); - publish.kernel.integer_parameters.emplace("src1_element_count", rhs_count); - publish.kernel.compile_parameters = plan.dispatches.front().kernel.compile_parameters; - } else { - publish.kernel.compile_parameters = plan.dispatches.front().kernel.compile_parameters; - } - publish.bindings.push_back({ lhs_value, 0, ggml_nbytes(lhs) }); - publish.bindings.push_back({ rhs_value, 0, ggml_nbytes(rhs) }); - publish.bindings.push_back({ published_f32, 0, f32_bytes }); - publish.bindings.push_back({ published_q8, 0, q8_bytes }); - plan.dispatches.push_back(std::move(publish)); - - const ggml::hrx::CommandProgram commands = ggml::hrx::build_command_program(graph, plan, corpus, target); - REQUIRE(commands.valid()); - REQUIRE(ggml::hrx::verify_command_program(commands, corpus, target).valid()); - const ggml::hrx::CommandProgramBindings bindings = - ggml::hrx::CommandProgramBindings::from_value_map(graph.values()); - REQUIRE(bindings.valid()); - const ggml::hrx::CommandProgramExecutionContext execution_context = { - backend_context->device->device, - backend_context->stream, - target, - &corpus, - &backend_context->kernel_executables, - &backend_context->transient_arena, - &backend_context->host_transfers, - &backend_context->host_weights, - }; - REQUIRE(ggml::hrx::execute_command_program(execution_context, commands, bindings)); - ggml_backend_synchronize(backend); - - const auto arena = backend_context->transient_arena.current_allocation(); - const auto * reference_f32_allocation = - ggml::hrx::find_transient_allocation(commands.transients, reference_f32); - const auto * reference_q8_allocation = - ggml::hrx::find_transient_allocation(commands.transients, reference_q8); - const auto * published_f32_allocation = - ggml::hrx::find_transient_allocation(commands.transients, published_f32); - const auto * published_q8_allocation = - ggml::hrx::find_transient_allocation(commands.transients, published_q8); - REQUIRE(reference_f32_allocation != nullptr && reference_q8_allocation != nullptr); - REQUIRE(published_f32_allocation != nullptr && published_q8_allocation != nullptr); - const auto reference_f32_bytes = read_transient_bytes(backend_context, arena, *reference_f32_allocation); - const auto reference_q8_bytes = read_transient_bytes(backend_context, arena, *reference_q8_allocation); - const auto published_f32_bytes = read_transient_bytes(backend_context, arena, *published_f32_allocation); - const auto published_q8_bytes = read_transient_bytes(backend_context, arena, *published_q8_allocation); - REQUIRE(reference_f32_bytes == published_f32_bytes); - REQUIRE(reference_q8_bytes == published_q8_bytes); - REQUIRE(reference_f32_bytes.size() == expected.size() * sizeof(float)); - REQUIRE(std::memcmp(reference_f32_bytes.data(), expected.data(), reference_f32_bytes.size()) == 0); - - ggml_backend_buffer_free(buffer); - ggml_free(ctx); - } - ggml_backend_free(backend); -} - static void run_qwen_expert_table_partition_prefill_512_execution() { ggml_backend_t backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); @@ -12874,158 +9653,6 @@ static void run_add_f32() { ggml_backend_free(backend); } -static std::vector run_rope_scale_numerical_case(ggml_backend_t backend, bool per_token, int mode) { - ggml_init_params params = {}; - params.mem_size = 512 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - constexpr int64_t head_size = 128; - constexpr int64_t n_dims = 96; - constexpr int64_t head_count = 2; - constexpr int64_t token_count = 3; - constexpr int64_t elements = head_size * head_count * token_count; - - ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); - ggml_tensor * pos = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); - ggml_tensor * freqs = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_dims / 2); - ggml_tensor * rope = - ggml_rope_ext(ctx, input, pos, freqs, n_dims, mode, 0, 10000.0f, 1.0f, 0.0f, 1.19024f, 32.0f, 1.0f); - ggml_tensor * scales = per_token ? ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, 1, token_count) : nullptr; - ggml_tensor * output = per_token ? ggml_mul(ctx, rope, scales) : ggml_scale(ctx, rope, 0.0883883461f); - REQUIRE(input != nullptr); - REQUIRE(pos != nullptr); - REQUIRE(freqs != nullptr); - REQUIRE(rope != nullptr); - REQUIRE(output != nullptr); - - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - - ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); - REQUIRE(buffer != nullptr); - std::vector input_data(elements); - std::vector freq_data(n_dims / 2); - const std::vector positions = { 3, 7, 11 }; - const std::vector scale_data = { 0.125f, -0.25f, 0.5f }; - for (int64_t i = 0; i < elements; ++i) { - input_data[i] = static_cast((i * 17) % 113) * 0.03125f - 1.5f; - } - for (int64_t i = 0; i < n_dims / 2; ++i) { - freq_data[i] = 0.75f + static_cast(i % 7) * 0.125f; - } - ggml_backend_tensor_set(input, input_data.data(), 0, input_data.size() * sizeof(float)); - ggml_backend_tensor_set(pos, positions.data(), 0, positions.size() * sizeof(int32_t)); - ggml_backend_tensor_set(freqs, freq_data.data(), 0, freq_data.size() * sizeof(float)); - if (per_token) { - ggml_backend_tensor_set(scales, scale_data.data(), 0, scale_data.size() * sizeof(float)); - } - - REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); - std::vector actual(elements); - ggml_backend_tensor_get(output, actual.data(), 0, actual.size() * sizeof(float)); - - ggml_backend_buffer_free(buffer); - ggml_free(ctx); - return actual; -} - -static void run_rope_scale_f32_numerics() { - ggml_backend_t hrx = ggml_backend_hrx_init(0); - ggml_backend_t cpu = ggml_backend_cpu_init(); - REQUIRE(hrx != nullptr); - REQUIRE(cpu != nullptr); - - for (const bool per_token : { false, true }) { - for (const int mode : { static_cast(GGML_ROPE_TYPE_NORMAL), static_cast(GGML_ROPE_TYPE_NEOX) }) { - const std::vector expected = run_rope_scale_numerical_case(cpu, per_token, mode); - const std::vector actual = run_rope_scale_numerical_case(hrx, per_token, mode); - REQUIRE(actual.size() == expected.size()); - for (size_t i = 0; i < actual.size(); ++i) { - const float tolerance = 2.0e-4f + 2.0e-4f * std::fabs(expected[i]); - REQUIRE(std::fabs(actual[i] - expected[i]) <= tolerance); - } - } - } - - ggml_backend_free(cpu); - ggml_backend_free(hrx); -} - -struct GatherAddRmsNormOutputs { - std::vector raw; - std::vector normalized; -}; - -static GatherAddRmsNormOutputs run_gather_add_rmsnorm_numerical_case(ggml_backend_t backend) { - ggml_init_params params = {}; - params.mem_size = 512 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - constexpr int64_t hidden_size = 128; - constexpr int64_t source_token_count = 4; - constexpr int64_t output_token_count = 5; - GatherAddRmsNormGraph tensors = - build_gather_add_rmsnorm_graph(ctx, hidden_size, source_token_count, output_token_count); - ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); - REQUIRE(buffer != nullptr); - - std::vector attention(hidden_size * source_token_count); - std::vector residual(hidden_size * source_token_count); - std::vector weight(hidden_size); - const std::vector row_ids = { 3, 1, 3, 0, 2 }; - for (size_t i = 0; i < attention.size(); ++i) { - attention[i] = static_cast((i * 13) % 47) * 0.03125f - 0.5f; - residual[i] = static_cast((i * 7) % 31) * -0.015625f + 0.25f; - } - for (size_t i = 0; i < weight.size(); ++i) { - weight[i] = 0.5f + static_cast(i % 11) * 0.0625f; - } - ggml_backend_tensor_set(tensors.attention, attention.data(), 0, attention.size() * sizeof(float)); - ggml_backend_tensor_set(tensors.residual, residual.data(), 0, residual.size() * sizeof(float)); - ggml_backend_tensor_set(tensors.row_ids, row_ids.data(), 0, row_ids.size() * sizeof(int32_t)); - ggml_backend_tensor_set(tensors.weight, weight.data(), 0, weight.size() * sizeof(float)); - - REQUIRE(ggml_backend_graph_compute(backend, tensors.graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(backend); - - GatherAddRmsNormOutputs outputs; - outputs.raw.resize(hidden_size * output_token_count); - outputs.normalized.resize(hidden_size * output_token_count); - ggml_backend_tensor_get(tensors.raw_output, outputs.raw.data(), 0, outputs.raw.size() * sizeof(float)); - ggml_backend_tensor_get(tensors.normalized_output, outputs.normalized.data(), 0, - outputs.normalized.size() * sizeof(float)); - - ggml_backend_buffer_free(buffer); - ggml_free(ctx); - return outputs; -} - -static void run_gather_add_rmsnorm_f32_numerics() { - ggml_backend_t hrx = ggml_backend_hrx_init(0); - ggml_backend_t cpu = ggml_backend_cpu_init(); - REQUIRE(hrx != nullptr); - REQUIRE(cpu != nullptr); - - const GatherAddRmsNormOutputs expected = run_gather_add_rmsnorm_numerical_case(cpu); - const GatherAddRmsNormOutputs actual = run_gather_add_rmsnorm_numerical_case(hrx); - REQUIRE(actual.raw.size() == expected.raw.size()); - REQUIRE(actual.normalized.size() == expected.normalized.size()); - for (size_t i = 0; i < actual.raw.size(); ++i) { - REQUIRE(actual.raw[i] == expected.raw[i]); - const float tolerance = 3.0e-4f + 3.0e-4f * std::fabs(expected.normalized[i]); - REQUIRE(std::fabs(actual.normalized[i] - expected.normalized[i]) <= tolerance); - } - - ggml_backend_free(cpu); - ggml_backend_free(hrx); -} - static void run_scale_f32() { ggml_backend_t backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); @@ -13471,7 +10098,6 @@ static void run_lowtoken_residual_dispatch_checks() { bool second_add; bool fused; }; - const Case cases[] = { { GGML_TYPE_Q4_K, 6144, 5120, 1, 0, 0, false, false, true }, { GGML_TYPE_Q4_K, 17408, 5120, 1, 0, 0, false, false, true }, @@ -13495,12 +10121,12 @@ static void run_lowtoken_residual_dispatch_checks() { { GGML_TYPE_Q4_K, 6144, 5120, 5, 2, 0, false, false, false }, { GGML_TYPE_Q4_K, 6144, 5120, 1, 3, 0, false, false, false }, { GGML_TYPE_Q4_K, 6144, 5120, 5, 3, 0, false, false, false }, - { GGML_TYPE_Q4_K, 6144, 5120, 1, 0, 0, true, false, true }, + { GGML_TYPE_Q4_K, 6144, 5120, 1, 0, 0, true, false, false }, { GGML_TYPE_Q4_K, 6144, 5120, 5, 1, 0, true, false, false }, - { GGML_TYPE_Q4_K, 3840, 4096, 1, 0, 0, false, false, true }, - { GGML_TYPE_Q4_K, 6144, 4032, 1, 0, 0, false, false, true }, - { GGML_TYPE_Q4_K, 4096, 8192, 1, 0, 0, false, false, true }, - { GGML_TYPE_F16, 6144, 5120, 1, 0, 0, false, false, true }, + { GGML_TYPE_Q4_K, 3840, 4096, 1, 0, 0, false, false, false }, + { GGML_TYPE_Q4_K, 6144, 4032, 1, 0, 0, false, false, false }, + { GGML_TYPE_Q4_K, 4096, 8192, 1, 0, 0, false, false, false }, + { GGML_TYPE_F16, 6144, 5120, 1, 0, 0, false, false, false }, { GGML_TYPE_Q4_K, 6144, 5120, 6, 1, 0, false, false, false }, { GGML_TYPE_Q4_K, 6144, 5120, 256, 1, 0, false, false, false }, }; @@ -13540,99 +10166,37 @@ static void run_lowtoken_residual_dispatch_checks() { REQUIRE(program.valid()); REQUIRE(command_program_verifies(program)); size_t fused = 0, binary = 0; - bool fused_second_add = false; for (const auto & command : program.commands) { const std::string name = kernel_name_for_id(command.kernel.kernel_id); binary += name == "loom_libs:ggml_binary_f32"; - const bool tiled_postops = name == "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32"; - const bool skinny_postops = name == "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32"; - const bool vector_postops = name == "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32"; - if (!tiled_postops && !skinny_postops && !vector_postops) { + if (name != "loom_libs:ggml_mul_mat_bias_f32_f32_wmma" && name != "loom_libs:ggml_mul_mat_add_f32_f32_wmma") { continue; } ++fused; - if (vector_postops) { - fused_second_add = - fused_second_add || - (command.kernel.compile_parameters.at("ggml.matmul.vector_postops.apply_bias") == "1" && - command.kernel.compile_parameters.at("ggml.matmul.vector_postops.apply_residual") == "1"); - } else { - fused_second_add = - fused_second_add || command.kernel.compile_parameters.at("ggml.mul_mat_postops.epilogue") == "3"; - } - REQUIRE(command.bindings.size() >= 4); + REQUIRE(command.bindings.size() == 4); REQUIRE(command.kernel.integer_parameters.at("token_count") == c.tokens); - if (vector_postops) { - REQUIRE(command.kernel.compile_parameters.count("ggml.matmul.vector_postops.weight_format") == 1); - } else { - REQUIRE(command.kernel.compile_parameters.count("ggml.mul_mat_postops.weight_format") == 1); - } + REQUIRE(command.kernel.compile_parameters.at("ggml.mul_mat_postops.weight_format") == + (c.type == GGML_TYPE_Q4_K ? "44" : "46")); + REQUIRE(command.kernel.compile_parameters.at("ggml.mul_mat.activation_format") == "9"); } REQUIRE(fused == (c.fused ? 1 : 0)); - const size_t expected_binary = c.fused ? (c.second_add && !fused_second_add ? 1 : 0) : 1; - REQUIRE(binary == expected_binary); + REQUIRE(binary == (c.fused ? (c.second_add ? 1 : 0) : 1)); ggml_free(ctx); } } -static void run_lowtoken_bias_dispatch_checks() { - ggml_context * ctx = ggml_init({ 16 * 1024 * 1024, nullptr, true }); - REQUIRE(ctx != nullptr); - - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 3584, 1024); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 3584, 2); - ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); - ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1024); - ggml_tensor * output = ggml_add(ctx, projection, bias); - - ggml_cgraph * graph = ggml_new_graph_custom(ctx, 64, false); - ggml_build_forward_expand(graph, output); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); - const ggml::hrx::CommandProgram program = ggml::hrx::build_command_program( - imported.graph, scheduler.plan(), ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); - REQUIRE(program.valid()); - REQUIRE(command_program_verifies(program)); - - size_t skinny_bias = 0; - size_t tiled_bias = 0; - for (const auto & command : program.commands) { - const std::string name = kernel_name_for_id(command.kernel.kernel_id); - skinny_bias += name == "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32"; - tiled_bias += name == "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32"; - } - REQUIRE(skinny_bias == 1); - REQUIRE(tiled_bias == 0); - - ggml_free(ctx); -} - static void run_gdn_selected_rms_q8_dispatch_checks() { ggml_init_params params = {}; params.mem_size = 32 * 1024 * 1024; params.no_alloc = true; ggml_context * ctx = ggml_init(params); REQUIRE(ctx != nullptr); - enum { - UnreadyInput = 1, - UnreadyGate = 2, - GatherSide = 4, - AttentionSide = 8, - GdnObserved = 16, - GatherObserved = 32, - Unaligned = 64, - CacheRead = 128, - OutputSide = 256, - CachedGate = 512, - TanhGate = 1024, - AttentionReshape = 2048, + UnreadyInput = 1, UnreadyGate = 2, GatherSide = 4, AttentionSide = 8, + GdnObserved = 16, GatherObserved = 32, Unaligned = 64, CacheRead = 128, + OutputSide = 256, CachedGate = 512, TanhGate = 1024, AttentionReshape = 2048, AttentionObserved = 4096, }; - const struct { int64_t tokens, heads, sequences; int flags; @@ -13690,8 +10254,8 @@ static void run_gdn_selected_rms_q8_dispatch_checks() { } ggml_tensor * state = ggml_reshape_4d(ctx, gathered, width, width, heads, test.sequences); const auto view = [&](int64_t nheads, size_t offset) { - return ggml_view_4d(ctx, raw, width, nheads, test.tokens, test.sequences, width * sizeof(float), - hidden * sizeof(float), hidden * test.tokens * sizeof(float), offset); + return ggml_view_4d(ctx, raw, width, nheads, test.tokens, test.sequences, + width * sizeof(float), hidden * sizeof(float), hidden * test.tokens * sizeof(float), offset); }; ggml_tensor * q = ggml_l2_norm(ctx, view(qheads, 0), 1.e-6f); ggml_tensor * k = ggml_l2_norm(ctx, view(qheads, width * qheads * sizeof(float)), 1.e-6f); @@ -13701,32 +10265,26 @@ static void run_gdn_selected_rms_q8_dispatch_checks() { ggml_set_output(gdn); } const size_t attention_bytes = channels * test.tokens * test.sequences * sizeof(float); - ggml_tensor * attention = - ggml_view_4d(ctx, gdn, width, heads, test.tokens, test.sequences, width * sizeof(float), - channels * sizeof(float), channels * test.tokens * sizeof(float), 0); + ggml_tensor * attention = ggml_view_4d(ctx, gdn, width, heads, test.tokens, test.sequences, + width * sizeof(float), channels * sizeof(float), channels * test.tokens * sizeof(float), 0); if (test.flags & AttentionObserved) { ggml_set_output(attention); } - ggml_tensor * snapshots = ggml_view_3d(ctx, gdn, state_elements, test.sequences, written, state_bytes, - state_bytes * test.sequences, attention_bytes); - ggml_tensor * target = ggml_view_3d(ctx, cache, state_elements, test.sequences, written, state_bytes, - 4 * state_bytes, (test.flags & Unaligned) ? sizeof(float) : state_bytes); + ggml_tensor * snapshots = ggml_view_3d(ctx, gdn, state_elements, test.sequences, written, + state_bytes, state_bytes * test.sequences, attention_bytes); + ggml_tensor * target = ggml_view_3d(ctx, cache, state_elements, test.sequences, written, + state_bytes, 4 * state_bytes, (test.flags & Unaligned) ? sizeof(float) : state_bytes); ggml_tensor * norm_input = (test.flags & AttentionReshape) ? - ggml_reshape_2d(ctx, attention, width, heads * test.tokens * test.sequences) : - attention; + ggml_reshape_2d(ctx, attention, width, heads * test.tokens * test.sequences) : attention; ggml_tensor * rms_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, width); - ggml_tensor * raw_gate = - (test.flags & CachedGate) ? + ggml_tensor * raw_gate = (test.flags & CachedGate) ? ggml_view_4d(ctx, cache, width, heads, test.tokens, test.sequences, width * sizeof(float), channels * sizeof(float), channels * test.tokens * sizeof(float), 0) : - ggml_scale(ctx, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, width, heads, test.tokens, test.sequences), - 0.5f); + ggml_scale(ctx, ggml_new_tensor_4d(ctx, GGML_TYPE_F32, width, heads, test.tokens, test.sequences), 0.5f); ggml_tensor * shaped_gate = (test.flags & AttentionReshape) ? - ggml_reshape_2d(ctx, raw_gate, width, heads * test.tokens * test.sequences) : - raw_gate; + ggml_reshape_2d(ctx, raw_gate, width, heads * test.tokens * test.sequences) : raw_gate; ggml_tensor * norm = ggml_mul(ctx, ggml_rms_norm(ctx, norm_input, 1.e-6f), rms_weight); - ggml_tensor * gated = - ggml_mul(ctx, norm, (test.flags & TanhGate) ? ggml_tanh(ctx, shaped_gate) : ggml_silu(ctx, shaped_gate)); + ggml_tensor * gated = ggml_mul(ctx, norm, (test.flags & TanhGate) ? ggml_tanh(ctx, shaped_gate) : ggml_silu(ctx, shaped_gate)); ggml_tensor * projection_input = ggml_reshape_2d(ctx, gated, channels, test.tokens * test.sequences); ggml_tensor * projection_weight = ggml_new_tensor_2d(ctx, test.type, channels, 4096); ggml_tensor * projection = ggml_mul_mat(ctx, projection_weight, projection_input); @@ -13777,149 +10335,93 @@ static void run_gdn_selected_rms_q8_dispatch_checks() { REQUIRE(std::none_of(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { const auto name = kernel_name_for_id(dispatch.kernel.kernel_id); return name == "loom_libs:ggml_get_rows_f32" || name == "loom_libs:ggml_copy_f32" || - name == "loom_libs:ggml_rmsnorm_gate_f32_publish"; + name == "loom_libs:ggml_rmsnorm_gate_f32_q8_1_x4"; })); } - const auto commands = - ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); + const auto commands = ggml::hrx::build_command_program(imported.graph, plan, ggml::hrx::get_qwen_kernel_corpus(), "gfx1151"); REQUIRE(commands.valid()); - if (!(test.flags & UnreadyGate)) { - REQUIRE(command_program_verifies(commands)); - } + REQUIRE(command_program_verifies(commands)); } ggml_free(ctx); } -static void register_hrx_backend_host_cases(test_runner::Suite & suite) { - suite.host_case("status", [] { run_status_checks(); }); - suite.host_case("command_plan_metadata", [] { run_command_plan_metadata_checks(); }); - suite.host_case("dispatch_registry", [] { run_dispatch_registry_checks(); }); - suite.host_case("graph_import", [] { run_graph_import_checks(); }); - suite.host_case("graph_import_mixed_backend_boundary", [] { run_graph_import_mixed_backend_boundary_checks(); }); - suite.host_case("graph_view_external_use", [] { run_graph_view_external_use_checks(); }); - suite.host_case("scale_f32_dispatch", [] { run_scale_f32_dispatch_checks(); }); - suite.host_case("scale_add_f32_dispatch", [] { run_scale_add_f32_dispatch_checks(); }); - suite.host_case("rmsnorm_two_binary_dispatch", [] { run_rmsnorm_two_binary_dispatch_checks(); }); - suite.host_case("rmsnorm_lowtoken_striped_dispatch", [] { run_rmsnorm_lowtoken_striped_dispatch_checks(); }); - suite.host_case("cont_f32_dispatch", [] { run_cont_f32_dispatch_checks(); }); - suite.host_case("binary_f32_broadcast_dispatch", [] { run_binary_f32_broadcast_dispatch_checks(); }); - suite.host_case("binary_q8_publication", [] { run_binary_q8_publication_checks(); }); - suite.host_case("graph_snapshot_diagnostics", [] { run_graph_snapshot_diagnostics_checks(); }); - suite.host_case("unmatched_graph_diagnostics", [] { run_unmatched_graph_diagnostics_checks(); }); - suite.host_case("completion_counter_plan", [] { run_completion_counter_plan_checks(); }); - suite.host_case("graph_index", [] { run_graph_index_checks(); }); - suite.host_case("rope_set_rows_dispatch", [] { run_rope_set_rows_dispatch_checks(); }); - suite.host_case("graph_traversal", [] { run_graph_traversal_checks(); }); - suite.host_case("qwen_token_embedding_dispatch", [] { run_qwen_token_embedding_dispatch_checks(); }); - suite.host_case("get_rows_rmsnorm_binary_dispatch", [] { run_get_rows_rmsnorm_binary_dispatch_checks(); }); - suite.host_case("get_rows_scale_dispatch", [] { run_get_rows_scale_dispatch_checks(); }); - suite.host_case("gather_add_dispatch", [] { run_gather_add_dispatch_checks(); }); - suite.host_case("qwen_flash_attention_dispatch", [] { run_qwen_flash_attention_dispatch_checks(); }); - suite.host_case("tiled_pair_matmul_postops_dispatch", [] { run_tiled_pair_matmul_postops_dispatch_checks(); }); - suite.host_case("qwen_attention_postprocess_dispatch", [] { run_qwen_attention_postprocess_dispatch_checks(); }); - suite.host_case("qwen_matmul_dispatch", [] { run_qwen_matmul_dispatch_checks(); }); - suite.host_case("iq1_codebook_matmul_dispatch", [] { run_iq1_codebook_matmul_dispatch_checks(); }); - suite.host_case("iq3_xxs_codebook_matmul_dispatch", [] { run_iq3_xxs_codebook_matmul_dispatch_checks(); }); - suite.host_case("iq2_codebook_resource_reuse", [] { run_iq2_codebook_resource_reuse_checks(); }); - suite.host_case("iq4_matmul_dispatch", [] { run_iq4_matmul_dispatch_checks(); }); - suite.host_case("packed_f16_producer_consumer", [] { run_packed_f16_producer_consumer_checks(); }); - suite.host_case("generic_swiglu_k16_publish", [] { run_generic_swiglu_k16_publish_checks(); }); - suite.host_case("generic_swiglu_k16_publication_edges", [] { run_generic_swiglu_k16_publication_edge_checks(); }); - suite.host_case("tiled_matmul_alternate_publish", [] { run_tiled_matmul_alternate_publish_checks(); }); - suite.host_case("k16_major_preparation_reuse", [] { run_k16_major_preparation_reuse_checks(); }); - suite.host_case("rmsnorm_binary_k16_output", [] { run_rmsnorm_binary_k16_output_checks(); }); - suite.host_case("rmsnorm_binary_f16_output", [] { run_rmsnorm_binary_f16_output_checks(); }); - suite.host_case("swiglu_q8_output", [] { run_swiglu_q8_output_checks(); }); - suite.host_case("rmsnorm_gate_packed_output", [] { run_rmsnorm_gate_packed_output_checks(); }); - suite.host_case("rmsnorm_gate_q8_output", [] { run_rmsnorm_gate_q8_output_checks(); }); - suite.host_case("symmetric_i4_consumer_qualification", - [] { run_symmetric_i4_consumer_qualification_checks(); }); - suite.host_case("lowtoken_residual_dispatch", [] { run_lowtoken_residual_dispatch_checks(); }); - suite.host_case("lowtoken_bias_dispatch", [] { run_lowtoken_bias_dispatch_checks(); }); - suite.host_case("generic_ssm_conv_binary_dispatch", [] { run_generic_ssm_conv_binary_dispatch_checks(); }); - suite.host_case("quantized_conv4_dispatch", [] { run_quantized_conv4_dispatch_checks(); }); - suite.host_case("lfm_dconv3_dispatch", [] { run_lfm_dconv3_dispatch_checks(); }); - suite.host_case("gdn_rmsnorm_gate_dispatch", [] { run_gdn_rmsnorm_gate_dispatch_checks(); }); - suite.host_case("gdn_native_projection_pair_dispatch", [] { run_gdn_native_projection_pair_dispatch_checks(); }); - suite.host_case("gdn_selected_snapshot_dispatch", [] { run_gdn_selected_snapshot_dispatch_checks(); }); - suite.host_case("gdn_selected_rms_q8_dispatch", [] { run_gdn_selected_rms_q8_dispatch_checks(); }); - suite.host_case("llama_attention_matmul_dispatch", [] { run_llama_attention_matmul_dispatch_checks(); }); - suite.host_case("quantized_value_projection_dispatch", [] { run_quantized_value_projection_dispatch_checks(); }); - suite.host_case("prefill_value_projection_cache_dispatch", - [] { run_prefill_value_projection_cache_dispatch_checks(); }); - suite.host_case("qwen_terminal_q6k_q8_command.tokens1", [] { schedule_qwen_terminal_q6k_q8_command(1); }); - suite.host_case("qwen_terminal_q6k_q8_command.tokens18", [] { schedule_qwen_terminal_q6k_q8_command(18); }); - suite.host_case("qwen_terminal_q6k_vector_command", [] { schedule_qwen_terminal_q6k_vector_command(); }); - suite.host_case("qwen_decode_rmsnorm_publication", [] { run_qwen_decode_rmsnorm_publication_checks(); }); - suite.host_case("qwen_decode_publication_consumer_qualification", [] { - run_qwen_decode_publication_consumer_qualification_checks(); - }); - suite.host_case("get_rows_q8_1_alternate_command.q4_k", - [] { schedule_get_rows_q8_1_alternate_command(GGML_TYPE_Q4_K); }); - suite.host_case("get_rows_q8_1_alternate_command.bf16", - [] { schedule_get_rows_q8_1_alternate_command(GGML_TYPE_BF16); }); - suite.host_case("qwen_router_top8_dispatch", [] { run_qwen_router_top8_dispatch_checks(); }); - suite.host_case("common_mul_mat_id_swiglu_dispatch", [] { run_common_mul_mat_id_swiglu_dispatch_checks(); }); - suite.host_case("qwen_routed_gate_up_dispatch", [] { run_qwen_routed_gate_up_dispatch_checks(); }); - suite.host_case("qwen_weighted_reduce_next_rmsnorm_q8_publication", [] { - run_qwen_weighted_reduce_next_rmsnorm_q8_publication_checks(); - }); - suite.host_case("common_mul_mat_id_postops_dispatch", [] { run_common_mul_mat_id_postops_dispatch_checks(); }); - suite.host_case("alias_value_import", [] { run_alias_value_import_checks(); }); - suite.host_case("multi_dispatch", [] { run_multi_dispatch_checks(); }); - suite.host_case("layout_alias_scheduler_elision", [] { run_layout_alias_scheduler_elision_checks(); }); - suite.host_case("zero_output_scheduler_elision", [] { run_zero_output_scheduler_elision_checks(); }); - suite.host_case("transient_import", [] { run_transient_import_checks(); }); - suite.host_case("graph_view_preserves_shared_output_storage", - [] { run_graph_view_preserves_shared_output_storage_checks(); }); - suite.host_case("chained_dispatch_requires_transients", [] { run_chained_dispatch_requires_transients(); }); - suite.host_case("graph_replay_host_staging_is_not_ineligible", - [] { run_graph_replay_host_staging_is_not_ineligible(); }); - suite.host_case("multiple_transient_plan", [] { run_multiple_transient_plan_checks(); }); - suite.host_case("disjoint_transient_plan_packing", [] { run_disjoint_transient_plan_packing_checks(); }); - suite.host_case("command_program_kernel_dump", [] { run_command_program_kernel_dump_checks(); }); - suite.host_case("command_shape_hash", [] { run_command_shape_hash_checks(); }); - suite.host_case("graph_program_cache_uid_mismatch", [] { run_graph_program_cache_uid_mismatch_checks(); }); - suite.host_case("graph_match_empty_expert_views", [] { run_graph_match_empty_expert_views_checks(); }); - suite.host_case("graph_match_hash_collision", [] { run_graph_match_hash_collision_checks(); }); - suite.host_case("graph_match_bijection", [] { run_graph_match_bijection_checks(); }); - suite.host_case("validated_graph_uid_match", [] { run_validated_graph_uid_match_checks(); }); - suite.host_case("graph_executor_contract", [] { run_graph_executor_contract_checks(); }); - suite.host_case("loom_async_jit_environment", [] { - REQUIRE(ggml::hrx::loom_async_jit_enabled_from_environment() == async_jit_expected_from_environment()); - }); -} - -static void register_hrx_backend_device_cases(test_runner::Suite & suite) { - suite.device_case("zero_output_device_support", [] { run_zero_output_device_support_checks(); }); - suite.device_case("scale_f32_device_support", [] { run_scale_f32_device_support_checks(); }); - suite.device_case("binary_q8_publication_numerics", [] { run_binary_q8_publication_numerics(); }); - suite.device_case("qwen_expert_table_partition_prefill_512_execution", - [] { run_qwen_expert_table_partition_prefill_512_execution(); }); - suite.device_case("add_f32", [] { run_add_f32(); }); - suite.device_case("rope_scale_f32_numerics", [] { run_rope_scale_f32_numerics(); }); - suite.device_case("gather_add_rmsnorm_f32_numerics", [] { run_gather_add_rmsnorm_f32_numerics(); }); - suite.device_case("scale_f32", [] { run_scale_f32(); }); - suite.device_case("scale_f32_inplace", [] { run_scale_f32_inplace(); }); - suite.device_case("two_independent_add_f32", [] { run_two_independent_add_f32(); }); - suite.device_case("chained_add_f32", [] { run_chained_add_f32(); }); - suite.device_case("same_uid_distinct_graph_reuses_graph_program.default", - [] { run_same_uid_distinct_graph_reuses_graph_program(); }); - suite.device_case("same_uid_distinct_graph_reuses_graph_program.distinct_uid", - [] { run_same_uid_distinct_graph_reuses_graph_program(true); }); - suite.device_case("unsupported_op_fails", [] { run_unsupported_op_fails(); }); -} - -static void register_hrx_backend_cases(test_runner::Suite & suite) { - register_hrx_backend_host_cases(suite); - register_hrx_backend_device_cases(suite); -} - -int main(int argc, char ** argv) { - test_runner::Suite suite( - test_runner::Config::with_prefix("HRX backend test", "hrx-backend", "GGML_HRX_BACKEND_TEST")); - register_hrx_backend_cases(suite); - - const bool has_device = ggml_backend_hrx_get_device_count() != 0; - return suite.run(argc, argv, has_device); +int main() { + run_status_checks(); + run_command_plan_metadata_checks(); + run_dispatch_registry_checks(); + run_graph_import_checks(); + run_graph_import_mixed_backend_boundary_checks(); + run_graph_view_external_use_checks(); + run_scale_f32_dispatch_checks(); + run_cont_f32_dispatch_checks(); + run_binary_f32_broadcast_dispatch_checks(); + run_graph_snapshot_diagnostics_checks(); + run_unmatched_graph_diagnostics_checks(); + run_completion_counter_plan_checks(); + run_graph_index_checks(); + run_rope_set_rows_dispatch_checks(); + run_graph_traversal_checks(); + run_qwen_token_embedding_dispatch_checks(); + run_gather_add_dispatch_checks(); + run_qwen_flash_attention_dispatch_checks(); + run_qwen_attention_postprocess_dispatch_checks(); + run_qwen_matmul_dispatch_checks(); + run_packed_f16_producer_consumer_checks(); + run_k16_major_preparation_reuse_checks(); + run_rmsnorm_binary_k16_output_checks(); + run_swiglu_q8_output_checks(); + run_rmsnorm_gate_packed_output_checks(); + run_rmsnorm_gate_q8_output_checks(); + run_lowtoken_residual_dispatch_checks(); + run_quantized_conv4_dispatch_checks(); + run_gdn_rmsnorm_gate_dispatch_checks(); + run_gdn_native_projection_pair_dispatch_checks(); + run_gdn_selected_snapshot_dispatch_checks(); + run_gdn_selected_rms_q8_dispatch_checks(); + run_llama_attention_matmul_dispatch_checks(); + run_quantized_value_projection_dispatch_checks(); + schedule_qwen_terminal_q6k_q8_command(1); + schedule_qwen_terminal_q6k_q8_command(18); + schedule_get_rows_q8_1_alternate_command(GGML_TYPE_Q4_K); + schedule_get_rows_q8_1_alternate_command(GGML_TYPE_BF16); + run_qwen_router_top8_dispatch_checks(); + run_common_mul_mat_id_swiglu_dispatch_checks(); + run_qwen_routed_gate_up_dispatch_checks(); + run_common_mul_mat_id_postops_dispatch_checks(); + run_alias_value_import_checks(); + run_multi_dispatch_checks(); + run_layout_alias_scheduler_elision_checks(); + run_zero_output_scheduler_elision_checks(); + run_transient_import_checks(); + run_graph_view_preserves_shared_output_storage_checks(); + run_chained_dispatch_requires_transients(); + run_graph_replay_host_staging_is_not_ineligible(); + run_multiple_transient_plan_checks(); + run_disjoint_transient_plan_packing_checks(); + run_command_program_kernel_dump_checks(); + run_command_shape_hash_checks(); + run_graph_program_cache_uid_mismatch_checks(); + run_graph_match_hash_collision_checks(); + run_graph_match_bijection_checks(); + run_validated_graph_uid_match_checks(); + run_graph_executor_contract_checks(); + REQUIRE(ggml::hrx::loom_async_jit_enabled_from_environment() == async_jit_expected_from_environment()); + + if (ggml_backend_hrx_get_device_count() == 0) { + std::fprintf(stderr, "test skipped: no HRX devices available\n"); + return 0; + } + + run_zero_output_device_support_checks(); + run_scale_f32_device_support_checks(); + run_qwen_expert_table_partition_prefill_512_execution(); + run_add_f32(); + run_scale_f32(); + run_scale_f32_inplace(); + run_two_independent_add_f32(); + run_chained_add_f32(); + run_same_uid_distinct_graph_reuses_graph_program(); + run_same_uid_distinct_graph_reuses_graph_program(true); + run_unsupported_op_fails(); + return 0; } diff --git a/tests/test-hrx-loom-jit.cpp b/tests/test-hrx-loom-jit.cpp index d90ae6a4ee98..880ce1f56b86 100644 --- a/tests/test-hrx-loom-jit.cpp +++ b/tests/test-hrx-loom-jit.cpp @@ -3,7 +3,6 @@ #include "kernel-corpus/kernel-corpus.h" #include "runtime/kernel-executable-cache.h" #include "runtime/loom-kernel-jit.h" -#include "testing_suite.h" #include #include @@ -204,18 +203,6 @@ static std::map token_count_workload(int64_t token_count) }; } -static std::map moe_routing_config() { - return { - { "ggml.moe_routing.route_count", "8" }, - { "ggml.moe_routing.expert_count", "128" }, - { "ggml.moe_routing.descriptor_expert_mask", "127" }, - { "ggml.moe_routing.descriptor_partition_shift", "7" }, - { "ggml.moe_routing.descriptor_row_count_shift", "13" }, - { "ggml.moe_routing.partition_workgroup_size", "128" }, - { "ggml.workload.token_capacity", "1" }, - }; -} - static std::map flash_attention_workload(int64_t query_token_count, int64_t key_value_token_count) { return { @@ -401,29 +388,25 @@ static void run_cache_materialize_case(HrxTestDevice & devic ggml::hrx::Dispatch dispatch = make_binary_add_dispatch(binary, 64); std::vector constants; - ggml_hrx_loom_jit_launch_config launch; ggml::hrx::KernelExecutableCache cache(mode); - ggml::hrx::KernelExecutableRef ref = cache.get_or_compile(context, binary, dispatch); + ggml::hrx::KernelExecutableRef ref = cache.get_or_compile(context, binary, dispatch, constants); REQUIRE(ref.valid()); - std::shared_ptr executable = - cache.materialize(context, ref, dispatch.kernel, constants, launch); + std::shared_ptr executable = cache.materialize(context, ref, constants); REQUIRE(executable != nullptr); REQUIRE(executable->executable != nullptr); REQUIRE(executable->export_info.binding_count == dispatch.bindings.size()); REQUIRE(executable->export_info.constant_byte_length == constants.size()); - REQUIRE(launch.workgroup_count[0] > 0); - REQUIRE(launch.workgroup_size[0] > 0); + REQUIRE(executable->launch.workgroup_count[0] > 0); + REQUIRE(executable->launch.workgroup_size[0] > 0); - std::shared_ptr loaded_hit = - cache.materialize(context, ref, dispatch.kernel, constants, launch); + std::shared_ptr loaded_hit = cache.materialize(context, ref, constants); REQUIRE(loaded_hit == executable); std::vector constants_for_prepare; - ggml_hrx_loom_jit_launch_config launch_for_prepare; std::shared_ptr prepared = - cache.prepare(context, binary, dispatch, constants_for_prepare, launch_for_prepare); + cache.prepare(context, binary, dispatch, constants_for_prepare); REQUIRE(prepared == executable); const char * mode_name = mode == ggml::hrx::LoomJitMode::Async ? "async" : "sync"; @@ -431,96 +414,6 @@ static void run_cache_materialize_case(HrxTestDevice & devic device.architecture.c_str()); } -static void run_dynamic_cache_materialize_case(HrxTestDevice & device, - ggml::hrx::LoomJitMode mode, - const ggml::hrx::KernelDefinition & binary) { - ggml::hrx::KernelExecutablePrepareContext context = {}; - context.device = device.device; - context.target = device.architecture.c_str(); - - ggml::hrx::Dispatch first = make_binary_add_dispatch(binary, 64); - ggml::hrx::Dispatch second = make_binary_add_dispatch(binary, 257); - first.kernel.compile_parameters = binary_f32_exact_config("0", 512); - second.kernel.compile_parameters = first.kernel.compile_parameters; - first.kernel.workload_specialization = ggml::hrx::WorkloadSpecialization::Dynamic; - second.kernel.workload_specialization = ggml::hrx::WorkloadSpecialization::Dynamic; - - ggml::hrx::KernelExecutableCache cache(mode); - ggml::hrx::KernelExecutableRef first_ref = cache.get_or_compile(context, binary, first); - ggml::hrx::KernelExecutableRef second_ref = cache.get_or_compile(context, binary, second); - REQUIRE(first_ref.valid()); - REQUIRE(second_ref.valid()); - REQUIRE(first_ref.entry == second_ref.entry); - - std::vector first_constants; - std::vector second_constants; - ggml_hrx_loom_jit_launch_config first_launch; - ggml_hrx_loom_jit_launch_config second_launch; - std::shared_ptr first_executable = - cache.materialize(context, first_ref, first.kernel, first_constants, first_launch); - std::shared_ptr second_executable = - cache.materialize(context, second_ref, second.kernel, second_constants, second_launch); - REQUIRE(first_executable != nullptr); - REQUIRE(second_executable == first_executable); - REQUIRE(first_launch.workgroup_count[0] == 1); - REQUIRE(second_launch.workgroup_count[0] == 2); - REQUIRE(first_constants.size() == second_constants.size()); - REQUIRE(first_constants != second_constants); - - ggml::hrx::Dispatch exact_second = second; - exact_second.kernel.workload_specialization = ggml::hrx::WorkloadSpecialization::Exact; - ggml::hrx::KernelExecutableRef exact_ref = cache.get_or_compile(context, binary, exact_second); - REQUIRE(exact_ref.valid()); - REQUIRE(exact_ref.entry != first_ref.entry); -} - -static void run_dynamic_export_reuse_case( - HrxTestDevice & device, - const char * kernel_name, - const std::map & first_workload, - const std::map & second_workload, - const std::map & compile_config = {}) { - const ggml::hrx::KernelDefinition & definition = find_kernel(kernel_name); - ggml::hrx::KernelExecutablePrepareContext context = {}; - context.device = device.device; - context.target = device.architecture.c_str(); - - auto make_dispatch = [&](const std::map & workload) { - ggml::hrx::Dispatch dispatch; - dispatch.kernel.kernel_id = definition.id; - dispatch.kernel.integer_parameters = workload; - dispatch.kernel.compile_parameters = compile_config; - dispatch.kernel.workload_specialization = ggml::hrx::WorkloadSpecialization::Dynamic; - dispatch.bindings.resize(definition.bindings.size()); - return dispatch; - }; - - ggml::hrx::Dispatch first = make_dispatch(first_workload); - ggml::hrx::Dispatch second = make_dispatch(second_workload); - ggml::hrx::KernelExecutableCache cache(ggml::hrx::LoomJitMode::Sync); - ggml::hrx::KernelExecutableRef first_ref = cache.get_or_compile(context, definition, first); - ggml::hrx::KernelExecutableRef second_ref = cache.get_or_compile(context, definition, second); - REQUIRE(first_ref.valid()); - REQUIRE(second_ref.valid()); - REQUIRE(first_ref.entry == second_ref.entry); - - std::vector first_constants; - std::vector second_constants; - ggml_hrx_loom_jit_launch_config first_launch; - ggml_hrx_loom_jit_launch_config second_launch; - std::shared_ptr first_executable = - cache.materialize(context, first_ref, first.kernel, first_constants, first_launch); - std::shared_ptr second_executable = - cache.materialize(context, second_ref, second.kernel, second_constants, second_launch); - REQUIRE(first_executable != nullptr); - REQUIRE(second_executable == first_executable); - REQUIRE(first_launch.workgroup_count[0] > 0); - REQUIRE(second_launch.workgroup_count[0] > 0); - REQUIRE(first_constants.size() == second_constants.size()); - REQUIRE(first_constants != second_constants); - std::printf("dynamic executable reuse materialized %s for two workloads\n", kernel_name); -} - static void run_targeted_export_materialize_case(HrxTestDevice & device) { const ggml::hrx::KernelDefinition * definition = find_targeted_kernel("ggml_linear_q6k_q8_1_x4", device.architecture.c_str()); @@ -543,28 +436,26 @@ static void run_targeted_export_materialize_case(HrxTestDevice & device) { ggml::hrx::KernelExecutableCache cache(ggml::hrx::LoomJitMode::Sync); std::vector constants; - ggml_hrx_loom_jit_launch_config launch; - std::shared_ptr executable = - cache.prepare(context, *definition, dispatch, constants, launch); + std::shared_ptr executable = cache.prepare(context, *definition, dispatch, constants); REQUIRE(executable != nullptr); REQUIRE(executable->executable != nullptr); - REQUIRE(launch.workgroup_count[0] == 151936); + REQUIRE(executable->launch.workgroup_count[0] == 151936); std::printf("KernelExecutableCache materialized targeted %s as export %s for %s\n", definition->symbol, definition->name, device.architecture.c_str()); } } // namespace -static void run_loom_jit_compile_checks() { +int main() { static constexpr const char * kTarget = "gfx1100"; const ggml::hrx::KernelDefinition & binary = find_kernel("ggml_binary_f32"); const ggml::hrx::KernelDefinition & binary_bc = find_kernel("ggml_binary_bc_f32"); const ggml::hrx::KernelDefinition & gather_add = find_kernel("ggml_gather_add_f32"); - const ggml::hrx::KernelDefinition & rmsnorm = find_kernel("ggml_rmsnorm_f32"); + const ggml::hrx::KernelDefinition & rmsnorm = find_kernel("qwen3_moe_rmsnorm_f32"); const ggml::hrx::KernelDefinition & router_top8 = find_kernel("qwen3_moe_router_top8_f32"); - const ggml::hrx::KernelDefinition & expert_table = find_kernel("ggml_moe_build_expert_table"); - const ggml::hrx::KernelDefinition & partition_table = find_kernel("ggml_moe_build_expert_partition_table"); + const ggml::hrx::KernelDefinition & expert_table = find_kernel("qwen3_moe_build_expert_table"); + const ggml::hrx::KernelDefinition & partition_table = find_kernel("qwen3_moe_build_expert_partition_table"); const ggml::hrx::KernelDefinition & mul_mat_id_f16 = find_kernel("ggml_mul_mat_id_f16_f16_wmma"); const ggml::hrx::KernelDefinition & swiglu_f16 = find_kernel("ggml_mul_mat_id_swiglu_f16_f16_wmma"); const ggml::hrx::KernelDefinition & flash_prefill = find_kernel("ggml_flash_attention_f32_f16_wmma"); @@ -618,9 +509,10 @@ static void run_loom_jit_compile_checks() { { "token_count", 1 } }, { - { "ggml.rmsnorm_f32.hidden_size", "2048" }, - { "ggml.rmsnorm_f32.input_stride", "2048" }, - { "ggml.rmsnorm_f32.rms_epsilon", "0.000001" }, + { "qwen3_moe.model.hidden_size", "2048" }, + { "qwen3_moe.model.rms_epsilon", "0.000001" }, + { "qwen3_moe.workload.token_capacity", "1" }, + { "ggml.quantize_q8_1_x4.group_capacity", "256" }, })); refs.push_back(compile_kernel(*async_jit, "async-router-top8-1", router_top8, { @@ -639,23 +531,28 @@ static void run_loom_jit_compile_checks() { { "route_stride", 8 }, { "expert_count", 128 }, }, - moe_routing_config())); + { + { "qwen3_moe.routed_gate_up.expert_count", "128" }, + { "qwen3_moe.routed_gate_up.route_count", "8" }, + { "qwen3_moe.workload.token_capacity", "1" }, + })); refs.push_back(compile_kernel(*async_jit, "async-partition-table-1", partition_table, { { "token_count", 1 }, { "route_count", 8 }, { "expert_count", 128 }, }, - moe_routing_config())); + { + { "qwen3_moe.routed_gate_up.expert_count", "128" }, + { "qwen3_moe.routed_gate_up.route_count", "8" }, + { "qwen3_moe.workload.token_capacity", "1" }, + })); // Exercise gather-add coverage in the async JIT path. refs.push_back(compile_kernel(*async_jit, "async-gather-add-2-to-1", gather_add, { { "source_token_count", 2 }, { "output_token_count", 1 }, { "hidden_size", 2048 }, - }, - { - { "ggml.gather_add_f32.binary_op", "3" }, })); refs.push_back(compile_kernel(*async_jit, "async-mul-mat-id-f16-q4", mul_mat_id_f16, token_count_workload(4), mul_mat_id_f16_f16_config("4"))); @@ -711,7 +608,6 @@ static void run_loom_jit_compile_checks() { std::printf("async Loom JIT compiled %zu kernels\n", refs.size()); const auto & q6_prefill = find_kernel("ggml_mul_mat_q6_k_f16_wmma_prefill_wave32"); - const struct { int64_t tokens; int64_t inputs; @@ -738,13 +634,11 @@ static void run_loom_jit_compile_checks() { {1024, 5120, 1024, 0, 512, 2, 16 }, {2048, 5120, 1024, 0, 512, 4, 16 }, }; - for (const auto & test : q6_launch_cases) { - const std::string key = "q6-prefill-" + std::to_string(test.tokens) + "-" + std::to_string(test.inputs) + "-" + - std::to_string(test.outputs) + "-" + std::to_string(test.accumulate); - auto ref = - compile_kernel(*sync_jit, key, q6_prefill, token_count_workload(test.tokens), - { + const std::string key = "q6-prefill-" + std::to_string(test.tokens) + "-" + + std::to_string(test.inputs) + "-" + std::to_string(test.outputs) + "-" + + std::to_string(test.accumulate); + auto ref = compile_kernel(*sync_jit, key, q6_prefill, token_count_workload(test.tokens), { { "ggml.mul_mat_q6_k_packed.input_size", std::to_string(test.inputs) }, { "ggml.mul_mat_q6_k_packed.output_size", std::to_string(test.outputs) }, { "ggml.mul_mat_q6_k_packed.output_accumulation", std::to_string(test.accumulate) }, @@ -762,7 +656,6 @@ static void run_loom_jit_compile_checks() { } const auto & q4_prefill = find_kernel("ggml_mul_mat_q4_k_f16_wmma_prefill_wave32"); - const struct { int64_t tokens; int64_t inputs; @@ -788,12 +681,11 @@ static void run_loom_jit_compile_checks() { {1024, 5120, 4096, 1024, 4, 16, 0 }, { 512, 17408, 5120, 512, 1, 80, 1 }, }; - for (const auto & test : q4_launch_cases) { - const std::string key = "q4-prefill-" + std::to_string(test.tokens) + "-" + std::to_string(test.inputs) + "-" + - std::to_string(test.outputs) + "-" + std::to_string(test.input_layout); - auto ref = compile_kernel(*sync_jit, key, q4_prefill, token_count_workload(test.tokens), - { + const std::string key = "q4-prefill-" + std::to_string(test.tokens) + "-" + + std::to_string(test.inputs) + "-" + std::to_string(test.outputs) + "-" + + std::to_string(test.input_layout); + auto ref = compile_kernel(*sync_jit, key, q4_prefill, token_count_workload(test.tokens), { { "ggml.mul_mat.input_size", std::to_string(test.inputs) }, { "ggml.mul_mat.output_size", std::to_string(test.outputs) }, { "ggml.workload.token_capacity", std::to_string(test.tokens) }, @@ -811,9 +703,8 @@ static void run_loom_jit_compile_checks() { const auto & packed_f16 = find_kernel("ggml_copy_f16_k16_major"); for (const auto * input_is_f16 : { "0", "1" }) { - auto ref = - compile_kernel(*sync_jit, std::string("packed-f16-") + input_is_f16, packed_f16, token_count_workload(512), - { + auto ref = compile_kernel(*sync_jit, std::string("packed-f16-") + input_is_f16, + packed_f16, token_count_workload(512), { { "ggml.copy_f16_k16_major.input_size", "256" }, { "ggml.copy_f16_k16_major.token_count", "512" }, { "ggml.copy_f16_k16_major.input_is_f16", input_is_f16 }, @@ -822,27 +713,19 @@ static void run_loom_jit_compile_checks() { } const auto & transpose_f16 = find_kernel("ggml_copy_transpose_f16"); - auto transpose_ref = compile_kernel(*sync_jit, "transpose-f16-8192x1024", transpose_f16, - { - { "row_count", 8192 }, - { "column_count", 1024 }, - }, - { + auto transpose_ref = compile_kernel(*sync_jit, "transpose-f16-8192x1024", transpose_f16, {}, { + { "ggml.copy_transpose_f16.row_count", "8192" }, + { "ggml.copy_transpose_f16.column_count", "1024" }, }); require_compiled_kernel(transpose_ref); auto transposed_attention_config = flash_attention_config("24", "4", "256", "0.0625"); transposed_attention_config.emplace("ggml.flash_attention.value_layout", "1"); - auto transposed_attention_ref = - compile_kernel(*sync_jit, "flash-transposed-value", flash_prefill, - { - { "query_token_count", 512 }, - { "key_value_token_count", 8192 } - }, + auto transposed_attention_ref = compile_kernel(*sync_jit, "flash-transposed-value", flash_prefill, + { { "query_token_count", 512 }, { "key_value_token_count", 8192 } }, transposed_attention_config); require_compiled_kernel(transposed_attention_ref); const auto & paired_q4 = find_kernel("ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32"); - const struct { int64_t tokens; int64_t outputs; @@ -856,11 +739,9 @@ static void run_loom_jit_compile_checks() { {1024, 2048, 512, 2 }, {2048,17408, 512, 4 }, }; - for (const auto & test : paired_q4_cases) { const std::string key = "paired-q4-" + std::to_string(test.tokens) + "-" + std::to_string(test.outputs); - auto ref = compile_kernel(*sync_jit, key, paired_q4, token_count_workload(test.tokens), - { + auto ref = compile_kernel(*sync_jit, key, paired_q4, token_count_workload(test.tokens), { { "ggml.mul_mat_swiglu.input_size", "5120" }, { "ggml.mul_mat_swiglu.output_size", std::to_string(test.outputs) }, { "ggml.mul_mat_swiglu.op", "4" }, @@ -879,9 +760,8 @@ static void run_loom_jit_compile_checks() { } for (const int64_t tokens : { 256, 512, 1024, 2048 }) { - auto ref = compile_kernel(*sync_jit, "paired-q4-packed-output-" + std::to_string(tokens), paired_q4, - token_count_workload(tokens), - { + auto ref = compile_kernel(*sync_jit, "paired-q4-packed-output-" + std::to_string(tokens), + paired_q4, token_count_workload(tokens), { { "ggml.mul_mat_swiglu.input_size", "5120" }, { "ggml.mul_mat_swiglu.output_size", "17408" }, { "ggml.mul_mat_swiglu.op", "4" }, @@ -891,8 +771,8 @@ static void run_loom_jit_compile_checks() { require_compiled_kernel(ref); } - const auto & paired_q4_decode = find_kernel("ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot"); + const auto & paired_q4_decode = find_kernel("ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot"); const struct { int64_t tokens; int64_t inputs; @@ -920,12 +800,10 @@ static void run_loom_jit_compile_checks() { { 5, 5632, 4096, 512, 128 }, { 5, 5888, 4096, 128, 512 }, }; - for (const auto & test : paired_q4_decode_cases) { - const std::string key = "paired-q4-decode-" + std::to_string(test.tokens) + "-" + std::to_string(test.inputs) + - "-" + std::to_string(test.outputs); - auto ref = compile_kernel(*sync_jit, key, paired_q4_decode, token_count_workload(test.tokens), - { + const std::string key = "paired-q4-decode-" + std::to_string(test.tokens) + "-" + + std::to_string(test.inputs) + "-" + std::to_string(test.outputs); + auto ref = compile_kernel(*sync_jit, key, paired_q4_decode, token_count_workload(test.tokens), { { "ggml.mul_mat_swiglu.input_size", std::to_string(test.inputs) }, { "ggml.mul_mat_swiglu.output_size", std::to_string(test.outputs) }, { "ggml.mul_mat_swiglu.op", "4" }, @@ -941,8 +819,9 @@ static void run_loom_jit_compile_checks() { compiled.reset(); } - const auto & quant_decode = find_kernel("ggml_mul_mat_f32_f32_wmma"); + + const auto & quant_decode = find_kernel("ggml_mul_mat_f32_f32_wmma"); const struct { int64_t tokens; int64_t inputs; @@ -998,13 +877,11 @@ static void run_loom_jit_compile_checks() { { 5, 8192, 4032, 44, 128, 504 }, { 5, 8192, 8256, 44, 128, 1032 }, }; - for (const auto & test : quant_decode_cases) { - const std::string key = "quant-decode-" + std::to_string(test.tokens) + "-" + std::to_string(test.inputs) + - "-" + std::to_string(test.outputs) + "-" + std::to_string(test.format); - auto ref = - compile_kernel(*sync_jit, key, quant_decode, token_count_workload(test.tokens), - { + const std::string key = "quant-decode-" + std::to_string(test.tokens) + "-" + + std::to_string(test.inputs) + "-" + std::to_string(test.outputs) + "-" + + std::to_string(test.format); + auto ref = compile_kernel(*sync_jit, key, quant_decode, token_count_workload(test.tokens), { { "ggml.mul_mat.input_size", std::to_string(test.inputs) }, { "ggml.mul_mat.output_size", std::to_string(test.outputs) }, { "ggml.mul_mat.weight_format", std::to_string(test.format) }, @@ -1051,10 +928,7 @@ static void run_loom_jit_compile_checks() { const auto & conv_finish = find_kernel("llm_ssm_conv_dconv4_silu_prefill_finish_f32"); REQUIRE(conv_finish.bindings.size() == 5); for (const int64_t channels : { 2112, 8192, 8256, 10240 }) { - auto ref = compile_kernel(*sync_jit, "conv-finish-" + std::to_string(channels), conv_finish, - { - }, - { + auto ref = compile_kernel(*sync_jit, "conv-finish-" + std::to_string(channels), conv_finish, {}, { { "llm.ssm_conv.generic.d_inner", std::to_string(channels) }, }); REQUIRE(ref != nullptr); @@ -1067,10 +941,7 @@ static void run_loom_jit_compile_checks() { const auto & ssm_prefill = find_kernel("llm_ssm_conv_dconv4_silu_prefill_512_wg1024"); for (const int64_t channels : { 8192, 8224, 8256, 10208, 10240 }) { - auto ref = compile_kernel(*sync_jit, "ssm-prefill-" + std::to_string(channels), ssm_prefill, - { - }, - { + auto ref = compile_kernel(*sync_jit, "ssm-prefill-" + std::to_string(channels), ssm_prefill, {}, { { "llm.ssm_conv.prefill.d_conv", "4" }, { "llm.ssm_conv.prefill.d_inner", std::to_string(channels) }, { "llm.ssm_conv.prefill.n_t", "512" }, @@ -1101,14 +972,10 @@ static void run_loom_jit_compile_checks() { { "llm_gated_delta_net_f32_wmma_head128_snapshot", 5, 8 }, { "llm_gated_delta_net_f32_wmma_head128_snapshot_projection_epilogue", 5, 10 }, }; - for (const auto & test : gdn_cases) { const auto & definition = find_kernel(test.name); REQUIRE(definition.bindings.size() == test.bindings); - auto ref = compile_kernel(*sync_jit, test.name, definition, - { - }, - { + auto ref = compile_kernel(*sync_jit, test.name, definition, {}, { { "llm.gated_delta_net.head_width", "128" }, { "llm.gated_delta_net.head_count", "3" }, { "llm.gated_delta_net.token_count", std::to_string(test.tokens) }, @@ -1138,433 +1005,15 @@ static void run_loom_jit_compile_checks() { REQUIRE(compiled.launch_config.workgroup_count[1] == 1); compiled.reset(); } -} - -static void run_moe_routing_compile_checks() { - static constexpr const char * kTarget = "gfx1100"; - - std::string error; - std::unique_ptr jit = ggml::hrx::create_loom_jit(kTarget, ggml::hrx::LoomJitMode::Sync, error); - REQUIRE(jit != nullptr); - - const std::map expert_table_workload = { - { "token_count", 1 }, - { "route_count", 8 }, - { "route_stride", 8 }, - { "expert_count", 128 }, - }; - const auto & expert_table = find_kernel("ggml_moe_build_expert_table"); - require_compiled_kernel( - compile_kernel(*jit, "moe-expert-table", expert_table, expert_table_workload, moe_routing_config())); - - const std::map partition_table_workload = { - { "token_count", 1 }, - { "route_count", 8 }, - { "expert_count", 128 }, - }; - const auto & partition_table = find_kernel("ggml_moe_build_expert_partition_table"); - require_compiled_kernel(compile_kernel(*jit, "moe-expert-partition-table", partition_table, - partition_table_workload, moe_routing_config())); - - const auto & weighted_reduce_publish_q8 = - find_kernel("qwen3_moe_routed_down_weighted_reduce_next_rmsnorm_f32_publish_q8"); - REQUIRE(weighted_reduce_publish_q8.bindings.size() == 6); - require_compiled_kernel(compile_kernel(*jit, "moe-weighted-reduce-next-rmsnorm-publish-q8", - weighted_reduce_publish_q8, token_count_workload(2), { - { "qwen3_moe.model.hidden_size", "2048" }, - { "qwen3_moe.model.rms_epsilon", "0.000001" }, - { "qwen3_moe.routed_down.output_size", "2048" }, - { "qwen3_moe.routed_down.route_count", "8" }, - { "qwen3_moe.workload.token_capacity", "2" }, - { "ggml.quantize_q8_1_x4.group_capacity", "32" }, - })); -} - -static void run_prefill_fusion_compile_checks() { - static constexpr const char * kTarget = "gfx1100"; - - std::string error; - std::unique_ptr jit = ggml::hrx::create_loom_jit(kTarget, ggml::hrx::LoomJitMode::Sync, error); - REQUIRE(jit != nullptr); - - const auto & rmsnorm_f32_f16 = find_kernel("ggml_rmsnorm_binary_f32_f16"); - require_compiled_kernel(compile_kernel(*jit, "prefill-rmsnorm-f32-f16", rmsnorm_f32_f16, token_count_workload(512), - { - { "ggml.rmsnorm_binary_f32.hidden_size", "2048" }, - { "ggml.rmsnorm_binary_f32.rms_epsilon", "0.00001" }, - { "ggml.rmsnorm_binary_f32.op", "2" }, - })); - - const auto & rmsnorm_lowtoken_striped = find_kernel("ggml_rmsnorm_binary_f32_lowtoken_striped"); - require_compiled_kernel(compile_kernel(*jit, "rmsnorm-lowtoken-striped", rmsnorm_lowtoken_striped, - token_count_workload(1), { - { "ggml.rmsnorm_binary_f32.hidden_size", "5120" }, - { "ggml.rmsnorm_binary_f32.rms_epsilon", "0.000001" }, - { "ggml.rmsnorm_binary_f32.op", "2" }, - })); - - const auto & rmsnorm_strided = find_kernel("ggml_rmsnorm_binary_strided_f32"); - require_compiled_kernel(compile_kernel(*jit, "prefill-rmsnorm-strided", rmsnorm_strided, token_count_workload(64), - { - { "ggml.rmsnorm_binary_f32.hidden_size", "256" }, - { "ggml.rmsnorm_binary_f32.input_stride", "288" }, - { "ggml.rmsnorm_binary_f32.rms_epsilon", "0.00001" }, - { "ggml.rmsnorm_binary_f32.op", "2" }, - })); - - const auto & rmsnorm_q8_f16 = find_kernel("ggml_rmsnorm_binary_q8_1_x4_publish"); - require_compiled_kernel(compile_kernel(*jit, "prefill-rmsnorm-q8-f16", rmsnorm_q8_f16, token_count_workload(512), - { - { "ggml.rmsnorm_binary_q8_1_x4.hidden_size", "2048" }, - { "ggml.rmsnorm_binary_q8_1_x4.rms_epsilon", "0.00001" }, - { "ggml.rmsnorm_binary_q8_1_x4.op", "2" }, - { "ggml.rmsnorm_binary_q8_1_x4.publish_f16", "1" }, - })); - - const auto & flash_f16 = find_kernel("ggml_flash_attention_f32_f16_wmma_publish_f16"); - require_compiled_kernel(compile_kernel(*jit, "prefill-flash-f16", flash_f16, flash_attention_workload(512, 512), - flash_attention_config("32", "8", "64", "0.125"))); - - const auto & swiglu_k16 = find_kernel("ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_k16"); - require_compiled_kernel(compile_kernel(*jit, "prefill-swiglu-k16", swiglu_k16, token_count_workload(512), - { - { "ggml.mul_mat_swiglu.input_size", "2048" }, - { "ggml.mul_mat_swiglu.output_size", "8192" }, - { "ggml.mul_mat_swiglu.gate_weight_format", "4" }, - { "ggml.mul_mat_swiglu.up_weight_format", "4" }, - { "ggml.mul_mat_swiglu.op", "4" }, - { "ggml.workload.token_capacity", "512" }, - })); - - const auto & get_rows_rmsnorm = find_kernel("ggml_get_rows_rmsnorm_binary_q8_1_x4_f16"); - require_compiled_kernel(compile_kernel(*jit, "prefill-get-rows-rmsnorm", get_rows_rmsnorm, - { - { "token_count", 512 }, - { "row_count", 128256 }, - { "hidden_size", 2048 }, - }, - { - { "ggml.get_rows_f32.token_capacity", "512" }, - { "ggml.get_rows_f32.hidden_capacity", "2048" }, - { "ggml.get_rows_f32.weight_format", "6" }, - { "ggml.get_rows_rmsnorm.rms_epsilon", "0.00001" }, - })); - - const auto & get_rows_scale = find_kernel("ggml_get_rows_scale_f32"); - require_compiled_kernel(compile_kernel(*jit, "prefill-get-rows-scale", get_rows_scale, - { - { "token_count", 64 }, - { "row_count", 122753 }, - { "hidden_size", 2560 }, - }, - { - { "ggml.get_rows_f32.token_capacity", "64" }, - { "ggml.get_rows_f32.hidden_capacity", "2560" }, - { "ggml.get_rows_f32.weight_format", "32" }, - { "ggml.get_rows_scale_f32.scale", "0.177800179" }, - { "ggml.get_rows_scale_f32.bias", "0" }, - })); - require_compiled_kernel(compile_kernel(*jit, "gemma-prefill-get-rows-scale-q5_1", get_rows_scale, - { - { "token_count", 64 }, - { "row_count", 262144 }, - { "hidden_size", 640 }, - }, - { - { "ggml.get_rows_f32.token_capacity", "64" }, - { "ggml.get_rows_f32.hidden_capacity", "640" }, - { "ggml.get_rows_f32.weight_format", "51" }, - { "ggml.get_rows_scale_f32.scale", "25.2982216" }, - { "ggml.get_rows_scale_f32.bias", "0" }, - })); - require_compiled_kernel(compile_kernel(*jit, "gemma-decode-get-rows-scale-q5_1", get_rows_scale, - { - { "token_count", 1 }, - { "row_count", 262144 }, - { "hidden_size", 640 }, - }, - { - { "ggml.get_rows_f32.token_capacity", "1" }, - { "ggml.get_rows_f32.hidden_capacity", "640" }, - { "ggml.get_rows_f32.weight_format", "51" }, - { "ggml.get_rows_scale_f32.scale", "25.2982216" }, - { "ggml.get_rows_scale_f32.bias", "0" }, - })); - - const auto & lfm_dconv3 = find_kernel("llm_ssm_conv_lfm_dconv3_state_f32"); - REQUIRE(lfm_dconv3.bindings.size() == 5); - for (const int64_t tokens : {1, 64}) { - auto ref = compile_kernel(*jit, "lfm-dconv3-" + std::to_string(tokens), lfm_dconv3, {}, { - {"llm.ssm_conv.lfm_dconv3.n_t", std::to_string(tokens)}, - }); - REQUIRE(ref != nullptr); - REQUIRE(resolve_with_timeout(ref, std::chrono::seconds(120))); - auto compiled = ref->take_result(); - REQUIRE(compiled.hsaco_data != nullptr); - REQUIRE(compiled.launch_config.workgroup_size[0] == 256); - REQUIRE(compiled.launch_config.workgroup_count[0] == static_cast(8 * tokens)); - } -} -static void run_binary_fusion_compile_checks() { - std::string error; - std::unique_ptr jit = - ggml::hrx::create_loom_jit("gfx1100", ggml::hrx::LoomJitMode::Sync, error); - REQUIRE(jit != nullptr); - - const auto & gather = find_kernel("ggml_gather_add_f32"); - require_compiled_kernel(compile_kernel(*jit, "binary-fusion-gather-div", gather, - { - { "source_token_count", 2 }, - { "output_token_count", 1 }, - { "hidden_size", 2048 }, - }, - { - { "ggml.gather_add_f32.binary_op", "3" }, - })); - - const auto & scale = find_kernel("ggml_scale_add_f32"); - require_compiled_kernel(compile_kernel(*jit, "binary-fusion-scale-sub", scale, - { - { "element_count", 128 }, - }, - { - { "ggml.scale_add_f32.ne0", "128" }, - { "ggml.scale_add_f32.ne1", "1" }, - { "ggml.scale_add_f32.ne2", "1" }, - { "ggml.scale_add_f32.input_stride1", "128" }, - { "ggml.scale_add_f32.input_stride2", "128" }, - { "ggml.scale_add_f32.input_stride3", "128" }, - { "ggml.scale_add_f32.input_span", "128" }, - { "ggml.scale_add_f32.residual_stride1", "128" }, - { "ggml.scale_add_f32.residual_stride2", "128" }, - { "ggml.scale_add_f32.residual_stride3", "128" }, - { "ggml.scale_add_f32.residual_span", "128" }, - { "ggml.scale_add_f32.scale", "0.5" }, - { "ggml.scale_add_f32.bias", "0" }, - { "ggml.scale_add_f32.binary_op", "1" }, - { "ggml.scale_add_f32.scaled_lhs", "0" }, - })); - - const auto & rms = find_kernel("ggml_rmsnorm_mul_add_f32"); - require_compiled_kernel(compile_kernel(*jit, "binary-fusion-rms-div-sub", rms, - { - { "token_count", 1 }, - }, - { - { "ggml.rmsnorm_mul_add_f32.hidden_size", "128" }, - { "ggml.rmsnorm_mul_add_f32.rms_epsilon", "0.000001" }, - { "ggml.rmsnorm_mul_add_f32.normalized_op", "3" }, - { "ggml.rmsnorm_mul_add_f32.output_op", "1" }, - { "ggml.rmsnorm_mul_add_f32.binary_lhs", "0" }, - })); - - const auto & gather_rms = find_kernel("ggml_gather_add_rmsnorm_binary_f32"); - require_compiled_kernel(compile_kernel(*jit, "binary-fusion-gather-rms-mul", gather_rms, - { - { "source_token_count", 2 }, - { "output_token_count", 1 }, - }, - { - { "ggml.gather_add_rmsnorm_binary_f32.hidden_size", "128" }, - { "ggml.gather_add_rmsnorm_binary_f32.rms_epsilon", "0.000001" }, - { "ggml.gather_add_rmsnorm_binary_f32.gather_op", "2" }, - })); - - const auto & ssm = find_kernel("llm_ssm_conv_binary_f32"); - require_compiled_kernel(compile_kernel(*jit, "binary-fusion-ssm-rhs-div", ssm, - { - { "n_t", 8 }, - { "n_s", 1 }, - }, - { - { "llm.ssm_conv.generic.d_conv", "4" }, - { "llm.ssm_conv.generic.d_inner", "64" }, - { "llm.ssm_conv.generic.unary_op", "23" }, - { "llm.ssm_conv.generic.binary_op", "3" }, - { "llm.ssm_conv.generic.binary_lhs", "0" }, - { "llm.ssm_conv.generic.workgroup_size", "256" }, - })); - - const auto & vector_postops = find_kernel("ggml_mul_mat_vector_bias_residual_f32_f32"); - for (const int64_t epilogue : { int64_t{ 1 }, int64_t{ 2 }, int64_t{ 3 } }) { - require_compiled_kernel(compile_kernel( - *jit, "binary-fusion-vector-postops-" + std::to_string(epilogue), vector_postops, token_count_workload(1), - { - { "ggml.matmul.vector_postops.input_size", "2048" }, - { "ggml.matmul.vector_postops.output_size", "2048" }, - { "ggml.matmul.vector_postops.weight_format", "4" }, - { "ggml.matmul.vector_postops.apply_bias", epilogue & 1 ? "1" : "0" }, - { "ggml.matmul.vector_postops.apply_residual", epilogue & 2 ? "1" : "0" }, - { "ggml.workload.token_capacity", "1" }, - })); - } - - const auto compile_postops = [&](const char * kernel_name, const char * key_prefix, int64_t token_count, - int64_t token_capacity) { - const auto & kernel = find_kernel(kernel_name); - for (const int64_t epilogue : { int64_t{ 1 }, int64_t{ 2 }, int64_t{ 3 } }) { - require_compiled_kernel(compile_kernel( - *jit, std::string(key_prefix) + std::to_string(epilogue), kernel, token_count_workload(token_count), - { - { "ggml.mul_mat_postops.input_size", "2048" }, - { "ggml.mul_mat_postops.output_size", "2048" }, - { "ggml.mul_mat_postops.weight_format", "4" }, - { "ggml.mul_mat_postops.epilogue", std::to_string(epilogue) }, - { "ggml.workload.token_capacity", std::to_string(token_capacity) }, - })); - } - }; - compile_postops("ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", "binary-fusion-skinny-postops-", 2, 2); - compile_postops("ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", "binary-fusion-tiled-postops-", 32, 32); - - const auto & pair_postops = find_kernel("ggml_mul_mat_tiled_pair_input_f32_binary_bias_residual_publish_f32"); - for (const int64_t epilogue : { int64_t{ 1 }, int64_t{ 2 }, int64_t{ 3 } }) { - require_compiled_kernel(compile_kernel(*jit, "binary-fusion-pair-postops-" + std::to_string(epilogue), - pair_postops, token_count_workload(32), - { - { "ggml.mul_mat_swiglu.input_size", "2048" }, - { "ggml.mul_mat_swiglu.output_size", "2048" }, - { "ggml.mul_mat_swiglu.gate_weight_format", "4" }, - { "ggml.mul_mat_swiglu.up_weight_format", "4" }, - { "ggml.mul_mat_swiglu.op", "4" }, - { "ggml.mul_mat_swiglu.epilogue", std::to_string(epilogue) }, - { "ggml.workload.token_capacity", "32" }, - })); - } -} - -static void run_loom_jit_materialize_checks() { - const ggml::hrx::KernelDefinition & binary = find_kernel("ggml_binary_f32"); - - HrxTestDevice device; - REQUIRE(device.open()); - run_cache_materialize_case(device, ggml::hrx::LoomJitMode::Sync, binary); - run_cache_materialize_case(device, ggml::hrx::LoomJitMode::Async, binary); - run_dynamic_cache_materialize_case(device, ggml::hrx::LoomJitMode::Sync, binary); - run_dynamic_cache_materialize_case(device, ggml::hrx::LoomJitMode::Async, binary); - run_dynamic_export_reuse_case(device, "ggml_unary_f32", { { "element_count", 64 } }, - { { "element_count", 257 } }, { { "ggml.unary_f32.op", "4" } }); - run_dynamic_export_reuse_case(device, "ggml_copy_f32", { { "element_count", 64 } }, - { { "element_count", 257 } }); - run_dynamic_export_reuse_case( - device, "ggml_concat_dim0_f32", - { { "row_count", 3 }, { "lhs_span", 192 }, { "rhs_span", 96 } }, - { { "row_count", 7 }, { "lhs_span", 448 }, { "rhs_span", 224 } }, - { { "ggml.concat_dim0_f32.lhs_width", "64" }, { "ggml.concat_dim0_f32.rhs_width", "32" } }); - run_dynamic_export_reuse_case( - device, "ggml_rmsnorm_f32", { { "token_count", 1 } }, { { "token_count", 4 } }, - { { "ggml.rmsnorm_f32.hidden_size", "256" }, { "ggml.rmsnorm_f32.input_stride", "256" }, - { "ggml.rmsnorm_f32.rms_epsilon", "1e-06" } }); - run_dynamic_export_reuse_case( - device, "ggml_get_rows_f32", - { { "token_count", 3 }, { "row_count", 64 }, { "hidden_size", 256 } }, - { { "token_count", 7 }, { "row_count", 64 }, { "hidden_size", 256 } }, - { { "ggml.get_rows_f32.token_capacity", "8" }, { "ggml.get_rows_f32.hidden_capacity", "256" }, - { "ggml.get_rows_f32.weight_format", "32" } }); - run_dynamic_export_reuse_case( - device, "ggml_rope_f32", { { "token_count", 3 }, { "input_span", 768 } }, - { { "token_count", 7 }, { "input_span", 1792 } }, - { { "ggml.rope_f32.head_size", "64" }, { "ggml.rope_f32.n_dims", "64" }, - { "ggml.rope_f32.head_count", "4" }, { "ggml.rope_f32.token_capacity", "8" }, - { "ggml.rope_f32.input_stride1", "64" }, { "ggml.rope_f32.input_stride2", "256" }, - { "ggml.rope_f32.mscale", "1" }, { "ggml.rope_f32.mode", "0" } }); - run_dynamic_export_reuse_case( - device, "ggml_mul_mat_tiled_input_f32_publish_f32_aligned", { { "token_count", 32 } }, - { { "token_count", 64 } }, - { { "ggml.mul_mat.input_size", "256" }, { "ggml.mul_mat.output_size", "64" }, - { "ggml.mul_mat.output_accumulation", "0" }, { "ggml.mul_mat.output_unary_op", "23" }, - { "ggml.mul_mat.weight_format", "4" } }); - run_dynamic_export_reuse_case( - device, "ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_aligned", { { "token_count", 32 } }, - { { "token_count", 64 } }, - { { "ggml.mul_mat_swiglu.input_size", "256" }, { "ggml.mul_mat_swiglu.output_size", "64" }, - { "ggml.mul_mat_swiglu.gate_weight_format", "4" }, - { "ggml.mul_mat_swiglu.up_weight_format", "4" }, { "ggml.mul_mat_swiglu.op", "4" } }); - run_dynamic_export_reuse_case( - device, "llm_ssm_conv_f32", { { "n_t", 8 }, { "n_s", 1 } }, { { "n_t", 16 }, { "n_s", 2 } }, - { { "llm.ssm_conv.generic.d_conv", "4" }, { "llm.ssm_conv.generic.d_inner", "64" }, - { "llm.ssm_conv.generic.unary_op", "23" }, { "llm.ssm_conv.generic.workgroup_size", "256" } }); - run_dynamic_export_reuse_case( - device, "llm_ssm_conv_binary_f32", { { "n_t", 8 }, { "n_s", 1 } }, - { { "n_t", 16 }, { "n_s", 2 } }, - { { "llm.ssm_conv.generic.d_conv", "4" }, { "llm.ssm_conv.generic.d_inner", "64" }, - { "llm.ssm_conv.generic.unary_op", "23" }, { "llm.ssm_conv.generic.binary_op", "3" }, - { "llm.ssm_conv.generic.binary_lhs", "0" }, - { "llm.ssm_conv.generic.workgroup_size", "256" } }); - run_dynamic_export_reuse_case(device, "ggml_copy_transpose_f16", - { { "row_count", 512 }, { "column_count", 256 } }, - { { "row_count", 1024 }, { "column_count", 512 } }); - run_dynamic_export_reuse_case( - device, "ggml_flash_attention_f32_f16_wmma", flash_attention_workload(16, 64), - flash_attention_workload(17, 65), flash_attention_config("32", "4", "128", "0.0883883461")); - run_dynamic_export_reuse_case( - device, "ggml_flash_attention_f32_f16_wmma_publish_f16", flash_attention_workload(16, 64), - flash_attention_workload(17, 65), flash_attention_config("32", "4", "128", "0.0883883461")); - run_dynamic_export_reuse_case( - device, "ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", - flash_attention_decode_workload(65), flash_attention_decode_workload(127), - flash_attention_decode_config("32", "4", "128", "0.0883883461", "128")); - run_dynamic_export_reuse_case( - device, "llm_gated_delta_net_projection_epilogue_f32", { { "element_count", 48 } }, - { { "element_count", 192 } }, - { { "llm.gated_delta_net.epilogue_head_count", "48" }, - { "llm.gated_delta_net.epilogue_workgroup_size", "256" } }); - run_targeted_export_materialize_case(device); -} - -static void run_binary_publication_compile_checks() { - std::string error; - std::unique_ptr jit = - ggml::hrx::create_loom_jit("gfx1100", ggml::hrx::LoomJitMode::Sync, error); - REQUIRE(jit != nullptr); - - const auto & binary_q8 = find_kernel("ggml_binary_f32_publish_q8_1_x4"); - require_compiled_kernel(compile_kernel(*jit, "binary-mul-q8", binary_q8, { - { "token_count", 2 }, - }, { - { "ggml.binary_f32.op", "2" }, - { "ggml.binary_f32.ne0", "2048" }, - { "ggml.binary_f32.src0_stride1", "2048" }, - { "ggml.binary_f32.src1_stride1", "2048" }, - { "ggml.binary_f32.src0_span", "4096" }, - { "ggml.binary_f32.src1_span", "4096" }, - })); - - const auto & binary_bc_q8 = find_kernel("ggml_binary_bc_f32_publish_q8_1_x4"); - require_compiled_kernel(compile_kernel(*jit, "binary-bc-mul-q8", binary_bc_q8, { - { "token_count", 2 }, - { "hidden_size", 256 }, - { "src0_element_count", 512 }, - { "src1_element_count", 256 }, - }, { - { "ggml.binary_bc_f32.op", "2" }, - { "ggml.binary_bc_f32.src0_broadcast_dim0", "0" }, - { "ggml.binary_bc_f32.src0_broadcast_dim1", "0" }, - { "ggml.binary_bc_f32.src1_broadcast_dim0", "0" }, - { "ggml.binary_bc_f32.src1_broadcast_dim1", "1" }, - })); -} - -static bool has_hrx_test_device() { HrxTestDevice device; - return device.open(); -} - -static void register_hrx_loom_jit_cases(test_runner::Suite & suite) { - suite.host_case("compile", [] { run_loom_jit_compile_checks(); }); - suite.host_case("binary-publication", [] { run_binary_publication_compile_checks(); }); - suite.host_case("moe-routing", [] { run_moe_routing_compile_checks(); }); - suite.host_case("prefill-fusions", [] { run_prefill_fusion_compile_checks(); }); - suite.host_case("binary-fusions", [] { run_binary_fusion_compile_checks(); }); - suite.device_case("materialize", [] { run_loom_jit_materialize_checks(); }); -} - -int main(int argc, char ** argv) { - test_runner::Suite suite( - test_runner::Config::with_prefix("HRX Loom JIT test", "hrx-loom-jit", "GGML_HRX_LOOM_JIT_TEST", 2)); - register_hrx_loom_jit_cases(suite); + if (device.open()) { + run_cache_materialize_case(device, ggml::hrx::LoomJitMode::Sync, binary); + run_cache_materialize_case(device, ggml::hrx::LoomJitMode::Async, binary); + run_targeted_export_materialize_case(device); + } else { + std::printf("skipping KernelExecutableCache materialization checks: no HRX device available\n"); + } - const bool has_device = has_hrx_test_device(); - return suite.run(argc, argv, has_device); + return 0; } diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp index 6cc2f14900c8..51162048f255 100644 --- a/tests/test-hrx-ops.cpp +++ b/tests/test-hrx-ops.cpp @@ -1,24 +1,19 @@ #include "backend-context.h" #include "dispatch/dispatch-scheduler.h" -#include "dispatch_registration/common/dispatch-activation-publication.h" #include "ggml-alloc.h" #include "ggml-backend.h" #include "ggml-hrx.h" -#include "ggml-quants.h" #include "ggml.h" #include "graph/graph.h" #include "kernel-corpus/kernel-corpus.h" #include "runtime/graph-executor.h" -#include "testing_suite.h" #include -#include #include #include #include #include #include -#include #include #include #include @@ -103,12 +98,6 @@ static std::vector make_weight(int64_t hidden_size) { return data; } -static const char * expected_plain_rmsnorm_binary_kernel(int64_t hidden_size, int64_t token_count) { - return token_count < 6 && hidden_size >= 2048 && hidden_size <= 5120 ? - "loom_libs:ggml_rmsnorm_binary_f32_lowtoken_striped" : - "loom_libs:ggml_rmsnorm_binary_f32"; -} - static std::vector make_router_input(int64_t hidden_size, int64_t token_count) { std::vector data(hidden_size * token_count); for (int64_t i = 0; i < static_cast(data.size()); ++i) { @@ -397,22 +386,6 @@ static std::vector scheduled_kernel_sequence(ggml_cgraph * graph) { return names; } -static ggml::hrx::KernelSpecialization scheduled_kernel_specialization(ggml_cgraph * graph, - const char * kernel_name) { - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" })); - REQUIRE(scheduler.plan().valid()); - for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { - if (kernel_name_for_id(dispatch.kernel.kernel_id) == kernel_name) { - return dispatch.kernel; - } - } - REQUIRE(false); - return {}; -} - static size_t producer_index_for_tensor(const ggml::hrx::Graph & graph, const ggml_tensor * tensor) { const ggml::hrx::Value * value = graph.values().find_tensor(tensor); REQUIRE(value != nullptr); @@ -463,7 +436,7 @@ static void match_dispatch_at_index(const ggml::hrx::Graph & graph, ggml::hrx::DispatchMatch & match) { REQUIRE(root_index < graph.nodes().size()); const ggml::hrx::DispatchMatchContext context = { - graph, &graph.nodes()[root_index], root_index, covered_nodes, plan, next_plan_value(graph, plan), ®istry, + graph, &graph.nodes()[root_index], root_index, covered_nodes, plan, next_plan_value(graph, plan), }; ggml::hrx::DispatchMatchDiagnostics diagnostics; if (!registry.match(context, match, &diagnostics)) { @@ -524,11 +497,10 @@ static void require_kernel_subsequence(const std::vector & sequence } static void run_alternate_value_alias_lookup_checks() { - constexpr int64_t element_count = 2048; - const size_t full_bytes = ggml_row_size(GGML_TYPE_F32, element_count); - const size_t f16_bytes = ggml_row_size(GGML_TYPE_F16, element_count); - const size_t q8_bytes = ggml_row_size(GGML_TYPE_Q8_1, element_count); - const size_t symmetric_bytes = 1536; + constexpr int64_t element_count = 2048; + const size_t full_bytes = ggml_row_size(GGML_TYPE_F32, element_count); + const size_t f16_bytes = ggml_row_size(GGML_TYPE_F16, element_count); + const size_t q8_bytes = ggml_row_size(GGML_TYPE_Q8_1, element_count); ggml::hrx::Graph graph; ggml::hrx::Status status; @@ -537,10 +509,6 @@ static void run_alternate_value_alias_lookup_checks() { const ggml::hrx::ValueId root(0); const ggml::hrx::ValueId full_alias(1); const ggml::hrx::ValueId partial_alias(2); - const ggml::hrx::ValueId reordered_alias(3); - const ggml::hrx::ValueId unrelated_same_storage(4); - const ggml::hrx::ValueId different_shape_alias(5); - const ggml::hrx::ValueId incompatible_q8_shape_alias(6); const ggml::hrx::ValueId f16_alternate(99); const ggml::hrx::ValueId q8_alternate(100); const ggml::hrx::ValueStorageId storage(0); @@ -556,41 +524,9 @@ static void run_alternate_value_alias_lookup_checks() { status = graph.values().add_snapshot_value( make_test_value(partial_alias, storage, root, root, 0, full_bytes, GGML_TYPE_F32, element_count / 2)); REQUIRE(status.success()); - ggml::hrx::Value reordered = - make_test_value(reordered_alias, storage, root, root, 0, full_bytes, GGML_TYPE_F32, element_count); - reordered.contiguous = false; - status = graph.values().add_snapshot_value(std::move(reordered)); - REQUIRE(status.success()); - status = graph.values().add_snapshot_value( - make_test_value(unrelated_same_storage, storage, root, root, 0, full_bytes, GGML_TYPE_F32, element_count)); - REQUIRE(status.success()); - ggml::hrx::Value different_shape = - make_test_value(different_shape_alias, storage, root, root, 0, full_bytes, GGML_TYPE_F32, element_count); - different_shape.ne = { element_count / 2, 2, 1, 1 }; - different_shape.nb = { sizeof(float), static_cast(element_count / 2) * sizeof(float), full_bytes, - full_bytes }; - status = graph.values().add_snapshot_value(std::move(different_shape)); - REQUIRE(status.success()); - ggml::hrx::Value incompatible_q8_shape = - make_test_value(incompatible_q8_shape_alias, storage, root, root, 0, full_bytes, GGML_TYPE_F32, element_count); - incompatible_q8_shape.ne = { 64, element_count / 64, 1, 1 }; - incompatible_q8_shape.nb = { sizeof(float), 64 * sizeof(float), full_bytes, full_bytes }; - status = graph.values().add_snapshot_value(std::move(incompatible_q8_shape)); - REQUIRE(status.success()); - - graph.add_node(GGML_OP_RESHAPE, full_alias, { root }); - graph.add_node(GGML_OP_VIEW, partial_alias, { root }); - graph.add_node(GGML_OP_TRANSPOSE, reordered_alias, { root }); - graph.add_node(GGML_OP_CPY, unrelated_same_storage, { root }); - graph.add_node(GGML_OP_RESHAPE, different_shape_alias, { root }); - graph.add_node(GGML_OP_RESHAPE, incompatible_q8_shape_alias, { root }); - status = graph.build_index(); - REQUIRE(status.success()); REQUIRE(plan.metadata.append_alternate_value({ root, q8_alternate, GGML_TYPE_Q8_1, q8_bytes, "q8" }, status)); REQUIRE(plan.metadata.append_alternate_value({ root, f16_alternate, GGML_TYPE_F16, f16_bytes, "f16" }, status)); - REQUIRE(plan.metadata.append_alternate_value( - { root, ggml::hrx::ValueId(104), GGML_TYPE_COUNT, symmetric_bytes, "symmetric-i4" }, status)); const ggml::hrx::CommandPlanAlternateValue * exact = ggml::hrx::find_alternate_value(graph, plan, root, GGML_TYPE_Q8_1, q8_bytes); @@ -606,412 +542,14 @@ static void run_alternate_value_alias_lookup_checks() { ggml::hrx::find_alternate_value(graph, plan, partial_alias, GGML_TYPE_Q8_1, q8_bytes); REQUIRE(through_partial_alias == nullptr); - const ggml::hrx::CommandPlanAlternateValue * through_reordered_alias = - ggml::hrx::find_alternate_value(graph, plan, reordered_alias, GGML_TYPE_Q8_1, q8_bytes); - REQUIRE(through_reordered_alias == nullptr); - - REQUIRE(ggml::hrx::find_alternate_value(graph, plan, unrelated_same_storage, GGML_TYPE_Q8_1, q8_bytes) == nullptr); - REQUIRE(ggml::hrx::find_alternate_value(graph, plan, different_shape_alias, GGML_TYPE_Q8_1, q8_bytes) != nullptr); - REQUIRE(ggml::hrx::find_alternate_value(graph, plan, incompatible_q8_shape_alias, GGML_TYPE_Q8_1, q8_bytes) == - nullptr); - const ggml::hrx::CommandPlanAlternateValue * f16_exact = ggml::hrx::find_alternate_value(graph, plan, root, GGML_TYPE_F16, f16_bytes); REQUIRE(f16_exact != nullptr); REQUIRE(f16_exact->alternate_value == f16_alternate); - REQUIRE(ggml::hrx::find_alternate_value(graph, plan, different_shape_alias, GGML_TYPE_F16, f16_bytes) != nullptr); - REQUIRE(ggml::hrx::find_alternate_value(graph, plan, full_alias, GGML_TYPE_COUNT, symmetric_bytes) != nullptr); - REQUIRE(ggml::hrx::find_alternate_value(graph, plan, different_shape_alias, GGML_TYPE_COUNT, symmetric_bytes) == - nullptr); ggml::hrx::Status conflict_status; REQUIRE(!plan.metadata.append_alternate_value( { root, ggml::hrx::ValueId(101), GGML_TYPE_Q8_1, q8_bytes, "q8-conflict" }, conflict_status)); - - ggml::hrx::Status generated_status; - const ggml::hrx::ValueId k16_value(102); - REQUIRE(plan.metadata.append_generated_resource( - { root, ggml::hrx::GeneratedResourceRole::F16K16Major, k16_value, f16_bytes, {} }, generated_status)); - const ggml::hrx::CommandPlanGeneratedResource * generated_through_alias = ggml::hrx::find_generated_resource( - graph, plan, full_alias, ggml::hrx::GeneratedResourceRole::F16K16Major, f16_bytes); - REQUIRE(generated_through_alias != nullptr); - REQUIRE(generated_through_alias->generated_value == k16_value); - REQUIRE(ggml::hrx::find_generated_resource(graph, plan, partial_alias, - ggml::hrx::GeneratedResourceRole::F16K16Major, f16_bytes) == nullptr); - REQUIRE(ggml::hrx::find_generated_resource(graph, plan, reordered_alias, - ggml::hrx::GeneratedResourceRole::F16K16Major, f16_bytes) == nullptr); - REQUIRE(ggml::hrx::find_generated_resource(graph, plan, different_shape_alias, - ggml::hrx::GeneratedResourceRole::F16K16Major, f16_bytes) == nullptr); -} - -static bool publication_accepts_add(const ggml::hrx::DispatchMatchContext &, - const ggml::hrx::GraphNode & consumer, - const ggml::hrx::Value & input) { - return consumer.op == GGML_OP_ADD && - std::find(consumer.inputs.begin(), consumer.inputs.end(), input.id) != consumer.inputs.end(); -} - -static bool publication_accepts_mul(const ggml::hrx::DispatchMatchContext &, - const ggml::hrx::GraphNode & consumer, - const ggml::hrx::Value & input) { - return consumer.op == GGML_OP_MUL && - std::find(consumer.inputs.begin(), consumer.inputs.end(), input.id) != consumer.inputs.end(); -} - -static bool publication_accepts_mul_output_13(const ggml::hrx::DispatchMatchContext & context, - const ggml::hrx::GraphNode & consumer, - const ggml::hrx::Value & input) { - return consumer.output == ggml::hrx::ValueId(13) && publication_accepts_mul(context, consumer, input); -} - -static bool publication_accepts_mul_output_15(const ggml::hrx::DispatchMatchContext & context, - const ggml::hrx::GraphNode & consumer, - const ggml::hrx::Value & input) { - return consumer.output == ggml::hrx::ValueId(15) && publication_accepts_mul(context, consumer, input); -} - -static void run_activation_publication_contract_checks() { - constexpr int64_t row_size = 2048; - constexpr int64_t row_count = 2; - constexpr int64_t elements = row_size * row_count; - const size_t f32_bytes = ggml_row_size(GGML_TYPE_F32, elements); - - auto make_activation = [&](ggml::hrx::ValueId id, ggml::hrx::ValueStorageId storage, ggml::hrx::ValueId root, - ggml::hrx::ValueId alias, int64_t element_count = elements) { - ggml::hrx::Value value = make_test_value(id, storage, root, alias, 0, f32_bytes, GGML_TYPE_F32, element_count); - if (element_count == elements) { - value.ne = { row_size, row_count, 1, 1 }; - value.nb = { sizeof(float), static_cast(row_size) * sizeof(float), f32_bytes, f32_bytes }; - } - return value; - }; - - ggml::hrx::Graph graph; - ggml::hrx::Status status; - const ggml::hrx::ValueId root(0); - const ggml::hrx::ValueId reshaped(1); - const ggml::hrx::ValueId partial(2); - const ggml::hrx::ValueId transposed(3); - const ggml::hrx::ValueId different_shape(8); - const ggml::hrx::ValueId incompatible_q8_shape(10); - const ggml::hrx::ValueId shape_back(12); - const ggml::hrx::ValueId mismatched_alias_source(14); - status = graph.values().add_snapshot_storage({ ggml::hrx::ValueStorageId(0), root, f32_bytes }); - REQUIRE(status.success()); - status = graph.values().add_snapshot_value( - make_activation(root, ggml::hrx::ValueStorageId(0), root, ggml::hrx::ValueId())); - REQUIRE(status.success()); - status = graph.values().add_snapshot_value(make_activation(reshaped, ggml::hrx::ValueStorageId(0), root, root)); - REQUIRE(status.success()); - status = graph.values().add_snapshot_value( - make_activation(partial, ggml::hrx::ValueStorageId(0), root, root, elements / 2)); - REQUIRE(status.success()); - ggml::hrx::Value reordered = make_activation(transposed, ggml::hrx::ValueStorageId(0), root, root); - reordered.contiguous = false; - status = graph.values().add_snapshot_value(std::move(reordered)); - REQUIRE(status.success()); - - for (int32_t value = 4; value < 8; ++value) { - const ggml::hrx::ValueStorageId storage(value - 3); - const ggml::hrx::ValueId id(value); - status = graph.values().add_snapshot_storage({ storage, id, f32_bytes }); - REQUIRE(status.success()); - status = graph.values().add_snapshot_value(make_activation(id, storage, id, ggml::hrx::ValueId())); - REQUIRE(status.success()); - } - ggml::hrx::Value different = make_activation(different_shape, ggml::hrx::ValueStorageId(0), root, root); - different.ne = { row_size / 2, row_count * 2, 1, 1 }; - different.nb = { sizeof(float), static_cast(row_size / 2) * sizeof(float), f32_bytes, f32_bytes }; - status = graph.values().add_snapshot_value(std::move(different)); - REQUIRE(status.success()); - status = graph.values().add_snapshot_storage({ ggml::hrx::ValueStorageId(5), ggml::hrx::ValueId(9), f32_bytes }); - REQUIRE(status.success()); - status = graph.values().add_snapshot_value(make_activation(ggml::hrx::ValueId(9), ggml::hrx::ValueStorageId(5), - ggml::hrx::ValueId(9), ggml::hrx::ValueId())); - REQUIRE(status.success()); - ggml::hrx::Value incompatible = make_activation(incompatible_q8_shape, ggml::hrx::ValueStorageId(0), root, root); - incompatible.ne = { 64, elements / 64, 1, 1 }; - incompatible.nb = { sizeof(float), 64 * sizeof(float), f32_bytes, f32_bytes }; - status = graph.values().add_snapshot_value(std::move(incompatible)); - REQUIRE(status.success()); - status = graph.values().add_snapshot_storage({ ggml::hrx::ValueStorageId(6), ggml::hrx::ValueId(11), f32_bytes }); - REQUIRE(status.success()); - status = graph.values().add_snapshot_value(make_activation(ggml::hrx::ValueId(11), ggml::hrx::ValueStorageId(6), - ggml::hrx::ValueId(11), ggml::hrx::ValueId())); - REQUIRE(status.success()); - status = graph.values().add_snapshot_value( - make_activation(shape_back, ggml::hrx::ValueStorageId(0), root, different_shape)); - REQUIRE(status.success()); - status = graph.values().add_snapshot_storage({ ggml::hrx::ValueStorageId(7), ggml::hrx::ValueId(13), f32_bytes }); - REQUIRE(status.success()); - status = graph.values().add_snapshot_value(make_activation(ggml::hrx::ValueId(13), ggml::hrx::ValueStorageId(7), - ggml::hrx::ValueId(13), ggml::hrx::ValueId())); - REQUIRE(status.success()); - ggml::hrx::Value mismatched = make_activation( - mismatched_alias_source, ggml::hrx::ValueStorageId(0), root, root); - mismatched.ne = { row_size / 2, row_count * 2, 1, 1 }; - mismatched.nb = { sizeof(float), static_cast(row_size / 2) * sizeof(float), f32_bytes, f32_bytes }; - status = graph.values().add_snapshot_value(std::move(mismatched)); - REQUIRE(status.success()); - status = graph.values().add_snapshot_storage({ ggml::hrx::ValueStorageId(8), ggml::hrx::ValueId(15), f32_bytes }); - REQUIRE(status.success()); - status = graph.values().add_snapshot_value(make_activation(ggml::hrx::ValueId(15), ggml::hrx::ValueStorageId(8), - ggml::hrx::ValueId(15), ggml::hrx::ValueId())); - REQUIRE(status.success()); - - graph.add_node(GGML_OP_RESHAPE, reshaped, { root }); - graph.add_node(GGML_OP_VIEW, partial, { root }); - graph.add_node(GGML_OP_TRANSPOSE, transposed, { root }); - graph.add_node(GGML_OP_ADD, ggml::hrx::ValueId(4), { root, root }); - graph.add_node(GGML_OP_MUL, ggml::hrx::ValueId(5), { reshaped, reshaped }); - graph.add_node(GGML_OP_MUL, ggml::hrx::ValueId(6), { partial, partial }); - graph.add_node(GGML_OP_MUL, ggml::hrx::ValueId(7), { transposed, transposed }); - graph.add_node(GGML_OP_RESHAPE, different_shape, { root }); - graph.add_node(GGML_OP_MUL, ggml::hrx::ValueId(9), { different_shape, different_shape }); - graph.add_node(GGML_OP_RESHAPE, incompatible_q8_shape, { root }); - graph.add_node(GGML_OP_MUL, ggml::hrx::ValueId(11), { incompatible_q8_shape, incompatible_q8_shape }); - graph.add_node(GGML_OP_RESHAPE, shape_back, { different_shape }); - graph.add_node(GGML_OP_MUL, ggml::hrx::ValueId(13), { shape_back, shape_back }); - graph.add_node(GGML_OP_RESHAPE, mismatched_alias_source, { different_shape }); - graph.add_node(GGML_OP_MUL, ggml::hrx::ValueId(15), { mismatched_alias_source, mismatched_alias_source }); - status = graph.build_index(); - REQUIRE(status.success()); - - ggml::hrx::DispatchRegistryBuilder publication_registry_builder; - publication_registry_builder.add_activation_consumer({ - "test.q8.mul", - GGML_OP_MUL, - ggml::hrx::DispatchActivationInputQ8_1X4, - ggml::hrx::DispatchActivationConsumerUse::BandwidthLimited, - ggml::hrx::DispatchActivationProducerRequirementNone, - ggml::hrx::DispatchSource::Common, - publication_accepts_mul, - }); - publication_registry_builder.add_activation_consumer({ - "test.f16-row.add", - GGML_OP_ADD, - ggml::hrx::DispatchActivationInputF16Row, - ggml::hrx::DispatchActivationConsumerUse::Row, - ggml::hrx::DispatchActivationProducerRequirementNone, - ggml::hrx::DispatchSource::Common, - publication_accepts_add, - }); - const ggml::hrx::DispatchRegistry publication_registry = publication_registry_builder.build(); - const std::vector covered(graph.nodes().size(), false); - ggml::hrx::CommandPlan plan; - const ggml::hrx::DispatchMatchContext context = { - graph, &graph.nodes().front(), 0, covered, plan, ggml::hrx::ValueId(100), &publication_registry, - }; - const ggml::hrx::Value * produced = graph.values().find(root); - REQUIRE(produced != nullptr); - - const ggml::hrx::CommonActivationPublicationCapabilities direct_only = { - ggml::hrx::DispatchActivationInputQ8_1X4, - }; - const ggml::hrx::CommonActivationPublicationPlan direct_only_plan = - ggml::hrx::common_select_activation_publication_plan(context, *produced, direct_only); - REQUIRE(!direct_only_plan.matched()); - REQUIRE(direct_only_plan.fallback_reason == - ggml::hrx::CommonActivationPublicationFallbackReason::NoQualifiedConsumer); - - const ggml::hrx::CommonActivationPublicationCapabilities q8_capabilities = { - ggml::hrx::DispatchActivationInputQ8_1X4, - ggml::hrx::DispatchActivationInputQ8_1X4, - }; - const ggml::hrx::CommonActivationPublicationPlan q8_plan = - ggml::hrx::common_select_activation_publication_plan(context, *produced, q8_capabilities); - REQUIRE(q8_plan.matched()); - const ggml::hrx::CommonActivationPublicationDemand q8_demand = q8_plan.preferred(); - REQUIRE(q8_demand.format == ggml::hrx::CommonActivationPublicationFormat::Q8_1X4); - REQUIRE(q8_demand.consumer_value->id == reshaped); - REQUIRE(q8_demand.alias_depth == 1); - REQUIRE(q8_demand.mixed_fanout); - - const ggml::hrx::CommonActivationPublicationCapabilities f16_capabilities = { - ggml::hrx::DispatchActivationInputF16Row, - }; - const ggml::hrx::CommonActivationPublicationPlan f16_plan = - ggml::hrx::common_select_activation_publication_plan(context, *produced, f16_capabilities); - REQUIRE(f16_plan.matched()); - const ggml::hrx::CommonActivationPublicationDemand f16_demand = f16_plan.preferred(); - REQUIRE(f16_demand.format == ggml::hrx::CommonActivationPublicationFormat::F16Row); - REQUIRE(f16_demand.consumer_value->id == root); - REQUIRE(f16_demand.alias_depth == 0); - REQUIRE(f16_demand.mixed_fanout); - - ggml::hrx::DispatchRegistryBuilder direct_k16_registry_builder; - direct_k16_registry_builder.add_activation_consumer({ - "test.k16.mul", - GGML_OP_MUL, - ggml::hrx::DispatchActivationInputF16K16Major, - ggml::hrx::DispatchActivationConsumerUse::Tiled, - ggml::hrx::DispatchActivationProducerRequirementNone, - ggml::hrx::DispatchSource::Common, - publication_accepts_mul, - }); - const ggml::hrx::DispatchRegistry direct_k16_registry = direct_k16_registry_builder.build(); - const ggml::hrx::DispatchMatchContext direct_k16_context = { - graph, &graph.nodes().front(), 0, covered, plan, ggml::hrx::ValueId(100), &direct_k16_registry, - }; - const ggml::hrx::CommonActivationPublicationCapabilities mixed_alias_capabilities = { - ggml::hrx::DispatchActivationInputQ8_1X4 | ggml::hrx::DispatchActivationInputF16K16Major, - ggml::hrx::DispatchActivationInputQ8_1X4, - ggml::hrx::DispatchActivationProducerRequirementNone, - }; - const ggml::hrx::CommonActivationPublicationDemand direct_k16_through_q8_alias = - ggml::hrx::common_select_activation_publication_plan( - direct_k16_context, *produced, mixed_alias_capabilities).publications.front(); - REQUIRE(!direct_k16_through_q8_alias.matched()); - - ggml::hrx::DispatchRegistryBuilder shape_back_registry_builder; - shape_back_registry_builder.add_activation_consumer({ - "test.k16.shape-back", - GGML_OP_MUL, - ggml::hrx::DispatchActivationInputF16K16Major, - ggml::hrx::DispatchActivationConsumerUse::Tiled, - ggml::hrx::DispatchActivationProducerRequirementNone, - ggml::hrx::DispatchSource::Common, - publication_accepts_mul_output_13, - }); - const ggml::hrx::DispatchRegistry shape_back_registry = shape_back_registry_builder.build(); - const ggml::hrx::DispatchMatchContext shape_back_context = { - graph, &graph.nodes().front(), 0, covered, plan, ggml::hrx::ValueId(100), &shape_back_registry, - }; - REQUIRE(!ggml::hrx::common_select_activation_publication_plan( - shape_back_context, *produced, mixed_alias_capabilities).matched()); - - ggml::hrx::CommonActivationPublicationCapabilities multiple_capabilities = { - ggml::hrx::DispatchActivationInputQ8_1X4 | ggml::hrx::DispatchActivationInputF16Row, - ggml::hrx::DispatchActivationInputQ8_1X4, - ggml::hrx::DispatchActivationProducerRequirementNone, - ggml::hrx::DispatchActivationInputQ8_1X4 | ggml::hrx::DispatchActivationInputF16Row, - }; - ggml::hrx::DispatchRegistryBuilder mismatched_alias_registry_builder; - mismatched_alias_registry_builder.add_activation_consumer({ - "test.q8.mismatched-alias-source", - GGML_OP_MUL, - ggml::hrx::DispatchActivationInputQ8_1X4, - ggml::hrx::DispatchActivationConsumerUse::BandwidthLimited, - ggml::hrx::DispatchActivationProducerRequirementNone, - ggml::hrx::DispatchSource::Common, - publication_accepts_mul_output_15, - }); - const ggml::hrx::DispatchRegistry mismatched_alias_registry = mismatched_alias_registry_builder.build(); - const ggml::hrx::DispatchMatchContext mismatched_alias_context = { - graph, &graph.nodes().front(), 0, covered, plan, ggml::hrx::ValueId(100), &mismatched_alias_registry, - }; - REQUIRE(!ggml::hrx::common_select_activation_publication_plan( - mismatched_alias_context, *produced, multiple_capabilities).matched()); - - const ggml::hrx::CommonActivationPublicationPlan multiple_plan = - ggml::hrx::common_select_activation_publication_plan(context, *produced, multiple_capabilities); - REQUIRE(multiple_plan.publication_count == 2); - REQUIRE(multiple_plan.publications[0].format == ggml::hrx::CommonActivationPublicationFormat::Q8_1X4); - REQUIRE(multiple_plan.publications[1].format == ggml::hrx::CommonActivationPublicationFormat::F16Row); - ggml::hrx::DispatchRegistryBuilder q8_only_registry_builder; - q8_only_registry_builder.add_activation_consumer({ - "test.q8.only", - GGML_OP_MUL, - ggml::hrx::DispatchActivationInputQ8_1X4, - ggml::hrx::DispatchActivationConsumerUse::BandwidthLimited, - ggml::hrx::DispatchActivationProducerRequirementNone, - ggml::hrx::DispatchSource::Common, - publication_accepts_mul, - }); - const ggml::hrx::DispatchRegistry q8_only_registry = q8_only_registry_builder.build(); - const ggml::hrx::DispatchMatchContext q8_only_context = { - graph, &graph.nodes().front(), 0, covered, plan, ggml::hrx::ValueId(100), &q8_only_registry, - }; - const ggml::hrx::CommonActivationPublicationPlan companion_plan = - ggml::hrx::common_select_activation_publication_plan(q8_only_context, *produced, multiple_capabilities); - REQUIRE(companion_plan.publication_count == 2); - REQUIRE(companion_plan.publications[0].format == ggml::hrx::CommonActivationPublicationFormat::Q8_1X4); - REQUIRE(companion_plan.publications[1].format == ggml::hrx::CommonActivationPublicationFormat::F16Row); - REQUIRE(companion_plan.publications[1].consumer == companion_plan.publications[0].consumer); - multiple_capabilities.co_publication_formats = ggml::hrx::DispatchActivationInputNone; - REQUIRE(ggml::hrx::common_select_activation_publication_plan(context, *produced, multiple_capabilities) - .publication_count == 1); - - const ggml::hrx::CommonActivationPublicationCapabilities disabled_capabilities; - const ggml::hrx::CommonActivationPublicationPlan disabled_plan = - ggml::hrx::common_select_activation_publication_plan(context, *produced, disabled_capabilities); - REQUIRE(!disabled_plan.matched()); - REQUIRE(disabled_plan.fallback_reason == - ggml::hrx::CommonActivationPublicationFallbackReason::NoEnabledCandidates); - - ggml::hrx::DispatchMatch match; - match.completion_counter_requests.push_back({ ggml::hrx::ValueId(100), "existing.counter", 1 }); - const ggml::hrx::CommonActivationPublicationDemand symmetric_demand = { - ggml::hrx::CommonActivationPublicationFormat::SymmetricI4K32, - produced, - &graph.nodes()[3], - }; - ggml::hrx::CommonActivationPublication symmetric; - ggml::hrx::CommonActivationPublication q8; - ggml::hrx::CommonActivationPublication f16; - ggml::hrx::CommonActivationPublication k16; - REQUIRE(ggml::hrx::common_reserve_activation_publication(context, match, *produced, symmetric_demand, - "test.symmetric_i4", symmetric)); - REQUIRE(ggml::hrx::common_append_activation_publication(match, symmetric, "test.symmetric_i4")); - REQUIRE(ggml::hrx::common_reserve_activation_publication(context, match, *produced, q8_demand, "test.q8", q8)); - REQUIRE(ggml::hrx::common_append_activation_publication(match, q8, "test.q8")); - REQUIRE(ggml::hrx::common_reserve_activation_publication(context, match, *produced, f16_demand, "test.f16", f16)); - REQUIRE(ggml::hrx::common_append_activation_publication(match, f16, "test.f16")); - const ggml::hrx::CommonActivationPublicationDemand k16_demand = { - ggml::hrx::CommonActivationPublicationFormat::F16K16Major, - q8_demand.consumer_value, - q8_demand.consumer, - }; - REQUIRE(ggml::hrx::common_reserve_activation_publication(context, match, *produced, k16_demand, "test.k16", k16)); - REQUIRE(ggml::hrx::common_append_activation_publication(match, k16, "test.k16")); - - ggml::hrx::DispatchMatch invalid_id_match; - const ggml::hrx::DispatchMatchContext invalid_id_context = { - graph, &graph.nodes().front(), 0, covered, plan, ggml::hrx::ValueId(-1), - }; - ggml::hrx::CommonActivationPublication invalid_id_publication; - REQUIRE(!ggml::hrx::common_reserve_activation_publication(invalid_id_context, invalid_id_match, *produced, - f16_demand, "test.invalid_id", invalid_id_publication)); - REQUIRE(invalid_id_match.transients.empty()); - - ggml::hrx::DispatchMatch conflict_match; - REQUIRE(ggml::hrx::common_append_activation_publication(conflict_match, q8, "test.q8")); - ggml::hrx::CommonActivationPublication conflicting_q8 = q8; - conflicting_q8.alternate_value = ggml::hrx::ValueId(200); - REQUIRE(!ggml::hrx::common_append_activation_publication(conflict_match, conflicting_q8, "test.q8.conflict")); - REQUIRE(!conflict_match.status.success()); - - REQUIRE(match.status.success()); - REQUIRE(match.transients.size() == 4); - REQUIRE(symmetric.alternate_value == ggml::hrx::ValueId(101)); - REQUIRE(q8.alternate_value == ggml::hrx::ValueId(102)); - REQUIRE(f16.alternate_value == ggml::hrx::ValueId(103)); - REQUIRE(k16.alternate_value == ggml::hrx::ValueId(104)); - REQUIRE(symmetric.payload_bytes == static_cast(elements) / 2); - REQUIRE(symmetric.scales_offset % 256 == 0); - REQUIRE(symmetric.sums_offset % 256 == 0); - REQUIRE(q8.byte_count == static_cast(row_count) * ggml_row_size(GGML_TYPE_Q8_1, row_size)); - REQUIRE(f16.byte_count == static_cast(elements) * sizeof(ggml_fp16_t)); - REQUIRE(symmetric.binding(symmetric.scales_offset, symmetric.scales_bytes).value == symmetric.alternate_value); - - ggml::hrx::CommandPlan recorded; - recorded.metadata = std::move(match.metadata); - const std::vector & diagnostics = - recorded.metadata.activation_publication_diagnostics(); - REQUIRE(diagnostics.size() == 4); - REQUIRE(diagnostics[1].source_value == root); - REQUIRE(diagnostics[1].alternate_value == q8.alternate_value); - REQUIRE(diagnostics[1].consumer_value == reshaped); - REQUIRE(diagnostics[1].requested_format == "q8-1-x4"); - REQUIRE(diagnostics[1].publication_name == "test.q8"); - REQUIRE(diagnostics[1].publication_stage == "published"); - REQUIRE(diagnostics[1].fallback_reason == "none"); - REQUIRE(diagnostics[1].alias_depth == 1); - REQUIRE(diagnostics[1].mixed_fanout); - REQUIRE(ggml::hrx::find_alternate_value(graph, recorded, reshaped, GGML_TYPE_Q8_1, q8.byte_count) != nullptr); - REQUIRE(ggml::hrx::find_alternate_value(graph, recorded, reshaped, GGML_TYPE_F16, f16.byte_count) != nullptr); - REQUIRE(ggml::hrx::find_generated_resource(graph, recorded, reshaped, ggml::hrx::GeneratedResourceRole::F16K16Major, - k16.byte_count) != nullptr); } static std::vector make_pattern_f32(size_t element_count, int seed, float scale = 0.01f) { @@ -1060,27 +598,10 @@ static std::vector make_pattern_f16(size_t element_count, int seed, static std::vector make_quantized_rows(ggml_type type, int64_t row_length, int64_t row_count, int seed) { const ggml_type_traits * traits = ggml_get_type_traits(type); REQUIRE(traits != nullptr); + REQUIRE(traits->from_float_ref != nullptr); ggml_quantize_init(type); const size_t row_size = ggml_row_size(type, row_length); std::vector data(static_cast(row_count) * row_size); - if (type == GGML_TYPE_IQ2_XXS || type == GGML_TYPE_IQ2_XS) { - std::vector source(static_cast(row_length * row_count)); - std::vector importance(static_cast(row_length), 1.0f); - for (int64_t r = 0; r < row_count; ++r) { - for (int64_t c = 0; c < row_length; ++c) { - const int value = static_cast((r * 13 + c * 7 + seed * 31) % 101) - 50; - source[static_cast(r * row_length + c)] = static_cast(value) * 0.005f; - } - } - const size_t bytes_written = type == GGML_TYPE_IQ2_XXS ? - quantize_iq2_xxs(source.data(), data.data(), row_count, row_length, - importance.data()) : - quantize_iq2_xs(source.data(), data.data(), row_count, row_length, - importance.data()); - REQUIRE(bytes_written == data.size()); - return data; - } - REQUIRE(traits->from_float_ref != nullptr); std::vector row(static_cast(row_length)); for (int64_t r = 0; r < row_count; ++r) { for (int64_t c = 0; c < row_length; ++c) { @@ -1092,59 +613,7 @@ static std::vector make_quantized_rows(ggml_type type, int64_t row_leng return data; } -static std::vector make_iq1_rows(ggml_type type, int64_t row_length, int64_t row_count, int seed) { - REQUIRE(type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M); - REQUIRE(row_length % 256 == 0); - const size_t row_size = ggml_row_size(type, row_length); - const size_t block_size = type == GGML_TYPE_IQ1_S ? 50 : 56; - REQUIRE(row_size == static_cast(row_length / 256) * block_size); - std::vector data(static_cast(row_count) * row_size); - for (int64_t row = 0; row < row_count; ++row) { - uint8_t * row_data = data.data() + static_cast(row) * row_size; - for (int64_t block = 0; block < row_length / 256; ++block) { - uint8_t * dst = row_data + static_cast(block) * block_size; - if (type == GGML_TYPE_IQ1_S) { - const ggml_fp16_t d = ggml_fp32_to_fp16(0.015625f); - std::memcpy(dst, &d, sizeof(d)); - for (int i = 0; i < 32; ++i) { - dst[2 + i] = static_cast((row * 13 + block * 29 + i * 37 + seed) & 0xff); - } - for (int group = 0; group < 8; ++group) { - uint16_t qh = 0; - for (int entry = 0; entry < 4; ++entry) { - qh |= static_cast(((row + block + group + entry + seed) & 7) << (3 * entry)); - } - qh |= static_cast(((row + block + group) & 7) << 12); - qh |= static_cast(((row + group) & 1) << 15); - std::memcpy(dst + 34 + 2 * group, &qh, sizeof(qh)); - } - } else { - for (int i = 0; i < 32; ++i) { - dst[i] = static_cast((row * 17 + block * 31 + i * 41 + seed) & 0xff); - } - for (int i = 0; i < 16; ++i) { - const uint8_t low = static_cast((row + block + i + seed) & 7); - const uint8_t high = static_cast((row + block + i + seed + 3) & 7); - dst[32 + i] = static_cast(low | ((i & 1) ? 0x08 : 0) | (high << 4) | - ((i & 2) ? 0x80 : 0)); - } - const uint16_t base_bits = ggml_fp32_to_fp16(0.015625f); - for (int word = 0; word < 4; ++word) { - const uint16_t local_scales = static_cast(1 | (2 << 3) | (3 << 6) | (4 << 9)); - const uint16_t base_nibble = static_cast((base_bits >> (4 * word)) & 0xf); - const uint16_t scales = static_cast(local_scales | (base_nibble << 12)); - std::memcpy(dst + 48 + 2 * word, &scales, sizeof(scales)); - } - } - } - } - return data; -} - static std::vector make_matmul_weight_bytes(ggml_type type, int64_t row_length, int64_t row_count, int seed) { - if (type == GGML_TYPE_IQ1_S || type == GGML_TYPE_IQ1_M) { - return make_iq1_rows(type, row_length, row_count, seed); - } if (type == GGML_TYPE_F32) { (void) seed; const std::vector weights(static_cast(row_length * row_count), 0.00390625f); @@ -1598,7 +1067,8 @@ static void run_rmsnorm_support_checks() { REQUIRE(q8_graph != nullptr); ggml_build_forward_expand(q8_graph, q8_output); require_kernel_subsequence(scheduled_kernel_sequence(q8_graph), - { "loom_libs:ggml_rmsnorm_binary_f32", "loom_libs:ggml_mul_mat_vector_q6_f32_f32" }); + { "loom_libs:ggml_rmsnorm_binary_f32", + "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma" }); ggml_tensor * wrong_type_input = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, 256, 1); ggml_tensor * wrong_type_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F16, 256); @@ -1643,8 +1113,7 @@ static void run_rmsnorm_mul_case(int64_t hidden_size, int64_t token_count) { REQUIRE(graph != nullptr); ggml_build_forward_expand(graph, output); - require_kernel_subsequence(scheduled_kernel_sequence(graph), - { expected_plain_rmsnorm_binary_kernel(hidden_size, token_count) }); + require_kernel_subsequence(scheduled_kernel_sequence(graph), { "loom_libs:ggml_rmsnorm_binary_f32" }); ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); REQUIRE(buffer != nullptr); @@ -1737,147 +1206,6 @@ static void run_rmsnorm_binary_add_cpu_reference_case() { ggml_backend_free(hrx_backend); } -static void run_gemma_post_proj_cpu_reference_case(int64_t token_count) { - constexpr int64_t hidden_size = 640; - constexpr float epsilon = 1.0e-6f; - ggml_backend_t cpu_backend = init_cpu_backend(); - ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); - REQUIRE(hrx_backend != nullptr); - - ggml_init_params params = {}; - params.mem_size = 1024 * 1024; - params.no_alloc = true; - ggml_context * cpu_ctx = ggml_init(params); - ggml_context * hrx_ctx = ggml_init(params); - REQUIRE(cpu_ctx != nullptr); - REQUIRE(hrx_ctx != nullptr); - - struct GraphTensors { - ggml_tensor * input; - ggml_tensor * weight; - ggml_tensor * residual; - ggml_tensor * raw; - ggml_tensor * raw_side; - ggml_tensor * output; - }; - - auto build = [&](ggml_context * ctx) { - GraphTensors g = {}; - g.input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, token_count); - g.weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - g.residual = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, token_count); - g.raw = ggml_add(ctx, g.input, g.residual); - ggml_tensor * normalized = ggml_rms_norm(ctx, g.raw, epsilon); - ggml_tensor * scaled = ggml_mul(ctx, normalized, g.weight); - g.output = ggml_add(ctx, scaled, g.residual); - g.raw_side = ggml_mul(ctx, g.raw, g.residual); - REQUIRE(g.output != nullptr); - REQUIRE(g.raw_side != nullptr); - return g; - }; - GraphTensors cpu = build(cpu_ctx); - GraphTensors hrx = build(hrx_ctx); - ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); - ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); - REQUIRE(cpu_graph != nullptr); - REQUIRE(hrx_graph != nullptr); - ggml_build_forward_expand(cpu_graph, cpu.output); - ggml_build_forward_expand(cpu_graph, cpu.raw_side); - ggml_build_forward_expand(hrx_graph, hrx.output); - ggml_build_forward_expand(hrx_graph, hrx.raw_side); - require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rmsnorm_mul_add_f32" }); - - ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); - ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); - REQUIRE(cpu_buffer != nullptr); - REQUIRE(hrx_buffer != nullptr); - const std::vector input = make_pattern_f32(hidden_size * token_count, 37, 0.02f); - const std::vector residual = make_pattern_f32(hidden_size * token_count, 43, 0.04f); - const std::vector weight = make_weight(hidden_size); - set_tensor_pair_bytes(cpu_backend, cpu.input, hrx_backend, hrx.input, input.data(), input.size() * sizeof(float)); - set_tensor_pair_bytes(cpu_backend, cpu.residual, hrx_backend, hrx.residual, residual.data(), - residual.size() * sizeof(float)); - set_tensor_pair_bytes(cpu_backend, cpu.weight, hrx_backend, hrx.weight, weight.data(), - weight.size() * sizeof(float)); - REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); - REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(cpu_backend); - ggml_backend_synchronize(hrx_backend); - require_close(get_f32_tensor(hrx_backend, hrx.output), get_f32_tensor(cpu_backend, cpu.output), 5.0e-4f); - require_close(get_f32_tensor(hrx_backend, hrx.raw_side), get_f32_tensor(cpu_backend, cpu.raw_side), 1.0e-5f); - - ggml_backend_buffer_free(cpu_buffer); - ggml_backend_buffer_free(hrx_buffer); - ggml_free(cpu_ctx); - ggml_free(hrx_ctx); - ggml_backend_free(cpu_backend); - ggml_backend_free(hrx_backend); -} - -static void run_gemma_post_proj_matcher_negatives() { - constexpr int64_t hidden_size = 640; - ggml_init_params params = {}; - params.mem_size = 1024 * 1024; - params.no_alloc = true; - const ggml::hrx::DispatchRegistry * registry = ggml::hrx::find_dispatch_registry({ "gfx1151" }); - REQUIRE(registry != nullptr); - - for (const int variant : { 0, 1, 2, 3, 4 }) { - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 3); - ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, hidden_size); - ggml_tensor * rhs = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 3); - ggml_tensor * raw = ggml_add(ctx, input, rhs); - ggml_tensor * residual = variant == 2 ? ggml_view_2d(ctx, raw, hidden_size, 3, raw->nb[1], 0) : - variant == 3 ? ggml_mul(ctx, raw, rhs) : - rhs; - ggml_tensor * rms = ggml_rms_norm(ctx, raw, 1.0e-6f); - ggml_tensor * scaled = variant == 4 ? ggml_add(ctx, rms, weight) : ggml_mul(ctx, rms, weight); - ggml_tensor * output = ggml_add(ctx, scaled, residual); - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - if (variant == 1) { - ggml_build_forward_expand(graph, ggml_mul(ctx, scaled, rhs)); - } - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - const size_t rms_index = producer_index_for_tensor(imported.graph, rms); - const size_t raw_index = producer_index_for_tensor(imported.graph, raw); - ggml::hrx::CommandPlan plan; - std::vector covered(imported.graph.nodes().size(), false); - auto kernel = [&]() { - ggml::hrx::DispatchMatch match; - const ggml::hrx::DispatchMatchContext context = { - imported.graph, &imported.graph.nodes()[rms_index], rms_index, covered, - plan, next_plan_value(imported.graph, plan), - }; - REQUIRE(registry->match(context, match)); - REQUIRE(match.dispatches.size() == 1); - return kernel_name_for_id(match.dispatches[0].kernel.kernel_id); - }; - REQUIRE(kernel() == "loom_libs:ggml_rmsnorm_binary_f32"); - covered[raw_index] = true; - if (variant == 3) { - REQUIRE(kernel() == "loom_libs:ggml_rmsnorm_binary_f32"); - covered[producer_index_for_tensor(imported.graph, residual)] = true; - } - if (variant == 2) { - const size_t view_index = producer_index_for_tensor(imported.graph, residual); - covered[view_index] = true; - const ggml::hrx::Value * raw_value = imported.graph.values().find_tensor(raw); - const ggml::hrx::Value * residual_value = imported.graph.values().find_tensor(residual); - REQUIRE(raw_value != nullptr); - REQUIRE(residual_value != nullptr); - REQUIRE(raw_value->storage_root == residual_value->storage_root); - } - REQUIRE(kernel() == (variant == 0 || variant == 3 ? "loom_libs:ggml_rmsnorm_mul_add_f32" : - "loom_libs:ggml_rmsnorm_binary_f32")); - ggml_free(ctx); - } -} - static void run_rmsnorm_gate_cpu_reference_case(int64_t hidden_size, int64_t head_count, int64_t token_count, @@ -1900,7 +1228,6 @@ static void run_rmsnorm_gate_cpu_reference_case(int64_t hidden_size, ggml_tensor * output; ggml_cgraph * graph; }; - auto build = [&](ggml_context * ctx) { Graph g; g.input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hidden_size, head_count, token_count); @@ -1915,7 +1242,7 @@ static void run_rmsnorm_gate_cpu_reference_case(int64_t hidden_size, }; Graph a = build(cpu_ctx); Graph b = build(hrx_ctx); - require_kernel_subsequence(scheduled_kernel_sequence(b.graph), { "loom_libs:ggml_rmsnorm_gate_f32_publish" }); + require_kernel_subsequence(scheduled_kernel_sequence(b.graph), { "loom_libs:ggml_rmsnorm_gate_f32_f16" }); ggml_backend_buffer_t a_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu); ggml_backend_buffer_t b_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx); REQUIRE(a_buffer != nullptr && b_buffer != nullptr); @@ -1987,8 +1314,6 @@ static void run_rmsnorm_cpu_reference_case(int64_t hidden_size = 256, ggml_build_forward_expand(hrx_graph, hrx_output); require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rmsnorm_f32" }); - REQUIRE(scheduled_kernel_specialization(hrx_graph, "loom_libs:ggml_rmsnorm_f32").workload_specialization == - ggml::hrx::WorkloadSpecialization::Dynamic); ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); @@ -2084,77 +1409,6 @@ static void run_strided_rmsnorm_cpu_reference_case(int64_t token_count) { ggml_backend_free(hrx_backend); } -static void run_strided_rmsnorm_binary_cpu_reference_case() { - ggml_backend_t cpu_backend = init_cpu_backend(); - ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); - REQUIRE(hrx_backend != nullptr); - - constexpr int64_t hidden_size = 256; - constexpr int64_t token_count = 64; - constexpr int64_t input_stride = 288; - constexpr float epsilon = 1.0e-5f; - const size_t storage_elements = static_cast((token_count - 1) * input_stride + hidden_size); - - ggml_init_params params = {}; - params.mem_size = 2 * 1024 * 1024; - params.no_alloc = true; - ggml_context * cpu_ctx = ggml_init(params); - ggml_context * hrx_ctx = ggml_init(params); - REQUIRE(cpu_ctx != nullptr); - REQUIRE(hrx_ctx != nullptr); - - ggml_tensor * cpu_storage = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, storage_elements); - ggml_tensor * cpu_input = - ggml_view_2d(cpu_ctx, cpu_storage, hidden_size, token_count, input_stride * sizeof(float), 0); - ggml_tensor * cpu_scale = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, hidden_size); - ggml_tensor * cpu_output = ggml_mul(cpu_ctx, ggml_rms_norm(cpu_ctx, cpu_input, epsilon), cpu_scale); - ggml_tensor * hrx_storage = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, storage_elements); - ggml_tensor * hrx_input = - ggml_view_2d(hrx_ctx, hrx_storage, hidden_size, token_count, input_stride * sizeof(float), 0); - ggml_tensor * hrx_scale = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, hidden_size); - ggml_tensor * hrx_output = ggml_mul(hrx_ctx, ggml_rms_norm(hrx_ctx, hrx_input, epsilon), hrx_scale); - REQUIRE(cpu_output != nullptr); - REQUIRE(hrx_output != nullptr); - - ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); - ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); - REQUIRE(cpu_graph != nullptr); - REQUIRE(hrx_graph != nullptr); - ggml_build_forward_expand(cpu_graph, cpu_output); - ggml_build_forward_expand(hrx_graph, hrx_output); - require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rmsnorm_binary_strided_f32" }); - - ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); - ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); - REQUIRE(cpu_buffer != nullptr); - REQUIRE(hrx_buffer != nullptr); - - const std::vector rows = make_pattern_f32(hidden_size * token_count, 54, 0.015f); - const std::vector scale = make_pattern_f32(hidden_size, 55, 0.03f); - std::vector storage(storage_elements, -19.0f); - for (int64_t token = 0; token < token_count; ++token) { - std::memcpy(storage.data() + static_cast(token * input_stride), - rows.data() + static_cast(token * hidden_size), - static_cast(hidden_size) * sizeof(float)); - } - set_tensor_pair_bytes(cpu_backend, cpu_storage, hrx_backend, hrx_storage, storage.data(), - storage.size() * sizeof(float)); - set_tensor_pair_bytes(cpu_backend, cpu_scale, hrx_backend, hrx_scale, scale.data(), scale.size() * sizeof(float)); - - REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); - REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(cpu_backend); - ggml_backend_synchronize(hrx_backend); - require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 5.0e-4f); - - ggml_backend_buffer_free(cpu_buffer); - ggml_backend_buffer_free(hrx_buffer); - ggml_free(cpu_ctx); - ggml_free(hrx_ctx); - ggml_backend_free(cpu_backend); - ggml_backend_free(hrx_backend); -} - static void run_strided_cont_cpu_reference_case(int64_t token_count) { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); @@ -2224,10 +1478,7 @@ static void run_strided_cont_cpu_reference_case(int64_t token_count) { ggml_backend_free(hrx_backend); } -static void run_rmsnorm_mul_cpu_reference_case(int64_t hidden_size, - int64_t token_count, - float epsilon, - bool adversarial = false) { +static void run_rmsnorm_mul_cpu_reference_case(int64_t hidden_size, int64_t token_count, float epsilon) { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); REQUIRE(hrx_backend != nullptr); @@ -2259,29 +1510,14 @@ static void run_rmsnorm_mul_cpu_reference_case(int64_t hidden_size, ggml_build_forward_expand(cpu_graph, cpu_output); ggml_build_forward_expand(hrx_graph, hrx_output); - require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), - { expected_plain_rmsnorm_binary_kernel(hidden_size, token_count) }); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rmsnorm_binary_f32" }); ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); REQUIRE(cpu_buffer != nullptr); REQUIRE(hrx_buffer != nullptr); - std::vector input = make_pattern_f32(hidden_size * token_count, 23, 0.02f); - if (adversarial) { - for (size_t i = 0; i < input.size(); ++i) { - switch (i % 8) { - case 0: input[i] = 0.0f; break; - case 1: input[i] = 1.0e-6f; break; - case 2: input[i] = -1.0e-6f; break; - case 3: input[i] = 8.0f; break; - case 4: input[i] = -8.0f; break; - case 5: input[i] = 0.125f; break; - case 6: input[i] = -0.25f; break; - default: input[i] = 1.5f; break; - } - } - } + const std::vector input = make_pattern_f32(hidden_size * token_count, 23, 0.02f); const std::vector weight = make_weight(hidden_size); set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), @@ -2442,10 +1678,6 @@ static void run_scheduled_hrx_scale_rmsnorm_case() { REQUIRE(graph != nullptr); ggml_build_forward_expand(graph, output); - const ggml::hrx::KernelSpecialization scale_specialization = - scheduled_kernel_specialization(graph, "loom_libs:ggml_scale_f32"); - REQUIRE(scale_specialization.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); - REQUIRE(ggml_backend_sched_alloc_graph(sched, graph)); REQUIRE(ggml_backend_sched_get_tensor_backend(sched, scaled) == hrx_backend); REQUIRE(ggml_backend_sched_get_tensor_backend(sched, output) == hrx_backend); @@ -2813,15 +2045,6 @@ static void run_qwen_flash_attention_case() { require_kernel_subsequence(scheduled_kernel_sequence(graph), { "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8" }); - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" })); - const ggml::hrx::Value * output_value = imported.graph.values().find_tensor(output); - REQUIRE(output_value != nullptr); - REQUIRE(scheduler.plan().metadata.find_alternate_value( - output_value->id, GGML_TYPE_Q8_1, - static_cast(kQueryTokenCount) * ggml_row_size(GGML_TYPE_Q8_1, kQwenFlashHeadSize)) == nullptr); static constexpr const char * kDisableQwenDispatchEnv = "GGML_HRX_DISABLE_QWEN_DISPATCH"; const char * original_env = std::getenv(kDisableQwenDispatchEnv); const bool had_original_env = original_env != nullptr; @@ -2913,11 +2136,6 @@ static void run_common_flash_attention_cpu_reference_case(int64_t head_size ggml_build_forward_expand(graph, output); require_kernel_subsequence(scheduled_kernel_sequence(graph), { expected_kernel }); - const ggml::hrx::WorkloadSpecialization expected_specialization = - query_token_count == 256 && key_value_token_count == 256 ? ggml::hrx::WorkloadSpecialization::Exact : - ggml::hrx::WorkloadSpecialization::Dynamic; - REQUIRE(scheduled_kernel_specialization(graph, expected_kernel).workload_specialization == - expected_specialization); ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); REQUIRE(buffer != nullptr); @@ -2988,8 +2206,6 @@ static void run_asymmetric_flash_attention_cpu_reference_case() { ggml_build_forward_expand(graph, output); require_kernel_subsequence(scheduled_kernel_sequence(graph), { "loom_libs:ggml_flash_attention_f32_f16_wmma" }); - REQUIRE(scheduled_kernel_specialization(graph, "loom_libs:ggml_flash_attention_f32_f16_wmma") - .workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); REQUIRE(buffer != nullptr); @@ -3119,10 +2335,7 @@ static void run_qwen_decode_split_flash_attention_scheduling_case(int64_t query_ const int64_t output_hidden_size = query_head_count * value_head_size; const size_t q8_row_bytes = ggml_row_size(GGML_TYPE_Q8_1, output_hidden_size); const size_t q8_output_bytes = static_cast(query_token_count) * q8_row_bytes; - int64_t key_value_capacity = 64; - while (key_value_capacity < key_value_token_count) { - key_value_capacity *= 2; - } + const int64_t key_value_capacity = (key_value_token_count + 63) / 64 * 64; const int64_t key_value_blocks = key_value_capacity / 64; const size_t partial_scalar_bytes = static_cast(key_value_head_count * key_value_blocks * 16) * sizeof(float); @@ -3137,18 +2350,10 @@ static void run_qwen_decode_split_flash_attention_scheduling_case(int64_t query_ QwenFlashAttentionLayoutGraph attention = build_qwen_flash_attention_layout_graph( ctx, query_token_count, key_value_token_count, qk_head_size, value_head_size, scale); - ggml_tensor * attention_output = - ggml_reshape_2d(ctx, attention.output, output_hidden_size, query_token_count); - ggml_tensor * consumer_weight = - ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, output_hidden_size, kQwenHiddenSize); - ggml_tensor * consumer_output = ggml_mul_mat(ctx, consumer_weight, attention_output); - REQUIRE(attention_output != nullptr); - REQUIRE(consumer_weight != nullptr); - REQUIRE(consumer_output != nullptr); ggml_cgraph * cgraph = ggml_new_graph(ctx); REQUIRE(cgraph != nullptr); - ggml_build_forward_expand(cgraph, consumer_output); + ggml_build_forward_expand(cgraph, attention.output); ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*cgraph); REQUIRE(imported.valid()); @@ -3157,7 +2362,7 @@ static void run_qwen_decode_split_flash_attention_scheduling_case(int64_t query_ REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics)); const ggml::hrx::CommandPlan & plan = scheduler.plan(); REQUIRE(plan.valid()); - REQUIRE(plan.dispatches.size() == static_cast(query_token_count + 1)); + REQUIRE(plan.dispatches.size() == static_cast(query_token_count)); REQUIRE(plan.transients.size() == 4); REQUIRE(plan.completion_counter_requests.size() == 1); @@ -3167,16 +2372,14 @@ static void run_qwen_decode_split_flash_attention_scheduling_case(int64_t query_ REQUIRE(plan.transients[3].size == q8_output_bytes); REQUIRE(plan.completion_counter_requests[0].count == key_value_head_count); - const ggml::hrx::Value * query_value = imported.graph.values().find_tensor(attention.query); - const ggml::hrx::Value * mask_value = imported.graph.values().find_tensor(attention.mask); - const ggml::hrx::Value * output_value = imported.graph.values().find_tensor(attention.output); - const ggml::hrx::Value * publication_value = imported.graph.values().find_tensor(attention_output); + const ggml::hrx::Value * query_value = imported.graph.values().find_tensor(attention.query); + const ggml::hrx::Value * mask_value = imported.graph.values().find_tensor(attention.mask); + const ggml::hrx::Value * output_value = imported.graph.values().find_tensor(attention.output); REQUIRE(query_value != nullptr); REQUIRE(mask_value != nullptr); REQUIRE(output_value != nullptr); - REQUIRE(publication_value != nullptr); const ggml::hrx::CommandPlanAlternateValue * q8_alternate = - plan.metadata.find_alternate_value(publication_value->id, GGML_TYPE_Q8_1, q8_output_bytes); + plan.metadata.find_alternate_value(output_value->id, GGML_TYPE_Q8_1, q8_output_bytes); REQUIRE(q8_alternate != nullptr); REQUIRE(q8_alternate->alternate_value == plan.transients[3].value); @@ -3187,7 +2390,6 @@ static void run_qwen_decode_split_flash_attention_scheduling_case(int64_t query_ const ggml::hrx::Dispatch & dispatch = plan.dispatches[static_cast(row)]; const char * expected_kernel = "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"; REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == expected_kernel); - REQUIRE(dispatch.kernel.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); REQUIRE(dispatch.kernel.integer_parameters.at("key_value_token_count") == key_value_token_count); REQUIRE(dispatch.kernel.compile_parameters.at("ggml.flash_attention.query_head_count") == std::to_string(query_head_count)); @@ -3254,8 +2456,6 @@ static void run_qwen_decode_attention_output_next_q8_scheduling_case(bool includ ggml_tensor * norm_weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, kQwenHiddenSize); ggml_tensor * rms = ggml_rms_norm(ctx, residual, kQwenRmsNormEps); ggml_tensor * normalized = ggml_mul(ctx, rms, norm_weight); - ggml_tensor * next_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, kQwenHiddenSize, kQwenHiddenSize); - ggml_tensor * next_output = ggml_mul_mat(ctx, next_weight, normalized); REQUIRE(attention_output != nullptr); REQUIRE(output_weight != nullptr); REQUIRE(projection != nullptr); @@ -3264,12 +2464,10 @@ static void run_qwen_decode_attention_output_next_q8_scheduling_case(bool includ REQUIRE(norm_weight != nullptr); REQUIRE(rms != nullptr); REQUIRE(normalized != nullptr); - REQUIRE(next_weight != nullptr); - REQUIRE(next_output != nullptr); ggml_cgraph * cgraph = ggml_new_graph(ctx); REQUIRE(cgraph != nullptr); - ggml_build_forward_expand(cgraph, next_output); + ggml_build_forward_expand(cgraph, normalized); ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*cgraph); REQUIRE(imported.valid()); @@ -3278,15 +2476,10 @@ static void run_qwen_decode_attention_output_next_q8_scheduling_case(bool includ REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics)); const ggml::hrx::CommandPlan & plan = scheduler.plan(); REQUIRE(plan.valid()); - REQUIRE(plan.dispatches.size() == 3); + REQUIRE(plan.dispatches.size() == 2); REQUIRE(plan.completion_counter_requests.size() == 2); - const auto projection_it = std::find_if(plan.dispatches.begin(), plan.dispatches.end(), [](const auto & dispatch) { - return kernel_name_for_id(dispatch.kernel.kernel_id) == - "qwen3_moe:qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8"; - }); - REQUIRE(projection_it != plan.dispatches.end()); - const ggml::hrx::Dispatch & projection_dispatch = *projection_it; + const ggml::hrx::Dispatch & projection_dispatch = plan.dispatches.back(); REQUIRE(kernel_name_for_id(projection_dispatch.kernel.kernel_id) == "qwen3_moe:qwen3_moe_dense_linear_q4k_q8_1_x4_next_q8"); REQUIRE(projection_dispatch.kernel.integer_parameters.at("token_count") == query_token_count); @@ -3297,7 +2490,7 @@ static void run_qwen_decode_attention_output_next_q8_scheduling_case(bool includ REQUIRE(projection_dispatch.kernel.compile_parameters.at("qwen3_moe.dense_quantized.output_accumulation") == "1"); REQUIRE(projection_dispatch.bindings.size() == 7); - const ggml::hrx::Value * flash_output_value = imported.graph.values().find_tensor(attention_output); + const ggml::hrx::Value * flash_output_value = imported.graph.values().find_tensor(attention.output); const ggml::hrx::Value * residual_input_value = imported.graph.values().find_tensor(residual_input); const ggml::hrx::Value * residual_value = imported.graph.values().find_tensor(residual); const ggml::hrx::Value * normalized_value = imported.graph.values().find_tensor(normalized); @@ -3306,10 +2499,9 @@ static void run_qwen_decode_attention_output_next_q8_scheduling_case(bool includ REQUIRE(residual_value != nullptr); REQUIRE(normalized_value != nullptr); const ggml::hrx::CommandPlanAlternateValue * attention_q8 = - ggml::hrx::find_alternate_value(imported.graph, plan, flash_output_value->id, GGML_TYPE_Q8_1, - attention_q8_bytes); + ggml::hrx::find_alternate_value(plan, flash_output_value->id, GGML_TYPE_Q8_1, attention_q8_bytes); const ggml::hrx::CommandPlanAlternateValue * next_q8 = - ggml::hrx::find_alternate_value(imported.graph, plan, normalized_value->id, GGML_TYPE_Q8_1, next_q8_bytes); + ggml::hrx::find_alternate_value(plan, normalized_value->id, GGML_TYPE_Q8_1, next_q8_bytes); REQUIRE(attention_q8 != nullptr); REQUIRE(next_q8 != nullptr); REQUIRE(projection_dispatch.bindings[0].value == attention_q8->alternate_value); @@ -3381,7 +2573,7 @@ static void run_add_f32_cpu_reference_case() { ggml_backend_free(hrx_backend); } -static void run_dynamic_generic_ops_cpu_reference_case() { +static void run_view_mul_f32_cpu_reference_case() { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); REQUIRE(hrx_backend != nullptr); @@ -3394,115 +2586,35 @@ static void run_dynamic_generic_ops_cpu_reference_case() { REQUIRE(cpu_ctx != nullptr); REQUIRE(hrx_ctx != nullptr); - constexpr int64_t element_count = 257; - ggml_tensor * cpu_input = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, element_count); - ggml_tensor * hrx_input = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, element_count); - ggml_tensor * cpu_unary = ggml_sqr(cpu_ctx, cpu_input); - ggml_tensor * hrx_unary = ggml_sqr(hrx_ctx, hrx_input); - ggml_tensor * cpu_copy_target = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, element_count); - ggml_tensor * hrx_copy_target = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, element_count); - ggml_tensor * cpu_copy = ggml_cpy(cpu_ctx, cpu_input, cpu_copy_target); - ggml_tensor * hrx_copy = ggml_cpy(hrx_ctx, hrx_input, hrx_copy_target); - ggml_tensor * cpu_concat_lhs = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, 64, 7); - ggml_tensor * cpu_concat_rhs = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, 32, 7); - ggml_tensor * hrx_concat_lhs = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, 64, 7); - ggml_tensor * hrx_concat_rhs = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, 32, 7); - ggml_tensor * cpu_concat = ggml_concat(cpu_ctx, cpu_concat_lhs, cpu_concat_rhs, 0); - ggml_tensor * hrx_concat = ggml_concat(hrx_ctx, hrx_concat_lhs, hrx_concat_rhs, 0); - REQUIRE(cpu_unary != nullptr && hrx_unary != nullptr); - REQUIRE(cpu_copy != nullptr && hrx_copy != nullptr); - REQUIRE(cpu_concat != nullptr && hrx_concat != nullptr); + constexpr int64_t hidden_size = 2048; + constexpr int64_t token_count = 2; + ggml_tensor * cpu_source = ggml_new_tensor_3d(cpu_ctx, GGML_TYPE_F32, hidden_size, 3, token_count); + ggml_tensor * hrx_source = ggml_new_tensor_3d(hrx_ctx, GGML_TYPE_F32, hidden_size, 3, token_count); + REQUIRE(cpu_source != nullptr); + REQUIRE(hrx_source != nullptr); + ggml_tensor * cpu_lhs = ggml_view_2d(cpu_ctx, cpu_source, hidden_size, token_count, cpu_source->nb[2], 0); + ggml_tensor * cpu_rhs = + ggml_view_2d(cpu_ctx, cpu_source, hidden_size, token_count, cpu_source->nb[2], 2 * cpu_source->nb[1]); + ggml_tensor * hrx_lhs = ggml_view_2d(hrx_ctx, hrx_source, hidden_size, token_count, hrx_source->nb[2], 0); + ggml_tensor * hrx_rhs = + ggml_view_2d(hrx_ctx, hrx_source, hidden_size, token_count, hrx_source->nb[2], 2 * hrx_source->nb[1]); + REQUIRE(cpu_lhs != nullptr); + REQUIRE(cpu_rhs != nullptr); + REQUIRE(hrx_lhs != nullptr); + REQUIRE(hrx_rhs != nullptr); + ggml_tensor * cpu_output = ggml_mul(cpu_ctx, cpu_lhs, cpu_rhs); + ggml_tensor * hrx_output = ggml_mul(hrx_ctx, hrx_lhs, hrx_rhs); + REQUIRE(cpu_output != nullptr); + REQUIRE(hrx_output != nullptr); ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); - REQUIRE(cpu_graph != nullptr && hrx_graph != nullptr); - ggml_build_forward_expand(cpu_graph, cpu_unary); - ggml_build_forward_expand(cpu_graph, cpu_copy); - ggml_build_forward_expand(cpu_graph, cpu_concat); - ggml_build_forward_expand(hrx_graph, hrx_unary); - ggml_build_forward_expand(hrx_graph, hrx_copy); - ggml_build_forward_expand(hrx_graph, hrx_concat); + REQUIRE(cpu_graph != nullptr); + REQUIRE(hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_output); + ggml_build_forward_expand(hrx_graph, hrx_output); - require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), - { "loom_libs:ggml_unary_f32", "loom_libs:ggml_copy_f32", - "loom_libs:ggml_concat_dim0_f32" }); - for (const char * kernel_name : - { "loom_libs:ggml_unary_f32", "loom_libs:ggml_copy_f32", "loom_libs:ggml_concat_dim0_f32" }) { - const ggml::hrx::KernelSpecialization specialization = - scheduled_kernel_specialization(hrx_graph, kernel_name); - REQUIRE(specialization.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); - } - - ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); - ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); - REQUIRE(cpu_buffer != nullptr && hrx_buffer != nullptr); - const std::vector input = make_pattern_f32(element_count, 73, 0.03125f); - set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); - const std::vector concat_lhs = make_pattern_f32(64 * 7, 79, 0.015625f); - const std::vector concat_rhs = make_pattern_f32(32 * 7, 83, 0.0234375f); - set_tensor_pair_bytes(cpu_backend, cpu_concat_lhs, hrx_backend, hrx_concat_lhs, concat_lhs.data(), - concat_lhs.size() * sizeof(float)); - set_tensor_pair_bytes(cpu_backend, cpu_concat_rhs, hrx_backend, hrx_concat_rhs, concat_rhs.data(), - concat_rhs.size() * sizeof(float)); - - REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); - REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(cpu_backend); - ggml_backend_synchronize(hrx_backend); - require_close(get_f32_tensor(hrx_backend, hrx_unary), get_f32_tensor(cpu_backend, cpu_unary), 0.0f); - require_close(get_f32_tensor(hrx_backend, hrx_copy), get_f32_tensor(cpu_backend, cpu_copy), 0.0f); - require_close(get_f32_tensor(hrx_backend, hrx_concat), get_f32_tensor(cpu_backend, cpu_concat), 0.0f); - - ggml_backend_buffer_free(cpu_buffer); - ggml_backend_buffer_free(hrx_buffer); - ggml_free(cpu_ctx); - ggml_free(hrx_ctx); - ggml_backend_free(cpu_backend); - ggml_backend_free(hrx_backend); -} - -static void run_view_mul_f32_cpu_reference_case() { - ggml_backend_t cpu_backend = init_cpu_backend(); - ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); - REQUIRE(hrx_backend != nullptr); - - ggml_init_params params = {}; - params.mem_size = 512 * 1024; - params.no_alloc = true; - ggml_context * cpu_ctx = ggml_init(params); - ggml_context * hrx_ctx = ggml_init(params); - REQUIRE(cpu_ctx != nullptr); - REQUIRE(hrx_ctx != nullptr); - - constexpr int64_t hidden_size = 2048; - constexpr int64_t token_count = 2; - ggml_tensor * cpu_source = ggml_new_tensor_3d(cpu_ctx, GGML_TYPE_F32, hidden_size, 3, token_count); - ggml_tensor * hrx_source = ggml_new_tensor_3d(hrx_ctx, GGML_TYPE_F32, hidden_size, 3, token_count); - REQUIRE(cpu_source != nullptr); - REQUIRE(hrx_source != nullptr); - ggml_tensor * cpu_lhs = ggml_view_2d(cpu_ctx, cpu_source, hidden_size, token_count, cpu_source->nb[2], 0); - ggml_tensor * cpu_rhs = - ggml_view_2d(cpu_ctx, cpu_source, hidden_size, token_count, cpu_source->nb[2], 2 * cpu_source->nb[1]); - ggml_tensor * hrx_lhs = ggml_view_2d(hrx_ctx, hrx_source, hidden_size, token_count, hrx_source->nb[2], 0); - ggml_tensor * hrx_rhs = - ggml_view_2d(hrx_ctx, hrx_source, hidden_size, token_count, hrx_source->nb[2], 2 * hrx_source->nb[1]); - REQUIRE(cpu_lhs != nullptr); - REQUIRE(cpu_rhs != nullptr); - REQUIRE(hrx_lhs != nullptr); - REQUIRE(hrx_rhs != nullptr); - ggml_tensor * cpu_output = ggml_mul(cpu_ctx, cpu_lhs, cpu_rhs); - ggml_tensor * hrx_output = ggml_mul(hrx_ctx, hrx_lhs, hrx_rhs); - REQUIRE(cpu_output != nullptr); - REQUIRE(hrx_output != nullptr); - - ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); - ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); - REQUIRE(cpu_graph != nullptr); - REQUIRE(hrx_graph != nullptr); - ggml_build_forward_expand(cpu_graph, cpu_output); - ggml_build_forward_expand(hrx_graph, hrx_output); - - require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_binary_f32" }); + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_binary_f32" }); ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); @@ -3706,75 +2818,6 @@ static void run_gather_add_f32_cpu_reference_case() { ggml_backend_free(hrx_backend); } -static void run_scale_add_f32_cpu_reference_case(bool strided_residual) { - ggml_backend_t cpu_backend = init_cpu_backend(); - ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); - REQUIRE(hrx_backend != nullptr); - - ggml_init_params params = {}; - params.mem_size = 512 * 1024; - params.no_alloc = true; - ggml_context * cpu_ctx = ggml_init(params); - ggml_context * hrx_ctx = ggml_init(params); - REQUIRE(cpu_ctx != nullptr); - REQUIRE(hrx_ctx != nullptr); - - constexpr int64_t width = 257; - constexpr int64_t token_count = 3; - constexpr int64_t residual_width = 320; - ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, width, token_count); - ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, width, token_count); - ggml_tensor * cpu_residual_storage = - ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, strided_residual ? residual_width : width, token_count); - ggml_tensor * hrx_residual_storage = - ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, strided_residual ? residual_width : width, token_count); - ggml_tensor * cpu_residual = strided_residual ? ggml_view_2d(cpu_ctx, cpu_residual_storage, width, token_count, - cpu_residual_storage->nb[1], 0) : - cpu_residual_storage; - ggml_tensor * hrx_residual = strided_residual ? ggml_view_2d(hrx_ctx, hrx_residual_storage, width, token_count, - hrx_residual_storage->nb[1], 0) : - hrx_residual_storage; - ggml_tensor * cpu_scaled = ggml_scale_bias(cpu_ctx, cpu_input, 0.177800179f, -0.125f); - ggml_tensor * hrx_scaled = ggml_scale_bias(hrx_ctx, hrx_input, 0.177800179f, -0.125f); - ggml_tensor * cpu_output = ggml_add(cpu_ctx, cpu_scaled, cpu_residual); - ggml_tensor * hrx_output = ggml_add(hrx_ctx, hrx_scaled, hrx_residual); - REQUIRE(cpu_output != nullptr); - REQUIRE(hrx_output != nullptr); - - ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); - ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); - REQUIRE(cpu_graph != nullptr); - REQUIRE(hrx_graph != nullptr); - ggml_build_forward_expand(cpu_graph, cpu_output); - ggml_build_forward_expand(hrx_graph, hrx_output); - require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_scale_add_f32" }); - - ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); - ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); - REQUIRE(cpu_buffer != nullptr); - REQUIRE(hrx_buffer != nullptr); - - const std::vector input = make_pattern_f32(width * token_count, 31, 0.03125f); - const std::vector residual = - make_pattern_f32((strided_residual ? residual_width : width) * token_count, 37, 0.046875f); - set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); - set_tensor_pair_bytes(cpu_backend, cpu_residual_storage, hrx_backend, hrx_residual_storage, residual.data(), - residual.size() * sizeof(float)); - - REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); - REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(cpu_backend); - ggml_backend_synchronize(hrx_backend); - require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 0.0f); - - ggml_backend_buffer_free(cpu_buffer); - ggml_backend_buffer_free(hrx_buffer); - ggml_free(cpu_ctx); - ggml_free(hrx_ctx); - ggml_backend_free(cpu_backend); - ggml_backend_free(hrx_backend); -} - static void run_get_rows_f32_cpu_reference_case(ggml_type weight_type, int64_t hidden_size = kQwenHiddenSize, int64_t vocabulary_count = 64, @@ -3808,15 +2851,6 @@ static void run_get_rows_f32_cpu_reference_case(ggml_type weight_type, ggml_build_forward_expand(hrx_graph, hrx_output); require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_get_rows_f32" }); - const ggml::hrx::KernelSpecialization specialization = - scheduled_kernel_specialization(hrx_graph, "loom_libs:ggml_get_rows_f32"); - REQUIRE(specialization.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); - int64_t expected_token_capacity = 1; - while (expected_token_capacity < token_count) { - expected_token_capacity *= 2; - } - REQUIRE(specialization.compile_parameters.at("ggml.get_rows_f32.token_capacity") == - std::to_string(expected_token_capacity)); ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); @@ -3887,77 +2921,13 @@ static void run_get_rows_q8_1_zero_weight_case() { ggml_backend_free(hrx_backend); } -static void run_get_rows_scale_f32_cpu_reference_case(ggml_type weight_type = GGML_TYPE_F32, - int64_t hidden_size = 256, - int64_t row_count = 8, - int64_t token_count = 6) { - ggml_backend_t cpu_backend = init_cpu_backend(); - ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); - REQUIRE(hrx_backend != nullptr); - - ggml_init_params params = {}; - params.mem_size = 4 * 1024 * 1024; - params.no_alloc = true; - ggml_context * cpu_ctx = ggml_init(params); - ggml_context * hrx_ctx = ggml_init(params); - REQUIRE(cpu_ctx != nullptr); - REQUIRE(hrx_ctx != nullptr); - - const float scale = weight_type == GGML_TYPE_Q5_1 ? 25.2982216f : 0.177800179f; - ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, hidden_size, row_count); - ggml_tensor * cpu_ids = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_I32, token_count); - ggml_tensor * cpu_rows = ggml_get_rows(cpu_ctx, cpu_weight, cpu_ids); - ggml_tensor * cpu_output = ggml_scale(cpu_ctx, cpu_rows, scale); - ggml_tensor * hrx_weight = ggml_new_tensor_2d(hrx_ctx, weight_type, hidden_size, row_count); - ggml_tensor * hrx_ids = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_I32, token_count); - ggml_tensor * hrx_rows = ggml_get_rows(hrx_ctx, hrx_weight, hrx_ids); - ggml_tensor * hrx_output = ggml_scale(hrx_ctx, hrx_rows, scale); - REQUIRE(cpu_output != nullptr); - REQUIRE(hrx_output != nullptr); - ggml_set_output(cpu_rows); - ggml_set_output(hrx_rows); - - ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); - ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); - REQUIRE(cpu_graph != nullptr); - REQUIRE(hrx_graph != nullptr); - ggml_build_forward_expand(cpu_graph, cpu_output); - ggml_build_forward_expand(hrx_graph, hrx_output); - require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_get_rows_scale_f32" }); - - ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); - ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); - REQUIRE(cpu_buffer != nullptr); - REQUIRE(hrx_buffer != nullptr); - - const std::vector weight = make_matmul_weight_bytes(weight_type, hidden_size, row_count, 9); - const std::vector ids = { 0, static_cast(row_count - 1), 2, - 2, static_cast(row_count - 2), 1 }; - set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); - set_tensor_pair_bytes(cpu_backend, cpu_ids, hrx_backend, hrx_ids, ids.data(), ids.size() * sizeof(int32_t)); - - REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); - REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(cpu_backend); - ggml_backend_synchronize(hrx_backend); - const float tolerance = weight_type == GGML_TYPE_F32 ? 0.0f : 5.0e-4f; - require_close(get_f32_tensor(hrx_backend, hrx_rows), get_f32_tensor(cpu_backend, cpu_rows), tolerance); - require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), tolerance); - - ggml_backend_buffer_free(cpu_buffer); - ggml_backend_buffer_free(hrx_buffer); - ggml_free(cpu_ctx); - ggml_free(hrx_ctx); - ggml_backend_free(cpu_backend); - ggml_backend_free(hrx_backend); -} - static auto allocate_test_weight(ggml_backend_t backend, ggml_tensor * weight) { const auto buft = ggml_backend_get_default_buffer_type(backend); const size_t alignment = ggml_backend_buft_get_alignment(buft); const size_t size = GGML_PAD(ggml_backend_buft_get_alloc_size(buft, weight), alignment); auto buffer = std::unique_ptr( - ggml_backend_buft_alloc_buffer(buft, size), ggml_backend_buffer_free); + ggml_backend_buft_alloc_buffer(buft, size), + ggml_backend_buffer_free); REQUIRE(buffer != nullptr); ggml_tallocr allocator = ggml_tallocr_new(buffer.get()); REQUIRE(ggml_tallocr_alloc(&allocator, weight) == GGML_STATUS_SUCCESS); @@ -4026,17 +2996,11 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, const auto sequence = scheduled_kernel_sequence(hrx_graph); require_kernel_subsequence(sequence, { expected_kernel }); - const bool packed_q4_input = weight_type == GGML_TYPE_Q4_K; - const bool iq4_xs_weight = weight_type == GGML_TYPE_IQ4_XS; - const bool prefill_tokens = token_count >= 256; - const bool packed_iq4_xs_input = iq4_xs_weight && prefill_tokens; - const bool packed_input = packed_q4_input || packed_iq4_xs_input || square_output; - const bool normalized_packed_input = normalize_input && packed_input; - if (normalized_packed_input) { - require_kernel_subsequence(sequence, { "loom_libs:ggml_rmsnorm_binary_q8_1_x4_publish", expected_kernel }); + if (normalize_input && (weight_type == GGML_TYPE_Q4_K || square_output)) { + require_kernel_subsequence(sequence, { "loom_libs:ggml_rmsnorm_binary_q8_1_x4", expected_kernel }); REQUIRE(std::find(sequence.begin(), sequence.end(), "qwen3_moe:ggml_quantize_q8_1_x4_f32") == sequence.end()); } else if (normalize_input) { - REQUIRE(std::find(sequence.begin(), sequence.end(), "loom_libs:ggml_rmsnorm_binary_q8_1_x4_publish") == sequence.end()); + REQUIRE(std::find(sequence.begin(), sequence.end(), "loom_libs:ggml_rmsnorm_binary_q8_1_x4") == sequence.end()); } ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); @@ -4049,8 +3013,8 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, const bool weight_is_bf16 = weight_type == GGML_TYPE_BF16; const bool weight_is_f32 = weight_type == GGML_TYPE_F32; const bool weight_is_dense_float = weight_is_f16 || weight_is_bf16 || weight_is_f32; - std::vector input = weight_is_dense_float ? - std::vector(static_cast(input_size * token_count), 0.00390625f) : + std::vector input = + weight_is_dense_float ? std::vector(static_cast(input_size * token_count), 0.00390625f) : make_pattern_f32(input_size * token_count, 7, 0.01f); if (cache_type != GGML_TYPE_COUNT) { for (size_t i = 0; i < input.size(); ++i) { @@ -4062,8 +3026,7 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); if (normalize_input) { const std::vector norm_w = make_weight(input_size); - set_tensor_pair_bytes(cpu_backend, cpu_norm_w, hrx_backend, hrx_norm_w, norm_w.data(), - norm_w.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_norm_w, hrx_backend, hrx_norm_w, norm_w.data(), norm_w.size() * sizeof(float)); } if (cache_type != GGML_TYPE_COUNT) { const int64_t index = 7; @@ -4076,8 +3039,8 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(cpu_backend); ggml_backend_synchronize(hrx_backend); - const float abs_tolerance = cache_type != GGML_TYPE_COUNT ? 2.0e-3f : 1.0f; - require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), abs_tolerance, + require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), + cache_type != GGML_TYPE_COUNT ? 1.0e-3f : 1.0f, cache_type != GGML_TYPE_COUNT ? 1.0e-4f : 3.0e-2f); ggml_backend_buffer_free(cpu_buffer); @@ -4088,115 +3051,10 @@ static void run_dense_matmul_cpu_reference_case(ggml_type weight_type, ggml_backend_free(hrx_backend); } -static void run_prefill_value_cache_repeated_indices_case() { - static constexpr const char * disable_env = "GGML_HRX_DISABLE_PREFILL_V_CACHE_FUSION"; - const char * old_env = std::getenv(disable_env); - const bool had_env = old_env != nullptr; - const std::string old_value = had_env ? old_env : ""; - REQUIRE(unsetenv(disable_env) == 0); - - constexpr int64_t input_size = 1536; - constexpr int64_t output_size = 256; - constexpr int64_t token_count = 64; - constexpr int64_t vocabulary_count = 128; - constexpr int64_t cache_row_count = 512; - - ggml_backend_t fused_backend = ggml_backend_hrx_init(0); - ggml_backend_t split_backend = ggml_backend_hrx_init(0); - REQUIRE(fused_backend != nullptr); - REQUIRE(split_backend != nullptr); - - ggml_init_params params = {}; - params.mem_size = static_cast(32 * 1024 * 1024); - params.no_alloc = true; - ggml_context * fused_ctx = ggml_init(params); - ggml_context * split_ctx = ggml_init(params); - REQUIRE(fused_ctx != nullptr); - REQUIRE(split_ctx != nullptr); - - auto build_graph = [](ggml_context * ctx) { - ggml_tensor * embedding = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, input_size, vocabulary_count); - ggml_tensor * tokens = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); - ggml_tensor * norm = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, input_size); - ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, input_size, output_size); - ggml_tensor * indices = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); - ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, output_size, cache_row_count); - ggml_tensor * rows = ggml_get_rows(ctx, embedding, tokens); - ggml_tensor * input = build_rmsnorm_mul_graph(ctx, rows, norm); - ggml_tensor * projected = ggml_mul_mat(ctx, weight, input); - ggml_tensor * output = ggml_set_rows(ctx, cache, projected, indices); - return std::array{ embedding, tokens, norm, weight, indices, cache, output }; - }; - - const auto fused = build_graph(fused_ctx); - const auto split = build_graph(split_ctx); - auto fused_embedding_buffer = allocate_test_weight(fused_backend, fused[0]); - auto fused_projection_buffer = allocate_test_weight(fused_backend, fused[3]); - auto split_embedding_buffer = allocate_test_weight(split_backend, split[0]); - auto split_projection_buffer = allocate_test_weight(split_backend, split[3]); - - ggml_cgraph * fused_graph = ggml_new_graph(fused_ctx); - ggml_cgraph * split_graph = ggml_new_graph(split_ctx); - REQUIRE(fused_graph != nullptr); - REQUIRE(split_graph != nullptr); - ggml_build_forward_expand(fused_graph, fused[6]); - ggml_build_forward_expand(split_graph, split[6]); - require_kernel_subsequence(scheduled_kernel_sequence(fused_graph), - { "loom_libs:ggml_get_rows_rmsnorm_binary_q8_1_x4_f16", - "loom_libs:llm_attention_v_matmul_set_rows_tiled_f32_f32" }); - REQUIRE(setenv(disable_env, "1", 1) == 0); - require_kernel_subsequence(scheduled_kernel_sequence(split_graph), - { "loom_libs:ggml_get_rows_rmsnorm_binary_q8_1_x4_f16", - "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_aligned", - "loom_libs:ggml_set_rows" }); - - ggml_backend_buffer_t fused_buffer = ggml_backend_alloc_ctx_tensors(fused_ctx, fused_backend); - ggml_backend_buffer_t split_buffer = ggml_backend_alloc_ctx_tensors(split_ctx, split_backend); - REQUIRE(fused_buffer != nullptr); - REQUIRE(split_buffer != nullptr); - - const std::vector embedding = make_matmul_weight_bytes(GGML_TYPE_Q4_K, input_size, vocabulary_count, 5); - const std::vector projection = make_matmul_weight_bytes(GGML_TYPE_Q6_K, input_size, output_size, 11); - const std::vector norm = make_weight(input_size); - const std::vector tokens(static_cast(token_count), 3); - std::vector indices(static_cast(token_count)); - for (int64_t token = 0; token < token_count; ++token) { - indices[static_cast(token)] = (token * 37 + 11) % 31; - } - const std::vector cache(static_cast(output_size * cache_row_count), ggml_fp32_to_fp16(-2.5f)); - - set_tensor_pair_bytes(fused_backend, fused[0], split_backend, split[0], embedding.data(), embedding.size()); - set_tensor_pair_bytes(fused_backend, fused[1], split_backend, split[1], tokens.data(), - tokens.size() * sizeof(int32_t)); - set_tensor_pair_bytes(fused_backend, fused[2], split_backend, split[2], norm.data(), norm.size() * sizeof(float)); - set_tensor_pair_bytes(fused_backend, fused[3], split_backend, split[3], projection.data(), projection.size()); - set_tensor_pair_bytes(fused_backend, fused[4], split_backend, split[4], indices.data(), - indices.size() * sizeof(int64_t)); - set_tensor_pair_bytes(fused_backend, fused[5], split_backend, split[5], cache.data(), - cache.size() * sizeof(ggml_fp16_t)); - - REQUIRE(unsetenv(disable_env) == 0); - REQUIRE(ggml_backend_graph_compute(fused_backend, fused_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(fused_backend); - REQUIRE(setenv(disable_env, "1", 1) == 0); - REQUIRE(ggml_backend_graph_compute(split_backend, split_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(split_backend); - require_close(get_f32_tensor(fused_backend, fused[6]), get_f32_tensor(split_backend, split[6]), 0.0f); - - ggml_backend_buffer_free(fused_buffer); - ggml_backend_buffer_free(split_buffer); - ggml_free(fused_ctx); - ggml_free(split_ctx); - ggml_backend_free(fused_backend); - ggml_backend_free(split_backend); - restore_environment_value(disable_env, had_env, old_value); -} - static void run_dense_matmul_unary_cpu_reference_case(ggml_type weight_type, const char * expected_kernel, int64_t token_count, - int64_t output_size, - int64_t input_size = kQwenHiddenSize) { + int64_t output_size) { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); REQUIRE(hrx_backend != nullptr); @@ -4209,6 +3067,7 @@ static void run_dense_matmul_unary_cpu_reference_case(ggml_type weight_type, REQUIRE(cpu_ctx != nullptr); REQUIRE(hrx_ctx != nullptr); + constexpr int64_t input_size = kQwenHiddenSize; ggml_tensor * cpu_weight = ggml_new_tensor_2d(cpu_ctx, weight_type, input_size, output_size); ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, token_count); ggml_tensor * cpu_matmul = ggml_mul_mat(cpu_ctx, cpu_weight, cpu_input); @@ -4462,112 +3321,6 @@ static void run_dense_matmul_binary_cpu_reference_case(ggml_type lhs_weight_t ggml_backend_free(hrx_backend); } -static void run_vector_q8_publish_cpu_reference_case() { - constexpr int64_t input_size = 256; - constexpr int64_t intermediate_size = 256; - constexpr int64_t output_size = 128; - constexpr int64_t token_count = 1; - - ggml_backend_t cpu_backend = init_cpu_backend(); - ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); - REQUIRE(hrx_backend != nullptr); - - ggml_init_params params = {}; - params.mem_size = static_cast(32 * 1024 * 1024); - params.no_alloc = true; - ggml_context * cpu_ctx = ggml_init(params); - ggml_context * hrx_ctx = ggml_init(params); - REQUIRE(cpu_ctx != nullptr); - REQUIRE(hrx_ctx != nullptr); - - ggml_tensor * cpu_first_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, intermediate_size); - ggml_tensor * cpu_second_weight = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_Q4_K, intermediate_size, output_size); - ggml_tensor * cpu_input = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, input_size, token_count); - ggml_tensor * cpu_hidden = ggml_mul_mat(cpu_ctx, cpu_first_weight, cpu_input); - ggml_tensor * cpu_output = ggml_mul_mat(cpu_ctx, cpu_second_weight, cpu_hidden); - ggml_tensor * hrx_first_weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, intermediate_size); - ggml_tensor * hrx_second_weight = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_Q4_K, intermediate_size, output_size); - ggml_tensor * hrx_input = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, input_size, token_count); - ggml_tensor * hrx_hidden = ggml_mul_mat(hrx_ctx, hrx_first_weight, hrx_input); - ggml_tensor * hrx_output = ggml_mul_mat(hrx_ctx, hrx_second_weight, hrx_hidden); - REQUIRE(cpu_output != nullptr); - REQUIRE(hrx_output != nullptr); - - auto hrx_first_weight_buffer = allocate_test_weight(hrx_backend, hrx_first_weight); - auto hrx_second_weight_buffer = allocate_test_weight(hrx_backend, hrx_second_weight); - - ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); - ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); - REQUIRE(cpu_graph != nullptr); - REQUIRE(hrx_graph != nullptr); - ggml_build_forward_expand(cpu_graph, cpu_output); - ggml_build_forward_expand(hrx_graph, hrx_output); - - require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), - { "loom_libs:ggml_mul_mat_vector_f32_f32", "qwen3_moe:ggml_quantize_q8_1_x4_f32", - "loom_libs:ggml_mul_mat_vector_f32_f32" }); - - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*hrx_graph); - REQUIRE(imported.valid()); - ggml::hrx::DispatchScheduler scheduler; - ggml::hrx::DispatchScheduleDiagnostics diagnostics; - REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151" }, &diagnostics)); - const ggml::hrx::CommandPlan & plan = scheduler.plan(); - REQUIRE(plan.valid()); - - const ggml::hrx::Value * hidden_value = imported.graph.values().find_tensor(hrx_hidden); - REQUIRE(hidden_value != nullptr); - const size_t q8_byte_count = static_cast(token_count) * ggml_row_size(GGML_TYPE_Q8_1, intermediate_size); - const ggml::hrx::CommandPlanAlternateValue * q8_alternate = - plan.metadata.find_alternate_value(hidden_value->id, GGML_TYPE_Q8_1, q8_byte_count); - REQUIRE(q8_alternate != nullptr); - - const ggml::hrx::Dispatch * second_vector = nullptr; - int vector_count = 0; - for (const ggml::hrx::Dispatch & dispatch : plan.dispatches) { - if (kernel_name_for_id(dispatch.kernel.kernel_id) == "loom_libs:ggml_mul_mat_vector_f32_f32") { - ++vector_count; - if (vector_count == 2) { - second_vector = &dispatch; - break; - } - } - } - REQUIRE(second_vector != nullptr); - REQUIRE(second_vector->bindings.size() == 4); - REQUIRE(second_vector->bindings[0].value == q8_alternate->alternate_value); - REQUIRE(second_vector->kernel.compile_parameters.at("ggml.matmul.vector.activation_format") == - std::to_string(GGML_TYPE_Q8_1)); - - ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); - ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); - REQUIRE(cpu_buffer != nullptr); - REQUIRE(hrx_buffer != nullptr); - - const std::vector first_weight = make_matmul_weight_bytes(GGML_TYPE_F32, input_size, intermediate_size, 6); - const std::vector second_weight = - make_matmul_weight_bytes(GGML_TYPE_Q4_K, intermediate_size, output_size, 11); - const std::vector input = make_pattern_f32(input_size * token_count, 7, 0.01f); - set_tensor_pair_bytes(cpu_backend, cpu_first_weight, hrx_backend, hrx_first_weight, first_weight.data(), - first_weight.size()); - set_tensor_pair_bytes(cpu_backend, cpu_second_weight, hrx_backend, hrx_second_weight, second_weight.data(), - second_weight.size()); - set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); - - REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); - REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(cpu_backend); - ggml_backend_synchronize(hrx_backend); - require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 1.0f, 5.0e-2f); - - ggml_backend_buffer_free(cpu_buffer); - ggml_backend_buffer_free(hrx_buffer); - ggml_free(cpu_ctx); - ggml_free(hrx_ctx); - ggml_backend_free(cpu_backend); - ggml_backend_free(hrx_backend); -} - static void run_dense_matmul_postops_cpu_reference_case(ggml_type weight_type, const char * expected_kernel, int64_t token_count, @@ -4799,157 +3552,6 @@ static void run_endpoint_rmsnorm_q6k_q8_cpu_reference_case() { ggml_backend_free(hrx_backend); } -static void run_lfm_head_rope_cpu_reference_case(int64_t head_count, int64_t token_count, bool cache_output) { - constexpr int64_t head_size = 64; - constexpr float epsilon = 1.0e-5f; - const int64_t row_size = head_size * head_count; - const int64_t cache_rows = token_count + 64; - - ggml_backend_t cpu_backend = init_cpu_backend(); - ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); - REQUIRE(hrx_backend != nullptr); - - ggml_init_params params = {}; - params.mem_size = 1024 * 1024; - params.no_alloc = true; - ggml_context * cpu_ctx = ggml_init(params); - ggml_context * hrx_ctx = ggml_init(params); - REQUIRE(cpu_ctx != nullptr); - REQUIRE(hrx_ctx != nullptr); - - struct HeadGraph { - ggml_tensor * input; - ggml_tensor * weight; - ggml_tensor * positions; - ggml_tensor * cache; - ggml_tensor * row_ids; - ggml_tensor * output; - }; - - auto build = [&](ggml_context * ctx) { - HeadGraph g = {}; - g.input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, head_size, head_count, token_count); - g.weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, head_size); - g.positions = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, token_count); - ggml_tensor * normalized = ggml_rms_norm(ctx, g.input, epsilon); - ggml_tensor * scaled = ggml_mul(ctx, normalized, g.weight); - ggml_tensor * rope = ggml_rope_ext(ctx, scaled, g.positions, nullptr, head_size, GGML_ROPE_TYPE_NEOX, 0, - 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); - g.output = rope; - if (cache_output) { - g.cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F16, row_size, cache_rows); - g.row_ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, token_count); - ggml_tensor * rows = ggml_reshape_2d(ctx, rope, row_size, token_count); - ggml_tensor * view = ggml_view_2d(ctx, rows, row_size, token_count, rows->nb[1], 0); - g.output = ggml_set_rows(ctx, g.cache, view, g.row_ids); - } - REQUIRE(g.output != nullptr); - return g; - }; - - HeadGraph cpu = build(cpu_ctx); - HeadGraph hrx = build(hrx_ctx); - ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); - ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); - REQUIRE(cpu_graph != nullptr); - REQUIRE(hrx_graph != nullptr); - ggml_build_forward_expand(cpu_graph, cpu.output); - ggml_build_forward_expand(hrx_graph, hrx.output); - const std::vector kernels = scheduled_kernel_sequence(hrx_graph); - if (cache_output) { - REQUIRE(std::find(kernels.begin(), kernels.end(), "loom_libs:ggml_rmsnorm_mul_rope_f32") == kernels.end()); - require_kernel_subsequence(kernels, { "loom_libs:ggml_rope_f32", "loom_libs:ggml_set_rows" }); - } else { - require_kernel_subsequence(kernels, { "loom_libs:ggml_rmsnorm_mul_rope_f32" }); - } - - ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); - ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); - REQUIRE(cpu_buffer != nullptr); - REQUIRE(hrx_buffer != nullptr); - - const std::vector input = make_pattern_f32(static_cast(row_size * token_count), 47, 0.03f); - const std::vector weight = make_weight(head_size); - std::vector positions(static_cast(token_count)); - std::vector row_ids(static_cast(token_count)); - for (int64_t i = 0; i < token_count; ++i) { - positions[static_cast(i)] = static_cast(3 + 2 * i); - row_ids[static_cast(i)] = (17 * i + 5) % cache_rows; - } - set_tensor_pair_bytes(cpu_backend, cpu.input, hrx_backend, hrx.input, input.data(), input.size() * sizeof(float)); - set_tensor_pair_bytes(cpu_backend, cpu.weight, hrx_backend, hrx.weight, weight.data(), - weight.size() * sizeof(float)); - set_tensor_pair_bytes(cpu_backend, cpu.positions, hrx_backend, hrx.positions, positions.data(), - positions.size() * sizeof(int32_t)); - if (cache_output) { - const std::vector cache(static_cast(row_size * cache_rows), ggml_fp32_to_fp16(-3.0f)); - set_tensor_pair_bytes(cpu_backend, cpu.cache, hrx_backend, hrx.cache, cache.data(), - cache.size() * sizeof(ggml_fp16_t)); - set_tensor_pair_bytes(cpu_backend, cpu.row_ids, hrx_backend, hrx.row_ids, row_ids.data(), - row_ids.size() * sizeof(int64_t)); - } - - REQUIRE(ggml_backend_graph_compute(cpu_backend, cpu_graph) == GGML_STATUS_SUCCESS); - REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); - ggml_backend_synchronize(cpu_backend); - ggml_backend_synchronize(hrx_backend); - const float tolerance = cache_output ? 1.0e-3f : 1.0e-4f; - require_close(get_f32_tensor(hrx_backend, hrx.output), get_f32_tensor(cpu_backend, cpu.output), tolerance, - tolerance); - - ggml_backend_buffer_free(cpu_buffer); - ggml_backend_buffer_free(hrx_buffer); - ggml_free(cpu_ctx); - ggml_free(hrx_ctx); - ggml_backend_free(cpu_backend); - ggml_backend_free(hrx_backend); -} - -static void run_lfm_head_rope_availability_checks() { - ggml_init_params params = {}; - params.mem_size = 1024 * 1024; - params.no_alloc = true; - ggml_context * ctx = ggml_init(params); - REQUIRE(ctx != nullptr); - - ggml_tensor * input = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 64, 8, 1); - ggml_tensor * weight = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 64); - ggml_tensor * pos_f32 = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); - ggml_tensor * positions = ggml_cast(ctx, pos_f32, GGML_TYPE_I32); - ggml_tensor * rms = ggml_rms_norm(ctx, input, 1.0e-5f); - ggml_tensor * scaled = ggml_mul(ctx, rms, weight); - ggml_tensor * output = ggml_rope_ext(ctx, scaled, positions, nullptr, 64, GGML_ROPE_TYPE_NEOX, 0, 10000.0f, 1.0f, - 0.0f, 1.0f, 0.0f, 0.0f); - REQUIRE(output != nullptr); - ggml_cgraph * graph = ggml_new_graph(ctx); - REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, output); - - ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); - REQUIRE(imported.valid()); - const ggml::hrx::DispatchRegistry * registry = ggml::hrx::find_dispatch_registry({ "gfx1151" }); - REQUIRE(registry != nullptr); - const size_t rms_index = producer_index_for_tensor(imported.graph, rms); - const size_t pos_index = producer_index_for_tensor(imported.graph, positions); - ggml::hrx::CommandPlan plan; - std::vector covered(imported.graph.nodes().size(), false); - auto matched_kernel = [&]() { - ggml::hrx::DispatchMatch match; - const ggml::hrx::DispatchMatchContext context = { - imported.graph, &imported.graph.nodes()[rms_index], rms_index, covered, - plan, next_plan_value(imported.graph, plan), - }; - REQUIRE(registry->match(context, match)); - REQUIRE(match.dispatches.size() == 1); - return kernel_name_for_id(match.dispatches[0].kernel.kernel_id); - }; - REQUIRE(matched_kernel() == "loom_libs:ggml_rmsnorm_f32"); - covered[pos_index] = true; - REQUIRE(matched_kernel() == "loom_libs:ggml_rmsnorm_mul_rope_f32"); - - ggml_free(ctx); -} - static void run_rope_set_rows_cpu_reference_case() { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); @@ -4990,12 +3592,6 @@ static void run_rope_set_rows_cpu_reference_case() { ggml_build_forward_expand(cpu_graph, cpu_out); ggml_build_forward_expand(hrx_graph, hrx_out); require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { "loom_libs:ggml_rope_f32" }); - const ggml::hrx::KernelSpecialization specialization = - scheduled_kernel_specialization(hrx_graph, "loom_libs:ggml_rope_f32"); - REQUIRE(specialization.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); - REQUIRE(specialization.compile_parameters.at("ggml.rope_f32.token_capacity") == "4"); - REQUIRE(specialization.compile_parameters.count("ggml.rope_f32.input_span") == 0); - REQUIRE(specialization.integer_parameters.at("input_span") == hidden_size * token_count); ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); @@ -5734,22 +4330,18 @@ static void run_decode_routed_moe_scheduling_case(ggml_type down_weight_type, bo REQUIRE(next_rms != nullptr); ggml_tensor * output = ggml_mul(ctx, next_rms, next_norm_weight); REQUIRE(output != nullptr); - ggml_tensor * next_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, kQwenHiddenSize, kQwenHiddenSize); - ggml_tensor * next_output = ggml_mul_mat(ctx, next_weight, output); - REQUIRE(next_weight != nullptr); - REQUIRE(next_output != nullptr); ggml_cgraph * graph = ggml_new_graph(ctx); REQUIRE(graph != nullptr); - ggml_build_forward_expand(graph, next_output); + ggml_build_forward_expand(graph, output); std::vector expected = { "qwen3_moe:qwen3_moe_rmsnorm_f32_quantize_q8_1_x4", "qwen3_moe:qwen3_moe_router_projection_top8_fused_decode_f32", - down_weight_type == GGML_TYPE_Q4_K ? "loom_libs:ggml_mul_mat_id_vector_pair_binary_publish_q8" : - "loom_libs:ggml_mul_mat_id_vector_pair_binary_publish_f32", - down_weight_type == GGML_TYPE_Q4_K ? "loom_libs:ggml_mul_mat_id_vector_weighted_wave32_publish_q8" : - "loom_libs:ggml_mul_mat_id_vector_weighted_wave64_publish_q8", + down_weight_type == GGML_TYPE_Q4_K ? "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_q8_1_x4_next_q8" : + "qwen3_moe:qwen3_moe_routed_gate_up_swiglu_q4k_q8", + down_weight_type == GGML_TYPE_Q4_K ? "qwen3_moe:qwen3_moe_routed_down_q4k_q8_1_x4_next_q8" : + "qwen3_moe:qwen3_moe_routed_down_q6k_f32_wave64_next_q8", }; require_kernel_subsequence(scheduled_kernel_sequence(graph), expected); ggml_free(ctx); @@ -5788,12 +4380,7 @@ struct SsmConvPrefillGraph { ggml_cgraph * graph = nullptr; }; -static SsmConvPrefillGraph make_ssm_conv_prefill_graph(ggml_type type, - int64_t k, - int64_t n, - bool expose, - bool gathered, - bool clear_state) { +static SsmConvPrefillGraph make_ssm_conv_prefill_graph(ggml_type type, int64_t k, int64_t n, bool expose, bool gathered, bool clear_state) { ggml_init_params params = {}; params.mem_size = 32 * 1024 * 1024; params.no_alloc = true; @@ -5835,8 +4422,8 @@ static SsmConvPrefillGraph make_ssm_conv_prefill_graph(ggml_type type, "loom_libs:ggml_mul_mat_quantized_f16_wmma_prefill_conv4"; REQUIRE((std::find(sequence.begin(), sequence.end(), fused) != sequence.end()) == !expose); if (gathered && !expose) { - REQUIRE(std::find(sequence.begin(), sequence.end(), "loom_libs:llm_ssm_conv_dconv4_silu_prefill_finish_f32") != - sequence.end()); + REQUIRE(std::find(sequence.begin(), sequence.end(), + "loom_libs:llm_ssm_conv_dconv4_silu_prefill_finish_f32") != sequence.end()); } return c; } @@ -5882,704 +4469,257 @@ static void run_ssm_conv_prefill_recurrent_case(ggml_type type, bool gathered, b ggml_backend_free(backend); } -static void run_vector_output_capacity_experiment() { - constexpr int64_t kSmallVectorInputSize = 256; - constexpr int64_t kAboveOldOutputSize = 262208; - constexpr int64_t kLiftedOutputSize = 524288; - - run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, kAboveOldOutputSize, - kSmallVectorInputSize); - run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, - kAboveOldOutputSize, kSmallVectorInputSize); - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32", 1, - kAboveOldOutputSize, true, false, false, kSmallVectorInputSize); - run_dense_matmul_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, kLiftedOutputSize, - kSmallVectorInputSize); -} - -static void run_vector_postops_cpu_reference_case() { - constexpr int64_t kVectorInputSize = 256; - constexpr int64_t kVectorOutputSize = 256; +int main() { + run_rmsnorm_support_checks(); + run_alternate_value_alias_lookup_checks(); - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32", 1, - kVectorOutputSize, true, false, false, kVectorInputSize); - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32", 1, - kVectorOutputSize, true, true, false, kVectorInputSize); -} - -static const char * glu_op_name(ggml_glu_op op) { - switch (op) { - case GGML_GLU_OP_REGLU: - return "reglu"; - case GGML_GLU_OP_GEGLU: - return "geglu"; - case GGML_GLU_OP_SWIGLU: - return "swiglu"; - case GGML_GLU_OP_SWIGLU_OAI: - return "swiglu_oai"; - case GGML_GLU_OP_GEGLU_ERF: - return "geglu_erf"; - case GGML_GLU_OP_GEGLU_QUICK: - return "geglu_quick"; - case GGML_GLU_OP_COUNT: - break; + if (ggml_backend_hrx_get_device_count() == 0) { + std::fprintf(stderr, "test skipped: no HRX devices available\n"); + return 0; } - return "unknown_glu"; -} - -static std::string type_name(ggml_type type) { - return ggml_type_name(type); -} -using test_runner::Suite; - -static void register_basic_ops_cases(Suite & suite) { - suite.host_case("support.rmsnorm", [] { run_rmsnorm_support_checks(); }); - suite.host_case("support.alternate_value_alias_lookup", [] { run_alternate_value_alias_lookup_checks(); }); - suite.host_case("support.activation_publication_contract", [] { run_activation_publication_contract_checks(); }); - suite.device_case("add_f32.cpu_reference", [] { run_add_f32_cpu_reference_case(); }); - suite.device_case("dynamic_generic_ops.cpu_reference", [] { run_dynamic_generic_ops_cpu_reference_case(); }); + run_add_f32_cpu_reference_case(); for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { - suite.device_case("ssm_conv_prefill_recurrent." + type_name(type) + ".plain", - [type] { run_ssm_conv_prefill_recurrent_case(type, false, false); }); - suite.device_case("ssm_conv_prefill_recurrent." + type_name(type) + ".gathered", - [type] { run_ssm_conv_prefill_recurrent_case(type, true, false); }); - suite.device_case("ssm_conv_prefill_recurrent." + type_name(type) + ".gathered_clear_state", - [type] { run_ssm_conv_prefill_recurrent_case(type, true, true); }); - } - suite.device_case("view_mul_f32.cpu_reference", [] { run_view_mul_f32_cpu_reference_case(); }); - for (const int64_t token_count : { 2, 23 }) { - suite.device_case("strided_cont.tokens" + std::to_string(token_count), - [token_count] { run_strided_cont_cpu_reference_case(token_count); }); - } - for (const ggml_glu_op op : - { GGML_GLU_OP_REGLU, GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_GEGLU_ERF, GGML_GLU_OP_GEGLU_QUICK }) { - suite.device_case(std::string("glu_split_f32.") + glu_op_name(op), - [op] { run_glu_split_f32_cpu_reference_case(op); }); - suite.device_case(std::string("glu_packed_f32.") + glu_op_name(op) + ".normal", - [op] { run_glu_packed_f32_cpu_reference_case(op, false); }); - } - suite.device_case("glu_packed_f32.swiglu.gated", - [] { run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, true); }); - suite.device_case("glu_packed_f32.swiglu.outputs8192.tokens2", - [] { run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, false, 8192, 2); }); - suite.device_case("glu_packed_f32.swiglu.outputs8192.tokens14", - [] { run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, false, 8192, 14); }); - suite.device_case("gather_add_f32.cpu_reference", [] { run_gather_add_f32_cpu_reference_case(); }); - suite.device_case("scale_add_f32.cpu_reference.packed", [] { run_scale_add_f32_cpu_reference_case(false); }); - suite.device_case("scale_add_f32.cpu_reference.strided", [] { run_scale_add_f32_cpu_reference_case(true); }); - for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_IQ4_XS, GGML_TYPE_F16, - GGML_TYPE_BF16, GGML_TYPE_F32 }) { - suite.device_case("get_rows_f32." + type_name(type), [type] { run_get_rows_f32_cpu_reference_case(type); }); - } - suite.device_case("get_rows_f32.q1_0.rows2048", [] { run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q1_0, 2048); }); - suite.device_case("get_rows_f32.q5_1.rows640", [] { run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q5_1, 640); }); - suite.device_case("get_rows_q8_1.zero_weight", [] { run_get_rows_q8_1_zero_weight_case(); }); - suite.device_case("get_rows_scale_f32.cpu_reference", [] { run_get_rows_scale_f32_cpu_reference_case(); }); - suite.device_case("get_rows_scale_f32.q5_1.cpu_reference", - [] { run_get_rows_scale_f32_cpu_reference_case(GGML_TYPE_Q5_1, 640); }); -} - -static void register_dense_matmul_cases(Suite & suite) { - // Distinct gate/up weights cover cached metadata, multiple K blocks, and tile tails. - for (const ggml_type gate : { GGML_TYPE_Q5_K, GGML_TYPE_IQ4_XS }) { - for (const ggml_type up : { GGML_TYPE_Q5_K, GGML_TYPE_IQ4_XS }) { - for (const int64_t input_size : { 256, 768 }) { - suite.device_case("dense_matmul.metadata_pair." + type_name(gate) + "." + type_name(up) + ".k" + std::to_string(input_size), [gate, up, input_size] { - run_dense_matmul_swiglu_cpu_reference_case(gate, up, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", 33, 65, GGML_GLU_OP_SWIGLU, input_size); - }); - } - } - } - - // Multiple blocks and output tails check the prepared state lifetime and row guards. - for (const ggml_type type : { GGML_TYPE_Q5_K, GGML_TYPE_IQ4_XS }) { - for (const int64_t input_size : { 256, 768 }) { - suite.device_case("dense_metadata_reuse." + type_name(type) + ".input" + std::to_string(input_size), [type, input_size] { - run_dense_matmul_postops_cpu_reference_case( - type, "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", - 33, 65, true, true, false, input_size); - }); - } - } - - // K=384 exercises a partial 256-weight tile; N=65 exercises an unpaired output row. - for (const int64_t input_size : { 256, 384, 768 }) { - suite.device_case("dense_matmul.q1_0.vector_tail.k" + std::to_string(input_size), [input_size] { + run_ssm_conv_prefill_recurrent_case(type, false, false); + run_ssm_conv_prefill_recurrent_case(type, true, false); + run_ssm_conv_prefill_recurrent_case(type, true, true); + } + run_view_mul_f32_cpu_reference_case(); + run_strided_cont_cpu_reference_case(2); + run_strided_cont_cpu_reference_case(23); + run_glu_split_f32_cpu_reference_case(GGML_GLU_OP_REGLU); + run_glu_split_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU); + run_glu_split_f32_cpu_reference_case(GGML_GLU_OP_GEGLU); + run_glu_split_f32_cpu_reference_case(GGML_GLU_OP_GEGLU_ERF); + run_glu_split_f32_cpu_reference_case(GGML_GLU_OP_GEGLU_QUICK); + run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_REGLU, false); + run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, false); + run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_GEGLU, false); + run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_GEGLU_ERF, false); + run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_GEGLU_QUICK, false); + run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, true); + run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, false, 8192, 2); + run_glu_packed_f32_cpu_reference_case(GGML_GLU_OP_SWIGLU, false, 8192, 14); + run_gather_add_f32_cpu_reference_case(); + run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q4_K); + run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q6_K); + run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q1_0, 2048); + run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q5_1, 640); + run_get_rows_f32_cpu_reference_case(GGML_TYPE_Q8_0); + run_get_rows_q8_1_zero_weight_case(); + run_get_rows_f32_cpu_reference_case(GGML_TYPE_F16); + run_get_rows_f32_cpu_reference_case(GGML_TYPE_BF16); + run_get_rows_f32_cpu_reference_case(GGML_TYPE_F32); + for (int64_t tokens : { 1, 3, 5 }) { + for (int64_t outputs : { 1, 4, 47, 48, 63 }) { run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q1_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 65, input_size); - }); - } - - for (const ggml_type type : { GGML_TYPE_IQ1_S, GGML_TYPE_IQ1_M }) { - const char * type_suffix = type == GGML_TYPE_IQ1_S ? "iq1_s" : "iq1_m"; - const char * vector_kernel = type == GGML_TYPE_IQ1_S ? - "loom_libs:ggml_mul_mat_vector_iq1_s_f32_f32" : - "loom_libs:ggml_mul_mat_vector_iq1_m_f32_f32"; - const char * tiled_kernel = type == GGML_TYPE_IQ1_S ? - "loom_libs:ggml_mul_mat_tiled_input_f32_iq1_s_publish_f32" : - "loom_libs:ggml_mul_mat_tiled_input_f32_iq1_m_publish_f32"; - suite.device_case(std::string("dense_matmul.") + type_suffix + ".vector.tokens1.outputs128", [=] { - run_dense_matmul_cpu_reference_case(type, vector_kernel, 1, 128, 2048); - }); - suite.device_case(std::string("dense_matmul.") + type_suffix + ".tiled.tokens64.outputs128", [=] { - run_dense_matmul_cpu_reference_case(type, tiled_kernel, 64, 128, 2048); - }); - } - for (const int64_t tokens : { 1, 3, 5 }) { - for (const int64_t outputs : { 1, 4, 47, 48, 63 }) { - suite.device_case( - "dense_matmul.q4_k.tokens" + std::to_string(tokens) + ".outputs" + std::to_string(outputs), - [tokens, outputs] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, - tokens == 1 ? - "loom_libs:ggml_mul_mat_vector_f32_f32" : - "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", - tokens, outputs); - }); + GGML_TYPE_Q4_K, + tokens == 1 ? "loom_libs:ggml_mul_mat_f32_f32_decode_wave64" : + "loom_libs:ggml_mul_mat_f32_f32_wmma", + tokens, outputs); } } - for (const ggml_type type : - { GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_F32 }) { - suite.device_case("dense_matmul." + type_name(type) + ".tokens3.outputs47", [type] { - run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 3, 47); - }); - } - for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { - suite.device_case("dense_matmul." + type_name(type) + ".tiled.tokens33.outputs65.input512", [type] { - run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 33, 65, - 512); - }); - } - suite.device_case("dense_matmul_unary.f32.tokens3.outputs47", [] { - run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", - 3, 47); - }); - suite.device_case("get_rows_f32.f32.rows262144.tokens4.outputs1", - [] { run_get_rows_f32_cpu_reference_case(GGML_TYPE_F32, 262144, 4, 1); }); - suite.device_case("dense_matmul.q4_k.skinny.tokens2.outputs128", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, - 128); - }); - suite.device_case("dense_matmul.q4_k.vector.tokens1.outputs128", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); - }); - for (const ggml_type type : { GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XS }) { - suite.device_case("dense_matmul." + type_name(type) + ".codebook.vector.tokens1.outputs128", [type] { - run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_vector_iq2_s_f32_f32", 1, 128, - 2048); - }); - suite.device_case("dense_matmul." + type_name(type) + ".codebook.tiled.tokens64.outputs128", [type] { - run_dense_matmul_cpu_reference_case( - type, "loom_libs:ggml_mul_mat_tiled_input_f32_iq2_s_publish_f32", 64, 128, 2048); - }); - } - suite.device_case("dense_matmul.iq3_xxs.codebook.vector.tokens1.outputs128", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ3_XXS, - "loom_libs:ggml_mul_mat_vector_iq3_xxs_f32_f32", 1, 128, 2048); - }); - suite.device_case("dense_matmul.iq3_xxs.codebook.tiled.tokens64.outputs128", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ3_XXS, - "loom_libs:ggml_mul_mat_tiled_input_f32_iq3_xxs_publish_f32", 64, - 128, 2048); - }); - for (const ggml_type type : { GGML_TYPE_IQ2_S, GGML_TYPE_IQ3_S }) { - const char * vector_kernel = type == GGML_TYPE_IQ2_S ? - "loom_libs:ggml_mul_mat_vector_iq2_s_f32_f32" : - "loom_libs:ggml_mul_mat_vector_iq3_s_f32_f32"; - const char * tiled_kernel = type == GGML_TYPE_IQ2_S ? - "loom_libs:ggml_mul_mat_tiled_input_f32_iq2_s_publish_f32" : - "loom_libs:ggml_mul_mat_tiled_input_f32_iq3_s_publish_f32"; - suite.device_case("dense_matmul." + type_name(type) + ".codebook_control.vector.tokens1.outputs128", - [type, vector_kernel] { - run_dense_matmul_cpu_reference_case(type, vector_kernel, 1, 128, 2048); - }); - suite.device_case("dense_matmul." + type_name(type) + ".codebook_control.tiled.tokens64.outputs128", - [type, tiled_kernel] { - run_dense_matmul_cpu_reference_case(type, tiled_kernel, 64, 128, 2048); - }); - } - suite.device_case("dense_matmul.iq4_xs.vector.tokens1.outputs128", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_XS, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 128); - }); - suite.device_case("dense_matmul.iq4_xs.tiled.tokens64.outputs128", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_XS, - "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_aligned", 64, 128); - }); - suite.device_case("dense_matmul.iq4_xs.prefill.tokens256.outputs128", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_XS, - "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 256, 128, - kQwenHiddenSize, true); - }); - // Exercise the production K dimension and a partial output tile. - suite.device_case("dense_matmul.iq4_xs.prefill.tokens512.outputs64", [] { - run_dense_matmul_cpu_reference_case( - GGML_TYPE_IQ4_XS, "loom_libs:ggml_mul_mat_q5_k_iq4_xs_q8_1_x4_wmma_token256", 512, 64, 3072, true); - }); - for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { - for (const ggml_type cache_type : { GGML_TYPE_F16, GGML_TYPE_F32 }) { - suite.device_case( - "attention_v_matmul." + type_name(type) + ".cache_" + type_name(cache_type), [type, cache_type] { - run_dense_matmul_cpu_reference_case(type, - "loom_libs:llm_attention_v_matmul_set_rows_vector_f32_f32", 1, - 1024, 2048, false, false, cache_type); - }); + for (ggml_type type : { GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, + GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_F32 }) { + run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_f32_f32_wmma", 3, 47); + } + run_dense_matmul_unary_cpu_reference_case( + GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 3, 47); + run_get_rows_f32_cpu_reference_case(GGML_TYPE_F32, 262144, 4, 1); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 1, 128); + for (ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { + for (ggml_type cache_type : { GGML_TYPE_F16, GGML_TYPE_F32 }) { + run_dense_matmul_cpu_reference_case(type, "loom_libs:llm_attention_v_matmul_set_rows_f32_f32_wmma", + 1, 1024, 2048, false, false, cache_type); } } - suite.device_case("attention_v_matmul.q6_k.prefill_f16.repeated_indices", - [] { run_prefill_value_cache_repeated_indices_case(); }); - for (const ggml_type type : - { GGML_TYPE_Q1_0, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, GGML_TYPE_IQ2_S, - GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_IQ4_NL, GGML_TYPE_F16, GGML_TYPE_BF16 }) { - const char * vector_kernel = type == GGML_TYPE_Q6_K ? "loom_libs:ggml_mul_mat_vector_q6_f32_f32" : - "loom_libs:ggml_mul_mat_vector_f32_f32"; - const int64_t outputs = type == GGML_TYPE_IQ2_S ? 640 : (type == GGML_TYPE_IQ4_NL ? 1024 : 128); - const int64_t input = type == GGML_TYPE_IQ4_NL ? 640 : kQwenHiddenSize; - const char * route = type == GGML_TYPE_IQ4_NL ? "tiled" : "skinny"; - const char * kernel = type == GGML_TYPE_IQ4_NL ? "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32" : - "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32"; - suite.device_case( - "dense_matmul." + type_name(type) + "." + route + ".tokens2.outputs" + std::to_string(outputs), - [type, outputs, input, kernel] { run_dense_matmul_cpu_reference_case(type, kernel, 2, outputs, input); }); - suite.device_case("dense_matmul." + type_name(type) + ".vector.tokens1.outputs" + std::to_string(outputs), - [type, vector_kernel, outputs, input] { - run_dense_matmul_cpu_reference_case(type, vector_kernel, 1, outputs, input); - }); - } - suite.device_case("dense_matmul.iq4_nl.tiled.tokens64.outputs1024", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_NL, - "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32_aligned", 64, 1024, - 640); - }); - suite.device_case("dense_matmul.q5_0.tiled.tokens2.outputs256.input640", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 2, - 256, 640); - }); - suite.device_case("dense_matmul.q5_0.tiled.tokens2.outputs128.input544", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32", 2, - 128, 544); - }); - suite.device_case("dense_matmul.q5_0.vector.tokens1.outputs256.input640", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 256, 640); - }); - suite.device_case("dense_matmul.q8_1.zero_weight.skinny", [] { - run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32"); - }); - suite.device_case("dense_matmul.q8_1.zero_weight.vector", [] { - run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_vector_f32_f32", 1); - }); - suite.device_case("dense_matmul.f32.skinny.tokens2.outputs256", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 2, - 256); - }); - suite.device_case("dense_matmul.f32.vector.tokens1.outputs256", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_vector_f32_f32", 1, 256); - }); - suite.device_case("dense_matmul_unary.f32.tokens2.outputs128", [] { - run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", - 2, 128); - }); - suite.device_case("dense_matmul_swiglu.q4_k.f16.swiglu.tokens2.outputs128", [] { - run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128); - }); - suite.device_case("dense_matmul_swiglu.q4_k.f16.geglu.tokens2.outputs128", [] { - run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, - GGML_GLU_OP_GEGLU); - }); - suite.device_case("dense_matmul_swiglu.iq4_nl.iq4_nl.geglu.tokens2.outputs2048.input640", [] { - run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL, - "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", 2, - 2048, GGML_GLU_OP_GEGLU, 640); - }); - suite.device_case("dense_matmul_swiglu.iq4_nl.iq4_nl.geglu.tokens64.outputs128.input640", [] { + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q1_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q1_0, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_0, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_1, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_1, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 256, 640); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128, 544); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_0, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 256, 640); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_1, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q5_1, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ2_S, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 640); + run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ2_S, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 640); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q8_0, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_Q8_0, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_NL, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 1024, 640); + run_dense_matmul_cpu_reference_case(GGML_TYPE_IQ4_NL, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 1024, 640); + run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_f32_f32_wmma"); + run_dense_matmul_zero_weight_case(GGML_TYPE_Q8_1, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1); + run_dense_matmul_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_BF16, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_BF16, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 128); + run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 256); + run_dense_matmul_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_decode_wave64", 1, 256); + run_dense_matmul_unary_cpu_reference_case(GGML_TYPE_F32, "loom_libs:ggml_mul_mat_f32_f32_wmma", 2, 128); + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128); + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, GGML_GLU_OP_GEGLU); + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma", 2, 2048, GGML_GLU_OP_GEGLU, + 640); + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, GGML_GLU_OP_REGLU); + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_F16, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, GGML_GLU_OP_GEGLU_ERF); + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_F16, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, GGML_GLU_OP_GEGLU_QUICK); + run_dense_matmul_packed_glu_cpu_reference_case(GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma", 2, 128); + run_dense_matmul_packed_glu_cpu_reference_case(GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_wmma", 2, 128, + GGML_GLU_OP_SWIGLU, true); + run_dense_matmul_binary_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, + "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128); + for (ggml_glu_op op : { GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_REGLU, + GGML_GLU_OP_GEGLU_ERF, GGML_GLU_OP_GEGLU_QUICK }) { run_dense_matmul_swiglu_cpu_reference_case( - GGML_TYPE_IQ4_NL, GGML_TYPE_IQ4_NL, - "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32_aligned", - 64, 128, GGML_GLU_OP_GEGLU, 640); - }); - for (const ggml_glu_op op : { GGML_GLU_OP_REGLU, GGML_GLU_OP_GEGLU_ERF, GGML_GLU_OP_GEGLU_QUICK }) { - suite.device_case(std::string("dense_matmul_swiglu.q4_k.f16.") + glu_op_name(op) + ".tokens2.outputs128", [op] { - run_dense_matmul_swiglu_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_F16, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, op); - }); - } - suite.device_case("dense_matmul_packed_glu.q4_k.tokens2.outputs128", [] { - run_dense_matmul_packed_glu_cpu_reference_case( - GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", 2, 128); - }); - suite.device_case("dense_matmul_packed_glu.q4_k.swiglu.tokens2.outputs128.gated", [] { - run_dense_matmul_packed_glu_cpu_reference_case(GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", - 2, 128, GGML_GLU_OP_SWIGLU, true); - }); - suite.device_case("dense_matmul_binary.q4_k.f16.tokens2.outputs128", [] { - run_dense_matmul_binary_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128); - }); - suite.device_case("vector_q8_publish.cpu_reference", [] { run_vector_q8_publish_cpu_reference_case(); }); - for (const ggml_glu_op op : - { GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_REGLU, GGML_GLU_OP_GEGLU_ERF, GGML_GLU_OP_GEGLU_QUICK }) { - for (const int64_t outputs : { 128, 4096 }) { - suite.device_case(std::string("dense_matmul_swiglu.q4_k.q4_k.lowtoken.") + glu_op_name(op) + - ".tokens4.outputs" + std::to_string(outputs), - [op, outputs] { - run_dense_matmul_swiglu_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 4, outputs, op); - }); - } - for (const int64_t tokens : { 256, 512 }) { - if (tokens == 512) { - suite.device_case( - std::string("dense_matmul_swiglu.q4_k.q4_k.prefill.") + glu_op_name(op) + ".tokens512.outputs4096", - [op] { - run_dense_matmul_swiglu_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 512, 4096, op); - }); - } else { - for (const int64_t outputs : { 128, 4096 }) { - suite.device_case(std::string("dense_matmul_swiglu.q4_k.q4_k.prefill.") + glu_op_name(op) + - ".tokens256.outputs" + std::to_string(outputs), - [op, outputs] { - run_dense_matmul_swiglu_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, - outputs, op); - }); - } - } - } + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 4, 128, op); + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 4, 4096, op); + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, 128, op); + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, 4096, op); + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 512, 4096, op); } for (const int64_t tokens : { 1, 2, 3, 5 }) { - suite.device_case( - "dense_matmul_swiglu.q4_k.q4_k.lowtoken.tokens" + std::to_string(tokens) + ".outputs4160", [tokens] { - run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", - tokens, 4160); - }); + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", tokens, 4160); } for (const int64_t tokens : { 1, 2, 5 }) { - suite.device_case( - "dense_matmul_swiglu.q4_k.q4_k.swiglu.tokens" + std::to_string(tokens) + ".outputs16384.input4096", - [tokens] { - run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", - tokens, 16384, GGML_GLU_OP_SWIGLU, 4096); - }); - } - suite.device_case("dense_matmul_swiglu.q4_k.q4_k.reglu.tokens3.outputs16384.input4352", [] { - run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 3, 16384, - GGML_GLU_OP_REGLU, 4352); - }); - suite.device_case("dense_matmul_binary.q4_k.q4_k.lowtoken.tokens4.outputs128", [] { - run_dense_matmul_binary_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 4, 128); - }); - suite.device_case("dense_matmul_binary.q4_k.q4_k.prefill.tokens256.outputs128", [] { - run_dense_matmul_binary_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, 128); - }); - suite.device_case("dense_matmul_binary.q4_k.q4_k.prefill.tokens256.outputs4096", [] { - run_dense_matmul_binary_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, 4096); - }); - suite.device_case("dense_matmul_binary.q4_k.q4_k.prefill.tokens1024.outputs2048", [] { - run_dense_matmul_binary_cpu_reference_case( - GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 1024, 2048); - }); - suite.device_case("dense_matmul_postops.f16.bias.tokens33.outputs256", [] { - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", - 33, 256, true, false, false); - }); - suite.device_case("dense_matmul_postops.f16.residual.tokens33.outputs256", [] { - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", - 33, 256, false, true, false); - }); - suite.device_case("dense_matmul_postops.f16.bias_residual.tokens33.outputs256", [] { - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", - 33, 256, true, true, false); - }); - suite.device_case("dense_matmul_postops.f16.residual.alias.tokens33.outputs256", [] { - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", - 33, 256, false, true, true); - }); - suite.device_case("dense_matmul_postops.f16.bias_residual.alias.tokens33.outputs256", [] { - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, - "loom_libs:ggml_mul_mat_tiled_input_f32_bias_residual_publish_f32", - 33, 256, true, true, true); - }); - for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + tokens, 16384, GGML_GLU_OP_SWIGLU, 4096); + } + run_dense_matmul_swiglu_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", + 3, 16384, GGML_GLU_OP_REGLU, 4352); + run_dense_matmul_binary_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_q8_1_x4_lowtoken_dot", 4, 128); + run_dense_matmul_binary_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, 128); + run_dense_matmul_binary_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 256, 4096); + run_dense_matmul_binary_cpu_reference_case( + GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_swiglu_q4_k_f16_wmma_prefill_wave32", 1024, 2048); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_bias_f32_f32_wmma", 33, 256, + true, false, false); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 33, 256, + false, true, false); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_bias_add_f32_f32_wmma", 33, 256, + true, true, false); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 33, 256, + false, true, true); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_bias_add_f32_f32_wmma", 33, 256, + true, true, true); + for (ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { const int64_t input_size = type == GGML_TYPE_Q4_K ? 20480 : 8192; - suite.device_case( - "dense_matmul." + type_name(type) + ".skinny.tokens5.outputs4096.input" + std::to_string(input_size), - [type, input_size] { - run_dense_matmul_cpu_reference_case(type, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, - 4096, input_size); - }); - suite.device_case("dense_matmul_postops." + type_name(type) + ".add.tokens5.outputs4160", [type, input_size] { - run_dense_matmul_postops_cpu_reference_case( - type, "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", 5, 4160, false, true, false, - input_size + 256); - }); - suite.device_case("dense_matmul_postops." + type_name(type) + ".add.alias.tokens5.outputs4224", - [type, input_size] { - run_dense_matmul_postops_cpu_reference_case( - type, "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", 5, 4224, - false, true, true, input_size + 256); - }); - } - suite.device_case("dense_matmul.q4_k.skinny.tokens5.outputs4096.input8192", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, - 4096, 8192); - }); + run_dense_matmul_cpu_reference_case( + type, "loom_libs:ggml_mul_mat_f32_f32_wmma", 5, 4096, input_size); + run_dense_matmul_postops_cpu_reference_case( + type, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4160, false, true, false, input_size + 256); + run_dense_matmul_postops_cpu_reference_case( + type, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4224, false, true, true, input_size + 256); + } + run_dense_matmul_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 5, 4096, 8192); for (const int64_t input_size : { 6144, 6400, 6656 }) { - suite.device_case( - "dense_matmul.q4_k.skinny.tokens5.outputs5120.input" + std::to_string(input_size), [input_size] { - run_dense_matmul_cpu_reference_case( - GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, 5120, input_size); - }); - suite.device_case("dense_matmul_postops.q4_k.add.alias.tokens5.outputs5120.input" + std::to_string(input_size), - [input_size] { - run_dense_matmul_postops_cpu_reference_case( - GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", - 5, 5120, false, true, true, input_size); - }); - } - suite.device_case("dense_matmul.q4_k.skinny.tokens4.outputs4096.input4096", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 4, - 4096, 4096); - }); - suite.device_case("dense_matmul_postops.q4_k.add.tokens4.outputs4160.input6400", [] { - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", - 4, 4160, false, true, false, 6400); - }); - suite.device_case("dense_matmul_postops.q4_k.add.alias.tokens4.outputs4224.input6656", [] { - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q4_K, - "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", - 4, 4224, false, true, true, 6656); - }); - suite.device_case("endpoint_rmsnorm_q6k_q8.cpu_reference", - [] { run_endpoint_rmsnorm_q6k_q8_cpu_reference_case(); }); - suite.device_case("dense_matmul.q6_k.skinny.tokens5.outputs4096.input16384", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", 5, - 4096, 16384); - }); - suite.device_case("dense_matmul_postops.q6_k.add.tokens5.outputs4160.input16640", [] { - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q6_K, - "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", - 5, 4160, false, true, false, 16640); - }); - suite.device_case("dense_matmul_postops.q6_k.add.alias.tokens5.outputs4224.input16640", [] { - run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_Q6_K, - "loom_libs:ggml_mul_mat_skinny_input_f32_bias_residual_publish_f32", - 5, 4224, false, true, true, 16640); - }); + run_dense_matmul_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 5, 5120, input_size); + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 5120, false, true, true, input_size); + } + run_dense_matmul_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 4, 4096, 4096); + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 4, 4160, false, true, false, 6400); + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 4, 4224, false, true, true, 6656); + run_endpoint_rmsnorm_q6k_q8_cpu_reference_case(); + run_dense_matmul_cpu_reference_case( + GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", 5, 4096, 16384); + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4160, false, true, false, 16640); + run_dense_matmul_postops_cpu_reference_case( + GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_add_f32_f32_wmma", 5, 4224, false, true, true, 16640); for (const int64_t tokens : { 1, 3, 5 }) { - suite.device_case( - "dense_matmul.q4_k.contiguous.tokens" + std::to_string(tokens) + ".outputs256.input5120", [tokens] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q4_K, - tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : - "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", - tokens, 256, 5120, true); - }); - suite.device_case( - "dense_matmul.q6_k.contiguous.tokens" + std::to_string(tokens) + ".outputs256.input5120", [tokens] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, - tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : - "loom_libs:ggml_mul_mat_skinny_input_f32_publish_f32", - tokens, 256, 5120, true, true); - }); - } - suite.device_case("dense_matmul.q6_k.vector_q6.tokens1.outputs256.input5120", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_vector_q6_f32_f32", 1, 256, 5120, - true); - }); - suite.device_case("dense_matmul.q6_k.vector_q6.tokens1.outputs320.input512", [] { - run_dense_matmul_cpu_reference_case(GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_vector_q6_f32_f32", 1, 320, 512, - true); - }); -} - -static void register_rmsnorm_and_scheduling_cases(Suite & suite) { - suite.host_case("support.lfm_head_rope.availability", [] { run_lfm_head_rope_availability_checks(); }); - suite.host_case("gemma_post_proj.matcher_negatives", [] { run_gemma_post_proj_matcher_negatives(); }); - for (const int64_t token_count : { 1, 3, 64 }) { - suite.device_case("gemma_post_proj.tokens" + std::to_string(token_count), - [token_count] { run_gemma_post_proj_cpu_reference_case(token_count); }); - } - suite.device_case("rmsnorm.default", [] { run_rmsnorm_cpu_reference_case(); }); - suite.device_case("rmsnorm.hidden3840.tokens18", [] { run_rmsnorm_cpu_reference_case(3840, 18, 1.0e-6f); }); - for (const int64_t token_count : { 2, 23 }) { - suite.device_case("strided_rmsnorm.tokens" + std::to_string(token_count), - [token_count] { run_strided_rmsnorm_cpu_reference_case(token_count); }); - } - suite.device_case("strided_rmsnorm_binary.minicpm_pp512_witness", - [] { run_strided_rmsnorm_binary_cpu_reference_case(); }); - suite.device_case("rmsnorm_binary_add.cpu_reference", [] { run_rmsnorm_binary_add_cpu_reference_case(); }); - suite.device_case("rmsnorm_gate.silu.hidden128.tokens1.stride7", - [] { run_rmsnorm_gate_cpu_reference_case(128, 1, 7, GGML_UNARY_OP_SILU); }); - suite.device_case("rmsnorm_gate.gelu.hidden256.tokens3.stride3", - [] { run_rmsnorm_gate_cpu_reference_case(256, 3, 3, GGML_UNARY_OP_GELU); }); - suite.device_case("rmsnorm_gate.relu.hidden512.tokens3.stride5", - [] { run_rmsnorm_gate_cpu_reference_case(512, 3, 5, GGML_UNARY_OP_RELU); }); - suite.device_case("rmsnorm_gate.gelu_erf.hidden1024.tokens1.stride4", - [] { run_rmsnorm_gate_cpu_reference_case(1024, 1, 4, GGML_UNARY_OP_GELU_ERF); }); - suite.device_case("rmsnorm_gate.silu.hidden128.tokens48.stride512", - [] { run_rmsnorm_gate_cpu_reference_case(128, 48, 512, GGML_UNARY_OP_SILU); }); - suite.device_case("rmsnorm_mul.hidden3840.tokens18", [] { run_rmsnorm_mul_cpu_reference_case(3840, 18, 1.0e-6f); }); - for (const int64_t hidden_size : { 1024, 2048, 2560, 3072, 4096, 5120 }) { - suite.device_case("rmsnorm_mul.lowtoken.random.hidden" + std::to_string(hidden_size) + ".tokens1", - [hidden_size] { run_rmsnorm_mul_cpu_reference_case(hidden_size, 1, 1.0e-6f); }); - suite.device_case("rmsnorm_mul.lowtoken.adversarial.hidden" + std::to_string(hidden_size) + ".tokens1", - [hidden_size] { run_rmsnorm_mul_cpu_reference_case(hidden_size, 1, 1.0e-5f, true); }); - } - for (const int64_t hidden_size : { 2048, 5120 }) { - for (const int64_t token_count : { 2, 4 }) { - suite.device_case("rmsnorm_mul.lowtoken.random.hidden" + std::to_string(hidden_size) + ".tokens" + - std::to_string(token_count), - [hidden_size, token_count] { - run_rmsnorm_mul_cpu_reference_case(hidden_size, token_count, 1.0e-6f); - }); - } - } - suite.device_case("gemma_scaled_rmsnorm.cpu_reference", [] { run_gemma_scaled_rmsnorm_cpu_reference_case(); }); - suite.device_case("gemma_scaled_rmsnorm_mul.cpu_reference", - [] { run_gemma_scaled_rmsnorm_mul_cpu_reference_case(); }); - suite.device_case("scheduled.scale_rmsnorm", [] { run_scheduled_hrx_scale_rmsnorm_case(); }); - suite.device_case("scheduled.rmsnorm_scale", [] { run_scheduled_hrx_rmsnorm_scale_case(); }); - suite.device_case("scheduled.rmsnorm_view_scale", [] { run_scheduled_hrx_rmsnorm_view_scale_case(); }); - suite.device_case("external_view_input_rmsnorm", [] { run_external_view_input_rmsnorm_case(); }); - suite.device_case("split_local_view_alias_import", [] { run_split_local_view_alias_import_case(); }); - suite.device_case("rope_set_rows.cpu_reference", [] { run_rope_set_rows_cpu_reference_case(); }); - for (const int64_t head_count : { 8, 32 }) { - for (const int64_t token_count : { 1, 64 }) { - suite.device_case( - "lfm_head_rope.query.heads" + std::to_string(head_count) + ".tokens" + std::to_string(token_count), - [head_count, token_count] { run_lfm_head_rope_cpu_reference_case(head_count, token_count, false); }); - } - } - suite.device_case("lfm_head_rope.query.heads7.tokens3", [] { run_lfm_head_rope_cpu_reference_case(7, 3, false); }); - for (const int64_t token_count : { 1, 64 }) { - suite.device_case("lfm_head_rope.cache.heads8.tokens" + std::to_string(token_count), - [token_count] { run_lfm_head_rope_cpu_reference_case(8, token_count, true); }); - } - suite.device_case("attention_postprocess.cpu_reference", [] { run_attention_postprocess_cpu_reference_case(); }); - suite.device_case("routed_moe.q4_k.next_rmsnorm", [] { run_routed_moe_cpu_reference_case(GGML_TYPE_Q4_K, true); }); - suite.device_case("flash_attention.asymmetric", [] { run_asymmetric_flash_attention_cpu_reference_case(); }); - suite.device_case("routed_moe.q6_k", [] { run_routed_moe_cpu_reference_case(GGML_TYPE_Q6_K, false); }); - suite.device_case("decode_attention_qkv.scheduling", [] { run_decode_attention_qkv_scheduling_case(); }); - suite.device_case("decode_routed_moe.q4_k", [] { run_decode_routed_moe_scheduling_case(GGML_TYPE_Q4_K); }); - suite.device_case("decode_routed_moe.q6_k", [] { run_decode_routed_moe_scheduling_case(GGML_TYPE_Q6_K); }); - suite.device_case("decode_routed_moe.q6_k.alias_gate_input", - [] { run_decode_routed_moe_scheduling_case(GGML_TYPE_Q6_K, true); }); - suite.device_case("rmsnorm_mul.hidden256.tokens1", [] { run_rmsnorm_mul_case(256, 1); }); - suite.device_case("rmsnorm_mul.hidden256.tokens4", [] { run_rmsnorm_mul_case(256, 4); }); - suite.device_case("rmsnorm_mul.hidden2048.tokens1", [] { run_rmsnorm_mul_case(2048, 1); }); - suite.device_case("router_projection.tokens4", [] { run_router_projection_case(4); }); - suite.device_case("router_top8.tokens4", [] { run_router_top8_case(4); }); - suite.device_case("qwen_flash_attention", [] { run_qwen_flash_attention_case(); }); -} - -static void register_flash_attention_cases(Suite & suite) { - suite.device_case("common_flash_attention.prefill.head128.query256.kv256.exact", [] { - run_common_flash_attention_cpu_reference_case(128, 256, 256, 1.0f / std::sqrt(128.0f), - "loom_libs:ggml_flash_attention_f32_f16_wmma"); - }); + run_dense_matmul_cpu_reference_case( + GGML_TYPE_Q4_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", tokens, 256, 5120, true); + run_dense_matmul_cpu_reference_case( + GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_f32_f32_wmma", tokens, 256, 5120, true, true); + } + run_dense_matmul_cpu_reference_case( + GGML_TYPE_Q6_K, "loom_libs:ggml_mul_mat_q6_k_packed_token1_f16_wmma", 1, 256, 5120, true); + run_rmsnorm_cpu_reference_case(); + run_rmsnorm_cpu_reference_case(3840, 18, 1.0e-6f); + run_strided_rmsnorm_cpu_reference_case(2); + run_strided_rmsnorm_cpu_reference_case(23); + run_rmsnorm_binary_add_cpu_reference_case(); + run_rmsnorm_gate_cpu_reference_case(128, 1, 7, GGML_UNARY_OP_SILU); + run_rmsnorm_gate_cpu_reference_case(256, 3, 3, GGML_UNARY_OP_GELU); + run_rmsnorm_gate_cpu_reference_case(512, 3, 5, GGML_UNARY_OP_RELU); + run_rmsnorm_gate_cpu_reference_case(1024, 1, 4, GGML_UNARY_OP_GELU_ERF); + run_rmsnorm_gate_cpu_reference_case(128, 48, 512, GGML_UNARY_OP_SILU); + run_rmsnorm_mul_cpu_reference_case(3840, 18, 1.0e-6f); + run_gemma_scaled_rmsnorm_cpu_reference_case(); + run_gemma_scaled_rmsnorm_mul_cpu_reference_case(); + run_scheduled_hrx_scale_rmsnorm_case(); + run_scheduled_hrx_rmsnorm_scale_case(); + run_scheduled_hrx_rmsnorm_view_scale_case(); + run_external_view_input_rmsnorm_case(); + run_split_local_view_alias_import_case(); + run_rope_set_rows_cpu_reference_case(); + run_attention_postprocess_cpu_reference_case(); + run_routed_moe_cpu_reference_case(GGML_TYPE_Q4_K, true); + run_asymmetric_flash_attention_cpu_reference_case(); + run_routed_moe_cpu_reference_case(GGML_TYPE_Q6_K, false); + run_decode_attention_qkv_scheduling_case(); + run_decode_routed_moe_scheduling_case(GGML_TYPE_Q4_K); + run_decode_routed_moe_scheduling_case(GGML_TYPE_Q6_K); + run_decode_routed_moe_scheduling_case(GGML_TYPE_Q6_K, true); + run_rmsnorm_mul_case(256, 1); + run_rmsnorm_mul_case(256, 4); + run_rmsnorm_mul_case(2048, 1); + run_router_projection_case(4); + run_router_top8_case(4); + run_qwen_flash_attention_case(); for (int64_t head_size = 64; head_size <= 512; head_size += 64) { - suite.device_case("common_flash_attention.decode.head" + std::to_string(head_size) + ".query2.kv65", - [head_size] { - run_common_flash_attention_cpu_reference_case( - head_size, 2, 65, 1.0f / std::sqrt(static_cast(head_size)), - "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); - }); - suite.device_case( - "common_flash_attention.prefill.head" + std::to_string(head_size) + ".query16.kv64", [head_size] { - run_common_flash_attention_cpu_reference_case(head_size, 16, 64, - 1.0f / std::sqrt(static_cast(head_size)), - "loom_libs:ggml_flash_attention_f32_f16_wmma"); - }); - suite.device_case( - "common_flash_attention.prefill.head" + std::to_string(head_size) + ".query17.kv65", [head_size] { - run_common_flash_attention_cpu_reference_case(head_size, 17, 65, - 1.0f / std::sqrt(static_cast(head_size)), - "loom_libs:ggml_flash_attention_f32_f16_wmma"); - }); - } - suite.device_case("common_flash_attention.decode.head96.query2.kv4.capacity320", [] { - run_common_flash_attention_cpu_reference_case( - 96, 2, 4, 1.0f / std::sqrt(96.0f), "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", 320); - }); - suite.device_case("common_flash_attention.decode.head320.query2.kv64.capacity64", [] { - run_common_flash_attention_cpu_reference_case( - 320, 2, 64, 1.0f / std::sqrt(320.0f), "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", - 64); - }); - suite.device_case("qwen_decode_split_flash_attention.query1.kv512", - [] { run_qwen_decode_split_flash_attention_scheduling_case(1, 512); }); - suite.device_case("qwen_decode_split_flash_attention.query4.kv513", - [] { run_qwen_decode_split_flash_attention_scheduling_case(4, 513); }); - suite.device_case("qwen_decode_split_flash_attention.query1.kv512.masked", - [] { run_qwen_decode_split_flash_attention_scheduling_case(1, 512, true); }); - suite.device_case("qwen_decode_split_flash_attention.query1.kv512.masked.head96.value64", [] { - run_qwen_decode_split_flash_attention_scheduling_case(1, 512, true, 96, 64, 1.0f / std::sqrt(96.0f)); - }); - suite.device_case("qwen_decode_split_flash_attention.query4.kv513.masked.head96.value64", [] { - run_qwen_decode_split_flash_attention_scheduling_case(4, 513, true, 96, 64, 1.0f / std::sqrt(96.0f)); - }); - suite.device_case("qwen_decode_attention_output_next_q8.no_selectors", - [] { run_qwen_decode_attention_output_next_q8_scheduling_case(false); }); - suite.device_case("qwen_decode_attention_output_next_q8.selectors", - [] { run_qwen_decode_attention_output_next_q8_scheduling_case(true); }); - suite.device_case("qwen_full_cache_prefill_flash_attention.tokens512", - [] { run_qwen_full_cache_prefill_flash_attention_scheduling_case(512); }); -} - -static void register_hrx_ops_cases(Suite & suite) { - register_basic_ops_cases(suite); - register_dense_matmul_cases(suite); - register_rmsnorm_and_scheduling_cases(suite); - register_flash_attention_cases(suite); -} - -int main(int argc, char ** argv) { - Suite suite(test_runner::Config::with_prefix("HRX test", "hrx", "GGML_HRX_TEST")); - register_hrx_ops_cases(suite); - - test_runner::Options options; - if (!suite.parse_options(argc, argv, options)) { - suite.print_usage(argc > 0 ? argv[0] : "test-hrx-ops"); - return 2; - } - if (options.help) { - return suite.run(options, false); - } - - const bool run_vector_cap_experiment = std::getenv("GGML_HRX_RUN_VECTOR_CAP_EXPERIMENT") != nullptr; - const bool run_vector_publish_test = std::getenv("GGML_HRX_RUN_VECTOR_PUBLISH_TEST") != nullptr; - const bool run_vector_postops_test = std::getenv("GGML_HRX_RUN_VECTOR_POSTOPS_TEST") != nullptr; - - const bool has_device = ggml_backend_hrx_get_device_count() != 0; - if (!options.runner_control && run_vector_cap_experiment) { - run_vector_output_capacity_experiment(); - return 0; - } - if (!options.runner_control && run_vector_publish_test) { - run_vector_q8_publish_cpu_reference_case(); - return 0; - } - if (!options.runner_control && run_vector_postops_test) { - run_vector_postops_cpu_reference_case(); - return 0; + run_common_flash_attention_cpu_reference_case(head_size, 2, 65, + 1.0f / std::sqrt(static_cast(head_size)), + "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8"); + run_common_flash_attention_cpu_reference_case(head_size, 16, 64, + 1.0f / std::sqrt(static_cast(head_size)), + "loom_libs:ggml_flash_attention_f32_f16_wmma"); + run_common_flash_attention_cpu_reference_case(head_size, 17, 65, + 1.0f / std::sqrt(static_cast(head_size)), + "loom_libs:ggml_flash_attention_f32_f16_wmma"); } - - return suite.run(options, has_device); + run_common_flash_attention_cpu_reference_case(96, 2, 4, 1.0f / std::sqrt(96.0f), + "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", 320); + run_common_flash_attention_cpu_reference_case(320, 2, 64, 1.0f / std::sqrt(320.0f), + "loom_libs:ggml_flash_attention_decode_split_f32_f16_wmma_next_q8", 64); + run_qwen_decode_split_flash_attention_scheduling_case(1, 512); + run_qwen_decode_split_flash_attention_scheduling_case(4, 513); + run_qwen_decode_split_flash_attention_scheduling_case(1, 512, true); + run_qwen_decode_split_flash_attention_scheduling_case(1, 512, true, 96, 64, 1.0f / std::sqrt(96.0f)); + run_qwen_decode_split_flash_attention_scheduling_case(4, 513, true, 96, 64, 1.0f / std::sqrt(96.0f)); + run_qwen_decode_attention_output_next_q8_scheduling_case(false); + run_qwen_decode_attention_output_next_q8_scheduling_case(true); + run_qwen_full_cache_prefill_flash_attention_scheduling_case(512); + return 0; } From b785c2daad1a02deb7fc7b85e36b8a978e6c9430 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Sun, 18 Oct 2026 02:11:59 -0300 Subject: [PATCH 2/2] tests: update two HRX test expectations to our later ggml-hrx changes (engine#422) The 522dab47 copies restored in the previous commit predate two changes of ours: - a4963bb29 lets GET_ROWS of IQ3_S and IQ4_NL rows run on HRX, so the token embedding graph for both types is now schedulable. hrx-backend-test still required it to be refused and aborted at line 3494; it now requires support. - 39588ef99 points ggml_flash_attention_f32_f16_wmma at a kernel that declares @ggml.flash_attention.value_stride. The dispatch passes value->nb[1] in f16 elements; the test's prefill config did not set it, so the async-flash-prefill compiles failed. It now sets key_value_head_count * head_size, the stride of a contiguous value row. Measured on strixhalo (performance mode), two runs each: test-hrx-loom-jit aborted at line 369 -> exit 0 (28 kernels compiled) hrx-backend-test aborts at line 3731 instead of 3494 (gather-add support expectation, not changed here) Co-Authored-By: Claude Opus 5.5 --- tests/hrx-backend-test.cpp | 8 ++++---- tests/test-hrx-loom-jit.cpp | 2 ++ 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 6573d16702f0..a508c706f6c9 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -3490,10 +3490,10 @@ static void run_qwen_token_embedding_dispatch_checks() { manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q5_1, GGML_TYPE_I32, GGML_TYPE_F32, 640, 262144, 1)); REQUIRE(manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_IQ4_XS, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, 1)); - REQUIRE(!manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_IQ3_S, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, - 1)); - REQUIRE(!manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_IQ4_NL, GGML_TYPE_I32, GGML_TYPE_F32, 2048, - 151936, 1)); + REQUIRE(manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_IQ3_S, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, + 1)); + REQUIRE(manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_IQ4_NL, GGML_TYPE_I32, GGML_TYPE_F32, 2048, + 151936, 1)); REQUIRE( manual_token_embedding_graph_is_supported(ctx, GGML_TYPE_Q8_0, GGML_TYPE_I32, GGML_TYPE_F32, 2048, 151936, 1)); REQUIRE( diff --git a/tests/test-hrx-loom-jit.cpp b/tests/test-hrx-loom-jit.cpp index 880ce1f56b86..cdcc326b6800 100644 --- a/tests/test-hrx-loom-jit.cpp +++ b/tests/test-hrx-loom-jit.cpp @@ -226,6 +226,8 @@ static std::map flash_attention_config(const char * qu { "ggml.flash_attention.key_value_head_count", key_value_head_count }, { "ggml.flash_attention.qk_head_size", head_size }, { "ggml.flash_attention.value_head_size", head_size }, + { "ggml.flash_attention.value_stride", + std::to_string(std::stoll(key_value_head_count) * std::stoll(head_size)) }, { "ggml.flash_attention.attention_scale", attention_scale }, { "ggml.flash_attention.apply_gate", "0" }, { "ggml.flash_attention.gate_stride_head", "1" },