diff --git a/tests/cpu/CMakeLists.txt b/tests/cpu/CMakeLists.txt index 70ba14ac4..2dc25ae73 100644 --- a/tests/cpu/CMakeLists.txt +++ b/tests/cpu/CMakeLists.txt @@ -6,22 +6,6 @@ add_executable(Mllm-Test-KaiW4A32Pack KaiW4A32PackTest.cpp) target_link_libraries(Mllm-Test-KaiW4A32Pack PRIVATE gtest_main MllmCPUBackend) target_include_directories(Mllm-Test-KaiW4A32Pack PRIVATE ${MLLM_INCLUDE_DIR}) -add_executable(Mllm-Test-MiniCPM5-Config MiniCPM5ConfigTest.cpp) -target_link_libraries(Mllm-Test-MiniCPM5-Config PRIVATE gtest_main MllmCPUBackend) -target_include_directories(Mllm-Test-MiniCPM5-Config PRIVATE ${MLLM_INCLUDE_DIR}) -target_compile_definitions(Mllm-Test-MiniCPM5-Config - PRIVATE MINICPM5_EXAMPLE_DIR="${CMAKE_CURRENT_SOURCE_DIR}/../../examples/minicpm5") - -add_executable(Mllm-Test-MiniCPM5-Tokenizer MiniCPM5TokenizerTest.cpp) -target_link_libraries(Mllm-Test-MiniCPM5-Tokenizer PRIVATE gtest_main MllmCPUBackend) -target_include_directories(Mllm-Test-MiniCPM5-Tokenizer PRIVATE ${MLLM_INCLUDE_DIR}) - -add_executable(Mllm-Test-MiniCPM5-Model MiniCPM5ModelTest.cpp) -target_link_libraries(Mllm-Test-MiniCPM5-Model PRIVATE gtest_main MllmCPUBackend) -target_include_directories(Mllm-Test-MiniCPM5-Model PRIVATE ${MLLM_INCLUDE_DIR}) -target_compile_definitions(Mllm-Test-MiniCPM5-Model - PRIVATE MINICPM5_EXAMPLE_DIR="${CMAKE_CURRENT_SOURCE_DIR}/../../examples/minicpm5") - add_executable(Mllm-Test-CPUContiguousOp ContiguousOpTest.cpp) target_link_libraries(Mllm-Test-CPUContiguousOp PRIVATE gtest_main MllmRT MllmCPUBackend) target_include_directories(Mllm-Test-CPUContiguousOp PRIVATE ${MLLM_INCLUDE_DIR}) @@ -31,5 +15,5 @@ add_test( NAME CPUKernelFocused COMMAND Mllm-Test-CPUKernel - --gtest_filter=CausalDepthwiseConvKernelTest.*:CausalDepthwiseConvCurrentFirstKernelTest.*:GatedDeltaRuleKernelTest.*) + --gtest_filter=CausalDepthwiseConvKernelTest.*:CausalDepthwiseConvCurrentFirstKernelTest.*:GatedDeltaRuleKernelTest.*:GqaDecodeKernelTest.*) set_tests_properties(CPUKernelFocused PROPERTIES LABELS cpu-kernel) diff --git a/tests/cpu/GqaDecodeKernelTest.hpp b/tests/cpu/GqaDecodeKernelTest.hpp new file mode 100644 index 000000000..9803ca022 --- /dev/null +++ b/tests/cpu/GqaDecodeKernelTest.hpp @@ -0,0 +1,399 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +// Focused oracle for the single-token grouped-query-attention decode kernel +// that reads native KV-head [B, H, S, D] cache views. +// +// The reference below is an independent scalar implementation that accumulates +// in double precision. It is deliberately not routed through the production +// kernel, so the vectorized QK dot product, the grouped P@V accumulation, and +// the strided cache addressing cannot validate themselves. Output tolerance is +// used for the reference comparison because the vector dot product reorders the +// float32 summation; slice equivalence and repeat stability are compared +// bitwise because the kernel promises a fixed accumulation order per head. + +#include + +#include +#include +#include +#include +#include +#include + +#include "mllm/backends/cpu/kernels/common/gqa_decode/fwd_bhsd.hpp" +#include "KernelTestHelper.hpp" + +namespace gqa_decode_kernel_test { + +using mllm::cpu::gqa_decode::BhsdStrides; +using mllm::cpu::gqa_decode::fwdBhsdFp32; + +// Deterministic index-derived fill. No RNG, so every host reproduces the same +// bytes without carrying a seed through the evidence record. +inline float patternValue(std::size_t index, int salt) { + const auto scaled = static_cast((index * 37U + static_cast(salt) * 11U) % 251U); + return (scaled - 125.0F) / 64.0F; +} + +inline std::vector makeBuffer(std::size_t count, int salt) { + std::vector buffer(count); + for (std::size_t index = 0; index < count; ++index) { buffer[index] = patternValue(index, salt); } + return buffer; +} + +struct Geometry { + int batch; + int query_heads; + int kv_heads; + int kv_sequence; + int qk_dim; + int value_dim; +}; + +inline std::string describe(const Geometry& geometry) { + return "B=" + std::to_string(geometry.batch) + " Hq=" + std::to_string(geometry.query_heads) + + " Hkv=" + std::to_string(geometry.kv_heads) + " S=" + std::to_string(geometry.kv_sequence) + + " D=" + std::to_string(geometry.qk_dim) + " Dv=" + std::to_string(geometry.value_dim); +} + +// A float buffer together with the [B, H, S, D] strides the kernel uses to +// address it. The strides are part of the contract under test: the production +// caller hands the kernel views into a larger static cache and transposed +// query/output tensors, never freshly packed contiguous arrays. +struct StridedView { + std::vector storage; + BhsdStrides strides; +}; + +inline std::size_t offset(const BhsdStrides& strides, int batch, int head, int row, int column) { + return static_cast(batch) * strides.batch + static_cast(head) * strides.head + + static_cast(row) * strides.sequence + static_cast(column) * strides.dimension; +} + +// Contiguous [B, H, rows, dim] strides. +inline BhsdStrides contiguousStrides(int heads, int rows, int dim) { return {heads * rows * dim, rows * dim, dim, 1}; } + +// Single-token query [B, Hq, 1, D] in contiguous memory. +inline StridedView makeContiguousQuery(const Geometry& geometry, int salt) { + return {makeBuffer(static_cast(geometry.batch) * geometry.query_heads * geometry.qk_dim, salt), + contiguousStrides(geometry.query_heads, 1, geometry.qk_dim)}; +} + +// KV cache [B, Hkv, S, dim] whose sequence extent is exactly kv_sequence. +inline StridedView makeContiguousCache(const Geometry& geometry, int dim, int salt) { + return {makeBuffer(static_cast(geometry.batch) * geometry.kv_heads * geometry.kv_sequence * dim, salt), + contiguousStrides(geometry.kv_heads, geometry.kv_sequence, dim)}; +} + +// Zero-filled single-token output [B, Hq, 1, Dv] in contiguous memory. +inline StridedView makeContiguousOutput(const Geometry& geometry) { + return {std::vector(static_cast(geometry.batch) * geometry.query_heads * geometry.value_dim, 0.0F), + contiguousStrides(geometry.query_heads, 1, geometry.value_dim)}; +} + +// The same cache rows placed inside a static cache [B, Hkv, capacity, dim] +// with capacity > kv_sequence, as nn::KVHeadStaticCache exposes them. Rows +// beyond kv_sequence hold unrelated pattern bytes that the kernel must never +// read. +inline StridedView makeStaticCacheView(const Geometry& geometry, int capacity, int dim, const StridedView& contiguous, + int salt) { + StridedView view{makeBuffer(static_cast(geometry.batch) * geometry.kv_heads * capacity * dim, salt), + contiguousStrides(geometry.kv_heads, capacity, dim)}; + for (int batch = 0; batch < geometry.batch; ++batch) { + for (int head = 0; head < geometry.kv_heads; ++head) { + for (int row = 0; row < geometry.kv_sequence; ++row) { + for (int column = 0; column < dim; ++column) { + view.storage[offset(view.strides, batch, head, row, column)] = + contiguous.storage[offset(contiguous.strides, batch, head, row, column)]; + } + } + } + } + return view; +} + +// The same query bytes described as a transposed [B, 1, Hq, D] tensor: head +// stride D, batch and sequence stride Hq * D. With one token the memory order +// coincides with the contiguous query, so only the stride bookkeeping differs. +inline StridedView makeTransposedQueryView(const Geometry& geometry, const StridedView& contiguous_query) { + return {contiguous_query.storage, + {geometry.query_heads * geometry.qk_dim, geometry.qk_dim, geometry.query_heads * geometry.qk_dim, 1}}; +} + +inline StridedView makeTransposedOutputView(const Geometry& geometry) { + return {std::vector(static_cast(geometry.batch) * geometry.query_heads * geometry.value_dim, 0.0F), + {geometry.query_heads * geometry.value_dim, geometry.value_dim, geometry.query_heads * geometry.value_dim, 1}}; +} + +// The kernel's grouped scratch: group_size * kv_sequence probabilities. +inline std::vector makeScratch(const Geometry& geometry) { + return std::vector(static_cast(geometry.query_heads / geometry.kv_heads) * geometry.kv_sequence, 0.0F); +} + +// Independent scalar reference: per query head, softmax(q . k / sqrt(D)) @ v +// with double accumulation, honoring the same strided addressing contract. +inline void referenceDecode(const Geometry& geometry, const StridedView& query, const StridedView& key, + const StridedView& value, std::vector& output) { + const int group_size = geometry.query_heads / geometry.kv_heads; + const double scale = 1.0 / std::sqrt(static_cast(geometry.qk_dim)); + output.assign(static_cast(geometry.batch) * geometry.query_heads * geometry.value_dim, 0.0); + std::vector scores(static_cast(geometry.kv_sequence)); + for (int batch = 0; batch < geometry.batch; ++batch) { + for (int head = 0; head < geometry.query_heads; ++head) { + const int kv_head = head / group_size; + double maximum = -std::numeric_limits::infinity(); + for (int row = 0; row < geometry.kv_sequence; ++row) { + double score = 0.0; + for (int column = 0; column < geometry.qk_dim; ++column) { + score += static_cast(query.storage[offset(query.strides, batch, head, 0, column)]) + * static_cast(key.storage[offset(key.strides, batch, kv_head, row, column)]); + } + scores[static_cast(row)] = score * scale; + maximum = std::max(maximum, scores[static_cast(row)]); + } + double denominator = 0.0; + for (int row = 0; row < geometry.kv_sequence; ++row) { + scores[static_cast(row)] = std::exp(scores[static_cast(row)] - maximum); + denominator += scores[static_cast(row)]; + } + auto* out = output.data() + (static_cast(batch) * geometry.query_heads + head) * geometry.value_dim; + for (int row = 0; row < geometry.kv_sequence; ++row) { + const double probability = scores[static_cast(row)] / denominator; + for (int column = 0; column < geometry.value_dim; ++column) { + out[column] += probability * static_cast(value.storage[offset(value.strides, batch, kv_head, row, column)]); + } + } + } + } +} + +constexpr float kReferenceTolerance = 1.0e-4F; + +inline void expectMatchesReference(const Geometry& geometry, const StridedView& output, const std::vector& reference, + const std::string& label) { + for (int batch = 0; batch < geometry.batch; ++batch) { + for (int head = 0; head < geometry.query_heads; ++head) { + for (int column = 0; column < geometry.value_dim; ++column) { + const auto expected = + reference[(static_cast(batch) * geometry.query_heads + head) * geometry.value_dim + column]; + const auto actual = output.storage[offset(output.strides, batch, head, 0, column)]; + ASSERT_NEAR(actual, expected, kReferenceTolerance) + << label << " " << describe(geometry) << " b=" << batch << " h=" << head << " d=" << column; + } + } + } +} + +inline void testMatchesScalarReference(const std::vector& geometries) { + for (const auto& geometry : geometries) { + const auto query = makeContiguousQuery(geometry, 1); + const auto key = makeContiguousCache(geometry, geometry.qk_dim, 2); + const auto value = makeContiguousCache(geometry, geometry.value_dim, 3); + auto output = makeContiguousOutput(geometry); + auto scratch = makeScratch(geometry); + + ASSERT_TRUE(fwdBhsdFp32(geometry.batch, geometry.query_heads, geometry.kv_heads, geometry.kv_sequence, geometry.qk_dim, + geometry.value_dim, query.storage.data(), query.strides, key.storage.data(), key.strides, + value.storage.data(), value.strides, output.storage.data(), output.strides, scratch.data())) + << describe(geometry); + + std::vector reference; + referenceDecode(geometry, query, key, value, reference); + expectMatchesReference(geometry, output, reference, "contiguous"); + } +} + +// The production consumer hands the kernel a KV view inside a larger static +// cache and a transposed query/output view. Both must be addressed through the +// strides, and the result must be bitwise identical to the contiguous +// computation because the same bytes are read in the same order. +inline void testNativeCacheAndTransposedQueryStrides(const Geometry& geometry, int cache_capacity) { + ASSERT_GT(cache_capacity, geometry.kv_sequence); + const auto contiguous_query = makeContiguousQuery(geometry, 1); + const auto contiguous_key = makeContiguousCache(geometry, geometry.qk_dim, 2); + const auto contiguous_value = makeContiguousCache(geometry, geometry.value_dim, 3); + auto contiguous_output = makeContiguousOutput(geometry); + auto scratch = makeScratch(geometry); + ASSERT_TRUE(fwdBhsdFp32(geometry.batch, geometry.query_heads, geometry.kv_heads, geometry.kv_sequence, geometry.qk_dim, + geometry.value_dim, contiguous_query.storage.data(), contiguous_query.strides, + contiguous_key.storage.data(), contiguous_key.strides, contiguous_value.storage.data(), + contiguous_value.strides, contiguous_output.storage.data(), contiguous_output.strides, + scratch.data())) + << describe(geometry); + + const auto query_view = makeTransposedQueryView(geometry, contiguous_query); + const auto key_view = makeStaticCacheView(geometry, cache_capacity, geometry.qk_dim, contiguous_key, 7); + const auto value_view = makeStaticCacheView(geometry, cache_capacity, geometry.value_dim, contiguous_value, 8); + auto output_view = makeTransposedOutputView(geometry); + ASSERT_TRUE(fwdBhsdFp32(geometry.batch, geometry.query_heads, geometry.kv_heads, geometry.kv_sequence, geometry.qk_dim, + geometry.value_dim, query_view.storage.data(), query_view.strides, key_view.storage.data(), + key_view.strides, value_view.storage.data(), value_view.strides, output_view.storage.data(), + output_view.strides, scratch.data())) + << describe(geometry); + + std::vector reference; + referenceDecode(geometry, query_view, key_view, value_view, reference); + expectMatchesReference(geometry, output_view, reference, "strided"); + for (int batch = 0; batch < geometry.batch; ++batch) { + for (int head = 0; head < geometry.query_heads; ++head) { + for (int column = 0; column < geometry.value_dim; ++column) { + ASSERT_EQ(output_view.storage[offset(output_view.strides, batch, head, 0, column)], + contiguous_output.storage[offset(contiguous_output.strides, batch, head, 0, column)]) + << describe(geometry) << " b=" << batch << " h=" << head << " d=" << column; + } + } + } +} + +// Every (batch, kv-head) group is scheduled independently and every query head +// accumulates its keys in increasing order, so a grouped call must reproduce +// the per-head single-KV-head call bitwise. This also proves the grouped +// scratch rows do not leak between heads or batches. +inline void testGroupedSlicesMatchSingleHeadCallsBitwise(const Geometry& geometry) { + const auto query = makeContiguousQuery(geometry, 1); + const auto key = makeContiguousCache(geometry, geometry.qk_dim, 2); + const auto value = makeContiguousCache(geometry, geometry.value_dim, 3); + auto grouped_output = makeContiguousOutput(geometry); + auto grouped_scratch = makeScratch(geometry); + ASSERT_TRUE(fwdBhsdFp32(geometry.batch, geometry.query_heads, geometry.kv_heads, geometry.kv_sequence, geometry.qk_dim, + geometry.value_dim, query.storage.data(), query.strides, key.storage.data(), key.strides, + value.storage.data(), value.strides, grouped_output.storage.data(), grouped_output.strides, + grouped_scratch.data())) + << describe(geometry); + + const int group_size = geometry.query_heads / geometry.kv_heads; + const Geometry single{1, 1, 1, geometry.kv_sequence, geometry.qk_dim, geometry.value_dim}; + const auto key_head_size = static_cast(geometry.kv_sequence) * geometry.qk_dim; + const auto value_head_size = static_cast(geometry.kv_sequence) * geometry.value_dim; + for (int batch = 0; batch < geometry.batch; ++batch) { + for (int head = 0; head < geometry.query_heads; ++head) { + const int kv_head = head / group_size; + // Slice the grouped operands down to one query head and its KV head. + const auto* query_head = query.storage.data() + offset(query.strides, batch, head, 0, 0); + const auto* key_head = key.storage.data() + offset(key.strides, batch, kv_head, 0, 0); + const auto* value_head = value.storage.data() + offset(value.strides, batch, kv_head, 0, 0); + const StridedView single_query{std::vector(query_head, query_head + geometry.qk_dim), + contiguousStrides(1, 1, geometry.qk_dim)}; + const StridedView single_key{std::vector(key_head, key_head + key_head_size), + contiguousStrides(1, geometry.kv_sequence, geometry.qk_dim)}; + const StridedView single_value{std::vector(value_head, value_head + value_head_size), + contiguousStrides(1, geometry.kv_sequence, geometry.value_dim)}; + auto single_output = makeContiguousOutput(single); + auto single_scratch = makeScratch(single); + ASSERT_TRUE(fwdBhsdFp32(single.batch, single.query_heads, single.kv_heads, single.kv_sequence, single.qk_dim, + single.value_dim, single_query.storage.data(), single_query.strides, single_key.storage.data(), + single_key.strides, single_value.storage.data(), single_value.strides, + single_output.storage.data(), single_output.strides, single_scratch.data())) + << describe(geometry); + for (int column = 0; column < geometry.value_dim; ++column) { + ASSERT_EQ(grouped_output.storage[offset(grouped_output.strides, batch, head, 0, column)], + single_output.storage[static_cast(column)]) + << describe(geometry) << " b=" << batch << " h=" << head << " d=" << column; + } + } + } +} + +inline void testRepeatedCallsAreBitwiseStable(const Geometry& geometry, int repeats) { + const auto query = makeContiguousQuery(geometry, 1); + const auto key = makeContiguousCache(geometry, geometry.qk_dim, 2); + const auto value = makeContiguousCache(geometry, geometry.value_dim, 3); + std::vector first; + for (int repeat = 0; repeat <= repeats; ++repeat) { + auto output = makeContiguousOutput(geometry); + auto scratch = makeScratch(geometry); + ASSERT_TRUE(fwdBhsdFp32(geometry.batch, geometry.query_heads, geometry.kv_heads, geometry.kv_sequence, geometry.qk_dim, + geometry.value_dim, query.storage.data(), query.strides, key.storage.data(), key.strides, + value.storage.data(), value.strides, output.storage.data(), output.strides, scratch.data())) + << describe(geometry); + if (repeat == 0) { + first = output.storage; + } else { + ASSERT_EQ(output.storage, first) << describe(geometry) << " repeat=" << repeat; + } + } +} + +// Invalid geometry, null buffers, and unsupported strides must be rejected +// before any byte of the output is touched; the backend op then takes the +// reference path. +inline void testRejectsInvalidGeometryAndStrides() { + const Geometry geometry{1, 4, 2, 3, 8, 8}; + const auto query = makeContiguousQuery(geometry, 1); + const auto key = makeContiguousCache(geometry, geometry.qk_dim, 2); + const auto value = makeContiguousCache(geometry, geometry.value_dim, 3); + auto output = makeContiguousOutput(geometry); + auto scratch = makeScratch(geometry); + ASSERT_TRUE(fwdBhsdFp32(geometry.batch, geometry.query_heads, geometry.kv_heads, geometry.kv_sequence, geometry.qk_dim, + geometry.value_dim, query.storage.data(), query.strides, key.storage.data(), key.strides, + value.storage.data(), value.strides, output.storage.data(), output.strides, scratch.data())); + + constexpr float kSentinel = 123.5F; + const auto expectRejected = [&](const std::string& label, int query_heads, int kv_sequence, int qk_dim, + const float* query_data, const float* key_data, BhsdStrides key_strides, + const float* value_data, BhsdStrides value_strides, BhsdStrides output_strides, + float* scratch_data) { + std::vector untouched(output.storage.size(), kSentinel); + EXPECT_FALSE(fwdBhsdFp32(geometry.batch, query_heads, geometry.kv_heads, kv_sequence, qk_dim, geometry.value_dim, + query_data, query.strides, key_data, key_strides, value_data, value_strides, untouched.data(), + output_strides, scratch_data)) + << label; + for (const auto element : untouched) { EXPECT_EQ(element, kSentinel) << label << " touched the output"; } + }; + const auto* query_data = query.storage.data(); + const auto* key_data = key.storage.data(); + const auto* value_data = value.storage.data(); + BhsdStrides non_unit_dimension = key.strides; + non_unit_dimension.dimension = 2; + BhsdStrides non_unit_output_dimension = output.strides; + non_unit_output_dimension.dimension = 2; + BhsdStrides zero_sequence = value.strides; + zero_sequence.sequence = 0; + + expectRejected("query heads not a multiple of kv heads", 3, geometry.kv_sequence, geometry.qk_dim, query_data, key_data, + key.strides, value_data, value.strides, output.strides, scratch.data()); + expectRejected("zero kv sequence", geometry.query_heads, 0, geometry.qk_dim, query_data, key_data, key.strides, value_data, + value.strides, output.strides, scratch.data()); + expectRejected("zero qk dim", geometry.query_heads, geometry.kv_sequence, 0, query_data, key_data, key.strides, value_data, + value.strides, output.strides, scratch.data()); + expectRejected("null query", geometry.query_heads, geometry.kv_sequence, geometry.qk_dim, nullptr, key_data, key.strides, + value_data, value.strides, output.strides, scratch.data()); + expectRejected("null key", geometry.query_heads, geometry.kv_sequence, geometry.qk_dim, query_data, nullptr, key.strides, + value_data, value.strides, output.strides, scratch.data()); + expectRejected("null value", geometry.query_heads, geometry.kv_sequence, geometry.qk_dim, query_data, key_data, key.strides, + nullptr, value.strides, output.strides, scratch.data()); + expectRejected("null scratch", geometry.query_heads, geometry.kv_sequence, geometry.qk_dim, query_data, key_data, key.strides, + value_data, value.strides, output.strides, nullptr); + expectRejected("non-unit key dimension stride", geometry.query_heads, geometry.kv_sequence, geometry.qk_dim, query_data, + key_data, non_unit_dimension, value_data, value.strides, output.strides, scratch.data()); + expectRejected("non-unit output dimension stride", geometry.query_heads, geometry.kv_sequence, geometry.qk_dim, query_data, + key_data, key.strides, value_data, value.strides, non_unit_output_dimension, scratch.data()); + expectRejected("non-positive value sequence stride", geometry.query_heads, geometry.kv_sequence, geometry.qk_dim, query_data, + key_data, key.strides, value_data, zero_sequence, output.strides, scratch.data()); +} + +} // namespace gqa_decode_kernel_test + +class GqaDecodeKernelTest : public KernelTest { + public: + void testMatchesScalarReference(const std::vector& geometries) { + gqa_decode_kernel_test::testMatchesScalarReference(geometries); + } + + void testNativeCacheAndTransposedQueryStrides(const gqa_decode_kernel_test::Geometry& geometry, int cache_capacity) { + gqa_decode_kernel_test::testNativeCacheAndTransposedQueryStrides(geometry, cache_capacity); + } + + void testGroupedSlicesMatchSingleHeadCallsBitwise(const gqa_decode_kernel_test::Geometry& geometry) { + gqa_decode_kernel_test::testGroupedSlicesMatchSingleHeadCallsBitwise(geometry); + } + + void testRepeatedCallsAreBitwiseStable(const gqa_decode_kernel_test::Geometry& geometry, int repeats) { + gqa_decode_kernel_test::testRepeatedCallsAreBitwiseStable(geometry, repeats); + } + + void testRejectsInvalidGeometryAndStrides() { gqa_decode_kernel_test::testRejectsInvalidGeometryAndStrides(); } +}; diff --git a/tests/cpu/KernelTest.cpp b/tests/cpu/KernelTest.cpp index b3f3d4d99..554a23161 100644 --- a/tests/cpu/KernelTest.cpp +++ b/tests/cpu/KernelTest.cpp @@ -630,6 +630,49 @@ TEST_F(GatedDeltaRuleKernelTest, ProductionGroupedHeadGeometry8LaneIsBitwiseStab /*repeats=*/24)); } +//===----------------------------------------------------------------------===// +// Grouped-query attention decode (native KV-head cache views) +//===----------------------------------------------------------------------===// +#include "GqaDecodeKernelTest.hpp" +TEST_F(GqaDecodeKernelTest, MatchesScalarReferenceAcrossFocusedMatrix) { + EXPECT_NO_FATAL_FAILURE(testMatchesScalarReference({ + // Scalar path and single KV head. + {1, 1, 1, 1, 1, 1}, + {1, 2, 1, 3, 3, 3}, + {2, 4, 1, 5, 4, 4}, + // Exact NEON vector blocks, group size 1 and 2. + {1, 2, 2, 7, 8, 8}, + {1, 4, 2, 16, 16, 16}, + // Vector loop with scalar tails in qk and value dims. + {1, 4, 2, 9, 5, 7}, + {2, 6, 3, 11, 13, 6}, + {1, 8, 2, 33, 127, 130}, + // Production 128-dim heads at several cache fills. + {1, 8, 2, 1, 128, 128}, + {1, 16, 2, 69, 128, 128}, + {1, 16, 2, 517, 128, 128}, + {2, 16, 2, 40, 128, 128}, + })); +} + +TEST_F(GqaDecodeKernelTest, NativeCacheViewAndTransposedQueryMatchContiguousBitwise) { + EXPECT_NO_FATAL_FAILURE(testNativeCacheAndTransposedQueryStrides({1, 16, 2, 69, 128, 128}, /*cache_capacity=*/2048)); + EXPECT_NO_FATAL_FAILURE(testNativeCacheAndTransposedQueryStrides({2, 6, 3, 11, 13, 6}, /*cache_capacity=*/32)); +} + +TEST_F(GqaDecodeKernelTest, GroupedSlicesMatchSingleHeadCallsBitwise) { + EXPECT_NO_FATAL_FAILURE(testGroupedSlicesMatchSingleHeadCallsBitwise({2, 16, 2, 69, 128, 128})); + EXPECT_NO_FATAL_FAILURE(testGroupedSlicesMatchSingleHeadCallsBitwise({1, 6, 3, 11, 13, 6})); +} + +TEST_F(GqaDecodeKernelTest, RepeatedCallsAreBitwiseStable) { + EXPECT_NO_FATAL_FAILURE(testRepeatedCallsAreBitwiseStable({1, 16, 2, 517, 128, 128}, /*repeats=*/8)); +} + +TEST_F(GqaDecodeKernelTest, RejectsInvalidGeometryAndStrides) { + EXPECT_NO_FATAL_FAILURE(testRejectsInvalidGeometryAndStrides()); +} + //===----------------------------------------------------------------------===// // Parallel linear //===----------------------------------------------------------------------===// diff --git a/tests/models/CMakeLists.txt b/tests/models/CMakeLists.txt index c427145b8..e023c0366 100644 --- a/tests/models/CMakeLists.txt +++ b/tests/models/CMakeLists.txt @@ -1 +1,2 @@ add_subdirectory(qwen3_5) +add_subdirectory(minicpm5) diff --git a/tests/models/minicpm5/CMakeLists.txt b/tests/models/minicpm5/CMakeLists.txt new file mode 100644 index 000000000..20cf320d1 --- /dev/null +++ b/tests/models/minicpm5/CMakeLists.txt @@ -0,0 +1,20 @@ +add_executable(Mllm-Test-MiniCPM5-Config MiniCPM5ConfigTest.cpp) +target_link_libraries(Mllm-Test-MiniCPM5-Config PRIVATE gtest_main MllmCPUBackend) +target_include_directories(Mllm-Test-MiniCPM5-Config PRIVATE ${MLLM_INCLUDE_DIR}) +target_compile_definitions(Mllm-Test-MiniCPM5-Config + PRIVATE MINICPM5_EXAMPLE_DIR="${CMAKE_CURRENT_SOURCE_DIR}/../../../examples/minicpm5") + +add_executable(Mllm-Test-MiniCPM5-Tokenizer MiniCPM5TokenizerTest.cpp) +target_link_libraries(Mllm-Test-MiniCPM5-Tokenizer PRIVATE gtest_main MllmCPUBackend) +target_include_directories(Mllm-Test-MiniCPM5-Tokenizer PRIVATE ${MLLM_INCLUDE_DIR}) + +add_executable(Mllm-Test-MiniCPM5-Model MiniCPM5ModelTest.cpp) +target_link_libraries(Mllm-Test-MiniCPM5-Model PRIVATE gtest_main MllmCPUBackend) +target_include_directories(Mllm-Test-MiniCPM5-Model PRIVATE ${MLLM_INCLUDE_DIR}) +target_compile_definitions(Mllm-Test-MiniCPM5-Model + PRIVATE MINICPM5_EXAMPLE_DIR="${CMAKE_CURRENT_SOURCE_DIR}/../../../examples/minicpm5") + +add_test(NAME MiniCPM5ConfigFocused COMMAND Mllm-Test-MiniCPM5-Config) +add_test(NAME MiniCPM5TokenizerFocused COMMAND Mllm-Test-MiniCPM5-Tokenizer) +add_test(NAME MiniCPM5ModelFocused COMMAND Mllm-Test-MiniCPM5-Model) +set_tests_properties(MiniCPM5ConfigFocused MiniCPM5TokenizerFocused MiniCPM5ModelFocused PROPERTIES LABELS minicpm5) diff --git a/tests/cpu/MiniCPM5ConfigTest.cpp b/tests/models/minicpm5/MiniCPM5ConfigTest.cpp similarity index 89% rename from tests/cpu/MiniCPM5ConfigTest.cpp rename to tests/models/minicpm5/MiniCPM5ConfigTest.cpp index 36e1321bd..ae95d7ee2 100644 --- a/tests/cpu/MiniCPM5ConfigTest.cpp +++ b/tests/models/minicpm5/MiniCPM5ConfigTest.cpp @@ -3,6 +3,7 @@ #include +#include #include #include @@ -10,8 +11,13 @@ namespace { +auto exampleDir() -> std::string { + const char* example_dir_override = std::getenv("MLLM_MINICPM5_EXAMPLE_DIR"); + return example_dir_override == nullptr ? std::string(MINICPM5_EXAMPLE_DIR) : std::string(example_dir_override); +} + auto loadConfig() -> mllm::models::minicpm5::MiniCPM5Config { - return mllm::models::minicpm5::MiniCPM5Config(std::string(MINICPM5_EXAMPLE_DIR) + "/config_1B_w4a32_kai.json"); + return mllm::models::minicpm5::MiniCPM5Config(exampleDir() + "/config_1B_w4a32_kai.json"); } auto parameterFile(mllm::ModelFileVersion version, const std::vector& embedding_shape, bool include_lm_head = true) diff --git a/tests/cpu/MiniCPM5ModelTest.cpp b/tests/models/minicpm5/MiniCPM5ModelTest.cpp similarity index 71% rename from tests/cpu/MiniCPM5ModelTest.cpp rename to tests/models/minicpm5/MiniCPM5ModelTest.cpp index 098e95d9b..7a3552ddb 100644 --- a/tests/cpu/MiniCPM5ModelTest.cpp +++ b/tests/models/minicpm5/MiniCPM5ModelTest.cpp @@ -3,14 +3,24 @@ #include +#include #include #include "mllm/mllm.hpp" #include "mllm/models/minicpm5/modeling_minicpm5.hpp" +namespace { + +auto exampleDir() -> std::string { + const char* example_dir_override = std::getenv("MLLM_MINICPM5_EXAMPLE_DIR"); + return example_dir_override == nullptr ? std::string(MINICPM5_EXAMPLE_DIR) : std::string(example_dir_override); +} + +} // namespace + TEST(MiniCPM5ModelTest, BuildsNativeKVHeadLogicalSlotsAndResetsThem) { mllm::initializeContext(); - const auto config = mllm::models::minicpm5::MiniCPM5Config(std::string(MINICPM5_EXAMPLE_DIR) + "/config_1B_w4a32_kai.json"); + const auto config = mllm::models::minicpm5::MiniCPM5Config(exampleDir() + "/config_1B_w4a32_kai.json"); auto model = mllm::models::minicpm5::MiniCPM5ForCausalLM(config); auto& cache = model.kvCache(); diff --git a/tests/cpu/MiniCPM5TokenizerTest.cpp b/tests/models/minicpm5/MiniCPM5TokenizerTest.cpp similarity index 100% rename from tests/cpu/MiniCPM5TokenizerTest.cpp rename to tests/models/minicpm5/MiniCPM5TokenizerTest.cpp diff --git a/tests/nn/CMakeLists.txt b/tests/nn/CMakeLists.txt index bd003caac..4d84b9399 100644 --- a/tests/nn/CMakeLists.txt +++ b/tests/nn/CMakeLists.txt @@ -28,3 +28,6 @@ target_include_directories(Mllm-Test-Nn-GatedDeltaRule PRIVATE ${MLLM_INCLUDE_DI include(GoogleTest) add_test(NAME GatedDeltaRuleFocused COMMAND Mllm-Test-Nn-GatedDeltaRule) +add_test(NAME GroupedQueryAttentionFocused COMMAND Mllm-Test-Nn-GroupedQueryAttention) +add_test(NAME KVHeadStaticCacheFocused COMMAND Mllm-Test-Nn-KVHeadStaticCache) +set_tests_properties(GroupedQueryAttentionFocused KVHeadStaticCacheFocused PROPERTIES LABELS nn-op)