diff --git a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net_f32_wmma.loom b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net_f32_wmma.loom index 2c2461fcfa72..03b43f21af24 100644 --- a/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net_f32_wmma.loom +++ b/ggml/src/ggml-hrx/kernel-corpus/kernels/loom-libs/ops/gated_delta_net_f32_wmma.loom @@ -3045,6 +3045,11 @@ template.def<@llm.gated_delta_net.f32_wmma_head128.body> device @llm_gated_delta kernel.barrier scope(workgroup) ordering(acq_rel) %rowh_r = index.rem %row0, %k64 : index %rowh = index.assume %rowh_r [range(%rowh_r, 0, 60)] : index + // Prefix factors can underflow to zero. Form decay ratios from log differences. + %snapshot_log0 = vector.extract %gq0[3] : vector<4xf32> -> f32 + %snapshot_log1 = vector.extract %gq1[3] : vector<4xf32> -> f32 + %snapshot_log2 = vector.extract %gq2[3] : vector<4xf32> -> f32 + %snapshot_log3 = vector.extract %gq3[3] : vector<4xf32> -> f32 %snapshot_next0, %snapshot_next1, %snapshot_next2, %snapshot_next3, %snapshot_next4, %snapshot_next5, %snapshot_next6, %snapshot_next7, %snapshot_next8, %snapshot_next9, %snapshot_next10, %snapshot_next11, %snapshot_next12, %snapshot_next13, %snapshot_next14, %snapshot_next15 = scf.if %reuse_snapshot_fragments -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<4xf32>) { scf.if %publish_snapshots { %snapshot_c0o = index.constant 0 : offset @@ -3238,7 +3243,7 @@ template.def<@llm.gated_delta_net.f32_wmma_head128.body> device @llm_gated_delta scf.if %has_prefix1 { %snapshot_index1 = index.sub %snapshot_last, %k1 : index %snapshot_slot1 = index.mul %snapshot_stride, %snapshot_index1 : index - %p10_log = scalar.subf %lc1, %lc0 : f32 + %p10_log = scalar.subf %snapshot_log1, %snapshot_log0 : f32 %p10 = scalar.expf %p10_log : f32 %p1d0f = scalar.mulf %dl0, %p10 : f32 %p1d0 = scalar.fptrunc %p1d0f : f32 to f16 @@ -3256,9 +3261,9 @@ template.def<@llm.gated_delta_net.f32_wmma_head128.body> device @llm_gated_delta scf.if %has_prefix2 { %snapshot_index2 = index.sub %snapshot_last, %k2 : index %snapshot_slot2 = index.mul %snapshot_stride, %snapshot_index2 : index - %p20_log = scalar.subf %lc2, %lc0 : f32 + %p20_log = scalar.subf %snapshot_log2, %snapshot_log0 : f32 %p20 = scalar.expf %p20_log : f32 - %p21_log = scalar.subf %lc2, %lc1 : f32 + %p21_log = scalar.subf %snapshot_log2, %snapshot_log1 : f32 %p21 = scalar.expf %p21_log : f32 %p2d0f = scalar.mulf %dl0, %p20 : f32 %p2d1f = scalar.mulf %dl1, %p21 : f32 @@ -3278,11 +3283,11 @@ template.def<@llm.gated_delta_net.f32_wmma_head128.body> device @llm_gated_delta scf.if %has_prefix3 { %snapshot_index3 = index.sub %snapshot_last, %k3 : index %snapshot_slot3 = index.mul %snapshot_stride, %snapshot_index3 : index - %p30_log = scalar.subf %lc3, %lc0 : f32 + %p30_log = scalar.subf %snapshot_log3, %snapshot_log0 : f32 %p30 = scalar.expf %p30_log : f32 - %p31_log = scalar.subf %lc3, %lc1 : f32 + %p31_log = scalar.subf %snapshot_log3, %snapshot_log1 : f32 %p31 = scalar.expf %p31_log : f32 - %p32_log = scalar.subf %lc3, %lc2 : f32 + %p32_log = scalar.subf %snapshot_log3, %snapshot_log2 : f32 %p32 = scalar.expf %p32_log : f32 %p3d0f = scalar.mulf %dl0, %p30 : f32 %p3d1f = scalar.mulf %dl1, %p31 : f32 @@ -3335,7 +3340,7 @@ template.def<@llm.gated_delta_net.f32_wmma_head128.body> device @llm_gated_delta scf.if %has_prefix1 { %snapshot_index1 = index.sub %snapshot_last, %k1 : index %snapshot_slot1 = index.mul %snapshot_stride, %snapshot_index1 : index - %p10_log = scalar.subf %lc1, %lc0 : f32 + %p10_log = scalar.subf %snapshot_log1, %snapshot_log0 : f32 %p10 = scalar.expf %p10_log : f32 %p1d0f = scalar.mulf %dl0, %p10 : f32 %p1d0 = scalar.fptrunc %p1d0f : f32 to f16 @@ -3353,9 +3358,9 @@ template.def<@llm.gated_delta_net.f32_wmma_head128.body> device @llm_gated_delta scf.if %has_prefix2 { %snapshot_index2 = index.sub %snapshot_last, %k2 : index %snapshot_slot2 = index.mul %snapshot_stride, %snapshot_index2 : index - %p20_log = scalar.subf %lc2, %lc0 : f32 + %p20_log = scalar.subf %snapshot_log2, %snapshot_log0 : f32 %p20 = scalar.expf %p20_log : f32 - %p21_log = scalar.subf %lc2, %lc1 : f32 + %p21_log = scalar.subf %snapshot_log2, %snapshot_log1 : f32 %p21 = scalar.expf %p21_log : f32 %p2d0f = scalar.mulf %dl0, %p20 : f32 %p2d1f = scalar.mulf %dl1, %p21 : f32 @@ -3375,11 +3380,11 @@ template.def<@llm.gated_delta_net.f32_wmma_head128.body> device @llm_gated_delta scf.if %has_prefix3 { %snapshot_index3 = index.sub %snapshot_last, %k3 : index %snapshot_slot3 = index.mul %snapshot_stride, %snapshot_index3 : index - %p30_log = scalar.subf %lc3, %lc0 : f32 + %p30_log = scalar.subf %snapshot_log3, %snapshot_log0 : f32 %p30 = scalar.expf %p30_log : f32 - %p31_log = scalar.subf %lc3, %lc1 : f32 + %p31_log = scalar.subf %snapshot_log3, %snapshot_log1 : f32 %p31 = scalar.expf %p31_log : f32 - %p32_log = scalar.subf %lc3, %lc2 : f32 + %p32_log = scalar.subf %snapshot_log3, %snapshot_log2 : f32 %p32 = scalar.expf %p32_log : f32 %p3d0f = scalar.mulf %dl0, %p30 : f32 %p3d1f = scalar.mulf %dl1, %p31 : f32 diff --git a/tests/test-hrx-ops.cpp b/tests/test-hrx-ops.cpp index be17e1e2fff4..a9b9839e1c50 100644 --- a/tests/test-hrx-ops.cpp +++ b/tests/test-hrx-ops.cpp @@ -2612,6 +2612,218 @@ static void run_add_f32_cpu_reference_case() { ggml_backend_free(hrx_backend); } +static void run_gdn_snapshot_cpu_reference_case(int64_t tokens) { + ggml_backend_t cpu = init_cpu_backend(); + ggml_backend_t hrx = ggml_backend_hrx_init(0); + REQUIRE(hrx != nullptr); + constexpr int64_t width = 128; + constexpr int64_t heads = 2; + constexpr int64_t state_size = width * width * heads; + constexpr size_t state_bytes = state_size * sizeof(float); + struct fixture { + ggml_context * ctx; + ggml_cgraph * graph; + ggml_tensor * raw; + ggml_tensor * bias; + ggml_tensor * decay; + ggml_tensor * q; + ggml_tensor * k; + ggml_tensor * v; + ggml_tensor * g; + ggml_tensor * beta; + ggml_tensor * state; + ggml_tensor * cache; + ggml_tensor * attention; + ggml_backend_buffer_t buffer; + }; + const auto build = [&](ggml_backend_t backend) { + fixture f{}; + f.ctx = ggml_init({ 2 * 1024 * 1024, nullptr, true }); + REQUIRE(f.ctx != nullptr); + constexpr int64_t hidden = 3 * width * heads; + f.raw = ggml_new_tensor_2d(f.ctx, GGML_TYPE_F32, hidden, tokens); + const auto view = [&](size_t offset) { + return ggml_view_4d(f.ctx, f.raw, width, heads, tokens, 1, + width * sizeof(float), hidden * sizeof(float), hidden * tokens * sizeof(float), offset); + }; + f.q = view(0); + f.k = view(width * heads * sizeof(float)); + f.v = view(2 * width * heads * sizeof(float)); + f.g = ggml_new_tensor_2d(f.ctx, GGML_TYPE_F32, heads, tokens); + f.beta = ggml_dup_tensor(f.ctx, f.g); + f.bias = ggml_new_tensor_1d(f.ctx, GGML_TYPE_F32, heads); + f.decay = ggml_dup_tensor(f.ctx, f.bias); + auto * gate = ggml_mul(f.ctx, ggml_softplus(f.ctx, ggml_add(f.ctx, + ggml_reshape_3d(f.ctx, f.g, heads, tokens, 1), f.bias)), f.decay); + gate = ggml_reshape_4d(f.ctx, gate, 1, heads, tokens, 1); + auto * beta = ggml_sigmoid(f.ctx, ggml_reshape_4d(f.ctx, f.beta, 1, heads, tokens, 1)); + f.state = ggml_new_tensor_2d(f.ctx, GGML_TYPE_F32, state_size, 1); + auto * state = ggml_reshape_4d(f.ctx, f.state, width, width, heads, 1); + f.cache = ggml_new_tensor_2d(f.ctx, GGML_TYPE_F32, state_size, tokens); + auto * gdn = ggml_gated_delta_net(f.ctx, ggml_l2_norm(f.ctx, f.q, 1.e-6f), + ggml_l2_norm(f.ctx, f.k, 1.e-6f), f.v, gate, beta, state, tokens); + f.attention = ggml_view_4d(f.ctx, gdn, width, heads, tokens, 1, + width * sizeof(float), width * heads * sizeof(float), width * heads * tokens * sizeof(float), 0); + auto * snapshots = ggml_view_3d(f.ctx, gdn, state_size, 1, tokens, + state_bytes, state_bytes, width * heads * tokens * sizeof(float)); + auto * target = ggml_view_3d(f.ctx, f.cache, state_size, 1, tokens, state_bytes, state_bytes, 0); + f.graph = ggml_new_graph(f.ctx); + ggml_build_forward_expand(f.graph, ggml_cpy(f.ctx, snapshots, target)); + f.attention = ggml_scale(f.ctx, f.attention, 0.5f); + ggml_build_forward_expand(f.graph, f.attention); + f.buffer = ggml_backend_alloc_ctx_tensors(f.ctx, backend); + REQUIRE(f.buffer != nullptr); + return f; + }; + fixture c = build(cpu); + fixture h = build(hrx); + require_kernel_subsequence(scheduled_kernel_sequence(h.graph), + { "loom_libs:llm_gated_delta_net_f32_wmma_head128_snapshot_projection_epilogue" }); + const auto q = make_pattern_f32(width * heads * tokens, 31, 0.03125f); + const auto k = make_pattern_f32(q.size(), 19, 0.046875f); + const auto v = make_pattern_f32(q.size(), 17, 0.03125f); + const auto state = make_pattern_f32(state_size, 23, 0.015625f); + std::vector gate(heads * tokens, -0.1f); + const std::vector beta(heads * tokens, 0.5f); + const auto set = [&](ggml_tensor * ct, ggml_tensor * ht, const std::vector & values) { + set_tensor_pair_bytes(cpu, ct, hrx, ht, values.data(), values.size() * sizeof(float)); + }; + std::vector raw(3 * q.size()); + const size_t row_width = width * heads; + for (int64_t t = 0; t < tokens; ++t) { + std::copy_n(q.data() + t * row_width, row_width, raw.data() + 3 * t * row_width); + std::copy_n(k.data() + t * row_width, row_width, raw.data() + (3 * t + 1) * row_width); + std::copy_n(v.data() + t * row_width, row_width, raw.data() + (3 * t + 2) * row_width); + } + set(c.raw, h.raw, raw); + set(c.g, h.g, gate); + set(c.beta, h.beta, beta); + set(c.bias, h.bias, std::vector(heads, 0.0f)); + // The second head underflows cumulative decay, while the first retains ordinary decay. + set(c.decay, h.decay, { -1.0f, -200.0f }); + set(c.state, h.state, state); + const auto compute_and_compare = [&] { + REQUIRE(ggml_backend_graph_compute(cpu, c.graph) == GGML_STATUS_SUCCESS); + REQUIRE(ggml_backend_graph_compute(hrx, h.graph) == GGML_STATUS_SUCCESS); + ggml_backend_synchronize(cpu); + ggml_backend_synchronize(hrx); + require_close(get_f32_tensor(hrx, h.attention), get_f32_tensor(cpu, c.attention), 5.0e-4f); + require_close(get_f32_tensor(hrx, h.cache), get_f32_tensor(cpu, c.cache), 5.0e-4f); + }; + compute_and_compare(); + + // Keep later contributions alive after the first cumulative decay underflows. + for (int64_t t = 1; t < tokens; ++t) { + gate[t * heads + 1] = -6.0f + 0.1f * t; + } + set(c.g, h.g, gate); + for (float & value : raw) { + value *= -0.75f; + } + set(c.raw, h.raw, raw); + compute_and_compare(); + const auto cpu_snapshots = get_f32_tensor(cpu, c.cache); + const auto hrx_snapshots = get_f32_tensor(hrx, h.cache); + + // Snapshots are newest first. Continue independently from every accepted prefix. + // Mild decay makes the resumed output sensitive to the restored state. + set(c.decay, h.decay, { -1.0f, -1.0f }); + for (int64_t accepted = 0; accepted <= tokens; ++accepted) { + const size_t offset = (tokens - accepted) * state_size; + const float * cpu_state = accepted == 0 ? state.data() : cpu_snapshots.data() + offset; + const float * hrx_state = accepted == 0 ? state.data() : hrx_snapshots.data() + offset; + ggml_backend_tensor_set(c.state, cpu_state, 0, state_bytes); + ggml_backend_tensor_set(h.state, hrx_state, 0, state_bytes); + compute_and_compare(); + } + for (auto * f : { &c, &h }) { + ggml_backend_buffer_free(f->buffer); + ggml_free(f->ctx); + } + ggml_backend_free(cpu); + ggml_backend_free(hrx); +} + +static void run_dynamic_generic_ops_cpu_reference_case() { + 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 = 512 * 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); + + constexpr int64_t element_count = 257; + ggml_tensor * cpu_input = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * hrx_input = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * cpu_unary = ggml_sqr(cpu_ctx, cpu_input); + ggml_tensor * hrx_unary = ggml_sqr(hrx_ctx, hrx_input); + ggml_tensor * cpu_copy_target = ggml_new_tensor_1d(cpu_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * hrx_copy_target = ggml_new_tensor_1d(hrx_ctx, GGML_TYPE_F32, element_count); + ggml_tensor * cpu_copy = ggml_cpy(cpu_ctx, cpu_input, cpu_copy_target); + ggml_tensor * hrx_copy = ggml_cpy(hrx_ctx, hrx_input, hrx_copy_target); + ggml_tensor * cpu_concat_lhs = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, 64, 7); + ggml_tensor * cpu_concat_rhs = ggml_new_tensor_2d(cpu_ctx, GGML_TYPE_F32, 32, 7); + ggml_tensor * hrx_concat_lhs = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, 64, 7); + ggml_tensor * hrx_concat_rhs = ggml_new_tensor_2d(hrx_ctx, GGML_TYPE_F32, 32, 7); + ggml_tensor * cpu_concat = ggml_concat(cpu_ctx, cpu_concat_lhs, cpu_concat_rhs, 0); + ggml_tensor * hrx_concat = ggml_concat(hrx_ctx, hrx_concat_lhs, hrx_concat_rhs, 0); + REQUIRE(cpu_unary != nullptr && hrx_unary != nullptr); + REQUIRE(cpu_copy != nullptr && hrx_copy != nullptr); + REQUIRE(cpu_concat != nullptr && hrx_concat != nullptr); + + ggml_cgraph * cpu_graph = ggml_new_graph(cpu_ctx); + ggml_cgraph * hrx_graph = ggml_new_graph(hrx_ctx); + REQUIRE(cpu_graph != nullptr && hrx_graph != nullptr); + ggml_build_forward_expand(cpu_graph, cpu_unary); + ggml_build_forward_expand(cpu_graph, cpu_copy); + ggml_build_forward_expand(cpu_graph, cpu_concat); + ggml_build_forward_expand(hrx_graph, hrx_unary); + ggml_build_forward_expand(hrx_graph, hrx_copy); + ggml_build_forward_expand(hrx_graph, hrx_concat); + + require_kernel_subsequence(scheduled_kernel_sequence(hrx_graph), + { "loom_libs:ggml_unary_f32", "loom_libs:ggml_copy_f32", + "loom_libs:ggml_concat_dim0_f32" }); + for (const char * kernel_name : + { "loom_libs:ggml_unary_f32", "loom_libs:ggml_copy_f32", "loom_libs:ggml_concat_dim0_f32" }) { + const ggml::hrx::KernelSpecialization specialization = + scheduled_kernel_specialization(hrx_graph, kernel_name); + REQUIRE(specialization.workload_specialization == ggml::hrx::WorkloadSpecialization::Dynamic); + } + + 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 && hrx_buffer != nullptr); + const std::vector input = make_pattern_f32(element_count, 73, 0.03125f); + set_tensor_pair_bytes(cpu_backend, cpu_input, hrx_backend, hrx_input, input.data(), input.size() * sizeof(float)); + const std::vector concat_lhs = make_pattern_f32(64 * 7, 79, 0.015625f); + const std::vector concat_rhs = make_pattern_f32(32 * 7, 83, 0.0234375f); + set_tensor_pair_bytes(cpu_backend, cpu_concat_lhs, hrx_backend, hrx_concat_lhs, concat_lhs.data(), + concat_lhs.size() * sizeof(float)); + set_tensor_pair_bytes(cpu_backend, cpu_concat_rhs, hrx_backend, hrx_concat_rhs, concat_rhs.data(), + concat_rhs.size() * sizeof(float)); + + 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); + require_close(get_f32_tensor(hrx_backend, hrx_unary), get_f32_tensor(cpu_backend, cpu_unary), 0.0f); + require_close(get_f32_tensor(hrx_backend, hrx_copy), get_f32_tensor(cpu_backend, cpu_copy), 0.0f); + require_close(get_f32_tensor(hrx_backend, hrx_concat), get_f32_tensor(cpu_backend, cpu_concat), 0.0f); + + 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); +} + static void run_view_mul_f32_cpu_reference_case() { ggml_backend_t cpu_backend = init_cpu_backend(); ggml_backend_t hrx_backend = ggml_backend_hrx_init(0); @@ -4649,6 +4861,14 @@ static void add_cases(test_runner::Suite & suite) { suite.host_case("rmsnorm_support", run_rmsnorm_support_checks); suite.host_case("alternate_value_alias_lookup", run_alternate_value_alias_lookup_checks); + // engine#1945 / upstream #155: the reference cases 44ee571 exists for -- the recurrent MTP snapshot path and the + // generic-op family around it. Registered HERE because upstream's parallel register_*_cases() layer was kept by the + // cherry-pick and never called by main(), which made #155's own reference case dead code. + add(suite, "dynamic_generic_ops.cpu_reference", [=] { run_dynamic_generic_ops_cpu_reference_case(); }); + for (int64_t tokens = 2; tokens <= 5; ++tokens) + add(suite, "gdn_snapshot.cpu_reference." + std::to_string(tokens), + [=] { run_gdn_snapshot_cpu_reference_case(tokens); }); + add(suite, "add_f32_cpu_reference", [=] { run_add_f32_cpu_reference_case(); }); add(suite, "ssm_conv.q4_k", @@ -5681,6 +5901,10 @@ static void add_cases(test_runner::Suite & suite) { [=] { run_qwen_decode_attention_output_next_q8_scheduling_case(true); }); add(suite, "fa_prefill_full_cache.512", [=] { run_qwen_full_cache_prefill_flash_attention_scheduling_case(512); }); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32", 1, + kVectorOutputSize, true, false, false, kVectorInputSize); + run_dense_matmul_postops_cpu_reference_case(GGML_TYPE_F16, "loom_libs:ggml_mul_mat_vector_bias_residual_f32_f32", 1, + kVectorOutputSize, true, true, false, kVectorInputSize); } int main(int argc, char ** argv) {