Skip to content

hrx: the Q4_K/Q6_K prefill reads its f16 activation copy at 2^-6, so an activation above 65504 stays finite (engine #1962) #298

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

Merged
bong-water-water-bong merged 1 commit into from
Nov 11, 2026
Merged

hrx: the Q4_K/Q6_K prefill reads its f16 activation copy at 2^-6, so an activation above 65504 stays finite (engine #1962) #298

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
15 changes: 9 additions & 6 deletions ggml/src/ggml-hrx/dispatch/command-plan-metadata.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,13 @@ static bool metadata_matches(const CommandPlanResourceMetadata & lhs, const Comm
static bool generated_resource_matches(const CommandPlanGeneratedResource & lhs,
const CommandPlanGeneratedResource & rhs) {
return lhs.source_value == rhs.source_value && lhs.role == rhs.role && lhs.generated_value == rhs.generated_value &&
lhs.byte_count == rhs.byte_count && metadata_matches(lhs.metadata, rhs.metadata);
lhs.byte_count == rhs.byte_count && lhs.scale_log2 == rhs.scale_log2 &&
metadata_matches(lhs.metadata, rhs.metadata);
}

static bool alternate_value_matches(const CommandPlanAlternateValue & lhs, const CommandPlanAlternateValue & rhs) {
return lhs.graph_value == rhs.graph_value && lhs.alternate_value == rhs.alternate_value && lhs.type == rhs.type &&
lhs.byte_count == rhs.byte_count && lhs.name == rhs.name;
lhs.byte_count == rhs.byte_count && lhs.name == rhs.name && lhs.scale_log2 == rhs.scale_log2;
}

static bool moe_routing_bundle_matches(const CommandPlanMoeRoutingBundle & lhs,
Expand Down Expand Up @@ -77,7 +78,7 @@ bool CommandPlanMetadata::append_generated_resource(CommandPlanGeneratedResource
bool CommandPlanMetadata::append_alternate_value(CommandPlanAlternateValue alternate, Status & status) {
for (const CommandPlanAlternateValue & existing : alternate_values_) {
if (existing.graph_value == alternate.graph_value && existing.type == alternate.type &&
existing.byte_count == alternate.byte_count) {
existing.byte_count == alternate.byte_count && existing.scale_log2 == alternate.scale_log2) {
if (alternate_value_matches(existing, alternate)) {
return true;
}
Expand Down Expand Up @@ -116,7 +117,7 @@ const CommandPlanGeneratedResource * CommandPlanMetadata::find_generated_resourc

const CommandPlanAlternateValue * CommandPlanMetadata::find_alternate_value(ValueId graph_value) const {
for (const CommandPlanAlternateValue & alternate : alternate_values_) {
if (alternate.graph_value == graph_value) {
if (alternate.graph_value == graph_value && alternate.scale_log2 == 0) {
return &alternate;
}
}
Expand All @@ -125,9 +126,11 @@ const CommandPlanAlternateValue * CommandPlanMetadata::find_alternate_value(Valu

const CommandPlanAlternateValue * CommandPlanMetadata::find_alternate_value(ValueId graph_value,
ggml_type type,
size_t byte_count) const {
size_t byte_count,
int scale_log2) const {
for (const CommandPlanAlternateValue & alternate : alternate_values_) {
if (alternate.graph_value == graph_value && alternate.type == type && alternate.byte_count == byte_count) {
if (alternate.graph_value == graph_value && alternate.type == type && alternate.byte_count == byte_count &&
alternate.scale_log2 == scale_log2) {
return &alternate;
}
}
Expand Down
8 changes: 7 additions & 1 deletion ggml/src/ggml-hrx/dispatch/command-plan-metadata.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,8 @@ struct CommandPlanGeneratedResource {
ValueId generated_value;
size_t byte_count = 0;
CommandPlanResourceMetadata metadata;
// The stored values are the source values * 2^-scale_log2 (0 = as they are). engine #1962
int scale_log2 = 0;
};

struct CommandPlanAlternateValue {
Expand All @@ -84,6 +86,9 @@ struct CommandPlanAlternateValue {
ggml_type type = GGML_TYPE_COUNT;
size_t byte_count = 0;
std::string name;
// The stored values are the graph values * 2^-scale_log2 (0 = as they are). A reader that does not ask for a
// scale only ever finds an unscaled alternate. engine #1962
int scale_log2 = 0;
};

struct CommandPlanMoeRoutingBundle {
Expand Down Expand Up @@ -118,7 +123,8 @@ class CommandPlanMetadata {

const CommandPlanAlternateValue * find_alternate_value(ValueId graph_value,
ggml_type type,
size_t byte_count) const;
size_t byte_count,
int scale_log2 = 0) const;

const CommandPlanMoeRoutingBundle * find_moe_routing_bundle(ValueId route_ids) const;

Expand Down
12 changes: 7 additions & 5 deletions ggml/src/ggml-hrx/dispatch/command-plan.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,8 +51,9 @@ inline const CommandPlanAlternateValue * find_alternate_value(const CommandPlan
inline const CommandPlanAlternateValue * find_alternate_value(const CommandPlan & plan,
ValueId graph_value,
ggml_type type,
size_t byte_count) {
return plan.metadata.find_alternate_value(graph_value, type, byte_count);
size_t byte_count,
int scale_log2 = 0) {
return plan.metadata.find_alternate_value(graph_value, type, byte_count, scale_log2);
}

inline bool same_full_value_range(const Value & lhs, const Value & rhs) {
Expand All @@ -63,8 +64,9 @@ inline const CommandPlanAlternateValue * find_alternate_value(const Graph &
const CommandPlan & plan,
ValueId graph_value,
ggml_type type,
size_t byte_count) {
const CommandPlanAlternateValue * exact = find_alternate_value(plan, graph_value, type, byte_count);
size_t byte_count,
int scale_log2 = 0) {
const CommandPlanAlternateValue * exact = find_alternate_value(plan, graph_value, type, byte_count, scale_log2);
if (exact != nullptr) {
return exact;
}
Expand All @@ -79,7 +81,7 @@ inline const CommandPlanAlternateValue * find_alternate_value(const Graph &
if (candidate == nullptr || !same_full_value_range(*value, *candidate)) {
return nullptr;
}
return find_alternate_value(plan, candidate_id, type, byte_count);
return find_alternate_value(plan, candidate_id, type, byte_count, scale_log2);
};

ValueId alias = value->alias_source;
Expand Down
14 changes: 10 additions & 4 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 @@ -586,11 +586,12 @@ static bool match_mul_mat_swiglu_q4_q8_prefill_dispatch(const DispatchMatchConte
const bool packed_input = match.token_count % 512 == 0 &&
(match.token_count / 512) * (match.output_size / 64) >= 64;
DispatchBinding activation;
int activation_scale_log2 = 0;
const bool prepared = packed_input ?
common_prepare_k16_major_f16_input(context, *match.input, match.input_size, match.token_count,
dispatch_match, activation) :
dispatch_match, activation, activation_scale_log2) :
common_prepare_f16_input(context, *match.input, match.input_size, match.token_count,
dispatch_match, activation);
dispatch_match, activation, activation_scale_log2);
if (!prepared) {
return false;
}
Expand All @@ -605,6 +606,7 @@ static bool match_mul_mat_swiglu_q4_q8_prefill_dispatch(const DispatchMatchConte

Dispatch dispatch;
dispatch.kernel = make_kernel_specialization(kMulMatSwiGLUQ4Q8PrefillKernel);
common_set_f16_prefill_activation_scale(dispatch, activation_scale_log2);
dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.f16_output_layout", packed_output ? "1" : "0");
dispatch.kernel.compile_parameters.emplace("ggml.mul_mat_swiglu.op",
common_to_config_value(static_cast<int64_t>(binary_kind_config_value(match.op))));
Expand All @@ -626,12 +628,16 @@ static bool match_mul_mat_swiglu_q4_q8_prefill_dispatch(const DispatchMatchConte
const char * f16_name = packed_output ? "common.mul_mat_swiglu.k16_major_f16" : "common.mul_mat_swiglu.f16";
dispatch_match.transients.push_back({ f16_output, f16_name, f16_bytes, 256 });
dispatch.bindings.push_back({ f16_output, 0, f16_bytes });
// The kernel publishes this f16 copy of the fused result as value * 2^-6 (the result is the next projection's
// activation, and a gated product is where a model leaves the f16 range: engine #1539, #1962).
Status status;
const bool recorded = packed_output ?
dispatch_match.metadata.append_generated_resource(
{ match.output->id, GeneratedResourceRole::F16K16Major, f16_output, f16_bytes, {} }, status) :
{ match.output->id, GeneratedResourceRole::F16K16Major, f16_output, f16_bytes, {},
kCommonF16ActivationScaleLog2 }, status) :
dispatch_match.metadata.append_alternate_value(
{ match.output->id, f16_output, GGML_TYPE_F16, f16_bytes, f16_name }, status);
{ match.output->id, f16_output, GGML_TYPE_F16, f16_bytes, f16_name, kCommonF16ActivationScaleLog2 },
status);
if (!recorded) {
dispatch_match.status.append(status);
return false;
Expand Down
61 changes: 46 additions & 15 deletions ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-common.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -552,19 +552,40 @@ inline CommonMulMatMatch common_match_mul_mat_any_format(const Graph & graph
return match;
}

// Share the conversion used by F16-operand matmuls.
// An f16 activation copy holds value * 2^-kCommonF16ActivationScaleLog2. A matmul narrows its f32 activation to f16
// for WMMA, and f16 ends at 65504: a larger activation became inf and the product NaN (engine #1962, the prefill
// form of #1539). Prescaled by 2^-6 the copy represents [3.9e-3, 4.19e6] at full precision; the kernel that reads it
// scales its f32 accumulator back by 2^6, which is exact. The scale travels with the alternate value / generated
// resource, so a reader always knows which form it holds.
constexpr int kCommonF16ActivationScaleLog2 = 6;
constexpr const char * kCommonF16PrefillActivationScaleParameter = "ggml.f16_prefill.activation_scale_log2";

// Tell a kernel built on ggml_mul_mat_quantized_f16_prefill_* which scale its f16 activation carries.
inline void common_set_f16_prefill_activation_scale(Dispatch & dispatch, int scale_log2) {
if (scale_log2 != 0) {
dispatch.kernel.compile_parameters.emplace(kCommonF16PrefillActivationScaleParameter,
common_to_config_value(static_cast<int64_t>(scale_log2)));
}
}

// Share the conversion used by F16-operand matmuls. `scale_log2` returns the scale of the bound copy: the copy made
// here (and shared by later readers) is prescaled; an unscaled F16 value another kernel published is used as it is.
inline bool common_prepare_f16_input(const DispatchMatchContext & context,
const Value & input,
int64_t input_size,
int64_t token_count,
DispatchMatch & match,
DispatchBinding & binding) {
DispatchBinding & binding,
int & scale_log2) {
const size_t bytes = static_cast<size_t>(token_count * input_size) * sizeof(ggml_fp16_t);
const CommandPlanAlternateValue * alternate =
find_alternate_value(context.graph, context.plan, input.id, GGML_TYPE_F16, bytes);
if (alternate != nullptr) {
binding = { alternate->alternate_value, 0, bytes };
return true;
for (const int candidate_scale : { kCommonF16ActivationScaleLog2, 0 }) {
const CommandPlanAlternateValue * alternate =
find_alternate_value(context.graph, context.plan, input.id, GGML_TYPE_F16, bytes, candidate_scale);
if (alternate != nullptr) {
binding = { alternate->alternate_value, 0, bytes };
scale_log2 = candidate_scale;
return true;
}
}

constexpr KernelCatalogRef kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_copy_f32_f16");
Expand All @@ -578,11 +599,13 @@ inline bool common_prepare_f16_input(const DispatchMatchContext & context,
convert.bindings.push_back({ activation, 0, bytes });
match.dispatches.push_back(std::move(convert));
Status status;
if (!match.metadata.append_alternate_value({ input.id, activation, GGML_TYPE_F16, bytes, name }, status)) {
if (!match.metadata.append_alternate_value(
{ input.id, activation, GGML_TYPE_F16, bytes, name, kCommonF16ActivationScaleLog2 }, status)) {
match.status.append(status);
return false;
}
binding = { activation, 0, bytes };
binding = { activation, 0, bytes };
scale_log2 = kCommonF16ActivationScaleLog2;
return true;
}

Expand All @@ -592,18 +615,25 @@ inline bool common_prepare_k16_major_f16_input(const DispatchMatchContext & cont
int64_t input_size,
int64_t token_count,
DispatchMatch & match,
DispatchBinding & binding) {
DispatchBinding & binding,
int & scale_log2) {
constexpr KernelCatalogRef kernel = GGML_HRX_KERNEL_REF("loom_libs", "ggml_copy_f16_k16_major");
constexpr const char * name = "common.mul_mat.k16_major_f16";
const size_t bytes = static_cast<size_t>(token_count * input_size) * sizeof(ggml_fp16_t);
const CommandPlanGeneratedResource * generated =
context.plan.metadata.find_generated_resource(input.id, GeneratedResourceRole::F16K16Major);
if (generated != nullptr && generated->byte_count == bytes) {
binding = { generated->generated_value, 0, bytes };
binding = { generated->generated_value, 0, bytes };
scale_log2 = generated->scale_log2;
return true;
}
const CommandPlanAlternateValue * alternate =
find_alternate_value(context.graph, context.plan, input.id, GGML_TYPE_F16, bytes);
// The pack keeps the scale of the F16 value it repacks; packed from the f32 value, the kernel prescales it.
const CommandPlanAlternateValue * alternate = find_alternate_value(
context.graph, context.plan, input.id, GGML_TYPE_F16, bytes, kCommonF16ActivationScaleLog2);
if (alternate == nullptr) {
alternate = find_alternate_value(context.graph, context.plan, input.id, GGML_TYPE_F16, bytes);
}
const int packed_scale_log2 = alternate != nullptr ? alternate->scale_log2 : kCommonF16ActivationScaleLog2;
const ValueId packed(context.next_plan_value.value + static_cast<int32_t>(match.transients.size()));
match.transients.push_back({ packed, name, bytes, 256 });
Dispatch copy;
Expand All @@ -618,11 +648,12 @@ inline bool common_prepare_k16_major_f16_input(const DispatchMatchContext & cont
match.dispatches.push_back(std::move(copy));
Status status;
if (!match.metadata.append_generated_resource(
{ input.id, GeneratedResourceRole::F16K16Major, packed, bytes, {} }, status)) {
{ input.id, GeneratedResourceRole::F16K16Major, packed, bytes, {}, packed_scale_log2 }, status)) {
match.status.append(status);
return false;
}
binding = { packed, 0, bytes };
binding = { packed, 0, bytes };
scale_log2 = packed_scale_log2;
return true;
}

Expand Down
Loading