diff --git a/ggml/src/ggml-hrx/backend-context.h b/ggml/src/ggml-hrx/backend-context.h index 43c27ba2bbe8..a629880cefb5 100644 --- a/ggml/src/ggml-hrx/backend-context.h +++ b/ggml/src/ggml-hrx/backend-context.h @@ -69,6 +69,24 @@ struct ggml_backend_hrx_context { ggml::hrx::HostWeightCache host_weights; ggml::hrx::GraphReplayStreamState graph_replay_state; std::string name; + std::mutex layout_policy_mutex; + bool row64_weight_policy = false; + bool layout_policy_locked = false; + + bool set_row64_weight_policy(bool enabled) { + std::lock_guard lock(layout_policy_mutex); + if (layout_policy_locked) { + return false; + } + row64_weight_policy = enabled; + return true; + } + + bool lock_layout_policy() { + std::lock_guard lock(layout_policy_mutex); + layout_policy_locked = true; + return row64_weight_policy; + } }; struct ggml_backend_hrx_reg_context { diff --git a/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp b/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp index 12d4f5fa338a..56dc2ed3d7ee 100644 --- a/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp +++ b/ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp @@ -31,11 +31,12 @@ static bool try_match_registration(const Graph & graph, const std::vector & covered_nodes, const CommandPlan & plan, const DispatchRegistry & registry, + const DispatchTarget & target, ValueId next_plan_value, DispatchMatch & match, DispatchMatchDiagnostics * diagnostics) { const DispatchMatchContext context = { - graph, node, node_index, covered_nodes, plan, next_plan_value, ®istry, + graph, node, node_index, covered_nodes, plan, next_plan_value, ®istry, target.row64_weight_policy, }; return registry.match(context, match, diagnostics); } @@ -197,7 +198,7 @@ bool DispatchScheduler::schedule_graph(Graph & graph, const ValueId next_plan_value(static_cast(graph.values().size() + plan_.transients.size() + plan_.completion_counter_requests.size())); DispatchMatchDiagnostics match_diagnostics; - if (!try_match_registration(graph, node, i, covered_nodes, plan_, *registry, next_plan_value, match, + if (!try_match_registration(graph, node, i, covered_nodes, plan_, *registry, target, next_plan_value, match, &match_diagnostics)) { if (can_elide_node(graph, *node, covered_nodes)) { pending_diagnostics.append(match.status); @@ -276,7 +277,7 @@ bool DispatchScheduler::schedule_graph(Graph & graph, const ValueId next_plan_value(static_cast(graph.values().size() + plan_.transients.size() + plan_.completion_counter_requests.size())); DispatchMatchDiagnostics match_diagnostics; - try_match_registration(graph, &nodes[i], i, covered_nodes, plan_, *registry, next_plan_value, match, + try_match_registration(graph, &nodes[i], i, covered_nodes, plan_, *registry, target, next_plan_value, match, &match_diagnostics); diagnostics->unsupported_node_index = i; diagnostics->unsupported_node = &nodes[i]; @@ -309,7 +310,7 @@ bool DispatchScheduler::supports_node(const Graph & graph, const GraphNode * nod DispatchMatch match; const ValueId next_plan_value(static_cast(graph.values().size())); const CommandPlan plan; - return try_match_registration(graph, node, node_index, covered_nodes, plan, *registry, next_plan_value, match, + return try_match_registration(graph, node, node_index, covered_nodes, plan, *registry, target, next_plan_value, match, nullptr); } diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp index cfeb31398a78..772aa3dd0822 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp @@ -137,6 +137,37 @@ struct PackedMulMatGluMatch { } }; +static CommonMulMatWeightFormat row64_swiglu_format(const DispatchMatchContext & context, + const Value & weight, + CommonMulMatWeightFormat format, + int64_t input_size, + int64_t output_size) { + const bool direct_weight = weight.alias_source.value < 0; + const bool aligned_input = input_size % 256 == 0; + const bool aligned_output = output_size % 64 == 0; + const bool q4_weight = format == CommonMulMatWeightFormat::Q4K; + const bool q6_weight = format == CommonMulMatWeightFormat::Q6K; + const bool supported_weight = q4_weight || q6_weight; + const bool use_row64 = context.row64_weight_policy && direct_weight && aligned_input && aligned_output && + supported_weight && dispatch_weight_layout_eligible_for_graph( + context.graph, weight, context.row64_weight_policy); + if (!use_row64) { + return format; + } + return q4_weight ? CommonMulMatWeightFormat::Q4KRow64 : CommonMulMatWeightFormat::Q6KRow64; +} + +static DispatchBinding row64_swiglu_binding(const Value & weight, CommonMulMatWeightFormat format, + int64_t input_size, int64_t output_size) { + const bool q4_row64 = format == CommonMulMatWeightFormat::Q4KRow64; + const bool q6_row64 = format == CommonMulMatWeightFormat::Q6KRow64; + if (!q4_row64 && !q6_row64) { + return { weight.id, 0, weight.byte_count }; + } + const char * layout = q4_row64 ? kQ4KPackedK256Row64Layout : kQ6KPackedK256Row64ScaleRowLayout; + return { weight.id, 0, weight.byte_count, layout, weight.type, input_size, output_size, weight.byte_count }; +} + static bool supported_symmetric_i4_pair(ggml_type gate_type, ggml_type up_type) { return (gate_type == GGML_TYPE_Q4_K && up_type == GGML_TYPE_Q4_K) || (gate_type == GGML_TYPE_Q5_K && up_type == GGML_TYPE_IQ4_XS) || @@ -718,7 +749,8 @@ static CommonActivationPublicationDemand select_swiglu_q8_publication(const Disp return publication_plan.matched() ? publication_plan.preferred() : CommonActivationPublicationDemand{}; } -static bool common_mul_mat_swiglu_is_generic_tiled_pair_route(const MulMatSwiGLUMatch & match) { +static bool common_mul_mat_swiglu_is_generic_tiled_pair_route(const DispatchMatchContext & context, + const MulMatSwiGLUMatch & match) { if (!match.matched()) { return false; } @@ -729,8 +761,10 @@ static bool common_mul_mat_swiglu_is_generic_tiled_pair_route(const MulMatSwiGLU const bool has_direct_up = match.up_weight->alias_source.value < 0; const bool has_q4_gate = match.gate_format == CommonMulMatWeightFormat::Q4K; const bool has_q4_up = match.up_format == CommonMulMatWeightFormat::Q4K; - const bool can_transform_gate = dispatch_weight_layout_eligible(*match.gate_weight); - const bool can_transform_up = dispatch_weight_layout_eligible(*match.up_weight); + const bool can_transform_gate = dispatch_weight_layout_eligible_for_graph( + context.graph, *match.gate_weight, context.row64_weight_policy); + const bool can_transform_up = dispatch_weight_layout_eligible_for_graph( + context.graph, *match.up_weight, context.row64_weight_policy); const bool pack_q4 = has_few_tokens && has_aligned_output && has_direct_gate && has_direct_up && has_q4_gate && has_q4_up && can_transform_gate && can_transform_up; const bool use_q8 = pack_q4; @@ -757,7 +791,7 @@ static CommonActivationPublicationDemand select_swiglu_k16_publication(const Dis static MulMatSwiGLUPostOpsMatch match_mul_mat_swiglu_postops(const DispatchMatchContext & context) { MulMatSwiGLUPostOpsMatch match; match.root = match_mul_mat_swiglu(context); - if (!common_mul_mat_swiglu_is_generic_tiled_pair_route(match.root) || !context.graph.has_index()) { + if (!common_mul_mat_swiglu_is_generic_tiled_pair_route(context, match.root) || !context.graph.has_index()) { return {}; } @@ -818,7 +852,7 @@ static bool build_tiled_pair_mul_mat_swiglu_dispatch(const DispatchMatchContext const MulMatSwiGLUPostOpsMatch * postops, DispatchMatch & dispatch_match) { const MulMatSwiGLUMatch & match = postops != nullptr ? postops->root : match_mul_mat_swiglu(context); - if (!common_mul_mat_swiglu_is_generic_tiled_pair_route(match)) { + if (!common_mul_mat_swiglu_is_generic_tiled_pair_route(context, match)) { return false; } @@ -847,12 +881,16 @@ static bool build_tiled_pair_mul_mat_swiglu_dispatch(const DispatchMatchContext common_to_config_value(match.input_size)); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.output_size", common_to_config_value(match.output_size)); + const CommonMulMatWeightFormat gate_format = row64_swiglu_format( + context, *match.gate_weight, match.gate_format, match.input_size, match.output_size); + const CommonMulMatWeightFormat up_format = row64_swiglu_format( + context, *match.up_weight, match.up_format, match.input_size, match.output_size); dispatch.kernel.compile_parameters.emplace( "ggml.mul_mat_swiglu.gate_weight_format", - common_to_config_value(common_mul_mat_format_config_value(match.gate_format))); + common_to_config_value(common_mul_mat_format_config_value(gate_format))); dispatch.kernel.compile_parameters.emplace( "ggml.mul_mat_swiglu.up_weight_format", - common_to_config_value(common_mul_mat_format_config_value(match.up_format))); + common_to_config_value(common_mul_mat_format_config_value(up_format))); dispatch.kernel.compile_parameters.emplace( "ggml.mul_mat_swiglu.op", common_to_config_value(static_cast(binary_kind_config_value(match.op)))); if (postops != nullptr) { @@ -862,8 +900,10 @@ static bool build_tiled_pair_mul_mat_swiglu_dispatch(const DispatchMatchContext (postops->has_residual ? int64_t{ 2 } : int64_t{ 0 }))); } dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); - dispatch.bindings.push_back({ match.gate_weight->id, 0, match.gate_weight->byte_count }); - dispatch.bindings.push_back({ match.up_weight->id, 0, match.up_weight->byte_count }); + dispatch.bindings.push_back(row64_swiglu_binding(*match.gate_weight, gate_format, + match.input_size, match.output_size)); + dispatch.bindings.push_back(row64_swiglu_binding(*match.up_weight, up_format, + match.input_size, match.output_size)); const Value * output = postops != nullptr ? postops->residual_output : match.output; if (postops != nullptr) { const Value * placeholder = postops->has_bias ? postops->bias : postops->residual_input; @@ -926,7 +966,7 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, if (!match.matched()) { return false; } - if (common_mul_mat_swiglu_is_generic_tiled_pair_route(match)) { + if (common_mul_mat_swiglu_is_generic_tiled_pair_route(context, match)) { return false; } @@ -970,10 +1010,16 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, const bool has_direct_up = match.up_weight->alias_source.value < 0; const bool has_q4_gate = match.gate_format == CommonMulMatWeightFormat::Q4K; const bool has_q4_up = match.up_format == CommonMulMatWeightFormat::Q4K; - const bool can_transform_gate = dispatch_weight_layout_eligible(*match.gate_weight); - const bool can_transform_up = dispatch_weight_layout_eligible(*match.up_weight); + const bool can_transform_gate = dispatch_weight_layout_eligible_for_graph( + context.graph, *match.gate_weight, context.row64_weight_policy); + const bool can_transform_up = dispatch_weight_layout_eligible_for_graph( + context.graph, *match.up_weight, context.row64_weight_policy); const bool pack_q4 = has_few_tokens && has_aligned_output && has_direct_gate && has_direct_up && has_q4_gate && has_q4_up && can_transform_gate && can_transform_up; + const CommonMulMatWeightFormat gate_format = pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : + row64_swiglu_format(context, *match.gate_weight, match.gate_format, match.input_size, match.output_size); + const CommonMulMatWeightFormat up_format = pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : + row64_swiglu_format(context, *match.up_weight, match.up_format, match.input_size, match.output_size); const bool use_q8 = pack_q4; const CommonActivationPublicationDemand q8_demand = use_q8 ? select_swiglu_q8_publication(context, match) : CommonActivationPublicationDemand{}; @@ -1003,11 +1049,9 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.output_size", common_to_config_value(match.output_size)); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.gate_weight_format", - common_to_config_value(common_mul_mat_format_config_value( - pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : match.gate_format))); + common_to_config_value(common_mul_mat_format_config_value(gate_format))); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.up_weight_format", - common_to_config_value(common_mul_mat_format_config_value( - pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : match.up_format))); + common_to_config_value(common_mul_mat_format_config_value(up_format))); dispatch.kernel.compile_parameters.emplace( "ggml.mul_mat_swiglu.op", common_to_config_value(static_cast(binary_kind_config_value(match.op)))); dispatch.bindings.push_back(activation); @@ -1017,8 +1061,10 @@ static bool match_mul_mat_swiglu_dispatch(const DispatchMatchContext & context, match.input_size, match.output_size, weight->byte_count }); } } else { - dispatch.bindings.push_back({ match.gate_weight->id, 0, match.gate_weight->byte_count }); - dispatch.bindings.push_back({ match.up_weight->id, 0, match.up_weight->byte_count }); + dispatch.bindings.push_back(row64_swiglu_binding(*match.gate_weight, gate_format, + match.input_size, match.output_size)); + dispatch.bindings.push_back(row64_swiglu_binding(*match.up_weight, up_format, + match.input_size, match.output_size)); } dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); @@ -1052,6 +1098,10 @@ static bool match_decode_mul_mat_swiglu_dispatch(const DispatchMatchContext & co } Dispatch dispatch; + const CommonMulMatWeightFormat gate_format = row64_swiglu_format( + context, *match.gate_weight, match.gate_format, match.input_size, match.output_size); + const CommonMulMatWeightFormat up_format = row64_swiglu_format( + context, *match.up_weight, match.up_format, match.input_size, match.output_size); dispatch.kernel = make_kernel_specialization(kMulMatSwiGLUF32F32DecodeWave64Kernel); dispatch.kernel.integer_parameters.emplace("token_count", match.token_count); dispatch.kernel.compile_parameters.emplace("ggml.workload.token_capacity", @@ -1062,15 +1112,17 @@ static bool match_decode_mul_mat_swiglu_dispatch(const DispatchMatchContext & co common_to_config_value(match.output_size)); dispatch.kernel.compile_parameters.emplace( "ggml.mul_mat_swiglu.gate_weight_format", - common_to_config_value(common_mul_mat_format_config_value(match.gate_format))); + common_to_config_value(common_mul_mat_format_config_value(gate_format))); dispatch.kernel.compile_parameters.emplace( "ggml.mul_mat_swiglu.up_weight_format", - common_to_config_value(common_mul_mat_format_config_value(match.up_format))); + common_to_config_value(common_mul_mat_format_config_value(up_format))); dispatch.kernel.compile_parameters.emplace( "ggml.mul_mat_swiglu.op", common_to_config_value(static_cast(binary_kind_config_value(match.op)))); dispatch.bindings.push_back({ match.input->id, 0, match.input->byte_count }); - dispatch.bindings.push_back({ match.gate_weight->id, 0, match.gate_weight->byte_count }); - dispatch.bindings.push_back({ match.up_weight->id, 0, match.up_weight->byte_count }); + dispatch.bindings.push_back(row64_swiglu_binding(*match.gate_weight, gate_format, + match.input_size, match.output_size)); + dispatch.bindings.push_back(row64_swiglu_binding(*match.up_weight, up_format, + match.input_size, match.output_size)); dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); if (!append_covered_node_index_once(context.graph, context.covered_nodes, match.gate_node, diff --git a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp index 4e785ac0fb79..96c4b8216940 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat.cpp @@ -272,6 +272,26 @@ static bool common_mul_mat_is_tiled_route(const CommonMulMatMatch & match) { return match.matched() && match.token_count > 1 && !common_mul_mat_is_skinny_route(match); } +static bool native_packed_glu_weight(const DispatchMatchContext & context, const Value & weight) { + const bool q4_weight = weight.type == GGML_TYPE_Q4_K; + const bool direct_weight = weight.alias_source.value < 0; + const bool has_graph_index = context.graph.has_index(); + const bool has_root = context.root_node != nullptr; + const bool inspect_consumer = context.row64_weight_policy && q4_weight && direct_weight && has_graph_index && has_root; + if (!inspect_consumer) { + return false; + } + const std::vector & consumers = context.graph.index().consumers(context.root_node->output); + if (consumers.size() != 1 || consumers.front() == nullptr) { + return false; + } + const GraphNode * consumer = consumers.front(); + const bool glu = consumer->op == GGML_OP_GLU; + const bool one_input = consumer->inputs.size() == 1; + const bool reads_projection = one_input && consumer->inputs[0] == context.root_node->output; + return glu && reads_projection; +} + static bool common_mul_mat_postops_is_skinny_route(const MulMatPostOpsMatch & match) { return match.matched() && common_mul_mat_is_low_token_batch(match.token_count) && common_mul_mat_has_skinny_reduction_shape(match.input_size); @@ -843,6 +863,7 @@ static bool build_q6_k_prefill_wave32_dispatch(const CommonMulMatMatch & match, ValueId input, size_t input_bytes, KernelCatalogRef kernel, + bool row64_weight, size_t root_index, DispatchMatch & dispatch_match) { Dispatch dispatch; @@ -853,10 +874,19 @@ static bool build_q6_k_prefill_wave32_dispatch(const CommonMulMatMatch & match, dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.output_size", common_to_config_value(match.output_size)); dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.output_accumulation", "0"); + if (row64_weight) { + dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.weight_format", "46"); + } dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_q6_k_packed.token_capacity", common_to_config_value(match.token_count)); dispatch.bindings.push_back({ input, 0, input_bytes }); - dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + if (row64_weight) { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count, + kQ6KPackedK256Row64ScaleRowLayout, match.weight->type, match.input_size, + match.output_size, match.weight->byte_count }); + } else { + dispatch.bindings.push_back({ match.weight->id, 0, match.weight->byte_count }); + } dispatch.bindings.push_back({ match.output->id, 0, match.output->byte_count }); dispatch_match.covered_nodes.push_back(root_index); @@ -867,12 +897,38 @@ static bool build_q6_k_prefill_wave32_dispatch(const CommonMulMatMatch & match, static bool match_q6_k_prefill_wave32_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { const CommonMulMatMatch match = common_match_mul_mat_any_format(context.graph, context.root_node, kMulMatQ6KF32WmmaPrefillWave32Kernel, false); - if (!match.matched() || match.weight->type != GGML_TYPE_Q6_K || match.weight->alias_source.value >= 0 || - match.token_count < 128 || match.token_count > 2048 || match.token_count % 128 != 0 || - match.input_size % 256 != 0 || match.output_size % 64 != 0) { + if (!match.matched()) { + return false; + } + const bool q6_weight = match.weight->type == GGML_TYPE_Q6_K; + const bool direct_weight = match.weight->alias_source.value < 0; + const bool aligned_input = match.input_size % 256 == 0; + const bool aligned_output = match.output_size % 64 == 0; + const bool supported_weight = q6_weight && direct_weight && aligned_input && aligned_output; + if (!supported_weight) { + return false; + } + const bool enough_tokens = match.token_count >= 64; + const bool short_prompt = match.token_count < 128; + const bool row64_weight = context.row64_weight_policy && dispatch_weight_layout_eligible_for_graph( + context.graph, *match.weight, context.row64_weight_policy); + const bool short_row64 = enough_tokens && short_prompt && row64_weight; + const bool ordinary_minimum = match.token_count >= 128; + const bool ordinary_maximum = match.token_count <= 2048; + const bool ordinary_aligned = match.token_count % 128 == 0; + const bool ordinary_prefill = ordinary_minimum && ordinary_maximum && ordinary_aligned; + const bool supported_tokens = short_row64 || ordinary_prefill; + if (!supported_tokens) { return false; } + // Short prompts use masked M128 tiles without a separate activation copy. + if (short_row64) { + return build_q6_k_prefill_wave32_dispatch(match, match.input->id, match.input->byte_count, + kMulMatQ6KF32WmmaPrefillWave32Kernel, row64_weight, + context.root_index, dispatch_match); + } + const bool packed_input = common_mul_mat_uses_k16_major_f16(match.weight_format, match.input_size, match.output_size, match.token_count); DispatchBinding activation; @@ -885,7 +941,8 @@ static bool match_q6_k_prefill_wave32_dispatch(const DispatchMatchContext & cont return false; } return build_q6_k_prefill_wave32_dispatch(match, activation.value, activation.length, - kMulMatQ6KF16WmmaPrefillWave32Kernel, context.root_index, dispatch_match); + kMulMatQ6KF16WmmaPrefillWave32Kernel, row64_weight, + context.root_index, dispatch_match); } static bool match_symmetric_i4_prefill_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { @@ -1037,8 +1094,17 @@ static bool match_packed_q8_1_x4_prefill_dispatch(const DispatchMatchContext & c const bool is_q4_weight = match.weight->type == GGML_TYPE_Q4_K; const bool has_direct_weight = match.weight->alias_source.value < 0; const bool has_wide_output = match.output_size >= match.input_size / 4; - const bool can_transform_weight = dispatch_weight_layout_eligible(*match.weight); - const bool use_f16 = is_q4_weight && has_direct_weight && has_wide_output && can_transform_weight; + const bool can_transform_weight = dispatch_weight_layout_eligible_for_graph( + context.graph, *match.weight, context.row64_weight_policy); + const bool native_packed_glu = native_packed_glu_weight(context, *match.weight); + const bool native_qwen30_attention = context.row64_weight_policy && + dispatch_weight_is_qwen30b_native_attention(context.graph, *match.weight); + const bool eligible_f16_weight = can_transform_weight || native_qwen30_attention; + const bool use_f16 = is_q4_weight && has_direct_weight && has_wide_output && eligible_f16_weight; + const bool unsupported_native_glu = native_packed_glu && !use_f16; + if (unsupported_native_glu) { + return false; + } DispatchBinding activation; const bool packed_input = use_f16 && @@ -1074,9 +1140,13 @@ static bool match_packed_q8_1_x4_prefill_dispatch(const DispatchMatchContext & c use_f16 ? "ggml.workload.token_capacity" : "ggml.mul_mat_q8_1_x4.token_capacity", common_to_config_value(match.token_count)); const bool pack_q4 = match.weight_format == CommonMulMatWeightFormat::Q4K && - has_direct_weight && can_transform_weight; + has_direct_weight && can_transform_weight && !native_packed_glu; if (use_f16) { dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.f16_input_layout", packed_input ? "1" : "0"); + dispatch.kernel.compile_parameters.emplace( + "ggml.mul_mat.weight_format", + common_to_config_value(common_mul_mat_format_config_value(pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : + match.weight_format))); } else { dispatch.kernel.compile_parameters.emplace( "ggml.mul_mat_q8_1_x4.weight_format", @@ -1316,15 +1386,21 @@ static bool build_mul_mat_dispatch( std::to_string(unary_kind_config_value(match.output_unary_op))); const bool has_few_tokens = match.token_count <= 5; const bool has_skinny_reduction = common_mul_mat_has_skinny_reduction_shape(match.input_size); + const bool has_aligned_input = match.input_size % 256 == 0; const bool has_aligned_output = match.output_size % 64 == 0; const bool has_direct_weight = match.weight->alias_source.value < 0; - const bool can_transform_weight = dispatch_weight_layout_eligible(*match.weight); - const bool pack_rows = has_few_tokens && has_skinny_reduction && has_aligned_output && - has_direct_weight && can_transform_weight; + const bool can_transform_weight = dispatch_weight_layout_eligible_for_graph( + context.graph, *match.weight, context.row64_weight_policy); + const bool use_native_packed_glu = native_packed_glu_weight(context, *match.weight); + const bool parent_pack_rows = has_few_tokens && has_skinny_reduction && has_aligned_output && + has_direct_weight && can_transform_weight && !use_native_packed_glu; + const bool target_pack_rows = context.row64_weight_policy && has_aligned_input && has_aligned_output && + has_direct_weight && can_transform_weight && !use_native_packed_glu; + const bool pack_rows = parent_pack_rows || target_pack_rows; const bool can_pack_q4 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q4K; const bool can_pack_q6 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q6K; - bool pack_q4 = false; - bool pack_q6 = false; + bool pack_q4 = target_pack_rows && can_pack_q4; + bool pack_q6 = target_pack_rows && can_pack_q6; DispatchBinding activation = { match.input->id, 0, match.input->byte_count }; const size_t f16_activation_bytes = static_cast(match.token_count * match.input_size) * sizeof(ggml_fp16_t); const CommandPlanAlternateValue * f16_activation = @@ -1334,14 +1410,14 @@ static bool build_mul_mat_dispatch( if (f16_activation != nullptr) { activation = { f16_activation->alternate_value, 0, f16_activation_bytes }; dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_F16)); - } else if (can_pack_q4 || can_pack_q6) { + } else if ((can_pack_q4 || can_pack_q6) && has_few_tokens) { if (common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, dispatch_match, activation, q8_policy)) { pack_q4 = can_pack_q4; pack_q6 = can_pack_q6; dispatch.kernel.compile_parameters.emplace("ggml.mul_mat.activation_format", std::to_string(GGML_TYPE_Q8_1)); - } else if (q8_policy == CommonQ8ActivationPolicy::AllowStandaloneQuantize) { + } else if (q8_policy == CommonQ8ActivationPolicy::AllowStandaloneQuantize && !target_pack_rows) { return false; } } @@ -1425,9 +1501,11 @@ static bool prepare_vector_activation_binding(const DispatchMatchContext & c const bool has_aligned_input = input_size % 256 == 0; const bool has_direct_weight = weight.alias_source.value < 0; const bool has_aligned_output = output_size % 64 == 0; - const bool can_transform_weight = dispatch_weight_layout_eligible(weight); + const bool can_transform_weight = dispatch_weight_layout_eligible_for_graph( + context.graph, weight, context.row64_weight_policy); const bool can_prepare_packed_vector = has_single_token && has_aligned_input && has_direct_weight && - has_aligned_output && can_transform_weight; + has_aligned_output && can_transform_weight && + !native_packed_glu_weight(context, weight); if (!can_prepare_packed_vector) { return true; } @@ -1438,6 +1516,12 @@ static bool prepare_vector_activation_binding(const DispatchMatchContext & c return true; } + // A row64 policy must choose the retained source before checking for an optional Q8 activation. + if (context.row64_weight_policy) { + kernel_weight_format = can_pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : CommonMulMatWeightFormat::Q6KRow64; + weight_layout = can_pack_q4 ? kQ4KPackedK256Row64Layout : kQ6KPackedK256Row64ScaleRowLayout; + } + DispatchBinding q8_activation = activation; if (!common_prepare_q8_1_x4_input(context, input, input_size, token_count, dispatch_match, q8_activation, CommonQ8ActivationPolicy::ExistingAlternateOnly)) { @@ -1445,8 +1529,10 @@ static bool prepare_vector_activation_binding(const DispatchMatchContext & c } activation = q8_activation; - kernel_weight_format = can_pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : CommonMulMatWeightFormat::Q6KRow64; - weight_layout = can_pack_q4 ? kQ4KPackedK256Row64Layout : kQ6KPackedK256Row64ScaleRowLayout; + if (!context.row64_weight_policy) { + kernel_weight_format = can_pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : CommonMulMatWeightFormat::Q6KRow64; + weight_layout = can_pack_q4 ? kQ4KPackedK256Row64Layout : kQ6KPackedK256Row64ScaleRowLayout; + } uses_q8_activation = true; return true; } @@ -1824,20 +1910,26 @@ static bool build_mul_mat_postops_dispatch(const DispatchMatchContext & context, common_to_config_value(match.output_size)); const bool has_few_tokens = match.token_count <= 5; const bool has_skinny_reduction = common_mul_mat_has_skinny_reduction_shape(match.input_size); + const bool has_aligned_input = match.input_size % 256 == 0; const bool has_aligned_output = match.output_size % 64 == 0; const bool has_direct_weight = match.weight->alias_source.value < 0; - const bool can_transform_weight = dispatch_weight_layout_eligible(*match.weight); - const bool pack_rows = has_few_tokens && has_skinny_reduction && has_aligned_output && - has_direct_weight && can_transform_weight; + const bool can_transform_weight = dispatch_weight_layout_eligible_for_graph( + context.graph, *match.weight, context.row64_weight_policy); + const bool use_native_packed_glu = native_packed_glu_weight(context, *match.weight); + const bool parent_pack_rows = has_few_tokens && has_skinny_reduction && has_aligned_output && + has_direct_weight && can_transform_weight && !use_native_packed_glu; + const bool target_pack_rows = context.row64_weight_policy && has_aligned_input && has_aligned_output && + has_direct_weight && can_transform_weight && !use_native_packed_glu; + const bool pack_rows = parent_pack_rows || target_pack_rows; const bool can_pack_q4 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q4K; const bool can_pack_q6 = pack_rows && match.weight_format == CommonMulMatWeightFormat::Q6K; - bool pack_q4 = false; - bool pack_q6 = false; + bool pack_q4 = target_pack_rows && can_pack_q4; + bool pack_q6 = target_pack_rows && can_pack_q6; DispatchBinding activation = { match.input->id, 0, match.input->byte_count }; if (match.requires_q8_activation && !can_pack_q4 && !can_pack_q6) { return false; } - if (can_pack_q4 || can_pack_q6) { + if ((can_pack_q4 || can_pack_q6) && has_few_tokens) { if (common_prepare_q8_1_x4_input(context, *match.input, match.input_size, match.token_count, dispatch_match, activation, CommonQ8ActivationPolicy::ExistingAlternateOnly)) { pack_q4 = can_pack_q4; @@ -1846,7 +1938,8 @@ static bool build_mul_mat_postops_dispatch(const DispatchMatchContext & context, std::to_string(GGML_TYPE_Q8_1)); } } - if (match.requires_q8_activation && !pack_q4 && !pack_q6) { + const bool has_q8_activation = activation.value != match.input->id; + if (match.requires_q8_activation && !has_q8_activation) { return false; } const CommonMulMatWeightFormat format = pack_q4 ? CommonMulMatWeightFormat::Q4KRow64 : diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp index 7555e200fa60..ef2551a2fda0 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.cpp @@ -6,11 +6,13 @@ #include "llm/dispatch-attention-qkv.h" #include "llm/dispatch-gated-delta-net.h" #include "llm/dispatch-ssm-conv.h" +#include "qwen/dispatch-llm-profiles.h" #include "qwen/dispatch-qwen.h" #include #include #include +#include #include namespace ggml::hrx { @@ -40,6 +42,103 @@ static bool qwen_dispatch_disabled_from_environment() { std::strcmp(value, "off") != 0 && std::strcmp(value, "OFF") != 0; } +static bool qwen30b_attention_needs_native_weight(const Graph & graph, const Value & value) { + if (qwen_dispatch_disabled_from_environment()) { + return false; + } + const ggml_tensor * tensor = value.tensor; + const bool has_tensor = tensor != nullptr; + const bool supported_type = value.type == GGML_TYPE_Q4_K || value.type == GGML_TYPE_Q6_K; + const bool matrix = value.ne[2] == 1 && value.ne[3] == 1; + const bool block_weight = has_tensor && std::strncmp(tensor->name, "blk.", 4) == 0; + const bool candidate = has_tensor && supported_type && matrix && block_weight; + if (!candidate) { + return false; + } + + // These four projections have native consumers in the measured Qwen3-30B + // prefill/decode graph. Packing them also materializes a second native copy. + const char * layer_start = tensor->name + 4; + const bool starts_with_digit = *layer_start >= '0' && *layer_start <= '9'; + if (!starts_with_digit) { + return false; + } + char * layer_end = nullptr; + const long layer = std::strtol(layer_start, &layer_end, 10); + const bool valid_layer = layer >= 0; + const bool valid_separator = layer_end != nullptr && *layer_end == '.'; + const bool named_layer = valid_layer && valid_separator; + if (!named_layer) { + return false; + } + + const char * role = layer_end + 1; + const int64_t hidden_size = kQwen30BMoeDispatchProfile.hidden_size; + const bool q4 = value.type == GGML_TYPE_Q4_K; + const bool query_role = std::strcmp(role, "attn_q.weight") == 0; + const bool key_role = std::strcmp(role, "attn_k.weight") == 0; + const bool value_role = std::strcmp(role, "attn_v.weight") == 0; + const bool output_role = std::strcmp(role, "attn_output.weight") == 0; + const bool query_shape = value.ne[0] == hidden_size && value.ne[1] == 4096; + const bool key_value_shape = value.ne[0] == hidden_size && value.ne[1] == 512; + const bool output_shape = value.ne[0] == 4096 && value.ne[1] == hidden_size; + const bool query = q4 && query_role && query_shape; + const bool key = q4 && key_role && key_value_shape; + const bool value_projection = value_role && key_value_shape; + const bool output = q4 && output_role && output_shape; + const bool attention_projection = query || key || value_projection || output; + if (!attention_projection) { + return false; + } + + // K/V dimensions alone also occur in other models. Require this layer's + // complete attention quartet and its Qwen30B MoE expert profile. + const std::string prefix(tensor->name, static_cast(layer_end - tensor->name + 1)); + bool has_query = false; + bool has_key = false; + bool has_value = false; + bool has_output = false; + bool has_experts = false; + for (const Value & peer : graph.values().values()) { + const ggml_tensor * peer_tensor = peer.tensor; + const bool has_peer_tensor = peer_tensor != nullptr; + const bool same_layer = has_peer_tensor && + std::strncmp(peer_tensor->name, prefix.c_str(), prefix.size()) == 0; + if (!same_layer) { + continue; + } + const char * peer_role = peer_tensor->name + prefix.size(); + const bool peer_matrix = peer.ne[2] == 1 && peer.ne[3] == 1; + const bool peer_q4 = peer.type == GGML_TYPE_Q4_K; + const bool peer_q6 = peer.type == GGML_TYPE_Q6_K; + const bool peer_hidden = peer.ne[0] == hidden_size; + const bool peer_kv = peer_matrix && peer_hidden && peer.ne[1] == 512; + const bool peer_query = peer_q4 && peer_matrix && peer_hidden && peer.ne[1] == 4096; + const bool peer_output = peer_q4 && peer_matrix && peer.ne[0] == 4096 && peer.ne[1] == hidden_size; + const bool peer_experts = peer_q4 && peer.ne[0] == hidden_size && + peer.ne[1] == kQwen30BMoeDispatchProfile.expert_hidden_size && + peer.ne[2] == kQwen30BMoeDispatchProfile.expert_count && peer.ne[3] == 1; + has_query = has_query || (peer_query && std::strcmp(peer_role, "attn_q.weight") == 0); + has_key = has_key || (peer_q4 && peer_kv && std::strcmp(peer_role, "attn_k.weight") == 0); + has_value = has_value || ((peer_q4 || peer_q6) && peer_kv && + std::strcmp(peer_role, "attn_v.weight") == 0); + has_output = has_output || (peer_output && std::strcmp(peer_role, "attn_output.weight") == 0); + has_experts = has_experts || (peer_experts && std::strcmp(peer_role, "ffn_gate_exps.weight") == 0); + } + return has_query && has_key && has_value && has_output && has_experts; +} + +static bool qwen17b_tied_token_embedding_needs_native_weight(const Value & value) { + const ggml_tensor * tensor = value.tensor; + const bool named_embedding = tensor != nullptr && std::strcmp(tensor->name, "token_embd.weight") == 0; + const bool q6_weight = value.type == GGML_TYPE_Q6_K; + const bool tied_source_shape = value.ne[0] == 2048 && value.ne[1] == kQwen3VocabularyCount && + value.ne[2] == 1 && value.ne[3] == 1; + // The measured Qwen3-1.7B GGUF omits output.weight: this exact source + // feeds both embedding lookup and the native output projection. + return named_embedding && q6_weight && tied_source_shape; +} + static bool dispatch_has_eligible_weight_layouts(const DispatchMatchContext & context, const DispatchMatch & match) { auto eligible = [&](const std::vector & dispatches) { for (const Dispatch & dispatch : dispatches) { @@ -49,7 +148,9 @@ static bool dispatch_has_eligible_weight_layouts(const DispatchMatchContext & co } const Value * value = context.graph.values().find(binding.value); const bool has_graph_value = value != nullptr; - const bool layout_is_eligible = !has_graph_value || dispatch_weight_layout_eligible(*value); + const bool layout_is_eligible = + !has_graph_value || + dispatch_weight_layout_eligible_for_graph(context.graph, *value, context.row64_weight_policy); const bool has_ineligible_weight = has_graph_value && !layout_is_eligible; if (has_ineligible_weight) { return false; @@ -79,6 +180,19 @@ bool dispatch_weight_layout_eligible(const Value & value) { return ggml_backend_hrx_weight_layout_eligible(value.tensor); } +bool dispatch_weight_layout_eligible_for_graph(const Graph & graph, const Value & value, bool row64_weight_policy) { + if (!row64_weight_policy) { + return dispatch_weight_layout_eligible(value); + } + const bool needs_native_embedding = qwen17b_tied_token_embedding_needs_native_weight(value); + const bool needs_native_attention = qwen30b_attention_needs_native_weight(graph, value); + return !needs_native_embedding && !needs_native_attention && dispatch_weight_layout_eligible(value); +} + +bool dispatch_weight_is_qwen30b_native_attention(const Graph & graph, const Value & value) { + return qwen30b_attention_needs_native_weight(graph, value); +} + bool DispatchRegistry::match(const DispatchMatchContext & context, DispatchMatch & match) const { return this->match(context, match, nullptr); } diff --git a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h index 2fe76c617983..d7f475075d13 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h +++ b/ggml/src/ggml-hrx/dispatch_registration/dispatch-registry.h @@ -16,6 +16,7 @@ class DispatchRegistry; struct DispatchTarget { std::string architecture; + bool row64_weight_policy = false; }; enum class DispatchMatchKind { @@ -46,10 +47,13 @@ struct DispatchMatchContext { const CommandPlan & plan; ValueId next_plan_value; const DispatchRegistry * registry = nullptr; + bool row64_weight_policy = false; }; // Non-native weight bindings require bytes that remain available after graph import. bool dispatch_weight_layout_eligible(const Value & value); +bool dispatch_weight_layout_eligible_for_graph(const Graph & graph, const Value & value, bool row64_weight_policy); +bool dispatch_weight_is_qwen30b_native_attention(const Graph & graph, const Value & value); enum DispatchActivationInputFormat : uint32_t { DispatchActivationInputNone = 0, diff --git a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp index 4e07e2933ad3..ff41195f677e 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-attention-qkv.cpp @@ -311,8 +311,9 @@ static bool can_fuse_q6_f16_prefill_v_cache_cell(const DispatchMatchContext & co nullptr; } -static bool can_use_packed_lowtoken_attention_matmul(const AttentionMatMulMatch & root) { - if (!dispatch_weight_layout_eligible(*root.weight)) { +static bool can_use_packed_lowtoken_attention_matmul(const DispatchMatchContext & context, const AttentionMatMulMatch & root) { + if (!dispatch_weight_layout_eligible_for_graph( + context.graph, *root.weight, context.row64_weight_policy)) { return false; } if (root.token_count < 1 || root.token_count > 5 || root.output_size % 64 != 0 || @@ -326,6 +327,21 @@ static bool can_use_packed_lowtoken_attention_matmul(const AttentionMatMulMatch return common_is_supported_dense_input_size(root.weight_format, root.input_size); } +static bool can_use_row64_attention(const DispatchMatchContext & context, + const AttentionMatMulMatch & root) { + const bool has_target_policy = context.row64_weight_policy; + const bool can_transform_weight = dispatch_weight_layout_eligible_for_graph( + context.graph, *root.weight, context.row64_weight_policy); + const bool has_direct_weight = root.weight->alias_source.value < 0; + const bool has_aligned_output = root.output_size % 64 == 0; + const bool has_q4_weight = root.weight_format == CommonMulMatWeightFormat::Q4K; + const bool has_q6_weight = root.weight_format == CommonMulMatWeightFormat::Q6K; + const bool has_supported_weight = has_q4_weight || has_q6_weight; + const bool has_supported_input = common_is_supported_dense_input_size(root.weight_format, root.input_size); + return has_target_policy && can_transform_weight && has_direct_weight && has_aligned_output && + has_supported_weight && has_supported_input; +} + static CommonMulMatWeightFormat packed_attention_lowtoken_format(CommonMulMatWeightFormat format) { return format == CommonMulMatWeightFormat::Q4K ? CommonMulMatWeightFormat::Q4KRow64 : CommonMulMatWeightFormat::Q6KRow64; @@ -534,9 +550,11 @@ static AttentionQkvMatch match_attention_qkv_projection(const DispatchMatchConte match.kernel = is_supported_decode_token_count(root.token_count) ? kAttentionKMatMulRopeSetRowsVectorF32F32Kernel : kAttentionKMatMulRopeSetRowsTiledF32F32Kernel; - if (can_use_packed_lowtoken_attention_matmul(root)) { + if (can_use_packed_lowtoken_attention_matmul(context, root) || can_use_row64_attention(context, root)) { match.root.weight_format = packed_attention_lowtoken_format(root.weight_format); - match.kernel = kAttentionKMatMulRopeSetRowsLegacyF32F32Kernel; + if (!context.row64_weight_policy) { + match.kernel = kAttentionKMatMulRopeSetRowsLegacyF32F32Kernel; + } } match.kind = AttentionQkvProjectionKind::Key; return match; @@ -545,9 +563,11 @@ static AttentionQkvMatch match_attention_qkv_projection(const DispatchMatchConte match.root = root; match.kernel = is_supported_decode_token_count(root.token_count) ? kAttentionQMatMulRopeVectorF32F32Kernel : kAttentionQMatMulRopeTiledF32F32Kernel; - if (can_use_packed_lowtoken_attention_matmul(root)) { + if (can_use_packed_lowtoken_attention_matmul(context, root) || can_use_row64_attention(context, root)) { match.root.weight_format = packed_attention_lowtoken_format(root.weight_format); - match.kernel = kAttentionQMatMulRopeLegacyF32F32Kernel; + if (!context.row64_weight_policy) { + match.kernel = kAttentionQMatMulRopeLegacyF32F32Kernel; + } } match.kind = AttentionQkvProjectionKind::Query; return match; @@ -567,9 +587,11 @@ static AttentionQkvMatch match_attention_qkv_projection(const DispatchMatchConte match.root = root; match.kernel = is_supported_decode_token_count(root.token_count) ? kAttentionVMatMulSetRowsVectorF32F32Kernel : kAttentionVMatMulSetRowsTiledF32F32Kernel; - if (can_use_packed_lowtoken_attention_matmul(root)) { + if (can_use_packed_lowtoken_attention_matmul(context, root) || can_use_row64_attention(context, root)) { match.root.weight_format = packed_attention_lowtoken_format(root.weight_format); - match.kernel = kAttentionVMatMulSetRowsLegacyF32F32Kernel; + if (!context.row64_weight_policy) { + match.kernel = kAttentionVMatMulSetRowsLegacyF32F32Kernel; + } } match.kind = AttentionQkvProjectionKind::Value; match.use_f16_activation = quantized_prefill; @@ -683,7 +705,20 @@ static bool match_attention_qkv_dispatch(const DispatchMatchContext & context, D dispatch_match, activation, CommonQ8ActivationPolicy::ExistingAlternateOnly)) { use_q8_activation = true; - } else { + if (context.row64_weight_policy) { + switch (match.kind) { + case AttentionQkvProjectionKind::Query: + match.kernel = kAttentionQMatMulRopeLegacyF32F32Kernel; + break; + case AttentionQkvProjectionKind::Key: + match.kernel = kAttentionKMatMulRopeSetRowsLegacyF32F32Kernel; + break; + case AttentionQkvProjectionKind::Value: + match.kernel = kAttentionVMatMulSetRowsLegacyF32F32Kernel; + break; + } + } + } else if (!context.row64_weight_policy) { CommonMulMatWeightFormat fallback_format = CommonMulMatWeightFormat::Q4K; if (!common_mul_mat_format_for_type(match.root.weight->type, fallback_format)) { return false; diff --git a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.cpp b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.cpp index a7e5a4ddb3eb..1fc967bde75b 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.cpp +++ b/ggml/src/ggml-hrx/dispatch_registration/llm/dispatch-ssm-conv.cpp @@ -807,11 +807,22 @@ static bool match_mul_mat_conv4_dispatch(const DispatchMatchContext & context, D dispatch.kernel.integer_parameters.emplace("token_count", mat.token_count); set_compile_parameter(dispatch.kernel, "ggml.mul_mat.input_size", mat.input_size); set_compile_parameter(dispatch.kernel, "ggml.mul_mat.output_size", mat.output_size); - set_compile_parameter(dispatch.kernel, "ggml.mul_mat.weight_format", mat.weight->type == GGML_TYPE_Q4_K ? 4 : 6); + const bool q4_weight = mat.weight->type == GGML_TYPE_Q4_K; + const bool q6_weight = mat.weight->type == GGML_TYPE_Q6_K; + const bool q6_row64 = context.row64_weight_policy && q6_weight && + dispatch_weight_layout_eligible_for_graph( + context.graph, *mat.weight, context.row64_weight_policy); + // Q4 conv4 binds row64 even when the model policy is off. Format 44 keeps + // that reader distinct from the native Q4 F16 prefill reader (format 4). + set_compile_parameter(dispatch.kernel, "ggml.mul_mat.weight_format", q4_weight ? 44 : (q6_row64 ? 46 : 6)); dispatch.bindings.push_back(activation); - if (mat.weight->type == GGML_TYPE_Q4_K) { + if (q4_weight) { dispatch.bindings.push_back({mat.weight->id, 0, mat.weight->byte_count, kQ4KPackedK256Row64Layout, mat.weight->type, mat.input_size, mat.output_size, mat.weight->byte_count}); + } else if (q6_row64) { + dispatch.bindings.push_back({mat.weight->id, 0, mat.weight->byte_count, + kQ6KPackedK256Row64ScaleRowLayout, mat.weight->type, mat.input_size, + mat.output_size, mat.weight->byte_count}); } else { dispatch.bindings.push_back({mat.weight->id, 0, mat.weight->byte_count}); } diff --git a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-llm-profiles.h b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-llm-profiles.h index ef9a06f9f8b0..84338e53e0e3 100644 --- a/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-llm-profiles.h +++ b/ggml/src/ggml-hrx/dispatch_registration/qwen/dispatch-llm-profiles.h @@ -20,6 +20,7 @@ constexpr LlmMoeDispatchProfile kLlmMoeQwen30BDispatchProfile = { static constexpr const LlmMoeDispatchProfile & kActiveLlmMoeDispatchProfile = kLlmMoeQwen30BDispatchProfile; static constexpr const LlmMoeDispatchProfile & kQwen30BMoeDispatchProfile = kLlmMoeQwen30BDispatchProfile; +static constexpr int64_t kQwen3VocabularyCount = 151936; constexpr bool is_llm_supported_query_length(const LlmMoeDispatchProfile & profile, int64_t query_length) { return query_length >= 1 && query_length <= profile.max_token_count; diff --git a/ggml/src/ggml-hrx/ggml-hrx.cpp b/ggml/src/ggml-hrx/ggml-hrx.cpp index 6be3802203f0..8073129c6ab1 100644 --- a/ggml/src/ggml-hrx/ggml-hrx.cpp +++ b/ggml/src/ggml-hrx/ggml-hrx.cpp @@ -1223,6 +1223,14 @@ static size_t registry_device_count(ggml_backend_reg_t registry) { return static_cast(registry->context)->devices.size(); } +static bool backend_set_row64_weight_policy(ggml_backend_t backend, bool enabled) { + if (!ggml_backend_is_hrx(backend)) { + return false; + } + auto * context = static_cast(backend->context); + return context->set_row64_weight_policy(enabled); +} + static ggml_backend_dev_t registry_device(ggml_backend_reg_t registry, size_t index) { auto * context = static_cast(registry->context); GGML_ASSERT(index < context->devices.size()); @@ -1231,7 +1239,9 @@ static ggml_backend_dev_t registry_device(ggml_backend_reg_t registry, size_t in static void * registry_proc(ggml_backend_reg_t registry, const char * name) { GGML_UNUSED(registry); - GGML_UNUSED(name); + if (std::strcmp(name, "ggml_backend_hrx_set_row64_weight_policy") == 0) { + return reinterpret_cast(backend_set_row64_weight_policy); + } return nullptr; } diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom index d4964444b253..2e87b871e5c3 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/dequant.loom @@ -30,7 +30,7 @@ func.def inline @ggml_paired_weight_words(%half_packet: i1, %paired: i1, %load_u func.return %words : vector<4xi32> } -func.def public inline @ggml_q4k_native_row64_f16_pair(%half_packet: i1, %paired: i1, %weight: buffer, %peer: buffer, %is_up: i1, %row_base: offset, %group0: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) { +func.def public inline @ggml_q4k_f16_pair(%native_weight: i1, %half_packet: i1, %paired: i1, %weight: buffer, %peer: buffer, %is_up: i1, %row_base: offset, %group0: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) { %c1 = index.constant 1 : index %c2 = index.constant 2 : index %c4 = index.constant 4 : index @@ -58,7 +58,16 @@ func.def public inline @ggml_q4k_native_row64_f16_pair(%half_packet: i1, %paired %word_bytes = index.constant 4 : offset %packet_word = index.rem %packet, %c4 : index %packet_add = index.scale %packet_word, %word_bytes : index, offset -> offset - %address = index.add %field_base, %packet_add : offset + %row64_address = index.add %field_base, %packet_add : offset + // Native Q4_K keeps the same header and stores each group pair in 32 code bytes. + %native_code_bytes = index.constant 16 : offset + %native_pair_bytes = index.constant 32 : offset + %native_code_base = index.add %row_base, %native_code_bytes : offset + %native_pair_add = index.scale %pair, %native_pair_bytes : index, offset -> offset + %native_packet_add = index.scale %packet, %word_bytes : index, offset -> offset + %native_pair_base = index.add %native_code_base, %native_pair_add : offset + %native_address = index.add %native_pair_base, %native_packet_add : offset + %address = scf.select %native_weight, %native_address, %row64_address : offset %codes = func.call @ggml_paired_weight_words(%half_packet, %paired, %is_up, %weight, %peer, %address) : (i1, i1, i1, buffer, buffer, offset) -> (vector<4xi32>) %mask = vector.constant 252645135 : vector<4xi32> %shift = vector.constant 4 : vector<4xi32> @@ -662,6 +671,120 @@ func.def inline @ggml_q4k_native_row64_f32_vector4(%weight: buffer, %row: index, func.return %result : vector<4xf32> } +// Keep the K256 header in registers across the eight K32 contractions. +func.def inline @ggml_q4k_row64_load_decode_state(%weight: buffer, %row: index, %input_size: index, %block: index) -> (vector<2xf16>, vector<3xi32>) { + %c0 = index.constant 0 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %record_bytes = index.constant 9216 : offset + %row_bytes = index.constant 16 : offset + %scale_offset = index.constant 4 : offset + %blocks = index.div %input_size, %c256 : index + %row_group = index.div %row, %c64 : index + %row_lane = index.rem %row, %c64 : index + %tile = index.madd %row_group, %blocks, %block : index + %tile_base = index.scale %tile, %record_bytes : index, offset -> offset + %row_add = index.scale %row_lane, %row_bytes : index, offset -> offset + %header_base = index.add %tile_base, %row_add : offset + %scale_base = index.add %header_base, %scale_offset : offset + %dm_view = buffer.view %weight[%header_base] : buffer -> view<2xf16> + %scale_view = buffer.view %weight[%scale_base] : buffer -> view<3xi32> + %dm = vector.load %dm_view[%c0] : view<2xf16> -> vector<2xf16> + %scales = vector.load %scale_view[%c0] : view<3xi32> -> vector<3xi32> + func.return %dm, %scales : vector<2xf16>, vector<3xi32> +} + +func.def inline @ggml_q4k_row64_preloaded_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q4_block: index, %q4_group: index, %packet: index, %dm: vector<2xf16>, %scales: vector<3xi32>) -> (vector<4xf32>) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c8 = index.constant 8 : index + %c2_i32 = scalar.constant 2 : i32 + %c4_i32 = scalar.constant 4 : i32 + %c15 = vector.constant 15 : vector<1xi32> + %c48 = vector.constant 48 : vector<1xi32> + %q4_mask = vector.constant 252645135 : vector<1xi32> + %bounded_group = index.assume %q4_group [range(%q4_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %record_bytes = index.constant 9216 : offset + %payload_offset = index.constant 1024 : offset + %blocks = index.div %input_size, %c256 : index + %row_group = index.div %row, %c64 : index + %row_lane = index.rem %row, %c64 : index + %group_block = index.madd %row_group, %blocks, %q4_block : index + %group_base = index.scale %group_block, %record_bytes : index, offset -> offset + %code_byte_base = index.add %group_base, %payload_offset : offset + %code_view = buffer.view %weight[%code_byte_base] : buffer -> view<8x64x4xi32> + %d_f16 = vector.extract %dm[0] : vector<2xf16> -> f16 + %dmin_f16 = vector.extract %dm[1] : vector<2xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %dmin = scalar.extf %dmin_f16 : f16 to f32 + %q_page0 = index.div %bounded_group, %c2 : index + %q_page = index.mul %q_page0, %c8 : index + %q_word_index0 = index.add %q_page, %bounded_packet : index + %q_word_index = index.assume %q_word_index0 [range(%q_word_index0, 0, 31)] : index + %is_low = index.cmp ult, %bounded_group, %c4 : index + %scale_lane = index.rem %bounded_group, %c4 : index + %scale_shift_index = index.mul %scale_lane, %c8 : index + %scale_shift_i32 = index.cast %scale_shift_index : index to i32 + %scale_shift = vector.splat %scale_shift_i32 : vector<1xi32> + %scale0_i32 = vector.extract %scales[0] : vector<3xi32> -> i32 + %scale1_i32 = vector.extract %scales[1] : vector<3xi32> -> i32 + %scale2_i32 = vector.extract %scales[2] : vector<3xi32> -> i32 + %scale0 = vector.splat %scale0_i32 : vector<1xi32> + %scale1 = vector.splat %scale1_i32 : vector<1xi32> + %scale2 = vector.splat %scale2_i32 : vector<1xi32> + %high_shift_i32 = scalar.addi %scale_shift_i32, %c2_i32 : i32 + %minimum_shift_i32 = scalar.addi %scale_shift_i32, %c4_i32 : i32 + %selected_scale_source = scf.select %is_low, %scale0, %scale2 : vector<1xi32> + %selected_minimum_source = scf.select %is_low, %scale1, %scale2 : vector<1xi32> + %selected_scale_high_shift_i32 = scf.select %is_low, %scale_shift_i32, %high_shift_i32 : i32 + %selected_minimum_low_shift_i32 = scf.select %is_low, %scale_shift_i32, %minimum_shift_i32 : i32 + %selected_scale_high_shift = vector.splat %selected_scale_high_shift_i32 : vector<1xi32> + %selected_minimum_low_shift = vector.splat %selected_minimum_low_shift_i32 : vector<1xi32> + %scale_low0 = vector.shrui %selected_scale_source, %scale_shift : vector<1xi32> + %scale_low = vector.andi %scale_low0, %c15 : vector<1xi32> + %scale_high0 = vector.shrui %scale0, %selected_scale_high_shift : vector<1xi32> + %scale_high = vector.andi %scale_high0, %c48 : vector<1xi32> + %scale = vector.ori %scale_low, %scale_high : vector<1xi32> + %minimum_low0 = vector.shrui %selected_minimum_source, %selected_minimum_low_shift : vector<1xi32> + %minimum_low = vector.andi %minimum_low0, %c15 : vector<1xi32> + %minimum_high0 = vector.shrui %scale1, %selected_scale_high_shift : vector<1xi32> + %minimum_high = vector.andi %minimum_high0, %c48 : vector<1xi32> + %minimum = vector.ori %minimum_low, %minimum_high : vector<1xi32> + %scale_f32 = vector.uitofp %scale : vector<1xi32> to vector<1xf32> + %minimum_f32 = vector.uitofp %minimum : vector<1xi32> to vector<1xf32> + %d_vector1 = vector.splat %d : vector<1xf32> + %dmin_vector1 = vector.splat %dmin : vector<1xf32> + %d_scale_vector1 = vector.mulf %d_vector1, %scale_f32 : vector<1xf32> + %minimum_scale_vector1 = vector.mulf %dmin_vector1, %minimum_f32 : vector<1xf32> + %d_scale = vector.extract %d_scale_vector1[0] : vector<1xf32> -> f32 + %minimum_scale = vector.extract %minimum_scale_vector1[0] : vector<1xf32> -> f32 + %payload_field = index.div %q_word_index, %c4 : index + %payload_word = index.rem %q_word_index, %c4 : index + %q_word = vector.load %code_view[%payload_field, %row_lane, %payload_word] : view<8x64x4xi32> -> vector<1xi32> + %q_half = index.rem %bounded_group, %c2 : index + %q_shift_index = index.mul %q_half, %c4 : index + %q_shift_i32 = index.cast %q_shift_index : index to i32 + %q_shift = vector.splat %q_shift_i32 : vector<1xi32> + %shifted_q = vector.shrui %q_word, %q_shift : vector<1xi32> + %masked_q = vector.andi %shifted_q, %q4_mask : vector<1xi32> + %q_i8 = vector.bitcast %masked_q : vector<1xi32> to vector<4xi8> + %q_f32 = vector.uitofp %q_i8 : vector<4xi8> to vector<4xf32> + %negative_minimum_scale = scalar.negf %minimum_scale : f32 + %q0 = vector.extract %q_f32[0] : vector<4xf32> -> f32 + %q1 = vector.extract %q_f32[1] : vector<4xf32> -> f32 + %q2 = vector.extract %q_f32[2] : vector<4xf32> -> f32 + %q3 = vector.extract %q_f32[3] : vector<4xf32> -> f32 + %value0 = scalar.fmaf %q0, %d_scale, %negative_minimum_scale : f32 + %value1 = scalar.fmaf %q1, %d_scale, %negative_minimum_scale : f32 + %value2 = scalar.fmaf %q2, %d_scale, %negative_minimum_scale : f32 + %value3 = scalar.fmaf %q3, %d_scale, %negative_minimum_scale : f32 + %result = vector.from_elements %value0, %value1, %value2, %value3 : vector<4xf32> + func.return %result : vector<4xf32> +} + // Q5_K uses the Q4_K affine header and a 32-byte fifth-bit plane. func.def inline @ggml_q5k_f32_vector4(%weight: buffer, %row_byte_base: offset, %q5_block: index, %q5_group: index, %packet: index) -> (vector<4xf32>) { %c0 = index.constant 0 : index @@ -2584,6 +2707,62 @@ func.def inline @ggml_q6k_f32_vector4(%weight: buffer, %weight_row_byte_base: of func.return %values : vector<4xf32> } +func.def inline @ggml_q6k_native_row64_f32_vector4(%weight: buffer, %row: index, %input_size: index, %q6_block: index, %q6_group: index, %packet: index) -> (vector<4xf32>) { + %c2 = index.constant 2 : index + %c4 = index.constant 4 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %tile_bytes = index.constant 13440 : offset + %scale_byte_add = index.constant 128 : offset + %raw_byte_add = index.constant 1152 : offset + %c4_i32v = vector.constant 4 : vector<1xi32> + %nibble_mask = vector.constant 252645135 : vector<1xi32> + %high_mask = vector.constant 50529027 : vector<1xi32> + %c32_f32v = vector.constant 32.0 : vector<4xf32> + %bounded_group = index.assume %q6_group [range(%q6_group, 0, 7)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 7)] : index + %blocks = index.div %input_size, %c256 : index + %row_group = index.div %row, %c64 : index + %row_lane = index.rem %row, %c64 : index + %tile = index.madd %row_group, %blocks, %q6_block : index + %tile_base = index.scale %tile, %tile_bytes : index, offset -> offset + %scale_byte_base = index.add %tile_base, %scale_byte_add : offset + %raw_byte_base = index.add %tile_base, %raw_byte_add : offset + %raw_view = buffer.view %weight[%raw_byte_base] : buffer -> view<2x3x64x8xi32> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<64x2x8xi8> + %d_view = buffer.view %weight[%tile_base] : buffer -> view<64xf16> + %group_in_half = index.rem %bounded_group, %c4 : index + %half = index.div %bounded_group, %c4 : index + %ql_side = index.rem %group_in_half, %c2 : index + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift_index = index.mul %nibble, %c4 : index + %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 + %nibble_shift = vector.splat %nibble_shift_i32 : vector<1xi32> + %qh_shift_index = index.mul %group_in_half, %c2 : index + %qh_shift_i32 = index.cast %qh_shift_index : index to i32 + %qh_shift = vector.splat %qh_shift_i32 : vector<1xi32> + %scale_packet_half = index.div %bounded_packet, %c4 : index + %ql_word = vector.load %raw_view[%half, %ql_side, %row_lane, %bounded_packet] : view<2x3x64x8xi32> -> vector<1xi32> + %qh_word = vector.load %raw_view[%half, %c2, %row_lane, %bounded_packet] : view<2x3x64x8xi32> -> vector<1xi32> + %ql_shifted = vector.shrui %ql_word, %nibble_shift : vector<1xi32> + %ql = vector.andi %ql_shifted, %nibble_mask : vector<1xi32> + %qh_shifted = vector.shrui %qh_word, %qh_shift : vector<1xi32> + %qh_low = vector.andi %qh_shifted, %high_mask : vector<1xi32> + %qh = vector.shli %qh_low, %c4_i32v : vector<1xi32> + %code = vector.ori %ql, %qh : vector<1xi32> + %code_i8 = vector.bitcast %code : vector<1xi32> to vector<4xi8> + %code_f32 = vector.uitofp %code_i8 : vector<4xi8> to vector<4xf32> + %centered = vector.subf %code_f32, %c32_f32v : vector<4xf32> + %scale_i8 = view.load %scale_view[%row_lane, %scale_packet_half, %bounded_group] : view<64x2x8xi8> -> i8 + %d_f16 = view.load %d_view[%row_lane] : view<64xf16> -> f16 + %scale = scalar.sitofp %scale_i8 : i8 to f32 + %d = scalar.extf %d_f16 : f16 to f32 + %combined_scale = scalar.mulf %scale, %d : f32 + %combined_scale_vector = vector.splat %combined_scale : vector<4xf32> + %values = vector.mulf %centered, %combined_scale_vector : vector<4xf32> + func.return %values : vector<4xf32> +} + func.def inline @ggml_q8_0_f32_vector4(%weight: buffer, %row_byte_base: offset, %q8_block: index, %packet: index) -> (vector<4xf32>) { %c0 = index.constant 0 : index %c4 = index.constant 4 : index @@ -2656,6 +2835,8 @@ func.def inline @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offse %c4 = index.constant 4 : index %c5 = index.constant 5 : index %c6 = index.constant 6 : index + %q4_row64_format = index.constant 44 : index + %q6_row64_format = index.constant 46 : index %c18 = index.constant 18 : index %c20 = index.constant 20 : index %c21 = index.constant 21 : index @@ -2699,9 +2880,13 @@ func.def inline @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offse %f32_tile_bytes = index.constant 1024 : offset %is_q1_0 = index.cmp eq, %weight_format, %q1_0_format : index %is_q3 = index.cmp eq, %weight_format, %c3 : index - %is_q4 = index.cmp eq, %weight_format, %c4 : index + %is_q4_native = index.cmp eq, %weight_format, %c4 : index + %is_q4_row64 = index.cmp eq, %weight_format, %q4_row64_format : index + %is_q4 = scalar.ori %is_q4_native, %is_q4_row64 : i1 %is_q5 = index.cmp eq, %weight_format, %c5 : index - %is_q6 = index.cmp eq, %weight_format, %c6 : index + %is_q6_native = index.cmp eq, %weight_format, %c6 : index + %is_q6_row64 = index.cmp eq, %weight_format, %q6_row64_format : index + %is_q6 = scalar.ori %is_q6_native, %is_q6_row64 : i1 %is_q4_0 = index.cmp eq, %weight_format, %q4_0_format : index %is_q4_1 = index.cmp eq, %weight_format, %q4_1_format : index %is_q5_0 = index.cmp eq, %weight_format, %q5_0_format : index @@ -2746,6 +2931,7 @@ func.def inline @ggml_dequant_weight_tile_bytes(%weight_format: index) -> (offse func.return %selected : offset } +// Row64 consumers pass logical native row offsets to the shared dequant API. func.def public inline @ggml_dequant_weight_row_bytes(%weight_format: index, %hidden_size: index) -> (offset) { %q1_0_format = index.constant 10 : index %iq1_s_format = index.constant 19 : index @@ -2753,6 +2939,8 @@ func.def public inline @ggml_dequant_weight_row_bytes(%weight_format: index, %hi %c4 = index.constant 4 : index %c5 = index.constant 5 : index %c6 = index.constant 6 : index + %q4_row64_format = index.constant 44 : index + %q6_row64_format = index.constant 46 : index %c18 = index.constant 18 : index %c20 = index.constant 20 : index %c21 = index.constant 21 : index @@ -2785,9 +2973,13 @@ func.def public inline @ggml_dequant_weight_row_bytes(%weight_format: index, %hi %f16_element_bytes = index.constant 2 : offset %f32_element_bytes = index.constant 4 : offset %uses_k_tile0 = index.cmp eq, %weight_format, %c3 : index - %uses_k_tile1 = index.cmp eq, %weight_format, %c4 : index + %is_q4_native = index.cmp eq, %weight_format, %c4 : index + %is_q4_row64 = index.cmp eq, %weight_format, %q4_row64_format : index + %uses_k_tile1 = scalar.ori %is_q4_native, %is_q4_row64 : i1 %uses_k_tile2 = index.cmp eq, %weight_format, %c5 : index - %uses_k_tile3 = index.cmp eq, %weight_format, %c6 : index + %is_q6_native = index.cmp eq, %weight_format, %c6 : index + %is_q6_row64 = index.cmp eq, %weight_format, %q6_row64_format : index + %uses_k_tile3 = scalar.ori %is_q6_native, %is_q6_row64 : i1 %uses_k_tile4 = index.cmp eq, %weight_format, %c20 : index %uses_k_tile5 = index.cmp eq, %weight_format, %c21 : index %uses_k_tile6 = index.cmp eq, %weight_format, %c23 : index @@ -2856,6 +3048,7 @@ func.def public inline @ggml_dequant_weight_row_bytes(%weight_format: index, %hi // Format ids: // 10: Q1_0 // 11/4/5/6: Q3_K/Q4_K/Q5_K/Q6_K +// 44/46: lossless Q4_K/Q6_K row64 // 40/41/50/51: Q4_0/Q4_1/Q5_0/Q5_1 // 22/20/21/23/24/25: IQ2_S/IQ4_NL/IQ3_S/IQ4_XS/IQ2_XXS/IQ2_XS // 80/81: Q8_0/Q8_1 @@ -2866,6 +3059,8 @@ func.def inline @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_gri %c4 = index.constant 4 : index %c5 = index.constant 5 : index %c6 = index.constant 6 : index + %q4_row64_format = index.constant 44 : index + %q6_row64_format = index.constant 46 : index %c8 = index.constant 8 : index %c20 = index.constant 20 : index %c21 = index.constant 21 : index @@ -2884,9 +3079,13 @@ func.def inline @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_gri // Decode the compact integer format id into one predicate per layout. %is_q1_0 = index.cmp eq, %weight_format, %q1_0_format : index %is_q3 = index.cmp eq, %weight_format, %c3 : index - %is_q4 = index.cmp eq, %weight_format, %c4 : index + %is_q4_native = index.cmp eq, %weight_format, %c4 : index + %is_q4_row64 = index.cmp eq, %weight_format, %q4_row64_format : index + %is_q4 = scalar.ori %is_q4_native, %is_q4_row64 : i1 %is_q5 = index.cmp eq, %weight_format, %c5 : index - %is_q6 = index.cmp eq, %weight_format, %c6 : index + %is_q6_native = index.cmp eq, %weight_format, %c6 : index + %is_q6_row64 = index.cmp eq, %weight_format, %q6_row64_format : index + %is_q6 = scalar.ori %is_q6_native, %is_q6_row64 : i1 %is_q4_0 = index.cmp eq, %weight_format, %q4_0_format : index %is_q4_1 = index.cmp eq, %weight_format, %q4_1_format : index %is_q5_0 = index.cmp eq, %weight_format, %q5_0_format : index @@ -2919,13 +3118,37 @@ func.def inline @ggml_dequant_f16_vector4(%iq4nl_table: vector<16xi8>, %iq3s_gri scf.yield %zero : vector<4xf16> } %q4_values = scf.if %is_q4 -> (vector<4xf16>) { - %values = func.call @ggml_q4k_f16_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) + %values = scf.if %is_q4_row64 -> (vector<4xf16>) { + %row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %input_size) : (index, index) -> (offset) + %row_byte_index = index.cast %row_byte_base : offset to i64 + %row_stride = index.cast %row_bytes : offset to i64 + %row_i64 = scalar.divui %row_byte_index, %row_stride : i64 + %row = index.cast %row_i64 : i64 to index + %wide = func.call @ggml_q4k_native_row64_f32_vector4(%weight, %row, %input_size, %quant_block, %quant_group, %packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + %values = vector.fptrunc %wide : vector<4xf32> to vector<4xf16> + scf.yield %values : vector<4xf16> + } else { + %values = func.call @ggml_q4k_f16_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } scf.yield %values : vector<4xf16> } else { scf.yield %zero : vector<4xf16> } %q6_values = scf.if %is_q6 -> (vector<4xf16>) { - %values = func.call @ggml_q6k_f16_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) + %values = scf.if %is_q6_row64 -> (vector<4xf16>) { + %row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %input_size) : (index, index) -> (offset) + %row_byte_index = index.cast %row_byte_base : offset to i64 + %row_stride = index.cast %row_bytes : offset to i64 + %row_i64 = scalar.divui %row_byte_index, %row_stride : i64 + %row = index.cast %row_i64 : i64 to index + %wide = func.call @ggml_q6k_native_row64_f32_vector4(%weight, %row, %input_size, %quant_block, %quant_group, %packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + %values = vector.fptrunc %wide : vector<4xf32> to vector<4xf16> + scf.yield %values : vector<4xf16> + } else { + %values = func.call @ggml_q6k_f16_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf16>) + scf.yield %values : vector<4xf16> + } scf.yield %values : vector<4xf16> } else { scf.yield %zero : vector<4xf16> @@ -3047,6 +3270,8 @@ func.def inline @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_gri %c4 = index.constant 4 : index %c5 = index.constant 5 : index %c6 = index.constant 6 : index + %q4_row64_format = index.constant 44 : index + %q6_row64_format = index.constant 46 : index %c8 = index.constant 8 : index %c20 = index.constant 20 : index %c21 = index.constant 21 : index @@ -3064,9 +3289,13 @@ func.def inline @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_gri %zero = vector.constant 0.0 : vector<4xf32> %is_q1_0 = index.cmp eq, %weight_format, %q1_0_format : index %is_q3 = index.cmp eq, %weight_format, %c3 : index - %is_q4 = index.cmp eq, %weight_format, %c4 : index + %is_q4_native = index.cmp eq, %weight_format, %c4 : index + %is_q4_row64 = index.cmp eq, %weight_format, %q4_row64_format : index + %is_q4 = scalar.ori %is_q4_native, %is_q4_row64 : i1 %is_q5 = index.cmp eq, %weight_format, %c5 : index - %is_q6 = index.cmp eq, %weight_format, %c6 : index + %is_q6_native = index.cmp eq, %weight_format, %c6 : index + %is_q6_row64 = index.cmp eq, %weight_format, %q6_row64_format : index + %is_q6 = scalar.ori %is_q6_native, %is_q6_row64 : i1 %is_q4_0 = index.cmp eq, %weight_format, %q4_0_format : index %is_q4_1 = index.cmp eq, %weight_format, %q4_1_format : index %is_q5_0 = index.cmp eq, %weight_format, %q5_0_format : index @@ -3095,13 +3324,35 @@ func.def inline @ggml_dequant_f32_vector4(%iq4nl_table: vector<16xi8>, %iq3s_gri scf.yield %zero : vector<4xf32> } %q4_values = scf.if %is_q4 -> (vector<4xf32>) { - %values = func.call @ggml_q4k_f32_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf32>) + %values = scf.if %is_q4_row64 -> (vector<4xf32>) { + %row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %input_size) : (index, index) -> (offset) + %row_byte_index = index.cast %row_byte_base : offset to i64 + %row_stride = index.cast %row_bytes : offset to i64 + %row_i64 = scalar.divui %row_byte_index, %row_stride : i64 + %row = index.cast %row_i64 : i64 to index + %values = func.call @ggml_q4k_native_row64_f32_vector4(%weight, %row, %input_size, %quant_block, %quant_group, %packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + %values = func.call @ggml_q4k_f32_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } scf.yield %values : vector<4xf32> } else { scf.yield %zero : vector<4xf32> } %q6_values = scf.if %is_q6 -> (vector<4xf32>) { - %values = func.call @ggml_q6k_f32_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf32>) + %values = scf.if %is_q6_row64 -> (vector<4xf32>) { + %row_bytes = func.call @ggml_dequant_weight_row_bytes(%weight_format, %input_size) : (index, index) -> (offset) + %row_byte_index = index.cast %row_byte_base : offset to i64 + %row_stride = index.cast %row_bytes : offset to i64 + %row_i64 = scalar.divui %row_byte_index, %row_stride : i64 + %row = index.cast %row_i64 : i64 to index + %values = func.call @ggml_q6k_native_row64_f32_vector4(%weight, %row, %input_size, %quant_block, %quant_group, %packet) : (buffer, index, index, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } else { + %values = func.call @ggml_q6k_f32_vector4(%weight, %row_byte_base, %quant_block, %quant_group, %packet) : (buffer, offset, index, index, index) -> (vector<4xf32>) + scf.yield %values : vector<4xf32> + } scf.yield %values : vector<4xf32> } else { scf.yield %zero : vector<4xf32> diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom index 4b4eb1c92e95..07a45aa51ec4 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_input_f32_dequant_accumulate_f32.loom @@ -1,3 +1,6 @@ +func.decl @ggml_q4k_row64_load_decode_state(%weight: buffer, %row: index, %input_size: index, %block: index) -> (vector<2xf16>, vector<3xi32>) +func.decl @ggml_q4k_row64_preloaded_f32_vector4(%weight: buffer, %row: index, %input_size: index, %block: index, %group: index, %packet: index, %dm: vector<2xf16>, %scales: vector<3xi32>) -> (vector<4xf32>) + func.decl @ggml_q5k_load_decode_state(%weight: buffer, %row_base: offset, %block: index, %packet: index) -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) func.decl @ggml_q5k_preloaded_f32_vector4(%weight: buffer, %row_byte_base: offset, %q5_block: index, %q5_group: index, %packet: index, %dm: vector<2xf16>, %scales: vector<3xi32>, %qh_bytes: vector<4xi8>) -> (vector<4xf32>) func.decl @ggml_iq4xs_load_decode_state(%weight: buffer, %row_base: offset, %block: index) -> (i32, i32) @@ -114,10 +117,13 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%full_token_ %reuse_q5_k = index.constant 5 : index %reuse_iq4_xs = index.constant 23 : index %weight_is_q5 = index.cmp eq, %weight_format, %reuse_q5_k : index + %weight_q4_row64_format = index.constant 44 : index + %weight_is_q4_row64 = index.cmp eq, %weight_format, %weight_q4_row64_format : index + %weight_reuses_header = scalar.ori %weight_is_q5, %weight_is_q4_row64 : i1 %weight_is_iq4 = index.cmp eq, %weight_format, %reuse_iq4_xs : index %result00, %result01, %result10, %result11 = scf.for %quant_block = [%quant_block_begin to %quant_block_end step %c1](%block_acc00 = %init00 : vector<8xf16>, %block_acc01 = %init01 : vector<8xf16>, %block_acc10 = %init10 : vector<8xf16>, %block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { - // Retain weight Q5_K headers and packet high bits for all eight K32 groups; payloads are still loaded per group. - %weight_dm, %weight_scales, %weight_qh = scf.if %weight_is_q5 -> (vector<4x2xf16>, vector<4x3xi32>, vector<4x4xi8>) { + // Retain Q4 row64 or Q5_K headers and Q5 packet high bits for all eight K32 groups; payloads are still loaded per group. + %weight_dm, %weight_scales, %weight_qh = scf.if %weight_reuses_header -> (vector<4x2xf16>, vector<4x3xi32>, vector<4x4xi8>) { %zero_dm = vector.constant 0.0 : vector<4x2xf16> %zero_scales = vector.constant 0 : vector<4x3xi32> %zero_qh = vector.constant 0 : vector<4x4xi8> @@ -127,7 +133,14 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%full_token_ %valid = index.cmp ult, %row, %bounded_output_size : index %dm, %scales, %qh = scf.if %valid -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) { %row_base = index.scale %row, %weight_row_bytes : index, offset -> offset - %d, %s, %q = func.call @ggml_q5k_load_decode_state(%weight_noalias, %row_base, %quant_block, %load_packet) : (buffer, offset, index, index) -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) + %d, %s, %q = scf.if %weight_is_q4_row64 -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) { + %d4, %s4 = func.call @ggml_q4k_row64_load_decode_state(%weight_noalias, %row, %bounded_input_size, %quant_block) : (buffer, index, index, index) -> (vector<2xf16>, vector<3xi32>) + %empty_qh = vector.constant 0 : vector<4xi8> + scf.yield %d4, %s4, %empty_qh : vector<2xf16>, vector<3xi32>, vector<4xi8> + } else { + %d5, %s5, %q5 = func.call @ggml_q5k_load_decode_state(%weight_noalias, %row_base, %quant_block, %load_packet) : (buffer, offset, index, index) -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) + scf.yield %d5, %s5, %q5 : vector<2xf16>, vector<3xi32>, vector<4xi8> + } scf.yield %d, %s, %q : vector<2xf16>, vector<3xi32>, vector<4xi8> } else { %empty_dm = vector.constant 0.0 : vector<2xf16> @@ -205,7 +218,16 @@ func.def inline @ggml_matmul_tiled_input_f32_dequant_accumulate_f32(%full_token_ %half = vector.fptrunc %values : vector<4xf32> to vector<4xf16> scf.yield %half : vector<4xf16> } else { - %fallback = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + %fallback = scf.if %weight_is_q4_row64 -> (vector<4xf16>) { + %dm = vector.extract %weight_dm[%slot] : vector<4x2xf16> -> vector<2xf16> + %scales = vector.extract %weight_scales[%slot] : vector<4x3xi32> -> vector<3xi32> + %wide = func.call @ggml_q4k_row64_preloaded_f32_vector4(%weight_noalias, %channel, %bounded_input_size, %quant_block, %quant_group, %load_packet, %dm, %scales) : (buffer, index, index, index, index, index, vector<2xf16>, vector<3xi32>) -> (vector<4xf32>) + %half = vector.fptrunc %wide : vector<4xf32> to vector<4xf16> + scf.yield %half : vector<4xf16> + } else { + %generic = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %weight_format, %weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %generic : vector<4xf16> + } scf.yield %fallback : vector<4xf16> } scf.yield %other : vector<4xf16> diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom index 8c0d57c746c1..02c0ec49af66 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/matmul/tiled_pair_input_f32_dequant_accumulate_f32.loom @@ -1,3 +1,6 @@ +func.decl @ggml_q4k_row64_load_decode_state(%weight: buffer, %row: index, %input_size: index, %block: index) -> (vector<2xf16>, vector<3xi32>) +func.decl @ggml_q4k_row64_preloaded_f32_vector4(%weight: buffer, %row: index, %input_size: index, %block: index, %group: index, %packet: index, %dm: vector<2xf16>, %scales: vector<3xi32>) -> (vector<4xf32>) + func.decl @ggml_q5k_load_decode_state(%weight: buffer, %row_base: offset, %block: index, %packet: index) -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) func.decl @ggml_q5k_preloaded_f32_vector4(%weight: buffer, %row_byte_base: offset, %q5_block: index, %q5_group: index, %packet: index, %dm: vector<2xf16>, %scales: vector<3xi32>, %qh_bytes: vector<4xi8>) -> (vector<4xf32>) func.decl @ggml_iq4xs_load_decode_state(%weight: buffer, %row_base: offset, %block: index) -> (i32, i32) @@ -97,12 +100,18 @@ func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%full_t %reuse_q5_k = index.constant 5 : index %reuse_iq4_xs = index.constant 23 : index %gate_is_q5 = index.cmp eq, %gate_weight_format, %reuse_q5_k : index + %gate_q4_row64_format = index.constant 44 : index + %gate_is_q4_row64 = index.cmp eq, %gate_weight_format, %gate_q4_row64_format : index + %gate_reuses_header = scalar.ori %gate_is_q5, %gate_is_q4_row64 : i1 %gate_is_iq4 = index.cmp eq, %gate_weight_format, %reuse_iq4_xs : index %up_is_q5 = index.cmp eq, %up_weight_format, %reuse_q5_k : index + %up_q4_row64_format = index.constant 44 : index + %up_is_q4_row64 = index.cmp eq, %up_weight_format, %up_q4_row64_format : index + %up_reuses_header = scalar.ori %up_is_q5, %up_is_q4_row64 : i1 %up_is_iq4 = index.cmp eq, %up_weight_format, %reuse_iq4_xs : index %gate_result00, %gate_result01, %gate_result10, %gate_result11, %up_result00, %up_result01, %up_result10, %up_result11 = scf.for %quant_block = [%c0 to %quant_block_count step %c1](%gate_block_acc00 = %init00 : vector<8xf16>, %gate_block_acc01 = %init01 : vector<8xf16>, %gate_block_acc10 = %init10 : vector<8xf16>, %gate_block_acc11 = %init11 : vector<8xf16>, %up_block_acc00 = %init00 : vector<8xf16>, %up_block_acc01 = %init01 : vector<8xf16>, %up_block_acc10 = %init10 : vector<8xf16>, %up_block_acc11 = %init11 : vector<8xf16>) -> (vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>, vector<8xf16>) { // Retain gate Q5_K headers and packet high bits for all eight K32 groups; payloads are still loaded per group. - %gate_dm, %gate_scales, %gate_qh = scf.if %gate_is_q5 -> (vector<4x2xf16>, vector<4x3xi32>, vector<4x4xi8>) { + %gate_dm, %gate_scales, %gate_qh = scf.if %gate_reuses_header -> (vector<4x2xf16>, vector<4x3xi32>, vector<4x4xi8>) { %zero_dm = vector.constant 0.0 : vector<4x2xf16> %zero_scales = vector.constant 0 : vector<4x3xi32> %zero_qh = vector.constant 0 : vector<4x4xi8> @@ -112,7 +121,14 @@ func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%full_t %valid = index.cmp ult, %row, %bounded_output_size : index %dm, %scales, %qh = scf.if %valid -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) { %row_base = index.scale %row, %gate_weight_row_bytes : index, offset -> offset - %d, %s, %q = func.call @ggml_q5k_load_decode_state(%gate_weight_noalias, %row_base, %quant_block, %load_packet) : (buffer, offset, index, index) -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) + %d, %s, %q = scf.if %gate_is_q4_row64 -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) { + %d4, %s4 = func.call @ggml_q4k_row64_load_decode_state(%gate_weight_noalias, %row, %bounded_input_size, %quant_block) : (buffer, index, index, index) -> (vector<2xf16>, vector<3xi32>) + %empty_qh = vector.constant 0 : vector<4xi8> + scf.yield %d4, %s4, %empty_qh : vector<2xf16>, vector<3xi32>, vector<4xi8> + } else { + %d5, %s5, %q5 = func.call @ggml_q5k_load_decode_state(%gate_weight_noalias, %row_base, %quant_block, %load_packet) : (buffer, offset, index, index) -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) + scf.yield %d5, %s5, %q5 : vector<2xf16>, vector<3xi32>, vector<4xi8> + } scf.yield %d, %s, %q : vector<2xf16>, vector<3xi32>, vector<4xi8> } else { %empty_dm = vector.constant 0.0 : vector<2xf16> @@ -160,7 +176,7 @@ func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%full_t scf.yield %empty_header, %empty_scales_l : vector<4xi32>, vector<4xi32> } // Retain up Q5_K headers and packet high bits for all eight K32 groups; payloads are still loaded per group. - %up_dm, %up_scales, %up_qh = scf.if %up_is_q5 -> (vector<4x2xf16>, vector<4x3xi32>, vector<4x4xi8>) { + %up_dm, %up_scales, %up_qh = scf.if %up_reuses_header -> (vector<4x2xf16>, vector<4x3xi32>, vector<4x4xi8>) { %zero_dm = vector.constant 0.0 : vector<4x2xf16> %zero_scales = vector.constant 0 : vector<4x3xi32> %zero_qh = vector.constant 0 : vector<4x4xi8> @@ -170,7 +186,14 @@ func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%full_t %valid = index.cmp ult, %row, %bounded_output_size : index %dm, %scales, %qh = scf.if %valid -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) { %row_base = index.scale %row, %up_weight_row_bytes : index, offset -> offset - %d, %s, %q = func.call @ggml_q5k_load_decode_state(%up_weight_noalias, %row_base, %quant_block, %load_packet) : (buffer, offset, index, index) -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) + %d, %s, %q = scf.if %up_is_q4_row64 -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) { + %d4, %s4 = func.call @ggml_q4k_row64_load_decode_state(%up_weight_noalias, %row, %bounded_input_size, %quant_block) : (buffer, index, index, index) -> (vector<2xf16>, vector<3xi32>) + %empty_qh = vector.constant 0 : vector<4xi8> + scf.yield %d4, %s4, %empty_qh : vector<2xf16>, vector<3xi32>, vector<4xi8> + } else { + %d5, %s5, %q5 = func.call @ggml_q5k_load_decode_state(%up_weight_noalias, %row_base, %quant_block, %load_packet) : (buffer, offset, index, index) -> (vector<2xf16>, vector<3xi32>, vector<4xi8>) + scf.yield %d5, %s5, %q5 : vector<2xf16>, vector<3xi32>, vector<4xi8> + } scf.yield %d, %s, %q : vector<2xf16>, vector<3xi32>, vector<4xi8> } else { %empty_dm = vector.constant 0.0 : vector<2xf16> @@ -248,7 +271,16 @@ func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%full_t %half = vector.fptrunc %values : vector<4xf32> to vector<4xf16> scf.yield %half : vector<4xf16> } else { - %fallback = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %gate_weight_format, %gate_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + %fallback = scf.if %gate_is_q4_row64 -> (vector<4xf16>) { + %dm = vector.extract %gate_dm[%slot] : vector<4x2xf16> -> vector<2xf16> + %scales = vector.extract %gate_scales[%slot] : vector<4x3xi32> -> vector<3xi32> + %wide = func.call @ggml_q4k_row64_preloaded_f32_vector4(%gate_weight_noalias, %channel, %bounded_input_size, %quant_block, %quant_group, %load_packet, %dm, %scales) : (buffer, index, index, index, index, index, vector<2xf16>, vector<3xi32>) -> (vector<4xf32>) + %half = vector.fptrunc %wide : vector<4xf32> to vector<4xf16> + scf.yield %half : vector<4xf16> + } else { + %generic = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %gate_weight_format, %gate_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %generic : vector<4xf16> + } scf.yield %fallback : vector<4xf16> } scf.yield %other : vector<4xf16> @@ -276,7 +308,16 @@ func.def inline @ggml_matmul_tiled_pair_input_f32_dequant_accumulate_f32(%full_t %half = vector.fptrunc %values : vector<4xf32> to vector<4xf16> scf.yield %half : vector<4xf16> } else { - %fallback = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %up_weight_format, %up_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + %fallback = scf.if %up_is_q4_row64 -> (vector<4xf16>) { + %dm = vector.extract %up_dm[%slot] : vector<4x2xf16> -> vector<2xf16> + %scales = vector.extract %up_scales[%slot] : vector<4x3xi32> -> vector<3xi32> + %wide = func.call @ggml_q4k_row64_preloaded_f32_vector4(%up_weight_noalias, %channel, %bounded_input_size, %quant_block, %quant_group, %load_packet, %dm, %scales) : (buffer, index, index, index, index, index, vector<2xf16>, vector<3xi32>) -> (vector<4xf32>) + %half = vector.fptrunc %wide : vector<4xf32> to vector<4xf16> + scf.yield %half : vector<4xf16> + } else { + %generic = func.call @ggml_dequant_f16_vector4(%iq4nl_table, %iq3s_grid_low0, %iq3s_grid_low1, %iq3s_grid_low2, %iq3s_grid_low3, %iq3s_grid_low4, %iq3s_grid_low5, %iq3s_grid_low6, %iq3s_grid_low7, %iq3s_grid_low8, %iq3s_grid_low9, %iq3s_grid_low10, %iq3s_grid_low11, %iq3s_grid_low12, %iq3s_grid_low13, %iq3s_grid_low14, %iq3s_grid_low15, %iq3s_grid_high0, %iq3s_grid_high1, %iq3s_grid_high2, %iq3s_grid_high3, %iq3s_grid_high4, %iq3s_grid_high5, %iq3s_grid_high6, %iq3s_grid_high7, %iq3s_grid_high8, %iq3s_grid_high9, %iq3s_grid_high10, %iq3s_grid_high11, %iq3s_grid_high12, %iq3s_grid_high13, %iq3s_grid_high14, %iq3s_grid_high15, %up_weight_format, %up_weight_noalias, %row_byte_base, %bounded_input_size, %quant_block, %quant_group, %load_packet, %weight_k) : (vector<16xi8>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, vector<32xf32>, index, buffer, offset, index, index, index, index, index) -> (vector<4xf16>) + scf.yield %generic : vector<4xf16> + } scf.yield %fallback : vector<4xf16> } scf.yield %other : vector<4xf16> diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_quantized_f16_prefill.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_quantized_f16_prefill.loom index 86a71706509a..faab82cc6072 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_quantized_f16_prefill.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/mul_mat_quantized_f16_prefill.loom @@ -3,8 +3,9 @@ template.decl @ggml.binary_f32.apply_vector8(%op: index, %lhs: vector<8xf32>, %rhs: vector<8xf32>) -> (vector<8xf32>) -func.decl @ggml_q4k_native_row64_f16_pair(%half_packet: i1, %paired: i1, %weight: buffer, %peer: buffer, %is_up: i1, %row_base: offset, %group0: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) +func.decl @ggml_q4k_f16_pair(%native_weight: i1, %half_packet: i1, %paired: i1, %weight: buffer, %peer: buffer, %is_up: i1, %row_base: offset, %group0: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) func.decl @ggml_q6k_f16_pair(%half_packet: i1, %weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) +func.decl @ggml_q6k_native_row64_f16_pair(%half_packet: i1, %weight: buffer, %row: index, %input_size: index, %q6_block: index, %q6_group: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) func.def inline @ggml_f16_prefill_load_pair(%input: buffer, %token_count: index, %input_size: index, %token_origin: index, %k_origin: index) -> (vector<8xf16>, vector<8xf16>) { %c0_offset = index.constant 0 : offset @@ -265,8 +266,13 @@ func.def inline @ggml_f16_prefill_stage_mma(%paired: i1, %packed_input: i1, %bou func.def inline @ggml_f16_prefill_publish_weights(%weight_format: index, %paired: i1, %packed_input: i1, %short_weight_packets: i1, %active: i1, %bounded_input_size: index, %bounded_output_size: index, %channel_tile: index, %weight_epoch: index, %weight_noalias: buffer, %up_weight: buffer, %weight_stage: buffer, %stage_offset: offset) { + %q4_format = index.constant 4 : index %q6_format = index.constant 6 : index - %q6_weights = index.cmp eq, %weight_format, %q6_format : index + %q6_row64_format = index.constant 46 : index + %q4_native_weights = index.cmp eq, %weight_format, %q4_format : index + %q6_native_weights = index.cmp eq, %weight_format, %q6_format : index + %q6_row64_weights = index.cmp eq, %weight_format, %q6_row64_format : index + %q6_weights = scalar.ori %q6_native_weights, %q6_row64_weights : i1 %layout_true = scalar.constant true : i1 %staged_input = scalar.xori %packed_input, %layout_true : i1 %workgroup_size = kernel.workgroup.size : index @@ -339,22 +345,33 @@ func.def inline @ggml_f16_prefill_publish_weights(%weight_format: index, %paired %valid_channel = index.cmp ult, %channel, %bounded_output_size : index %weight_values0, %weight_values1 = scf.if %valid_channel -> (vector<16xf16>, vector<16xf16>) { %decoded0, %decoded1 = scf.if %q6_weights -> (vector<16xf16>, vector<16xf16>) { - %native_block_bytes = index.constant 210 : offset - %native_row_bytes = index.scale %quant_block_count, %native_block_bytes : index, offset -> offset - %row_byte_base = index.scale %channel, %native_row_bytes : index, offset -> offset - %decoded0, %decoded1 = func.call @ggml_q6k_f16_pair(%short_weight_packets, %weight_noalias, %row_byte_base, %quant_block, %weight_group, %load_packet) : (i1, buffer, offset, index, index, index) -> (vector<16xf16>, vector<16xf16>) + %decoded0, %decoded1 = scf.if %q6_row64_weights -> (vector<16xf16>, vector<16xf16>) { + %row64_0, %row64_1 = func.call @ggml_q6k_native_row64_f16_pair(%short_weight_packets, %weight_noalias, %channel, %bounded_input_size, %quant_block, %weight_group, %load_packet) : (i1, buffer, index, index, index, index, index) -> (vector<16xf16>, vector<16xf16>) + scf.yield %row64_0, %row64_1 : vector<16xf16>, vector<16xf16> + } else { + %native_block_bytes = index.constant 210 : offset + %native_row_bytes = index.scale %quant_block_count, %native_block_bytes : index, offset -> offset + %row_byte_base = index.scale %channel, %native_row_bytes : index, offset -> offset + %native0, %native1 = func.call @ggml_q6k_f16_pair(%short_weight_packets, %weight_noalias, %row_byte_base, %quant_block, %weight_group, %load_packet) : (i1, buffer, offset, index, index, index) -> (vector<16xf16>, vector<16xf16>) + scf.yield %native0, %native1 : vector<16xf16>, vector<16xf16> + } scf.yield %decoded0, %decoded1 : vector<16xf16>, vector<16xf16> } else { %packed_record_bytes = index.constant 16 : offset %packed_block_bytes = index.constant 9216 : offset + %native_block_bytes = index.constant 144 : offset %packed_row_group = index.div %channel, %c64 : index %packed_row_lane = index.rem %channel, %c64 : index %packed_tile_block0 = index.mul %packed_row_group, %quant_block_count : index %packed_tile_block = index.add %packed_tile_block0, %quant_block : index %packed_tile_base = index.scale %packed_tile_block, %packed_block_bytes : index, offset -> offset %packed_row_add = index.scale %packed_row_lane, %packed_record_bytes : index, offset -> offset - %row_byte_base = index.add %packed_tile_base, %packed_row_add : offset - %decoded0, %decoded1 = func.call @ggml_q4k_native_row64_f16_pair(%short_weight_packets, %paired, %weight_noalias, %up_weight, %weight_is_up, %row_byte_base, %weight_group, %load_packet) : (i1, i1, buffer, buffer, i1, offset, index, index) -> (vector<16xf16>, vector<16xf16>) + %packed_row_base = index.add %packed_tile_base, %packed_row_add : offset + %native_row_block = index.mul %channel, %quant_block_count : index + %native_block = index.add %native_row_block, %quant_block : index + %native_block_base = index.scale %native_block, %native_block_bytes : index, offset -> offset + %row_byte_base = scf.select %q4_native_weights, %native_block_base, %packed_row_base : offset + %decoded0, %decoded1 = func.call @ggml_q4k_f16_pair(%q4_native_weights, %short_weight_packets, %paired, %weight_noalias, %up_weight, %weight_is_up, %row_byte_base, %weight_group, %load_packet) : (i1, i1, i1, buffer, buffer, i1, offset, index, index) -> (vector<16xf16>, vector<16xf16>) scf.yield %decoded0, %decoded1 : vector<16xf16>, vector<16xf16> } scf.yield %decoded0, %decoded1 : vector<16xf16>, vector<16xf16> @@ -521,8 +538,6 @@ func.def inline @ggml_prefill_conv4_strip(%deferred_state: i1, %width: index, %c } func.def inline @ggml_mul_mat_quantized_f16_prefill_wave32_body(%weight_format: index, %binary_op: index, %paired: i1, %packed_input: i1, %packed_output: i1, %token_count: index, %input_size0: index, %output_size0: index, %channel_tile: index, %token_tile: index, %input: buffer, %weight: buffer, %up_weight: buffer, %output: buffer, %f16_output: buffer, %conv_epilogue: i1, %conv_state: buffer, %conv_filter: buffer, %conv_output: buffer, %conv_cache: buffer, %deferred_state: i1) { - %q6_format = index.constant 6 : index - %q6_weights = index.cmp eq, %weight_format, %q6_format : index %layout_true = scalar.constant true : i1 %unpaired = scalar.xori %paired, %layout_true : i1 %packed_single = scalar.andi %unpaired, %packed_input : i1 diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q6_k_f16.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q6_k_f16.loom index 6c9bccfdff33..a13207cb2f14 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q6_k_f16.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/motifs/q6_k_f16.loom @@ -180,3 +180,115 @@ func.def public inline @ggml_q6k_f16_pair(%half_packet: i1, %weight: buffer, %we %values1 = vector.fptrunc %values_f32_1 : vector<16xf32> to vector<16xf16> func.return %values0, %values1 : vector<16xf16>, vector<16xf16> } + +// Row64 stores each raw field contiguously per row and transposes the scale pairs. +func.def public inline @ggml_q6k_native_row64_f16_pair(%half_packet: i1, %weight: buffer, %row: index, %input_size: index, %q6_block: index, %q6_group: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) { + %padding = vector.constant 0 : vector<2xi32> + %c1 = index.constant 1 : index + %c2 = index.constant 2 : index + %c3 = index.constant 3 : index + %c4 = index.constant 4 : index + %c64 = index.constant 64 : index + %c256 = index.constant 256 : index + %tile_bytes = index.constant 13440 : offset + %raw_field_bytes = index.constant 2048 : offset + %raw_row_bytes = index.constant 32 : offset + %scale_byte_add = index.constant 128 : offset + %raw_byte_add = index.constant 1152 : offset + %c4_i32v = vector.constant 4 : vector<4xi32> + %nibble_mask = vector.constant 252645135 : vector<4xi32> + %high_mask = vector.constant 50529027 : vector<4xi32> + %c32_f32v = vector.constant 32.0 : vector<16xf32> + %bounded_group = index.assume %q6_group [range(%q6_group, 0, 6), multiple_of(%q6_group, 2)] : index + %bounded_packet = index.assume %packet [range(%packet, 0, 6), multiple_of(%packet, 2)] : index + %blocks = index.div %input_size, %c256 : index + %row_group = index.div %row, %c64 : index + %row_lane = index.rem %row, %c64 : index + %tile = index.madd %row_group, %blocks, %q6_block : index + %tile_base = index.scale %tile, %tile_bytes : index, offset -> offset + %scale_byte_base = index.add %tile_base, %scale_byte_add : offset + %raw_byte_base = index.add %tile_base, %raw_byte_add : offset + %group_in_half = index.rem %bounded_group, %c4 : index + %half = index.div %bounded_group, %c4 : index + %ql_field = index.mul %half, %c3 : index + %ql_field_add = index.scale %ql_field, %raw_field_bytes : index, offset -> offset + %ql_row_add = index.scale %row_lane, %raw_row_bytes : index, offset -> offset + %ql_field_base = index.add %raw_byte_base, %ql_field_add : offset + %ql_base0 = index.add %ql_field_base, %ql_row_add : offset + %ql_base1 = index.add %ql_base0, %raw_field_bytes : offset + %qh_base = index.add %ql_base1, %raw_field_bytes : offset + %ql_view0 = buffer.view %weight[%ql_base0] : buffer -> view<8xi32> + %ql_view1 = buffer.view %weight[%ql_base1] : buffer -> view<8xi32> + %qh_view = buffer.view %weight[%qh_base] : buffer -> view<8xi32> + %scale_view = buffer.view %weight[%scale_byte_base] : buffer -> view<64x2x8xi8> + %d_view = buffer.view %weight[%tile_base] : buffer -> view<64xf16> + %nibble = index.div %group_in_half, %c2 : index + %nibble_shift_index = index.mul %nibble, %c4 : index + %nibble_shift_i32 = index.cast %nibble_shift_index : index to i32 + %nibble_shift = vector.splat %nibble_shift_i32 : vector<4xi32> + %qh_shift_index0 = index.mul %group_in_half, %c2 : index + %qh_shift_index1 = index.add %qh_shift_index0, %c2 : index + %qh_shift_i32_0 = index.cast %qh_shift_index0 : index to i32 + %qh_shift_i32_1 = index.cast %qh_shift_index1 : index to i32 + %qh_shift0 = vector.splat %qh_shift_i32_0 : vector<4xi32> + %qh_shift1 = vector.splat %qh_shift_i32_1 : vector<4xi32> + %scale_packet_half = index.div %bounded_packet, %c4 : index + %second_group = index.add %bounded_group, %c1 : index + %ql_word0 = scf.if %half_packet -> (vector<4xi32>) { + %words = vector.load %ql_view0[%bounded_packet] : view<8xi32> -> vector<2xi32> + %padded = vector.concat<0> %words, %padding : vector<2xi32>, vector<2xi32> -> vector<4xi32> + scf.yield %padded : vector<4xi32> + } else { + %words = vector.load %ql_view0[%bounded_packet] : view<8xi32> -> vector<4xi32> + scf.yield %words : vector<4xi32> + } + %ql_word1 = scf.if %half_packet -> (vector<4xi32>) { + %words = vector.load %ql_view1[%bounded_packet] : view<8xi32> -> vector<2xi32> + %padded = vector.concat<0> %words, %padding : vector<2xi32>, vector<2xi32> -> vector<4xi32> + scf.yield %padded : vector<4xi32> + } else { + %words = vector.load %ql_view1[%bounded_packet] : view<8xi32> -> vector<4xi32> + scf.yield %words : vector<4xi32> + } + %qh_word = scf.if %half_packet -> (vector<4xi32>) { + %words = vector.load %qh_view[%bounded_packet] : view<8xi32> -> vector<2xi32> + %padded = vector.concat<0> %words, %padding : vector<2xi32>, vector<2xi32> -> vector<4xi32> + scf.yield %padded : vector<4xi32> + } else { + %words = vector.load %qh_view[%bounded_packet] : view<8xi32> -> vector<4xi32> + scf.yield %words : vector<4xi32> + } + %d_f16 = view.load %d_view[%row_lane] : view<64xf16> -> f16 + %d = scalar.extf %d_f16 : f16 to f32 + %ql_shifted0 = vector.shrui %ql_word0, %nibble_shift : vector<4xi32> + %ql0 = vector.andi %ql_shifted0, %nibble_mask : vector<4xi32> + %qh_shifted0 = vector.shrui %qh_word, %qh_shift0 : vector<4xi32> + %qh_low0 = vector.andi %qh_shifted0, %high_mask : vector<4xi32> + %qh0 = vector.shli %qh_low0, %c4_i32v : vector<4xi32> + %code0 = vector.ori %ql0, %qh0 : vector<4xi32> + %code_i8_0 = vector.bitcast %code0 : vector<4xi32> to vector<16xi8> + %code_f32_0 = vector.uitofp %code_i8_0 : vector<16xi8> to vector<16xf32> + %centered0 = vector.subf %code_f32_0, %c32_f32v : vector<16xf32> + %scale_i8_0 = view.load %scale_view[%row_lane, %scale_packet_half, %bounded_group] : view<64x2x8xi8> -> i8 + %scale0 = scalar.sitofp %scale_i8_0 : i8 to f32 + %combined_scale0 = scalar.mulf %scale0, %d : f32 + %combined_scale_vector0 = vector.splat %combined_scale0 : vector<16xf32> + %values_f32_0 = vector.mulf %centered0, %combined_scale_vector0 : vector<16xf32> + %values0 = vector.fptrunc %values_f32_0 : vector<16xf32> to vector<16xf16> + %ql_shifted1 = vector.shrui %ql_word1, %nibble_shift : vector<4xi32> + %ql1 = vector.andi %ql_shifted1, %nibble_mask : vector<4xi32> + %qh_shifted1 = vector.shrui %qh_word, %qh_shift1 : vector<4xi32> + %qh_low1 = vector.andi %qh_shifted1, %high_mask : vector<4xi32> + %qh1 = vector.shli %qh_low1, %c4_i32v : vector<4xi32> + %code1 = vector.ori %ql1, %qh1 : vector<4xi32> + %code_i8_1 = vector.bitcast %code1 : vector<4xi32> to vector<16xi8> + %code_f32_1 = vector.uitofp %code_i8_1 : vector<16xi8> to vector<16xf32> + %centered1 = vector.subf %code_f32_1, %c32_f32v : vector<16xf32> + %scale_i8_1 = view.load %scale_view[%row_lane, %scale_packet_half, %second_group] : view<64x2x8xi8> -> i8 + %scale1 = scalar.sitofp %scale_i8_1 : i8 to f32 + %combined_scale1 = scalar.mulf %scale1, %d : f32 + %combined_scale_vector1 = vector.splat %combined_scale1 : vector<16xf32> + %values_f32_1 = vector.mulf %centered1, %combined_scale_vector1 : vector<16xf32> + %values1 = vector.fptrunc %values_f32_1 : vector<16xf32> to vector<16xf16> + func.return %values0, %values1 : vector<16xf16>, vector<16xf16> +} diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom index 2588f7d20cd3..f7dbc02b835c 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_f32_f32_wmma.loom @@ -175,7 +175,7 @@ kernel.def target(@ggml_mul_mat_gfx11_wave32) @ggml_mul_mat_q4_k_f16_wmma_prefil %input_layout = config.get @ggml.mul_mat.f16_input_layout : index %packed_layout = index.constant 1 : index %packed_input = index.cmp eq, %input_layout, %packed_layout : index - %prefill_weight_format = index.constant 44 : index + %prefill_weight_format = config.get @ggml.mul_mat.weight_format : index %prefill_binary_op = index.constant 0 : index func.call @ggml_mul_mat_quantized_f16_prefill_wave32(%prefill_weight_format, %prefill_binary_op, %paired, %packed_input, %paired, %bounded_token_count, %input_size, %bounded_output_size, %channel_tile, %token_tile, %input, %weight, %weight, %output, %output) : (index, index, i1, i1, i1, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) kernel.return diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q6_k_packed_f16_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q6_k_packed_f16_wmma.loom index 0f82a048db56..ace5d1c84816 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q6_k_packed_f16_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q6_k_packed_f16_wmma.loom @@ -6,9 +6,11 @@ config.decl @ggml.mul_mat_q6_k_packed.output_size : %value: index where [range(% config.decl @ggml.mul_mat_q6_k_packed.output_accumulation : %value: index where [range(%value, 0, 1)] +config.def @ggml.mul_mat_q6_k_packed.weight_format = 6 : index + config.decl @ggml.mul_mat_q6_k_packed.weight_offset : %value: index where [range(%value, 0, 2147483647), multiple_of(%value, 256)] -config.decl @ggml.mul_mat_q6_k_packed.token_capacity : %value: index where [range(%value, 128, 2048), multiple_of(%value, 128)] +config.decl @ggml.mul_mat_q6_k_packed.token_capacity : %value: index where [range(%value, 64, 2048)] config.decl @ggml.mul_mat_q6_k_shortlist.input_size : %value: index where [range(%value, 256, 32768), multiple_of(%value, 256)] @@ -32,6 +34,8 @@ amdgpu.target @ggml_mul_mat_q6_k_prefill_gfx11_wave32 {subgroup_s func.decl @ggml_q6k_f16_pair(%half_packet: i1, %weight: buffer, %weight_row_byte_base: offset, %q6_block: index, %q6_group: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) +func.decl @ggml_q6k_native_row64_f16_pair(%half_packet: i1, %weight: buffer, %row: index, %input_size: index, %q6_block: index, %q6_group: index, %packet: index) -> (vector<16xf16>, vector<16xf16>) + template.def<@ggml.mul_mat_q6_k_packed.publish_f16_wave32> device @ggml_mul_mat_q6_k_packed_publish_f16_wave32(%result: vector<16xf16>, %token_offset: index, %token_tile_base: index, %bounded_token_count: index, %bounded_output_size: index, %channel_tile_base: index, %subgroup_channel_add: index, %result_stage: buffer, %output_noalias: buffer, %output_accumulation: index) { %lane = kernel.subgroup.lane.id : index %subgroup0 = kernel.subgroup.id : index @@ -1298,7 +1302,6 @@ template.def<@ggml.mul_mat_q6_k_packed.prefill_wave32.body> device @ggml_mul_mat %c96 = index.constant 96 : index %c112 = index.constant 112 : index %c128 = index.constant 128 : index - %c210 = index.constant 210 : index %c256 = index.constant 256 : index %c0_offset = index.constant 0 : offset %weight_stage_bytes = index.constant 18432 : offset @@ -1310,7 +1313,6 @@ template.def<@ggml.mul_mat_q6_k_packed.prefill_wave32.body> device @ggml_mul_mat %n = index.constant 16 : index %k = index.constant 16 : index %quant_block_count = index.div %bounded_input_size, %c256 : index - %weight_row_bytes = index.mul %quant_block_count, %c210 : index %input_noalias, %weight_noalias, %output_noalias = buffer.assume.noalias %input, %weight, %output : buffer, buffer, buffer %input_f32_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf32> %input_f16_view = buffer.view %input_noalias[%c0_offset] : buffer -> view<[%bounded_token_count]x[%bounded_input_size]xf16> @@ -1354,10 +1356,20 @@ template.def<@ggml.mul_mat_q6_k_packed.prefill_wave32.body> device @ggml_mul_mat %channel = index.add %channel_tile_base, %weight_local_row : index %valid_channel = index.cmp ult, %channel, %bounded_output_size : index %weight_values0, %weight_values1 = scf.if %valid_channel -> (vector<16xf16>, vector<16xf16>) { - %channel_byte_index = index.mul %channel, %weight_row_bytes : index - %row_byte_base = index.cast %channel_byte_index : index to offset %half_packet = scalar.constant false : i1 - %decoded0, %decoded1 = func.call @ggml_q6k_f16_pair(%half_packet, %weight_noalias, %row_byte_base, %quant_block, %quant_group, %load_packet) : (i1, buffer, offset, index, index, index) -> (vector<16xf16>, vector<16xf16>) + %weight_format = config.get @ggml.mul_mat_q6_k_packed.weight_format : index + %row64_format = index.constant 46 : index + %row64_weight = index.cmp eq, %weight_format, %row64_format : index + %decoded0, %decoded1 = scf.if %row64_weight -> (vector<16xf16>, vector<16xf16>) { + %row64_0, %row64_1 = func.call @ggml_q6k_native_row64_f16_pair(%half_packet, %weight_noalias, %channel, %bounded_input_size, %quant_block, %quant_group, %load_packet) : (i1, buffer, index, index, index, index, index) -> (vector<16xf16>, vector<16xf16>) + scf.yield %row64_0, %row64_1 : vector<16xf16>, vector<16xf16> + } else { + %native_block_bytes = index.constant 210 : offset + %native_row_bytes = index.scale %quant_block_count, %native_block_bytes : index, offset -> offset + %row_byte_base = index.scale %channel, %native_row_bytes : index, offset -> offset + %native0, %native1 = func.call @ggml_q6k_f16_pair(%half_packet, %weight_noalias, %row_byte_base, %quant_block, %quant_group, %load_packet) : (i1, buffer, offset, index, index, index) -> (vector<16xf16>, vector<16xf16>) + scf.yield %native0, %native1 : vector<16xf16>, vector<16xf16> + } scf.yield %decoded0, %decoded1 : vector<16xf16>, vector<16xf16> } else { scf.yield %c0_f16x4, %c0_f16x4 : vector<16xf16>, vector<16xf16> @@ -1474,7 +1486,7 @@ kernel.def target(@ggml_mul_mat_q6_k_prefill_gfx11_wave32) export("ggml_mul_mat_ %input_size = config.get @ggml.mul_mat_q6_k_packed.input_size : index %output_size = config.get @ggml.mul_mat_q6_k_packed.output_size : index %token_capacity = config.get @ggml.mul_mat_q6_k_packed.token_capacity : index - %bounded_token_count = index.assume %token_count [range(%token_count, 128, 2048), le(%token_count, %token_capacity)] : index + %bounded_token_count = index.assume %token_count [range(%token_count, 64, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %token_tile = kernel.workgroup.id : index %input_is_f16 = scalar.constant false : i1 @@ -1561,14 +1573,14 @@ kernel.def target(@ggml_mul_mat_q6_k_prefill_gfx11_wave32) export("ggml_mul_mat_ %token_capacity = config.get @ggml.mul_mat_q6_k_packed.token_capacity : index %output_accumulation = config.get @ggml.mul_mat_q6_k_packed.output_accumulation : index %packed_input, %packed_n_tile = func.call pure @ggml_q6k_prefill_packed_tile(%token_capacity, %input_size, %output_size, %output_accumulation) : (index, index, index, index) -> (i1, index) - %bounded_token_count = index.assume %token_count [range(%token_count, 128, 2048), le(%token_count, %token_capacity)] : index + %bounded_token_count = index.assume %token_count [range(%token_count, 64, 2048), le(%token_count, %token_capacity)] : index %channel_tile = kernel.workgroup.id : index %token_tile = kernel.workgroup.id : index %workgroup_size = kernel.workgroup.size : index %c512 = index.constant 512 : index %wide_tile = index.cmp uge, %workgroup_size, %c512 : index scf.if %wide_tile { - %weight_format = index.constant 6 : index + %weight_format = config.get @ggml.mul_mat_q6_k_packed.weight_format : index %paired = scalar.constant false : i1 %prefill_binary_op = index.constant 0 : index func.call @ggml_mul_mat_quantized_f16_prefill_wave32(%weight_format, %prefill_binary_op, %paired, %packed_input, %paired, %bounded_token_count, %input_size, %output_size, %channel_tile, %token_tile, %input, %weight, %weight, %output, %output) : (index, index, i1, i1, i1, index, index, index, index, index, buffer, buffer, buffer, buffer, buffer) diff --git a/ggml/src/ggml-hrx/runtime/graph-executor.cpp b/ggml/src/ggml-hrx/runtime/graph-executor.cpp index d192393596a4..3040bb72d910 100644 --- a/ggml/src/ggml-hrx/runtime/graph-executor.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-executor.cpp @@ -41,7 +41,8 @@ GraphSupportResult GraphExecutor::can_execute(const ggml_cgraph & graph) const { } const KernelCorpus & corpus = get_qwen_kernel_corpus(); const GraphProgramSupportResult support = - context_.graph_programs.check_support(graph, corpus, context_.device->architecture); + context_.graph_programs.check_support(graph, corpus, context_.device->architecture, + context_.lock_layout_policy()); result.supported = support.supported; result.status.append(support.status); return result; @@ -77,7 +78,8 @@ GraphExecutionResult GraphExecutor::execute(const ggml_cgraph & graph) const { } const KernelCorpus & corpus = get_qwen_kernel_corpus(); - GraphProgramLookup lookup = context_.graph_programs.get_or_build(graph, corpus, context_.device->architecture); + GraphProgramLookup lookup = context_.graph_programs.get_or_build(graph, corpus, context_.device->architecture, + context_.lock_layout_policy()); if (!lookup.valid()) { result.status.append(lookup.status); result.status.append(lookup.match.status); diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp index d62b2be3ccfb..775fb7e9f408 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.cpp @@ -304,9 +304,11 @@ GraphProgram::GraphProgram(uint64_t uid, std::string target, std::unique_ptr graph, std::unique_ptr commands, - std::string command_shape) : + std::string command_shape, + bool row64_weight_policy) : uid_(uid), target_(std::move(target)), + row64_weight_policy_(row64_weight_policy), graph_(std::move(graph)), commands_(std::move(commands)), command_shape_(std::move(command_shape)), @@ -638,13 +640,15 @@ GraphProgramMatch GraphProgram::match_current_graph(const ggml_cgraph & current_ bool GraphProgramCache::can_execute(const ggml_cgraph & graph, const KernelCorpus & corpus, - const std::string & target) const { - return check_support(graph, corpus, target).supported; + const std::string & target, + bool row64_weight_policy) const { + return check_support(graph, corpus, target, row64_weight_policy).supported; } GraphProgramSupportResult GraphProgramCache::check_support(const ggml_cgraph & graph, const KernelCorpus & corpus, - const std::string & target) const { + const std::string & target, + bool row64_weight_policy) const { GraphProgramSupportResult result; if (graph.n_nodes == 0) { result.supported = true; @@ -656,7 +660,7 @@ GraphProgramSupportResult GraphProgramCache::check_support(const ggml_cgraph & return result; } std::unique_ptr program = - build_program_from_imported(graph.uid, std::move(imported.graph), corpus, target, result.status); + build_program_from_imported(graph.uid, std::move(imported.graph), corpus, target, row64_weight_policy, result.status); result.supported = program != nullptr && result.status.success(); return result; } @@ -664,10 +668,11 @@ GraphProgramSupportResult GraphProgramCache::check_support(const ggml_cgraph & GraphProgramLookup GraphProgramCache::build_from_imported(const ggml_cgraph & graph, Graph && imported_graph, const KernelCorpus & corpus, - const std::string & target) { + const std::string & target, + bool row64_weight_policy) { GraphProgramLookup result; std::unique_ptr program = - build_program_from_imported(graph.uid, std::move(imported_graph), corpus, target, result.status); + build_program_from_imported(graph.uid, std::move(imported_graph), corpus, target, row64_weight_policy, result.status); if (program == nullptr) { return result; } @@ -706,7 +711,8 @@ GraphProgramLookup GraphProgramCache::build_from_imported(const ggml_cgraph & g GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, const KernelCorpus & corpus, - const std::string & target) { + const std::string & target, + bool row64_weight_policy) { GraphProgramLookup result; if (graph.uid != 0) { const bool disable_fast_path = environment_flag_enabled("GGML_HRX_DISABLE_GRAPH_UID_FAST_PATH"); @@ -716,18 +722,20 @@ GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, { std::lock_guard lock(mutex_); if (!disable_fast_path && last_program_ != nullptr && last_program_->uid() == graph.uid && - last_program_->target() == target) { + last_program_->target() == target && last_program_->row64_weight_policy() == row64_weight_policy) { cached_program = last_program_; } else { const auto found = programs_.find(graph.uid); - if (found != programs_.end() && found->second->target() == target) { + if (found != programs_.end() && found->second->target() == target && + found->second->row64_weight_policy() == row64_weight_policy) { cached_program = found->second.get(); last_program_ = cached_program; } if (cached_program == nullptr && !disable_fast_path && allow_structural_reuse) { const auto matched = validated_matches_.find(graph.uid); if (matched != validated_matches_.end() && matched->second.nodes == graph.nodes && - matched->second.program->target() == target) { + matched->second.program->target() == target && + matched->second.program->row64_weight_policy() == row64_weight_policy) { cached_program = matched->second.program; } } @@ -768,7 +776,8 @@ GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, std::lock_guard lock(mutex_); for (const auto & entry : programs_) { GraphProgram * candidate = entry.second.get(); - if (candidate == cached_program || candidate->target() != target) { + if (candidate == cached_program || candidate->target() != target || + candidate->row64_weight_policy() != row64_weight_policy) { continue; } GraphProgramMatch match = candidate->match_current_graph(graph); @@ -795,7 +804,7 @@ GraphProgramLookup GraphProgramCache::get_or_build(const ggml_cgraph & graph, result.status.append(imported.status); return result; } - return build_from_imported(graph, std::move(imported.graph), corpus, target); + return build_from_imported(graph, std::move(imported.graph), corpus, target, row64_weight_policy); } GraphProgramCacheStats GraphProgramCache::stats() const { @@ -814,9 +823,10 @@ std::unique_ptr GraphProgramCache::build_program_from_imported(uin Graph && imported_graph, const KernelCorpus & corpus, const std::string & target, + bool row64_weight_policy, Status & errors) const { DispatchScheduler scheduler; - if (!scheduler.schedule_graph(imported_graph, { target })) { + if (!scheduler.schedule_graph(imported_graph, { target, row64_weight_policy })) { errors.append(scheduler.plan().status); return nullptr; } @@ -834,7 +844,7 @@ std::unique_ptr GraphProgramCache::build_program_from_imported(uin } return std::make_unique(uid, target, std::make_unique(std::move(imported_graph)), std::make_unique(std::move(commands)), - std::move(command_shape)); + std::move(command_shape), row64_weight_policy); } bool can_execute_standalone_op_as_graph(const ggml_tensor * op, const std::string & target) { diff --git a/ggml/src/ggml-hrx/runtime/graph-program-cache.h b/ggml/src/ggml-hrx/runtime/graph-program-cache.h index 54c41ca6873c..d9a8ab29518c 100644 --- a/ggml/src/ggml-hrx/runtime/graph-program-cache.h +++ b/ggml/src/ggml-hrx/runtime/graph-program-cache.h @@ -51,11 +51,13 @@ class GraphProgram { std::string target, std::unique_ptr graph, std::unique_ptr commands, - std::string command_shape); + std::string command_shape, + bool row64_weight_policy = false); uint64_t uid() const { return uid_; } const std::string & target() const { return target_; } + bool row64_weight_policy() const { return row64_weight_policy_; } const std::string & command_shape() const { return command_shape_; } @@ -87,6 +89,7 @@ class GraphProgram { uint64_t uid_ = 0; std::string target_; + bool row64_weight_policy_ = false; std::unique_ptr graph_; std::unique_ptr commands_; std::string command_shape_; @@ -130,13 +133,16 @@ struct GraphProgramSupportResult { class GraphProgramCache { public: - bool can_execute(const ggml_cgraph & graph, const KernelCorpus & corpus, const std::string & target) const; + bool can_execute(const ggml_cgraph & graph, const KernelCorpus & corpus, const std::string & target, + bool row64_weight_policy = false) const; GraphProgramSupportResult check_support(const ggml_cgraph & graph, const KernelCorpus & corpus, - const std::string & target) const; + const std::string & target, + bool row64_weight_policy = false) const; - GraphProgramLookup get_or_build(const ggml_cgraph & graph, const KernelCorpus & corpus, const std::string & target); + GraphProgramLookup get_or_build(const ggml_cgraph & graph, const KernelCorpus & corpus, const std::string & target, + bool row64_weight_policy = false); GraphProgramCacheStats stats() const; @@ -151,12 +157,14 @@ class GraphProgramCache { GraphProgramLookup build_from_imported(const ggml_cgraph & graph, Graph && imported_graph, const KernelCorpus & corpus, - const std::string & target); + const std::string & target, + bool row64_weight_policy); std::unique_ptr build_program_from_imported(uint64_t uid, Graph && imported_graph, const KernelCorpus & corpus, const std::string & target, + bool row64_weight_policy, Status & errors) const; mutable std::mutex mutex_; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 33ad7845f6bf..5f903155ea71 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -433,6 +433,17 @@ llama_context::llama_context( if (backend == nullptr) { throw std::runtime_error(format("failed to initialize %s backend", ggml_backend_dev_name(dev.dev))); } + const bool is_gemma4 = model.arch == LLM_ARCH_GEMMA4; + const bool row64_weight_policy = !is_gemma4; + if (row64_weight_policy) { + ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev.dev); + using set_row64_weight_policy_fn = bool (*)(ggml_backend_t, bool); + auto set_policy = reinterpret_cast( + ggml_backend_reg_get_proc_address(reg, "ggml_backend_hrx_set_row64_weight_policy")); + if (set_policy != nullptr && !set_policy(backend, true)) { + throw std::runtime_error("failed to set HRX row64 weight policy"); + } + } backends.emplace_back(backend); } diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index 292f389b313a..7b4e8627394b 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -14796,6 +14796,29 @@ static void run_registered_host_input_reuse_waits_before_mutation() { } static void register_hrx_backend_device_cases(test_runner::Suite & suite) { + suite.device_case("qwen27_policy_lifetime", [] { + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + ggml_backend_reg_t registry = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); + using set_policy_fn = bool (*)(ggml_backend_t, bool); + auto set_policy = reinterpret_cast( + ggml_backend_reg_get_proc_address(registry, "ggml_backend_hrx_set_row64_weight_policy")); + REQUIRE(set_policy != nullptr); + REQUIRE(set_policy(backend, true)); + ggml_context * ctx = ggml_init({ 256 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + ggml_tensor * input = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 128); + ggml_tensor * output = ggml_add(ctx, input, input); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + auto * state = static_cast(backend->context); + ggml::hrx::GraphExecutor executor(*state); + REQUIRE(executor.can_execute(*graph).supported); + REQUIRE(!set_policy(backend, false)); + REQUIRE(state->row64_weight_policy); + ggml_free(ctx); + ggml_backend_free(backend); + }); suite.device_case("registered_host_input_reuse_waits_before_mutation", [] { run_registered_host_input_reuse_waits_before_mutation(); }); diff --git a/tests/hrx-loom-jit-test.cpp b/tests/hrx-loom-jit-test.cpp index c335d1a08c86..d54c1d4d61ea 100644 --- a/tests/hrx-loom-jit-test.cpp +++ b/tests/hrx-loom-jit-test.cpp @@ -798,6 +798,7 @@ static void run_loom_jit_compile_checks() { { "ggml.mul_mat.output_size", std::to_string(test.outputs) }, { "ggml.workload.token_capacity", std::to_string(test.tokens) }, { "ggml.mul_mat.f16_input_layout", std::to_string(test.input_layout) }, + { "ggml.mul_mat.weight_format", "44" }, }); REQUIRE(ref != nullptr); REQUIRE(resolve_with_timeout(ref, std::chrono::seconds(120))); @@ -1028,7 +1029,7 @@ static void run_loom_jit_compile_checks() { "ggml_mul_mat_quantized_f16_wmma_prefill_conv4_interior" }) { const auto & quantized_conv4 = find_kernel(name); REQUIRE(quantized_conv4.bindings.size() == (std::string(name).find("interior") != std::string::npos ? 5 : 6)); - for (const int format : {4, 6}) { + for (const int format : {44, 6, 46}) { for (const int64_t channels : {2112, 8192, 10240}) { auto ref = compile_kernel(*sync_jit, std::string(name) + "-" + std::to_string(format) + "-" + std::to_string(channels), quantized_conv4, token_count_workload(512), { @@ -1553,6 +1554,28 @@ static void run_binary_publication_compile_checks() { })); } +static void run_native_q4_prefill_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 & kernel = find_kernel("ggml_mul_mat_q4_k_f16_wmma_prefill_wave32"); + for (const int64_t tokens : { int64_t{ 512 }, int64_t{ 1024 } }) { + for (const int64_t output_size : { int64_t{ 512 }, int64_t{ 4096 } }) { + const std::string key = "native-q4-prefill-" + std::to_string(tokens) + "-" + + std::to_string(output_size); + auto ref = compile_kernel(*jit, key, kernel, token_count_workload(tokens), { + { "ggml.mul_mat.input_size", "2048" }, + { "ggml.mul_mat.output_size", std::to_string(output_size) }, + { "ggml.workload.token_capacity", std::to_string(tokens) }, + { "ggml.mul_mat.f16_input_layout", "0" }, + { "ggml.mul_mat.weight_format", "4" }, + }); + require_compiled_kernel(ref); + } + } +} + static bool has_hrx_test_device() { HrxTestDevice device; return device.open(); @@ -1560,6 +1583,7 @@ static bool has_hrx_test_device() { static void register_hrx_loom_jit_cases(test_runner::Suite & suite) { suite.host_case("compile", [] { run_loom_jit_compile_checks(); }); + suite.host_case("native-q4-prefill", [] { run_native_q4_prefill_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(); }); diff --git a/tests/hrx-ops-test.cpp b/tests/hrx-ops-test.cpp index 5ccd092721e3..763fc60b8bf3 100644 --- a/tests/hrx-ops-test.cpp +++ b/tests/hrx-ops-test.cpp @@ -365,9 +365,9 @@ static ggml_tensor * build_qwen_router_top8_graph(ggml_context * ctx, return output; } -static std::string kernel_name_for_id(uint64_t kernel_id) { +static std::string kernel_name_for_id(uint64_t kernel_id, const std::string & architecture = "gfx1151") { const ggml::hrx::KernelResolveResult resolved = - ggml::hrx::resolve_kernel_definition(ggml::hrx::get_qwen_kernel_corpus(), "gfx1151", kernel_id); + ggml::hrx::resolve_kernel_definition(ggml::hrx::get_qwen_kernel_corpus(), architecture, kernel_id); REQUIRE(resolved.found()); return ggml::hrx::kernel_definition_name(*resolved.definition); } @@ -6259,6 +6259,46 @@ static void run_dense_matmul_unary_cpu_reference_case(ggml_type weight_type, ggml_backend_free(hrx_backend); } +// The scalar reference avoids the CPU backend's activation quantization. +static std::vector row64_matmul_reference(ggml_type type, const std::vector & bytes, + const std::vector & input, + int64_t inputs, int64_t outputs, int64_t tokens) { + const auto * traits = ggml_get_type_traits(type); + REQUIRE(traits != nullptr && traits->to_float != nullptr); + const size_t row_bytes = ggml_row_size(type, inputs); + std::vector row(static_cast(inputs)); + std::vector result(static_cast(outputs * tokens)); + for (int64_t output = 0; output < outputs; ++output) { + traits->to_float(bytes.data() + output * row_bytes, row.data(), inputs); + for (int64_t token = 0; token < tokens; ++token) { + double sum = 0.0; + for (int64_t k = 0; k < inputs; ++k) { + sum += static_cast(row[k]) * input[token * inputs + k]; + } + result[token * outputs + output] = static_cast(sum); + } + } + return result; +} + +static void require_row64_reference(const std::vector & actual, const std::vector & expected) { + double error = 0.0; + double magnitude = 0.0; + float maximum_error = 0.0f; + for (size_t i = 0; i < actual.size(); ++i) { + const double delta = static_cast(actual[i]) - expected[i]; + error += delta * delta; + magnitude += static_cast(expected[i]) * expected[i]; + maximum_error = std::max(maximum_error, static_cast(std::fabs(delta))); + } + REQUIRE(magnitude > 0.0); + const double relative_l2 = std::sqrt(error / magnitude); + std::fprintf(stderr, "row64 scalar reference: relative_l2=%g max_abs=%g elements=%zu\n", + relative_l2, maximum_error, actual.size()); + REQUIRE(relative_l2 <= 0.01); + require_close(actual, expected, 0.02f, 0.02f); +} + static void run_dense_matmul_swiglu_cpu_reference_case(ggml_type gate_weight_type, ggml_type up_weight_type, const char * expected_kernel, @@ -6267,10 +6307,20 @@ static void run_dense_matmul_swiglu_cpu_reference_case(ggml_type gate_weight_ ggml_glu_op glu_op = GGML_GLU_OP_SWIGLU, int64_t input_size = kQwenHiddenSize, bool verify_uploaded_native = false, - bool import_weights = false) { + bool import_weights = false, + bool row64_weight_policy = false, + bool varying_input = false) { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); REQUIRE(hrx_backend != nullptr); + if (row64_weight_policy) { + ggml_backend_reg_t registry = ggml_backend_dev_backend_reg(ggml_backend_get_device(hrx_backend)); + using set_policy_fn = bool (*)(ggml_backend_t, bool); + auto set_policy = reinterpret_cast( + ggml_backend_reg_get_proc_address(registry, "ggml_backend_hrx_set_row64_weight_policy")); + REQUIRE(set_policy != nullptr); + REQUIRE(set_policy(hrx_backend, true)); + } ggml_init_params params = {}; params.mem_size = static_cast(48 * 1024 * 1024); @@ -6309,7 +6359,34 @@ static void run_dense_matmul_swiglu_cpu_reference_case(ggml_type gate_weight_ 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_kernel }); + if (row64_weight_policy) { + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*hrx_graph); + REQUIRE(imported.valid()); + auto * backend_state = static_cast(hrx_backend->context); + REQUIRE(backend_state != nullptr); + REQUIRE(backend_state->device != nullptr); + const std::string architecture = backend_state->device->architecture; + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { architecture, true })); + for (ggml_tensor * weight : { hrx_gate_weight, hrx_up_weight }) { + const ggml::hrx::Value * source = imported.graph.values().find_tensor(weight); + REQUIRE(source != nullptr); + size_t matched = 0; + for (const auto & dispatch : scheduler.plan().dispatches) { + for (const auto & binding : dispatch.bindings) { + if (binding.value != source->id) { + continue; + } + REQUIRE(binding.layout == ggml::hrx::kQ4KPackedK256Row64Layout); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id, architecture) == expected_kernel); + ++matched; + } + } + REQUIRE(matched == 1); + } + } else { + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { expected_kernel }); + } 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); @@ -6318,11 +6395,13 @@ static void run_dense_matmul_swiglu_cpu_reference_case(ggml_type gate_weight_ if (verify_uploaded_native) { require_native_weight_bindings(hrx_graph, { hrx_gate_weight, hrx_up_weight }); } - if (import_weights) { + const bool validate_default_import = import_weights && !row64_weight_policy; + if (validate_default_import) { require_packed_weight_bindings(hrx_graph, { hrx_gate_weight, hrx_up_weight }); } - const std::vector input(static_cast(input_size * token_count), 0.00390625f); + const std::vector input = varying_input ? make_pattern_f32(input_size * token_count, 7, 0.01f) : + std::vector(static_cast(input_size * token_count), 0.00390625f); if (verify_uploaded_native) { ggml_backend_tensor_set(cpu_gate_weight, gate_weight.data(), 0, gate_weight.size()); ggml_backend_tensor_set(cpu_up_weight, up_weight.data(), 0, up_weight.size()); @@ -6348,7 +6427,21 @@ static void run_dense_matmul_swiglu_cpu_reference_case(ggml_type gate_weight_ REQUIRE(ggml_backend_graph_compute(hrx_backend, hrx_graph) == GGML_STATUS_SUCCESS); ggml_backend_synchronize(cpu_backend); ggml_backend_synchronize(hrx_backend); + if (row64_weight_policy) { + auto expected = row64_matmul_reference(gate_weight_type, gate_weight, input, input_size, output_size, token_count); + const auto up = row64_matmul_reference(up_weight_type, up_weight, input, input_size, output_size, token_count); + for (size_t i = 0; i < expected.size(); ++i) { + expected[i] = (expected[i] / (1.0f + std::exp(-expected[i]))) * up[i]; + } + require_row64_reference(get_f32_tensor(hrx_backend, hrx_output), expected); + } require_close(get_f32_tensor(hrx_backend, hrx_output), get_f32_tensor(cpu_backend, cpu_output), 1.0e-4f, 1.0e-3f); + if (row64_weight_policy) { + ggml_backend_hrx_cache_stats stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); + REQUIRE(stats.weight_duplicate_warnings == 0); + REQUIRE(stats.weight_materializations == 2); + } if (verify_uploaded_native) { const auto after = backend_state->host_weights.stats(); REQUIRE(after.misses == before.misses); @@ -6442,10 +6535,19 @@ static void run_dense_matmul_packed_glu_cpu_reference_case(ggml_type weight_t int64_t output_size, ggml_glu_op glu_op = GGML_GLU_OP_SWIGLU, bool swapped = false, - int64_t input_size = 256) { + int64_t input_size = 256, + bool import_weight = false) { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); REQUIRE(hrx_backend != nullptr); + if (import_weight) { + ggml_backend_reg_t registry = ggml_backend_dev_backend_reg(ggml_backend_get_device(hrx_backend)); + using set_policy_fn = bool (*)(ggml_backend_t, bool); + auto set_policy = reinterpret_cast( + ggml_backend_reg_get_proc_address(registry, "ggml_backend_hrx_set_row64_weight_policy")); + REQUIRE(set_policy != nullptr); + REQUIRE(set_policy(hrx_backend, true)); + } ggml_init_params params = {}; params.mem_size = static_cast(64 * 1024 * 1024); @@ -6473,26 +6575,84 @@ static void run_dense_matmul_packed_glu_cpu_reference_case(ggml_type weight_t ggml_build_forward_expand(cpu_graph, cpu_output); ggml_build_forward_expand(hrx_graph, hrx_output); + std::vector weight = make_matmul_weight_bytes(weight_type, input_size, 2 * output_size, 15); + ggml_backend_buffer_t imported_buffer = nullptr; + if (import_weight) { + imported_buffer = ggml_backend_dev_buffer_from_host_ptr(ggml_backend_get_device(hrx_backend), weight.data(), + weight.size(), weight.size()); + REQUIRE(imported_buffer != nullptr); + REQUIRE(ggml_backend_tensor_alloc(imported_buffer, hrx_weight, weight.data()) == GGML_STATUS_SUCCESS); + ggml_backend_buffer_set_usage(imported_buffer, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + } + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), { expected_kernel }); + if (import_weight) { + REQUIRE(ggml_backend_hrx_weight_has_repackable_source(hrx_weight)); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*hrx_graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151", true })); + const ggml::hrx::Value * value = imported.graph.values().find_tensor(hrx_weight); + REQUIRE(value != nullptr); + size_t weight_bindings = 0; + for (const auto & dispatch : scheduler.plan().dispatches) { + for (const auto & binding : dispatch.bindings) { + if (binding.value == value->id) { + REQUIRE(binding.layout == ggml::hrx::kNativeWeightLayout); + ++weight_bindings; + } + } + } + REQUIRE(weight_bindings > 0); + } + 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, input_size, 2 * output_size, 15); const std::vector input(static_cast(input_size * token_count), 0.00390625f); - set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), weight.size()); + if (import_weight) { + ggml_backend_tensor_set(cpu_weight, weight.data(), 0, weight.size()); + } else { + set_tensor_pair_bytes(cpu_backend, cpu_weight, hrx_backend, hrx_weight, weight.data(), 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, 3.0e-2f); + const std::vector actual = get_f32_tensor(hrx_backend, hrx_output); + const std::vector expected = get_f32_tensor(cpu_backend, cpu_output); + require_close(actual, expected, import_weight ? 0.02f : 1.0f, import_weight ? 0.02f : 3.0e-2f); + + if (import_weight) { + double squared_error = 0.0; + double squared_reference = 0.0; + for (size_t index = 0; index < actual.size(); ++index) { + const double error = static_cast(actual[index]) - expected[index]; + squared_error += error * error; + squared_reference += static_cast(expected[index]) * expected[index]; + } + REQUIRE(squared_reference > 0.0); + REQUIRE(std::sqrt(squared_error / squared_reference) <= 0.03); + } + + if (import_weight) { + ggml_backend_hrx_cache_stats stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(hrx_backend, &stats)); + REQUIRE(stats.weight_materializations == 1); + REQUIRE(stats.weight_live_allocations == 1); + REQUIRE(stats.weight_duplicate_warnings == 0); + } ggml_backend_buffer_free(cpu_buffer); ggml_backend_buffer_free(hrx_buffer); + if (imported_buffer != nullptr) { + ggml_backend_buffer_free(imported_buffer); + } ggml_free(cpu_ctx); ggml_free(hrx_ctx); ggml_backend_free(cpu_backend); @@ -8032,6 +8192,7 @@ struct SsmConvPrefillGraph { ggml_tensor * ids = nullptr; ggml_tensor * filter = nullptr; ggml_tensor * output = nullptr; + ggml_tensor * projection = nullptr; ggml_tensor * cache = nullptr; ggml_cgraph * graph = nullptr; }; @@ -8066,6 +8227,7 @@ static SsmConvPrefillGraph make_ssm_conv_prefill_graph(ggml_type type, state = ggml_reshape_3d(c.ctx, state, 3, n, 1); } auto raw = ggml_mul_mat(c.ctx, c.weight, c.input); + c.projection = raw; auto x = ggml_reshape_3d(c.ctx, raw, n, 512, 1); auto window = ggml_concat(c.ctx, state, ggml_transpose(c.ctx, x), 0); auto tail = ggml_view_3d(c.ctx, window, 3, n, 1, window->nb[1], window->nb[2], 512 * sizeof(float)); @@ -8089,22 +8251,307 @@ static SsmConvPrefillGraph make_ssm_conv_prefill_graph(ggml_type type, return c; } -static void run_ssm_conv_prefill_recurrent_case(ggml_type type, bool gathered, bool clear_state) { - const int64_t k = 256; +static void run_qwen27_ssm_policy_route_case(ggml_type type, bool gathered) { + auto graph = make_ssm_conv_prefill_graph(type, 256, 8192, false, gathered, false); + for (bool row64_weight_policy : { false, true }) { + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph.graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151", row64_weight_policy })); + const ggml::hrx::Value * weight = imported.graph.values().find_tensor(graph.weight); + REQUIRE(weight != nullptr); + size_t matched = 0; + for (const auto & dispatch : scheduler.plan().dispatches) { + for (const auto & binding : dispatch.bindings) { + if (binding.value != weight->id) { + continue; + } + const bool q4_weight = type == GGML_TYPE_Q4_K; + const bool q6_row64 = !q4_weight && row64_weight_policy; + const char * expected_layout = q4_weight ? ggml::hrx::kQ4KPackedK256Row64Layout : + q6_row64 ? ggml::hrx::kQ6KPackedK256Row64ScaleRowLayout : ggml::hrx::kNativeWeightLayout; + REQUIRE(binding.layout == expected_layout); + const auto format = dispatch.kernel.compile_parameters.find("ggml.mul_mat.weight_format"); + REQUIRE(format != dispatch.kernel.compile_parameters.end()); + REQUIRE(format->second == (q4_weight ? "44" : (q6_row64 ? "46" : "6"))); + ++matched; + } + } + REQUIRE(matched == 1); + } + ggml_free(graph.ctx); +} + +static void run_qwen27_matmul_policy_route_case(ggml_type type, int64_t tokens) { + ggml_context * ctx = ggml_init({ 4 * 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, type, 5120, 10240); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 5120, tokens); + ggml_tensor * output = ggml_silu(ctx, ggml_mul_mat(ctx, weight, input)); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + for (bool row64_weight_policy : { false, true }) { + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151", row64_weight_policy })); + const ggml::hrx::Value * source = imported.graph.values().find_tensor(weight); + REQUIRE(source != nullptr); + const char * row64 = type == GGML_TYPE_Q4_K ? ggml::hrx::kQ4KPackedK256Row64Layout : + ggml::hrx::kQ6KPackedK256Row64ScaleRowLayout; + const char * expected = row64_weight_policy ? row64 : ggml::hrx::kNativeWeightLayout; + size_t matched = 0; + for (const auto & dispatch : scheduler.plan().dispatches) { + for (const auto & binding : dispatch.bindings) { + if (binding.value == source->id) { + REQUIRE(binding.layout == expected); + const bool q6_weight = type == GGML_TYPE_Q6_K; + const bool short_prompt = tokens == 66; + const bool short_q6 = q6_weight && short_prompt; + if (short_q6) { + const char * expected_kernel = row64_weight_policy ? + "loom_libs:ggml_mul_mat_q6_k_f32_wmma_prefill_wave32" : + "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32"; + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id, "gfx1151") == expected_kernel); + } + ++matched; + } + } + } + REQUIRE(matched == 1); + } + ggml_free(ctx); +} + +static void run_gemma4_decode_parent_f32_route_case(ggml_type type) { + ggml_context * ctx = ggml_init({ 4 * 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, type, 15360, 3840); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 15360, 1); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, input); + ggml_tensor * output = ggml_scale(ctx, projection, 0.5f); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, output); + for (const char * target : { "gfx1100", "gfx1151" }) { + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.has_index()); + const ggml::hrx::Value * projection_value = imported.graph.values().find_tensor(projection); + const ggml::hrx::Value * source = imported.graph.values().find_tensor(weight); + REQUIRE(projection_value != nullptr); + REQUIRE(source != nullptr); + REQUIRE(!imported.graph.index().consumers(projection_value->id).empty()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { target, false })); + size_t matched = 0; + for (const auto & dispatch : scheduler.plan().dispatches) { + for (const auto & binding : dispatch.bindings) { + if (binding.value != source->id) { + continue; + } + REQUIRE(binding.layout == ggml::hrx::kNativeWeightLayout); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id, target) == + "loom_libs:ggml_mul_mat_vector_f32_f32"); + const auto format = dispatch.kernel.compile_parameters.find("ggml.matmul.vector.weight_format"); + REQUIRE(format != dispatch.kernel.compile_parameters.end()); + REQUIRE(format->second == (type == GGML_TYPE_Q4_K ? "4" : "6")); + REQUIRE(dispatch.kernel.compile_parameters.find("ggml.matmul.vector.activation_format") == + dispatch.kernel.compile_parameters.end()); + ++matched; + } + } + REQUIRE(matched == 1); + } + ggml_free(ctx); +} + +static void run_gemma4_decode_parent_terminal_q6_route_case() { + ggml_context * ctx = ggml_init({ 4 * 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, 15360, 3840); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 15360, 1); + ggml_cgraph * graph = ggml_new_graph(ctx); + ggml_build_forward_expand(graph, ggml_mul_mat(ctx, weight, input)); + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151", false })); + const ggml::hrx::Value * source = imported.graph.values().find_tensor(weight); + REQUIRE(source != nullptr); + size_t matched = 0; + for (const auto & dispatch : scheduler.plan().dispatches) { + for (const auto & binding : dispatch.bindings) { + if (binding.value != source->id) { + continue; + } + REQUIRE(binding.layout == ggml::hrx::kQ6KPackedK256Row64ScaleRowLayout); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:ggml_mul_mat_vector_q6_f32_f32"); + const auto storage = dispatch.kernel.compile_parameters.find("ggml.matmul.vector.q6_storage"); + REQUIRE(storage != dispatch.kernel.compile_parameters.end()); + REQUIRE(storage->second == "0"); + ++matched; + } + } + REQUIRE(matched == 1); + ggml_free(ctx); +} + +static void require_qwen27_native_dot_reference(ggml_type type, const std::vector & weights, + const std::vector & input, const std::vector & observed, + int64_t k, int64_t n, const std::vector & rows, + const std::vector & tokens, const char * label, + bool silu_output = false) { + const ggml_type_traits * traits = ggml_get_type_traits(type); + REQUIRE(traits != nullptr && traits->to_float != nullptr); + const size_t row_bytes = ggml_row_size(type, k); + std::vector dequantized(static_cast(k)); + std::vector observed_samples; + std::vector reference_samples; + for (int64_t row : rows) { + traits->to_float(weights.data() + static_cast(row) * row_bytes, dequantized.data(), k); + for (int64_t token : tokens) { + double reference = 0.0; + for (int64_t column = 0; column < k; ++column) { + reference += static_cast(dequantized[static_cast(column)]) * + static_cast(input[static_cast(token * k + column)]); + } + if (silu_output) { + reference = reference / (1.0 + std::exp(-reference)); + } + observed_samples.push_back(observed[static_cast(token * n + row)]); + reference_samples.push_back(static_cast(reference)); + } + } + double squared_error = 0.0; + double squared_reference = 0.0; + double maximum_error = 0.0; + for (size_t i = 0; i < observed_samples.size(); ++i) { + const double difference = static_cast(observed_samples[i]) - reference_samples[i]; + squared_error += difference * difference; + squared_reference += static_cast(reference_samples[i]) * reference_samples[i]; + maximum_error = std::max(maximum_error, std::fabs(difference)); + } + REQUIRE(squared_reference > 0.0); + const double relative_l2 = std::sqrt(squared_error / squared_reference); + std::fprintf(stderr, "qwen27 %s %s native-dequant reference: relative_l2=%g max_abs=%g samples=%zu\n", + label, ggml_type_name(type), relative_l2, maximum_error, observed_samples.size()); + REQUIRE(relative_l2 <= 0.01); + require_close(observed_samples, reference_samples, 0.02f, 0.02f); +} + +static void run_qwen27_dense_native_reference_case(ggml_type type, int64_t tokens, int64_t k = 512, + int64_t n = 128) { + auto backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + ggml_backend_reg_t registry = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); + using set_policy_fn = bool (*)(ggml_backend_t, bool); + auto set_policy = reinterpret_cast( + ggml_backend_reg_get_proc_address(registry, "ggml_backend_hrx_set_row64_weight_policy")); + REQUIRE(set_policy != nullptr); + REQUIRE(set_policy(backend, true)); + auto * backend_state = static_cast(backend->context); + REQUIRE(backend_state != nullptr && backend_state->device != nullptr); + const std::string architecture = backend_state->device->architecture; + + ggml_context * ctx = ggml_init({ 1024 * 1024, nullptr, true }); + REQUIRE(ctx != nullptr); + ggml_tensor * weight = ggml_new_tensor_2d(ctx, type, k, n); + ggml_tensor * activation = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, k, tokens); + ggml_tensor * projection = ggml_mul_mat(ctx, weight, activation); + ggml_tensor * output = ggml_silu(ctx, projection); + 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, { architecture, true })); + const ggml::hrx::Value * source = imported.graph.values().find_tensor(weight); + REQUIRE(source != nullptr); + const char * expected_layout = type == GGML_TYPE_Q4_K ? ggml::hrx::kQ4KPackedK256Row64Layout : + ggml::hrx::kQ6KPackedK256Row64ScaleRowLayout; + const bool q6_weight = type == GGML_TYPE_Q6_K; + const bool short_minimum = tokens >= 64; + const bool short_maximum = tokens < 128; + const bool short_q6 = q6_weight && short_minimum && short_maximum; + const bool prefill_minimum = tokens >= 128; + const bool prefill_aligned = tokens % 128 == 0; + const bool aligned_prefill = prefill_minimum && prefill_aligned; + const bool aligned_q6_prefill = q6_weight && aligned_prefill; + const char * expected_kernel = tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : + short_q6 ? "loom_libs:ggml_mul_mat_q6_k_f32_wmma_prefill_wave32" : + aligned_q6_prefill ? "loom_libs:ggml_mul_mat_q6_k_f16_wmma_prefill_wave32" : + "loom_libs:ggml_mul_mat_tiled_input_f32_publish_f32"; + size_t matched = 0; + for (const auto & dispatch : scheduler.plan().dispatches) { + for (const auto & binding : dispatch.bindings) { + if (binding.value != source->id) { + continue; + } + REQUIRE(binding.layout == expected_layout); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id, architecture) == expected_kernel); + ++matched; + } + } + REQUIRE(matched == 1); + + std::vector weights = make_matmul_weight_bytes(type, k, n, 6); + auto weight_buffer = import_test_weight(backend, weight, weights); + ggml_backend_buffer_t graph_buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(graph_buffer != nullptr); + const std::vector input = make_pattern_f32(k * tokens, 7, 0.0625f); + set_tensor_bytes(backend, activation, input.data(), input.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + const auto observed = get_f32_tensor(backend, output); + std::vector rows; + for (int64_t row = 0; row < n; ++row) { + rows.push_back(row); + } + const std::vector sample_tokens = tokens == 1 ? std::vector{ 0 } : + std::vector{ 0, 21, 43, tokens - 1 }; + require_qwen27_native_dot_reference(type, weights, input, observed, k, n, rows, sample_tokens, + tokens == 1 ? "vector" : short_q6 ? "short-prefill" : "tiled", true); + ggml_backend_hrx_cache_stats stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &stats)); + REQUIRE(stats.weight_duplicate_warnings == 0); + REQUIRE(stats.weight_materializations == 1); + ggml_backend_buffer_free(graph_buffer); + weight_buffer.reset(); + ggml_free(ctx); + ggml_backend_free(backend); +} + +static void run_ssm_conv_prefill_recurrent_case(ggml_type type, bool gathered, bool clear_state, + bool row64_weight_policy = false) { + const int64_t k = row64_weight_policy ? 512 : 256; const int64_t n = 8192; auto backend = ggml_backend_hrx_init(0); REQUIRE(backend != nullptr); + if (row64_weight_policy) { + ggml_backend_reg_t registry = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); + using set_policy_fn = bool (*)(ggml_backend_t, bool); + auto set_policy = reinterpret_cast( + ggml_backend_reg_get_proc_address(registry, "ggml_backend_hrx_set_row64_weight_policy")); + REQUIRE(set_policy != nullptr); + REQUIRE(set_policy(backend, true)); + } auto a = make_ssm_conv_prefill_graph(type, k, n, false, gathered, clear_state); auto b = make_ssm_conv_prefill_graph(type, k, n, true, gathered, clear_state); - auto aw = allocate_test_weight(backend, a.weight); - auto bw = allocate_test_weight(backend, b.weight); + auto weights = make_matmul_weight_bytes(type, k, n, 6); + auto weights_copy = weights; + auto aw = row64_weight_policy ? import_test_weight(backend, a.weight, weights) : + allocate_test_weight(backend, a.weight); + auto bw = row64_weight_policy ? import_test_weight(backend, b.weight, weights_copy) : + allocate_test_weight(backend, b.weight); auto ab = ggml_backend_alloc_ctx_tensors(a.ctx, backend); auto bb = ggml_backend_alloc_ctx_tensors(b.ctx, backend); REQUIRE(ab != nullptr && bb != nullptr); - const auto weights = make_matmul_weight_bytes(type, k, n, 6); const auto filters = make_pattern_f32(4 * n, 11, 0.0625f); const auto states = make_pattern_f32(6 * n, 17, 0.03125f); - set_tensor_pair_bytes(backend, a.weight, backend, b.weight, weights.data(), weights.size()); + if (!row64_weight_policy) { + set_tensor_pair_bytes(backend, a.weight, backend, b.weight, weights.data(), weights.size()); + } set_tensor_pair_bytes(backend, a.filter, backend, b.filter, filters.data(), filters.size() * sizeof(float)); set_tensor_pair_bytes(backend, a.pool, backend, b.pool, states.data(), states.size() * sizeof(float)); set_tensor_pair_bytes(backend, a.initial, backend, b.initial, states.data(), 3 * n * sizeof(float)); @@ -8120,6 +8567,23 @@ static void run_ssm_conv_prefill_recurrent_case(ggml_type type, bool gathered, b const auto expected = get_f32_tensor(backend, b.output); require_close(get_f32_tensor(backend, a.cache), get_f32_tensor(backend, b.cache), 0.0f, 0.0f); require_close(actual, expected, 0.0f, 0.0f); + if (row64_weight_policy && round == 0) { + const auto projection = get_f32_tensor(backend, b.projection); + std::vector rows; + for (int64_t row = 0; row < 64; ++row) { + rows.push_back(row); + } + rows.push_back(4095); + rows.push_back(8191); + require_qwen27_native_dot_reference(type, weights, input, projection, k, n, rows, + { 0, 127, 255, 511 }, "SSM"); + } + } + if (row64_weight_policy) { + ggml_backend_hrx_cache_stats stats = {}; + REQUIRE(ggml_backend_hrx_get_cache_stats(backend, &stats)); + REQUIRE(stats.weight_duplicate_warnings == 0); + REQUIRE(stats.weight_materializations == 2); } ggml_backend_buffer_free(ab); ggml_backend_buffer_free(bb); @@ -8181,7 +8645,345 @@ static std::string type_name(ggml_type type) { using test_runner::Suite; +static void run_qwen30b_attention_weight_layout_scheduling_case() { + 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; + const std::string original_env_value = had_original_env ? original_env : ""; + struct WeightCase { + const char * role; + ggml_type type; + int64_t input_size; + int64_t output_size; + }; + const WeightCase cases[] = { + { "attn_q.weight", GGML_TYPE_Q4_K, 2048, 4096 }, + { "attn_k.weight", GGML_TYPE_Q4_K, 2048, 512 }, + { "attn_v.weight", GGML_TYPE_Q4_K, 2048, 512 }, + { "attn_v.weight", GGML_TYPE_Q6_K, 2048, 512 }, + { "attn_output.weight", GGML_TYPE_Q4_K, 4096, 2048 }, + }; + + for (const WeightCase & test : cases) { + for (const int64_t token_count : { int64_t{ 1 }, int64_t{ 44 }, int64_t{ 512 }, int64_t{ 1024 } }) { + for (const bool qwen_profile : { true, false }) { + for (const bool qwen_dispatch_disabled : { false, true }) { + if (qwen_dispatch_disabled) { + REQUIRE(setenv(kDisableQwenDispatchEnv, "1", 1) == 0); + } else { + REQUIRE(unsetenv(kDisableQwenDispatchEnv) == 0); + } + 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, test.type, test.input_size, test.output_size); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, test.input_size, token_count); + REQUIRE(weight != nullptr); + REQUIRE(input != nullptr); + const std::string prefix = "blk.0."; + const std::string name = prefix + test.role; + ggml_set_name(weight, name.c_str()); + const auto peer = [&](const char * role, ggml_type type, int64_t ne0, int64_t ne1) { + if (std::strcmp(test.role, role) == 0) { + return static_cast(nullptr); + } + ggml_tensor * tensor = ggml_new_tensor_2d(ctx, type, ne0, ne1); + REQUIRE(tensor != nullptr); + const std::string peer_name = prefix + role; + ggml_set_name(tensor, peer_name.c_str()); + return tensor; + }; + ggml_tensor * peers[] = { + peer("attn_q.weight", GGML_TYPE_Q4_K, 2048, 4096), + peer("attn_k.weight", GGML_TYPE_Q4_K, 2048, 512), + peer("attn_v.weight", test.type == GGML_TYPE_Q6_K ? GGML_TYPE_Q6_K : GGML_TYPE_Q4_K, + 2048, 512), + peer("attn_output.weight", GGML_TYPE_Q4_K, 4096, 2048), + }; + ggml_tensor * experts = ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, 2048, 768, + qwen_profile ? 128 : 64); + REQUIRE(experts != nullptr); + ggml_set_name(experts, "blk.0.ffn_gate_exps.weight"); + 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()); + // Model graphs import these peer weights through their own layer ops. + // Add them directly here to keep this scheduling test focused on one matmul. + for (ggml_tensor * peer_weight : peers) { + if (peer_weight != nullptr) { + imported.graph.values().get_or_add_tensor_value(peer_weight, ggml::hrx::ValueKind::External); + } + } + imported.graph.values().get_or_add_tensor_value(experts, ggml::hrx::ValueKind::External); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151", true })); + const ggml::hrx::Value * value = imported.graph.values().find_tensor(weight); + REQUIRE(value != nullptr); + const bool native_profile = qwen_profile && !qwen_dispatch_disabled; + const char * expected_layout = native_profile ? ggml::hrx::kNativeWeightLayout : + test.type == GGML_TYPE_Q4_K ? ggml::hrx::kQ4KPackedK256Row64Layout : + ggml::hrx::kQ6KPackedK256Row64ScaleRowLayout; + size_t binding_count = 0; + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + for (const ggml::hrx::DispatchBinding & binding : dispatch.bindings) { + if (binding.value == value->id) { + REQUIRE(binding.layout == expected_layout); + const bool native_q4_prefill = native_profile && test.type == GGML_TYPE_Q4_K && + token_count >= 512; + if (native_q4_prefill) { + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:ggml_mul_mat_q4_k_f16_wmma_prefill_wave32"); + REQUIRE(dispatch.kernel.compile_parameters.at("ggml.mul_mat.weight_format") == "4"); + } + ++binding_count; + } + } + } + REQUIRE(binding_count > 0); + ggml_free(ctx); + } + } + } + } + + // LFM2.5-1.2B also has Q4 K/V weights [2048,512], but Q and O are + // [2048,2048] and the layer has no Qwen30B expert tensor. + REQUIRE(unsetenv(kDisableQwenDispatchEnv) == 0); + for (const int64_t token_count : { int64_t{ 1 }, int64_t{ 44 } }) { + 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_Q4_K, 2048, 512); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2048, token_count); + ggml_set_name(weight, "blk.2.attn_k.weight"); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + 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()); + for (const char * role : { "attn_q.weight", "attn_output.weight" }) { + ggml_tensor * peer = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, 2048, 2048); + const std::string name = std::string("blk.2.") + role; + ggml_set_name(peer, name.c_str()); + imported.graph.values().get_or_add_tensor_value(peer, ggml::hrx::ValueKind::External); + } + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1151", true })); + const ggml::hrx::Value * value = imported.graph.values().find_tensor(weight); + REQUIRE(value != nullptr); + size_t binding_count = 0; + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + for (const ggml::hrx::DispatchBinding & binding : dispatch.bindings) { + if (binding.value == value->id) { + REQUIRE(binding.layout == ggml::hrx::kQ4KPackedK256Row64Layout); + ++binding_count; + } + } + } + REQUIRE(binding_count > 0); + ggml_free(ctx); + } + restore_environment_value(kDisableQwenDispatchEnv, had_original_env, original_env_value); +} + +static void run_qwen30b_native_q4_prefill_reference_case_for(const char * role, + int64_t input_size, + int64_t output_size, + int64_t token_count) { + const std::vector weight_bytes = make_matmul_weight_bytes(GGML_TYPE_Q4_K, input_size, output_size, 53); + const std::vector input_bytes = make_pattern_f32(input_size * token_count, 29, 0.01f); + + 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_Q4_K, input_size, output_size); + ggml_tensor * input = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, input_size, token_count); + ggml_tensor * output = ggml_mul_mat(ctx, weight, input); + const std::string weight_name = std::string("blk.0.") + role; + ggml_set_name(weight, weight_name.c_str()); + const auto peer = [&](const char * peer_role, int64_t ne0, int64_t ne1) { + if (std::strcmp(role, peer_role) == 0) { + return static_cast(nullptr); + } + ggml_tensor * tensor = ggml_new_tensor_2d(ctx, GGML_TYPE_Q4_K, ne0, ne1); + const std::string name = std::string("blk.0.") + peer_role; + ggml_set_name(tensor, name.c_str()); + return tensor; + }; + ggml_tensor * query = peer("attn_q.weight", 2048, 4096); + ggml_tensor * key = peer("attn_k.weight", 2048, 512); + ggml_tensor * value = peer("attn_v.weight", 2048, 512); + ggml_tensor * projection = peer("attn_output.weight", 4096, 2048); + ggml_tensor * experts = ggml_new_tensor_3d(ctx, GGML_TYPE_Q4_K, 2048, 768, 128); + ggml_set_name(experts, "blk.0.ffn_gate_exps.weight"); + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, output); + // Active aliases keep peer weights in the imported graph without extra matmul work. + for (ggml_tensor * peer_weight : { query, key, value, projection, experts }) { + if (peer_weight != nullptr) { + ggml_tensor * marker = ggml_reshape_1d(ctx, peer_weight, ggml_nelements(peer_weight)); + ggml_build_forward_expand(graph, marker); + } + } + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + REQUIRE(imported.graph.values().find_tensor(experts) != nullptr); + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, { "gfx1100", true })); + const ggml::hrx::Value * weight_value = imported.graph.values().find_tensor(weight); + REQUIRE(weight_value != nullptr); + bool found_native_prefill = false; + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + for (const ggml::hrx::DispatchBinding & binding : dispatch.bindings) { + if (binding.value == weight_value->id) { + REQUIRE(binding.layout == ggml::hrx::kNativeWeightLayout); + REQUIRE(kernel_name_for_id(dispatch.kernel.kernel_id) == + "loom_libs:ggml_mul_mat_q4_k_f16_wmma_prefill_wave32"); + REQUIRE(dispatch.kernel.compile_parameters.at("ggml.mul_mat.weight_format") == "4"); + found_native_prefill = true; + } + } + } + REQUIRE(found_native_prefill); + + ggml_backend_t backend = ggml_backend_hrx_init(0); + REQUIRE(backend != nullptr); + ggml_backend_reg_t registry = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); + using set_policy_fn = bool (*)(ggml_backend_t, bool); + auto set_policy = reinterpret_cast( + ggml_backend_reg_get_proc_address(registry, "ggml_backend_hrx_set_row64_weight_policy")); + REQUIRE(set_policy != nullptr); + REQUIRE(set_policy(backend, true)); + ggml_backend_buffer_t buffer = ggml_backend_alloc_ctx_tensors(ctx, backend); + REQUIRE(buffer != nullptr); + ggml_backend_tensor_set(weight, weight_bytes.data(), 0, weight_bytes.size()); + ggml_backend_tensor_set(input, input_bytes.data(), 0, input_bytes.size() * sizeof(float)); + REQUIRE(ggml_backend_graph_compute(backend, graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(backend); + const std::vector actual = get_f32_tensor(backend, output); + + const auto * blocks = reinterpret_cast(weight_bytes.data()); + std::vector decoded(input_size); + const std::array selected_tokens = { 0, 127, token_count / 2 - 1, token_count / 2, token_count - 1 }; + std::array, 5> input_f16; + for (size_t sample = 0; sample < selected_tokens.size(); ++sample) { + input_f16[sample].resize(input_size); + for (int64_t k = 0; k < input_size; ++k) { + const float input_value = input_bytes[selected_tokens[sample] * input_size + k]; + input_f16[sample][k] = ggml_fp16_to_fp32(ggml_fp32_to_fp16(input_value)); + } + } + double reference_energy = 0.0; + double error_energy = 0.0; + for (int64_t channel = 0; channel < output_size; ++channel) { + dequantize_row_q4_K(blocks + channel * (input_size / QK_K), decoded.data(), input_size); + std::array expected = {}; + for (int64_t k = 0; k < input_size; ++k) { + const float weight_f16 = ggml_fp16_to_fp32(ggml_fp32_to_fp16(decoded[k])); + for (size_t sample = 0; sample < selected_tokens.size(); ++sample) { + expected[sample] = std::fma(weight_f16, input_f16[sample][k], expected[sample]); + } + } + for (size_t sample = 0; sample < selected_tokens.size(); ++sample) { + const float result = actual[selected_tokens[sample] * output_size + channel]; + const float error = result - expected[sample]; + REQUIRE(std::isfinite(result)); + REQUIRE(std::fabs(error) <= 0.03f + 0.01f * std::fabs(expected[sample])); + reference_energy += static_cast(expected[sample]) * expected[sample]; + error_energy += static_cast(error) * error; + } + } + REQUIRE(reference_energy > 1.0); + REQUIRE(error_energy / reference_energy < 0.0004); + + ggml_backend_buffer_free(buffer); + ggml_backend_free(backend); + ggml_free(ctx); +} + +static void run_qwen30b_native_q4_prefill_reference_case() { + run_qwen30b_native_q4_prefill_reference_case_for("attn_k.weight", 2048, 512, 512); + run_qwen30b_native_q4_prefill_reference_case_for("attn_q.weight", 2048, 4096, 512); + run_qwen30b_native_q4_prefill_reference_case_for("attn_output.weight", 4096, 2048, 1024); +} + +static void run_qwen17b_tied_embedding_layout_policy_case() { + struct WeightCase { + const char * name; + ggml_type type; + int64_t hidden_size; + int64_t vocabulary_count; + bool eligible; + }; + const WeightCase cases[] = { + { "token_embd.weight", GGML_TYPE_Q6_K, 2048, 151936, false }, + { "token_embd.weight", GGML_TYPE_Q4_K, 2048, 151936, true }, + { "output.weight", GGML_TYPE_Q6_K, 2048, 151936, true }, + { "token_embd.weight", GGML_TYPE_Q6_K, 2048, 151552, true }, + }; + ggml_init_params params = {}; + params.mem_size = 4 * 1024 * 1024; + params.no_alloc = true; + ggml_context * ctx = ggml_init(params); + REQUIRE(ctx != nullptr); + ggml::hrx::Graph graph; + for (const WeightCase & test : cases) { + ggml_tensor * tensor = ggml_new_tensor_2d(ctx, test.type, test.hidden_size, test.vocabulary_count); + REQUIRE(tensor != nullptr); + ggml_set_name(tensor, test.name); + const ggml::hrx::ValueId id = graph.values().get_or_add_tensor_value(tensor, ggml::hrx::ValueKind::External); + const ggml::hrx::Value * value = graph.values().find(id); + REQUIRE(value != nullptr); + REQUIRE(ggml::hrx::dispatch_weight_layout_eligible_for_graph(graph, *value, true) == test.eligible); + REQUIRE(ggml::hrx::dispatch_weight_layout_eligible_for_graph(graph, *value, false) == + ggml::hrx::dispatch_weight_layout_eligible(*value)); + } + ggml_free(ctx); +} + static void register_basic_ops_cases(Suite & suite) { + suite.host_case("qwen17b_tied_embedding.layout_policy", [] { + run_qwen17b_tied_embedding_layout_policy_case(); + }); + for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { + suite.host_case("gemma4_decode.parent_f32_route." + type_name(type), + [type] { run_gemma4_decode_parent_f32_route_case(type); }); + } + suite.host_case("gemma4_decode.parent_terminal_q6_route.gfx1151", [] { + run_gemma4_decode_parent_terminal_q6_route_case(); + }); + for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { + suite.host_case("qwen27.ssm_policy_route." + type_name(type) + ".plain", + [type] { run_qwen27_ssm_policy_route_case(type, false); }); + suite.host_case("qwen27.ssm_policy_route." + type_name(type) + ".interior", + [type] { run_qwen27_ssm_policy_route_case(type, true); }); + } + suite.host_case("qwen27.q4_decode_policy_route", [] { + run_qwen27_matmul_policy_route_case(GGML_TYPE_Q4_K, 1); + }); + suite.host_case("qwen27.q6_decode_policy_route", [] { + run_qwen27_matmul_policy_route_case(GGML_TYPE_Q6_K, 1); + }); + suite.host_case("qwen27.q6_prefill512_policy_route", [] { + run_qwen27_matmul_policy_route_case(GGML_TYPE_Q6_K, 512); + }); + suite.host_case("qwen27.q6_short66_policy_route", [] { + run_qwen27_matmul_policy_route_case(GGML_TYPE_Q6_K, 66); + }); 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(); }); @@ -8201,6 +9003,36 @@ static void register_basic_ops_cases(Suite & suite) { 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("qwen27.ssm_conv_q6_imported_numerics", [] { + run_ssm_conv_prefill_recurrent_case(GGML_TYPE_Q6_K, false, false, true); + }); + suite.device_case("qwen27.ssm_conv_q4_interior_imported_numerics", [] { + run_ssm_conv_prefill_recurrent_case(GGML_TYPE_Q4_K, true, false, true); + }); + for (const ggml_type type : { GGML_TYPE_Q4_K, GGML_TYPE_Q6_K }) { + suite.device_case("qwen27.native_dequant." + type_name(type) + ".vector", [type] { + run_qwen27_dense_native_reference_case(type, 1); + }); + suite.device_case("qwen27.native_dequant." + type_name(type) + ".tiled", [type] { + run_qwen27_dense_native_reference_case(type, 66); + }); + } + for (const int64_t tokens : { 63, 64, 127, 128, 129 }) { + suite.device_case("qwen27.q6_short_boundary.tokens" + std::to_string(tokens), [tokens] { + run_qwen27_dense_native_reference_case(GGML_TYPE_Q6_K, tokens); + }); + } + for (const bool varying_input : { false, true }) { + suite.device_case("qwen27.q4_metadata_paired.tokens66." + + std::string(varying_input ? "varied" : "constant"), [varying_input] { + run_dense_matmul_swiglu_cpu_reference_case(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, + "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", + 66, 192, GGML_GLU_OP_SWIGLU, 768, false, true, true, varying_input); + }); + } + suite.device_case("qwen27.q4_metadata_single.tokens66.input768.outputs192", [] { + run_qwen27_dense_native_reference_case(GGML_TYPE_Q4_K, 66, 768, 192); + }); 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), @@ -8235,6 +9067,12 @@ static void register_basic_ops_cases(Suite & suite) { } static void register_dense_matmul_cases(Suite & suite) { + suite.host_case("qwen30b_attention_weight_layout.scheduling", [] { + run_qwen30b_attention_weight_layout_scheduling_case(); + }); + suite.device_case("qwen30b_native_q4_prefill.reference", [] { + run_qwen30b_native_q4_prefill_reference_case(); + }); // 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 }) { @@ -8461,6 +9299,16 @@ static void register_dense_matmul_cases(Suite & suite) { GGML_TYPE_Q4_K, GGML_TYPE_F16, "loom_libs:ggml_mul_mat_swiglu_f32_f32_lowtoken_dot", 2, 128, op); }); } + for (const int64_t tokens : { int64_t{ 1 }, int64_t{ 35 }, int64_t{ 512 }, int64_t{ 1024 } }) { + suite.device_case("dense_matmul_packed_glu.q4_k.imported_native.tokens" + std::to_string(tokens), [tokens] { + run_dense_matmul_packed_glu_cpu_reference_case( + GGML_TYPE_Q4_K, + tokens >= 512 ? "loom_libs:ggml_mul_mat_q4_k_f16_wmma_prefill_wave32" : + tokens == 1 ? "loom_libs:ggml_mul_mat_vector_f32_f32" : + "loom_libs:ggml_mul_mat_tiled_pair_input_f32_binary_publish_f32", + tokens, 128, GGML_GLU_OP_SWIGLU, false, 512, true); + }); + } 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);