Skip to content

ggml-hrx: avoid duplicate Qwen weights with shared row64 layouts #181

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
AaronStGeorge wants to merge 11 commits into
base: hrx-graph-develop-v2
Choose a base branch
from
Draft

ggml-hrx: avoid duplicate Qwen weights with shared row64 layouts #181

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
6d36a01
ggml-hrx: keep Qwen3.8 27B quantized weights in one layout
AaronStGeorge Nov 4, 2026
45bca49
ggml-hrx: preserve Q4 conv reader and native packed GLU routes
AaronStGeorge Nov 4, 2026
2e9335b
test: cover Q4 SSM policy routes and two-block Q6 reference
AaronStGeorge Nov 4, 2026
5654df2
test: verify Qwen27 policy-on native dequant consumers
AaronStGeorge Nov 4, 2026
9e7a385
ggml-hrx: enable row64 weight policy outside Gemma4
AaronStGeorge Nov 4, 2026
f45fecf
test: seal Gemma routes and Q4 prefill format producers
AaronStGeorge Nov 4, 2026
21f9406
Pass layout policy context to paired SwiGLU route helper
AaronStGeorge Nov 4, 2026
28ebc27
Align Gemma and packed GLU controls with parent routes
AaronStGeorge Nov 4, 2026
c37cebc
Resolve Gemma route kernel against selected target
AaronStGeorge Nov 4, 2026
7485439
hrx: restore Qwen short Q6 prefill and Q4 metadata reuse
AaronStGeorge Nov 4, 2026
5d4d898
tests: clarify Qwen prefill route facts
AaronStGeorge Nov 4, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 18 additions & 0 deletions ggml/src/ggml-hrx/backend-context.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<std::mutex> lock(layout_policy_mutex);
if (layout_policy_locked) {
return false;
}
row64_weight_policy = enabled;
return true;
}

bool lock_layout_policy() {
std::lock_guard<std::mutex> lock(layout_policy_mutex);
layout_policy_locked = true;
return row64_weight_policy;
}
};

struct ggml_backend_hrx_reg_context {
Expand Down
9 changes: 5 additions & 4 deletions ggml/src/ggml-hrx/dispatch/dispatch-scheduler.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,11 +31,12 @@ static bool try_match_registration(const Graph & graph,
const std::vector<bool> & 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, &registry,
graph, node, node_index, covered_nodes, plan, next_plan_value, &registry, target.row64_weight_policy,
};
return registry.match(context, match, diagnostics);
}
Expand Down Expand Up @@ -197,7 +198,7 @@ bool DispatchScheduler::schedule_graph(Graph & graph,
const ValueId next_plan_value(static_cast<int32_t>(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);
Expand Down Expand Up @@ -276,7 +277,7 @@ bool DispatchScheduler::schedule_graph(Graph & graph,
const ValueId next_plan_value(static_cast<int32_t>(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];
Expand Down Expand Up @@ -309,7 +310,7 @@ bool DispatchScheduler::supports_node(const Graph & graph, const GraphNode * nod
DispatchMatch match;
const ValueId next_plan_value(static_cast<int32_t>(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);
}

Expand Down
96 changes: 74 additions & 22 deletions ggml/src/ggml-hrx/dispatch_registration/common/dispatch-gated-mul-mat.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) ||
Expand Down Expand Up @@ -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;
}
Expand All @@ -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;
Expand All @@ -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 {};
}

Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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<int64_t>(binary_kind_config_value(match.op))));
if (postops != nullptr) {
Expand All @@ -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;
Expand Down Expand Up @@ -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;
}

Expand Down Expand Up @@ -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{};
Expand Down Expand Up @@ -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<int64_t>(binary_kind_config_value(match.op))));
dispatch.bindings.push_back(activation);
Expand All @@ -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 });

Expand Down Expand Up @@ -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",
Expand All @@ -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<int64_t>(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,
Expand Down
Loading