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),