Skip to content

hrx: keep activations inside the f16 range in the two small-batch Q6_K kernels (engine #1962) #288

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 10, 2026
Merged

hrx: keep activations inside the f16 range in the two small-batch Q6_K kernels (engine #1962) #288

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
30 changes: 23 additions & 7 deletions ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/mul_mat_q6_k_packed_f16_wmma.loom
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,11 @@ template.def<@ggml.mul_mat_q6_k_packed.token1.body> device @ggml_mul_mat_q6_k_pa
%weight_stage_bytes = index.constant 9216 : offset
%activation_stage_bytes = index.constant 2304 : offset
%zero_accumulator = vector.constant 0.0 : vector<8xf32>
// engine #1962: the activations pass through f16 for the multiply and f16 ends at 65504, so a larger value became
// inf and the product NaN. They are narrowed at 2^-6 and the f32 accumulator is scaled back by 2^6 before it is
// published (exact while the scaled value is a normal f16; window [3.9e-3, 4.19e6]), as in the f32 matmul cores.
%f16_range_prescale = vector.constant 0.015625 : vector<4xf32>
%f16_range_unscale = vector.constant 64.0 : vector<8xf32>
%empty_words = vector.constant 0 : vector<4xi32>
%zero_scale_bytes = vector.constant 0 : vector<8xi8>
%nibble_mask = vector.constant 252645135 : vector<4xi32>
Expand Down Expand Up @@ -274,8 +279,10 @@ template.def<@ggml.mul_mat_q6_k_packed.token1.body> device @ggml_mul_mat_q6_k_pa
%input_k_next = index.add %input_k0, %activation_k_next : index
%loaded0 = vector.load %input_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32>
%loaded1 = vector.load %input_view[%bounded_token, %input_k_next] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32>
%narrow0 = vector.fptrunc %loaded0 : vector<4xf32> to vector<4xf16>
%narrow1 = vector.fptrunc %loaded1 : vector<4xf32> to vector<4xf16>
%prescaled0 = vector.mulf %loaded0, %f16_range_prescale : vector<4xf32>
%narrow0 = vector.fptrunc %prescaled0 : vector<4xf32> to vector<4xf16>
%prescaled1 = vector.mulf %loaded1, %f16_range_prescale : vector<4xf32>
%narrow1 = vector.fptrunc %prescaled1 : vector<4xf32> to vector<4xf16>
scf.yield %narrow0, %narrow1 : vector<4xf16>, vector<4xf16>
} else {
%zeros0 = vector.constant 0.0 : vector<4xf16>
Expand All @@ -297,13 +304,14 @@ template.def<@ggml.mul_mat_q6_k_packed.token1.body> device @ggml_mul_mat_q6_k_pa
scf.yield %block_result : vector<8xf32>
}

%result_unscaled = vector.mulf %result, %f16_range_unscale : vector<8xf32>
%overwrites_output = index.cmp eq, %output_accumulation, %c0 : index
%has_one_token = index.cmp eq, %bounded_token_count, %c1 : index
%use_token1_overwrite = scalar.andi %overwrites_output, %has_one_token : i1
scf.if %use_token1_overwrite {
template.apply<@ggml.mul_mat_q6_k_packed.publish_token1_overwrite_f32_wave32>(%result, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias) : (vector<8xf32>, index, index, index, buffer)
template.apply<@ggml.mul_mat_q6_k_packed.publish_token1_overwrite_f32_wave32>(%result_unscaled, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias) : (vector<8xf32>, index, index, index, buffer)
} else {
template.apply<@ggml.mul_mat_q6_k_packed.publish_f32_wave32>(%result, %c0, %c0, %bounded_token_count, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index)
template.apply<@ggml.mul_mat_q6_k_packed.publish_f32_wave32>(%result_unscaled, %c0, %c0, %bounded_token_count, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index)
}
template.return
}
Expand Down Expand Up @@ -383,6 +391,11 @@ template.def<@ggml.mul_mat_q6_k_i8_prepacked.body> device @ggml_mul_mat_q6_k_i8_
%weight_stage_bytes = index.constant 9216 : offset
%activation_stage_bytes = index.constant 2304 : offset
%zero_accumulator = vector.constant 0.0 : vector<8xf32>
// engine #1962: the activations pass through f16 for the multiply and f16 ends at 65504, so a larger value became
// inf and the product NaN. They are narrowed at 2^-6 and the f32 accumulator is scaled back by 2^6 before it is
// published (exact while the scaled value is a normal f16; window [3.9e-3, 4.19e6]), as in the f32 matmul cores.
%f16_range_prescale = vector.constant 0.015625 : vector<4xf32>
%f16_range_unscale = vector.constant 64.0 : vector<8xf32>
%m = index.constant 16 : index
%n = index.constant 16 : index
%k = index.constant 16 : index
Expand Down Expand Up @@ -468,8 +481,10 @@ template.def<@ggml.mul_mat_q6_k_i8_prepacked.body> device @ggml_mul_mat_q6_k_i8_
%input_k_next = index.add %input_k0, %activation_k_next : index
%loaded0 = vector.load %input_view[%bounded_token, %input_k] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32>
%loaded1 = vector.load %input_view[%bounded_token, %input_k_next] : view<[%bounded_token_count]x[%bounded_input_size]xf32> -> vector<4xf32>
%narrow0 = vector.fptrunc %loaded0 : vector<4xf32> to vector<4xf16>
%narrow1 = vector.fptrunc %loaded1 : vector<4xf32> to vector<4xf16>
%prescaled0 = vector.mulf %loaded0, %f16_range_prescale : vector<4xf32>
%narrow0 = vector.fptrunc %prescaled0 : vector<4xf32> to vector<4xf16>
%prescaled1 = vector.mulf %loaded1, %f16_range_prescale : vector<4xf32>
%narrow1 = vector.fptrunc %prescaled1 : vector<4xf32> to vector<4xf16>
scf.yield %narrow0, %narrow1 : vector<4xf16>, vector<4xf16>
} else {
%zeros0 = vector.constant 0.0 : vector<4xf16>
Expand All @@ -491,7 +506,8 @@ template.def<@ggml.mul_mat_q6_k_i8_prepacked.body> device @ggml_mul_mat_q6_k_i8_
scf.yield %block_result : vector<8xf32>
}

