From d91ecf9bdcf544b6b7040b46f7195a069498d01e Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Wed, 11 Nov 2026 00:50:25 -0400 Subject: [PATCH] hrx: the draft head's logits step reads only the group ids that were written, and stays finite above the f16 range (engine #1962) common.mul_mat.q6_k_token1_shortlist (a Q6_K projection onto a vocabulary of at least 65,536 at one token, downstream of the draft head's concat projection) had no test, and two defects: - The matcher asked for 96 selected activation groups at every width. The selection kernel writes one id per 32-wide group that exists and the scan reads as many ids as were asked for, so below a width of 3072 the tail of the id buffer was never written and was used as group indices: a device memory fault at width 256, 16 stray indices at 2560 (Qwen3.8-Flash-Next). The request is now min(96, input_size / 32). - The refine kernel narrowed the f32 activation to f16 unscaled and accumulated the logit in f16, so one activation above 65504 made every refined logit inf. The body now narrows at 2^-6 and the exported kernel scales the published logit back by 2^6. Tests: test-hrx-ops mul_mat.q6_k.token1_shortlist (ordinary data) and .large_activation (one activation above 65504, logits inside f16), comparing the refined candidates and the best token with the CPU; hrx-backend-test q6k_token1_shortlist_group_count.{256,2560,4096}, device-free. Co-Authored-By: Claude Fable 5.2 --- .../common/dispatch-mul-mat.cpp | 8 +- .../ops/mul_mat_q6_k_packed_f16_wmma.loom | 16 +- tests/hrx-backend-test.cpp | 64 ++++++ tests/test-hrx-ops.cpp | 198 ++++++++++++++++++ 4 files changed, 282 insertions(+), 4 deletions(-) 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 3785c64b441e..eb143331ff6d 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 @@ -422,7 +422,7 @@ static bool match_symmetric_i4_low_row_dispatch(const DispatchMatchContext & con } static bool match_q6_k_token1_shortlist_dispatch(const DispatchMatchContext & context, DispatchMatch & dispatch_match) { - constexpr int64_t selected_group_count = 96; + constexpr int64_t max_selected_groups = 96; constexpr int64_t candidate_count = 128; constexpr int64_t refine_tile_size = 64; constexpr size_t partition_entry_count = 1024; @@ -437,6 +437,12 @@ static bool match_q6_k_token1_shortlist_dispatch(const DispatchMatchContext & co return false; } + // engine #1962: the selection kernel writes one id per 32-wide activation group that exists, and the scan reads + // `selected_group_count` ids. Asking for 96 at an input below 3072 (Qwen3.8-Flash-Next is 2560: 80 groups) left the + // tail of the id buffer unwritten, and the scan used whatever was there as group indices: a read outside the weight + // (a device memory fault at 256) or stray terms in the coarse scores. Never ask for more groups than there are. + const int64_t selected_group_count = std::min(max_selected_groups, match.input_size / 32); + const CommonSymmetricI4ActivationLayout activation_layout = common_symmetric_i4_activation_layout(match.input_size, match.token_count); const CommandPlanAlternateValue * alternate = find_alternate_value(context.graph, context.plan, match.input->id, 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 52a683fa2b8c..4109a00c59c5 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 @@ -1171,8 +1171,15 @@ template.def<@ggml.mul_mat_q6_k_packed.selected_refine.body> device @ggml_mul_ma %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> + // engine #1962: the activation is narrowed to f16 and the logit is accumulated in f16, so an activation above + // 65504 became inf and so did every refined logit. Carry both at 2^-6 through the half-precision stage; the + // exported kernel scales the published logit back by 2^6. Both factors are powers of two, so a value that fit + // before is rounded as before; the usable range moves from 6.55e4 to 4.19e6 for the activation and the logit. + %activation_prescale = vector.constant 0.015625 : vector<4xf32> + %prescaled0 = vector.mulf %loaded0, %activation_prescale : vector<4xf32> + %prescaled1 = vector.mulf %loaded1, %activation_prescale : vector<4xf32> + %narrow0 = vector.fptrunc %prescaled0 : vector<4xf32> to vector<4xf16> + %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> @@ -1226,7 +1233,10 @@ kernel.def export("ggml_mul_mat_q6_k_packed_selected_refine_token1") @ggml_mul_m %candidate1 = index.assume %candidate0 [range(%candidate0, 0, 262143)] : index %candidate, %bounded_output_size = index.assume %candidate1, %output_size [lt(%candidate1, %output_size)] : index, index %output_view = buffer.view %output_noalias[%base] : buffer -> view<[%bounded_output_size]xf32> - %value = view.load %exact_view[%tid] : view<64xf32> -> f32 + // the body accumulated the logit at 2^-6 (engine #1962); undo it in f32 + %prescaled_value = view.load %exact_view[%tid] : view<64xf32> -> f32 + %logit_rescale = scalar.constant 64.0 : f32 + %value = scalar.mulf %prescaled_value, %logit_rescale : f32 view.store %value, %output_view[%candidate] : f32, view<[%bounded_output_size]xf32> } kernel.return diff --git a/tests/hrx-backend-test.cpp b/tests/hrx-backend-test.cpp index d6697b684cf3..dee609dcab95 100644 --- a/tests/hrx-backend-test.cpp +++ b/tests/hrx-backend-test.cpp @@ -3389,6 +3389,64 @@ static bool matmul_graph_is_supported(ggml_context * ctx, ggml_tensor * output) return ggml::hrx::DispatchScheduler::can_schedule_graph(imported.graph, test_dispatch_target()); } +// engine #1962: the draft head's logits step (common.mul_mat.q6_k_token1_shortlist) selects activation groups, and the +// scan reads as many group ids as it is told were selected. The selection kernel writes one id per 32-wide group that +// exists, so the request must never exceed input_size / 32: at 96 for every width, an input below 3072 left the tail +// of the id buffer unwritten and the scan used it as group indices. +static void schedule_q6k_token1_shortlist_group_count(int64_t hidden_size, int64_t vocab_size, int64_t expected_groups) { + 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 * table = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 8); + ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + ggml_tensor * external = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 1); + ggml_tensor * scale_a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 1); + ggml_tensor * scale_b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 1); + ggml_tensor * projection = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2 * hidden_size, hidden_size); + ggml_tensor * head = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, hidden_size, vocab_size); + ggml_tensor * branch_a = ggml_mul(ctx, ggml_rms_norm(ctx, ggml_get_rows(ctx, table, ids), 0.000001f), scale_a); + ggml_tensor * branch_b = ggml_mul(ctx, ggml_rms_norm(ctx, external, 0.000001f), scale_b); + ggml_tensor * projected = ggml_mul_mat(ctx, projection, ggml_concat(ctx, branch_a, branch_b, 0)); + ggml_tensor * logits = ggml_mul_mat(ctx, head, projected); + REQUIRE(logits != nullptr); + + ggml_cgraph * graph = ggml_new_graph(ctx); + REQUIRE(graph != nullptr); + ggml_build_forward_expand(graph, logits); + + ggml::hrx::GraphImportResult imported = ggml::hrx::import_ggml_graph(*graph); + REQUIRE(imported.valid()); + + ggml::hrx::DispatchScheduler scheduler; + REQUIRE(scheduler.schedule_graph(imported.graph, test_dispatch_target())); + REQUIRE(scheduler.plan().valid()); + + size_t select_count = 0; + size_t scan_count = 0; + for (const ggml::hrx::Dispatch & dispatch : scheduler.plan().dispatches) { + const std::string name = kernel_name_for_id(dispatch.kernel.kernel_id); + if (name == "loom_libs:ggml_select_symmetric_i4_k32_groups") { + ++select_count; + require_compile_parameter(dispatch, "ggml.mul_mat_q6_k_shortlist.selected_group_count", + std::to_string(expected_groups)); + REQUIRE(dispatch.bindings.size() == 2); + REQUIRE(dispatch.bindings[1].length == static_cast(expected_groups) * sizeof(int32_t)); + } else if (name == "loom_libs:ggml_mul_mat_q6_k_symmetric_i2_scan_token1") { + ++scan_count; + require_compile_parameter(dispatch, "ggml.mul_mat_q6_k_shortlist.selected_group_count", + std::to_string(expected_groups)); + } + } + // the case must reach the shortlist, or it asserts nothing + REQUIRE(select_count == 1); + REQUIRE(scan_count == 1); + + ggml_free(ctx); +} + static void schedule_qwen_terminal_q6k_q8_command(int64_t token_count) { ggml_init_params params = {}; params.mem_size = 4 * 1024 * 1024; @@ -11768,6 +11826,12 @@ int main(int argc, char ** argv) { { "gdn_selected_rms_q8_dispatch", [] { run_gdn_selected_rms_q8_dispatch_checks(); } }, { "llama_attention_matmul_dispatch", [] { run_llama_attention_matmul_dispatch_checks(); } }, { "quantized_value_projection_dispatch", [] { run_quantized_value_projection_dispatch_checks(); } }, + // widths 256 and 2560 (Qwen3.8-Flash-Next) have fewer than 96 groups; 4096 has more + { "q6k_token1_shortlist_group_count.256", [] { schedule_q6k_token1_shortlist_group_count(256, 65536, 8); } }, + { "q6k_token1_shortlist_group_count.2560", + [] { schedule_q6k_token1_shortlist_group_count(2560, 248320, 80); } }, + { "q6k_token1_shortlist_group_count.4096", + [] { schedule_q6k_token1_shortlist_group_count(4096, 65536, 96); } }, { "schedule_qwen_terminal_q6k_q8_command.tokens1", [] { schedule_qwen_terminal_q6k_q8_command(1); } }, { "schedule_qwen_terminal_q6k_q8_command.tokens18", [] { schedule_qwen_terminal_q6k_q8_command(18); } }, { "schedule_get_rows_q8_1_alternate_command.q4_k", [] { schedule_get_rows_q8_1_alternate_command(GGML_TYPE_Q4_K); } }, diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp index 9fdd1e325fbd..bb421100e8fc 100644 --- a/tests/test-hrx-ops.cpp +++ b/tests/test-hrx-ops.cpp @@ -3339,6 +3339,198 @@ static void run_dense_matmul_swiglu_cpu_reference_case(ggml_type gate_weight_ ggml_backend_free(hrx_backend); } +// engine #1962: the draft head's logits step (common.mul_mat.q6_k_token1_shortlist). A Q6_K projection onto a vocabulary +// of at least 65,536 at one token, whose input descends from a projection of concat(embedding row, external value), +// runs as a shortlist: a coarse scan, a top-128 selection, and an exact refine of those candidates; every other logit +// is filled negative. So the comparison with the CPU is over the refined candidates, not the whole output. +// `branch_scale` multiplies both normalised branches; it sets the size of the activation the logits step reads. +// Arms: Ordinary (everything inside f16); Spike (ONE activation above the f16 maximum, logits inside it, so the only +// value out of range is the one the refine kernel narrows); Large (every activation scaled up, so the activations and +// the logits are both above 65504 and inside the prescaled window). +enum class ShortlistArm { Ordinary, Spike, Large }; +static void run_q6_k_token1_shortlist_case(float branch_scale, ShortlistArm arm) { + const bool require_large_activation = arm != ShortlistArm::Ordinary; + constexpr int64_t hidden_size = 256; + constexpr int64_t vocab_size = 65536; + constexpr int64_t table_rows = 8; + ggml_backend_t cpu_backend = init_cpu_backend(); + ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); + REQUIRE(hrx_backend != nullptr); + + ggml_init_params params = {}; + params.mem_size = static_cast(48 * 1024 * 1024); + params.no_alloc = true; + ggml_context * cpu_ctx = ggml_init(params); + ggml_context * hrx_ctx = ggml_init(params); + REQUIRE(cpu_ctx != nullptr); + REQUIRE(hrx_ctx != nullptr); + + struct Graph { + ggml_tensor * table; + ggml_tensor * ids; + ggml_tensor * external; + ggml_tensor * scale_a; + ggml_tensor * scale_b; + ggml_tensor * projection_weight; + ggml_tensor * head_weight; + ggml_tensor * projected; + ggml_tensor * logits; + }; + auto build = [&](ggml_context * ctx) { + Graph g; + g.table = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, table_rows); + g.ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + g.external = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 1); + g.scale_a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 1); + g.scale_b = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, hidden_size, 1); + g.projection_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, 2 * hidden_size, hidden_size); + g.head_weight = ggml_new_tensor_2d(ctx, GGML_TYPE_Q6_K, hidden_size, vocab_size); + ggml_tensor * embedded = ggml_get_rows(ctx, g.table, g.ids); + ggml_tensor * branch_a = ggml_mul(ctx, ggml_rms_norm(ctx, embedded, 1.0e-6f), g.scale_a); + ggml_tensor * branch_b = ggml_mul(ctx, ggml_rms_norm(ctx, g.external, 1.0e-6f), g.scale_b); + g.projected = ggml_mul_mat(ctx, g.projection_weight, ggml_concat(ctx, branch_a, branch_b, 0)); + g.logits = ggml_mul_mat(ctx, g.head_weight, g.projected); + return g; + }; + const Graph cpu = build(cpu_ctx); + const Graph hrx = build(hrx_ctx); + + auto hrx_head_buffer = allocate_test_weight(hrx_backend, hrx.head_weight); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + ggml_build_forward_expand(cpu_graph, cpu.logits); + ggml_build_forward_expand(hrx_graph, hrx.logits); + + const std::vector sequence = scheduled_kernel_sequence(hrx_graph); + REQUIRE(test_sequence_has(sequence, "loom_libs:ggml_mul_mat_q6_k_symmetric_i2_scan_token1")); + REQUIRE(test_sequence_has(sequence, "loom_libs:ggml_mul_mat_q6_k_packed_selected_refine_token1")); + + ggml_backend_buffer_t cpu_buffer = ggml_backend_alloc_ctx_tensors(cpu_ctx, cpu_backend); + ggml_backend_buffer_t hrx_buffer = ggml_backend_alloc_ctx_tensors(hrx_ctx, hrx_backend); + REQUIRE(cpu_buffer != nullptr); + REQUIRE(hrx_buffer != nullptr); + + std::vector table(static_cast(hidden_size * table_rows)); + for (size_t i = 0; i < table.size(); ++i) { + table[i] = (static_cast((i * 29) % 23) - 11.0f) / 11.0f; + } + std::vector external(static_cast(hidden_size)); + for (size_t i = 0; i < external.size(); ++i) { + external[i] = (static_cast((i * 41) % 19) - 9.0f) / 9.0f; + } + std::vector scale_a(static_cast(hidden_size)); + std::vector scale_b(static_cast(hidden_size)); + for (size_t i = 0; i < scale_a.size(); ++i) { + scale_a[i] = branch_scale * (0.75f + 0.5f * static_cast(i % 5) / 4.0f); + scale_b[i] = branch_scale * (0.75f + 0.5f * static_cast(i % 7) / 6.0f); + } + std::vector projection(static_cast(2 * hidden_size * hidden_size)); + for (size_t i = 0; i < projection.size(); ++i) { + projection[i] = (static_cast((i * 53) % 31) - 15.0f) / 15.0f; + // The large arm puts ONE activation above the f16 maximum and keeps the rest small, so the logits (which this + // path publishes through f16) stay in range while the value the kernels narrow does not. + if (arm == ShortlistArm::Spike && i / static_cast(2 * hidden_size) != 7) { + projection[i] *= 1.0e-4f; + } + } + const std::vector head_weight = make_matmul_weight_bytes(GGML_TYPE_Q6_K, hidden_size, vocab_size, 23); + const int32_t id = 3; + set_tensor_pair_bytes(cpu_backend, cpu.table, hrx_backend, hrx.table, table.data(), table.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.ids, hrx_backend, hrx.ids, &id, sizeof(id)); + set_tensor_pair_bytes(cpu_backend, cpu.external, hrx_backend, hrx.external, external.data(), + external.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.scale_a, hrx_backend, hrx.scale_a, scale_a.data(), + scale_a.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.scale_b, hrx_backend, hrx.scale_b, scale_b.data(), + scale_b.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.projection_weight, hrx_backend, hrx.projection_weight, projection.data(), + projection.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu.head_weight, hrx_backend, hrx.head_weight, head_weight.data(), + head_weight.size()); + + 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); + + float activation_max = 0.0f; + for (const float value : get_f32_tensor(cpu_backend, cpu.projected)) { + activation_max = std::max(activation_max, std::fabs(value)); + } + const std::vector expected = get_f32_tensor(cpu_backend, cpu.logits); + const std::vector actual = get_f32_tensor(hrx_backend, hrx.logits); + REQUIRE(expected.size() == static_cast(vocab_size)); + REQUIRE(actual.size() == expected.size()); + + size_t expected_best = 0; + size_t actual_best = 0; + size_t non_finite = 0; + for (size_t i = 0; i < expected.size(); ++i) { + non_finite += std::isfinite(actual[i]) ? 0 : 1; + if (expected[i] > expected[expected_best]) { + expected_best = i; + } + if (actual[i] > actual[actual_best]) { + actual_best = i; + } + } + float logit_max = 0.0f; + for (const float value : expected) { + logit_max = std::max(logit_max, std::fabs(value)); + } + // the refined candidates are the logits the device did not fill negative; compare those with the CPU + double error = 0.0; + double norm = 0.0; + size_t refined = 0; + for (size_t i = 0; i < expected.size(); ++i) { + if (!std::isfinite(actual[i]) || actual[i] < -1.0e30f) { + continue; + } + ++refined; + const double diff = static_cast(actual[i]) - static_cast(expected[i]); + error += diff * diff; + norm += static_cast(expected[i]) * static_cast(expected[i]); + } + std::fprintf(stderr, + "q6_k shortlist: activation max=%g logit max=%g non-finite=%zu refined=%zu nmse(refined)=%g best " + "cpu=%zu (cpu %g, hrx %g) hrx=%zu (cpu %g, hrx %g)\n", + static_cast(activation_max), static_cast(logit_max), non_finite, refined, + norm > 0.0 ? error / norm : -1.0, expected_best, static_cast(expected[expected_best]), + static_cast(actual[expected_best]), actual_best, static_cast(expected[actual_best]), + static_cast(actual[actual_best])); + if (require_large_activation) { + REQUIRE(activation_max > 65504.0f); + } else { + REQUIRE(activation_max < 65504.0f); + } + if (arm == ShortlistArm::Large) { + REQUIRE(logit_max > 65504.0f); + REQUIRE(logit_max < 2.0e6f); + } else { + // the logits fit f16, so in the Spike arm the activation is the only value out of range + REQUIRE(logit_max < 65504.0f); + } + REQUIRE(non_finite == 0); + REQUIRE(refined > 0); + REQUIRE(error / norm < 1.0e-3); + // The device's best token must be as good as the CPU's. It may be a different token only when the two are a tie + // or a near tie: the generated weight repeats rows, and the refined logits are rounded to f16 (1 part in 2000). + // Not asserted in the Spike arm: the coarse scan reads the weight at 2 bits, so with one dominant activation it has + // four distinct scores and the 128 candidates are drawn from thousands of ties. That is the shortlist's recall on + // a degenerate input, not arithmetic; the arm checks that what is refined is finite and exact. + if (arm != ShortlistArm::Spike) { + REQUIRE(expected[actual_best] >= expected[expected_best] - 2.0e-3f * std::fabs(expected[expected_best])); + } + + ggml_backend_buffer_free(cpu_buffer); + ggml_backend_buffer_free(hrx_buffer); + ggml_free(cpu_ctx); + ggml_free(hrx_ctx); + ggml_backend_free(cpu_backend); + ggml_backend_free(hrx_backend); +} + // engine #1962: down(swiglu(gate(x), up(x))) at prefill, with a gated product above the f16 maximum. // The gate/up kernel hands the f16 form of its result to the down-projection. static void run_ffn_prefill_large_activation_case(int64_t token_count, float input_scale) { @@ -5099,6 +5291,12 @@ static void add_cases(test_runner::Suite & suite) { 4, 4096, GGML_GLU_OP_SWIGLU); }); + add(suite, "mul_mat.q6_k.token1_shortlist", + [=] { run_q6_k_token1_shortlist_case(1.0f, ShortlistArm::Ordinary); }); + add(suite, "mul_mat.q6_k.token1_shortlist.activation_spike", + [=] { run_q6_k_token1_shortlist_case(8000.0f, ShortlistArm::Spike); }); + add(suite, "mul_mat.q6_k.token1_shortlist.large_activation", + [=] { run_q6_k_token1_shortlist_case(8000.0f, ShortlistArm::Large); }); for (const int64_t tokens : { 256, 512 }) { // an input scale of 640 gives a gated maximum of 4.0e5 add(suite, "ffn.q4_k.prefill.large_activation." + std::to_string(tokens),