template.apply<@ggml.mul_mat_q6_k_packed.publish_f32_wave32>(%result, %c0, %c0, %bounded_token_count, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index)
%result_unscaled = vector.mulf %result, %f16_range_unscale : vector<8xf32>
template.apply<@ggml.mul_mat_q6_k_packed.publish_f32_wave32>(%result_unscaled, %c0, %c0, %bounded_token_count, %bounded_output_size, %channel_tile_base, %subgroup_channel_add, %output_noalias, %output_accumulation) : (vector<8xf32>, index, index, index, index, index, index, buffer, index)
template.return
}

Expand Down
44 changes: 44 additions & 0 deletions tests/test-backend-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4639,6 +4639,43 @@ struct test_mul_mat_resident_weight : public test_case {
}
};

// The two together (engine #1962): a resident weight, so the packed-layout registrations are reached, and activations
// outside [-1, 1]. A magnitude of 1 or more sets one activation in 64 of every fourth token to +-magnitude; a magnitude
// below 1 scales every activation by it (the lower edge of the same f16 window).
struct test_mul_mat_resident_weight_large_activation : public test_mul_mat_resident_weight {
const float magnitude;

std::string vars() override {
return test_mul_mat_resident_weight::vars() + ",magnitude=" + var_to_str(magnitude);
}

test_mul_mat_resident_weight_large_activation(ggml_type type_a, int64_t m, int64_t n, int64_t k, float magnitude)
: test_mul_mat_resident_weight(type_a, m, n, k), magnitude(magnitude) {}

void initialize_tensors(ggml_context * ctx) override {
test_case::initialize_tensors(ctx);
for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) {
if (strcmp(t->name, "b") != 0) {
continue;
}
std::vector<float> data(ggml_nelements(t));
ggml_backend_tensor_get(t, data.data(), 0, ggml_nbytes(t));
if (magnitude < 1.0f) {
for (float & value : data) {
value *= magnitude;
}
} else {
for (int64_t token = 0; token < t->ne[1]; token += 4) {
for (int64_t i = 0; i < t->ne[0]; i += 64) {
data[token * t->ne[0] + i] = ((i / 64 + token / 4) % 2 == 0) ? magnitude : -magnitude;
}
}
}
ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t));
}
}
};

// GGML_OP_MUL_MAT whose activations exceed the f16 range: the dense form of the engine #1539 case below.
struct test_mul_mat_large_activation : public test_mul_mat {
const float magnitude;
Expand Down Expand Up @@ -9489,6 +9526,13 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q6_K, 4096, 8, 512));
test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q6_K, 2048, 1, 4096));
test_cases.emplace_back(new test_mul_mat_resident_weight(GGML_TYPE_Q5_K, 2048, 512, 4096));
// engine #1962: the two small-batch Q6_K kernels with activations of the magnitude measured in flint-40b, and the
// lower edge (the smallest dense-input column maximum measured on four models). The 512-token prefill kernels for
// Q4_K, Q5_K and Q6_K still overflow above 65504 and are not listed until they pass.
test_cases.emplace_back(new test_mul_mat_resident_weight_large_activation(GGML_TYPE_Q6_K, 2048, 8, 4096, 9.1e4f));
test_cases.emplace_back(new test_mul_mat_resident_weight_large_activation(GGML_TYPE_Q6_K, 4096, 8, 512, 9.1e4f));
test_cases.emplace_back(new test_mul_mat_resident_weight_large_activation(GGML_TYPE_Q6_K, 2048, 8, 4096, 2.4e-2f));
test_cases.emplace_back(new test_mul_mat_resident_weight_large_activation(GGML_TYPE_Q6_K, 4096, 8, 512, 2.4e-2f));

test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 576, 512, 576, {1,1}, {1,1}));
test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 1, 2048, 8192, {1, 1}, {1, 1}));
Expand Down