diff --git a/README.md b/README.md index 4fa4ad85a..a511778c8 100644 --- a/README.md +++ b/README.md @@ -107,6 +107,7 @@ The mllm framework integrates seamlessly with popular community frameworks' chec | [Qwen3.5-0.8B](https://huggingface.co/Qwen/Qwen3.5-0.8B) | [✔️ w4a8](./examples/qwen3_5/README.md) | | | | [Qwen3.5-4B](https://huggingface.co/Qwen/Qwen3.5-4B) | [✔️ w4a8](./examples/qwen3_5/README.md) | | | | [MiniCPM5-1B](https://huggingface.co/openbmb/MiniCPM5-1B) | [✔️ w4a8](./examples/minicpm5/README.md) | | | +| [LFM2.5-2.6B](https://huggingface.co/LiquidAI/LFM2.5-2.6B) | [✔️ w4a8](./examples/lfm2/README.md) | | | | [DeepSeek-OCR](https://github.com/deepseek-ai/DeepSeek-OCR) | [✔️ w4a8](https://www.modelscope.cn/models/mllmTeam/DeepSeek-OCR-w4a8-i8mm-kai) | | | | [SmolLM3](https://huggingface.co/blog/smollm3)| [✔️ w4a8](https://www.modelscope.cn/models/mllmTeam/SmolLM3-3B-w4a8-i8mm-kai) | | | | [Qwen2-VL-2B-Instruct](https://qwenlm.github.io/zh/blog/qwen2-vl/)|[✔️ w4a8](https://www.modelscope.cn/models/mllmTeam/Qwen2-VL-2B-Instruct-w4a32kai) || | diff --git a/benchmarks/cpu/CMakeLists.txt b/benchmarks/cpu/CMakeLists.txt index 089eff6a1..065be03de 100644 --- a/benchmarks/cpu/CMakeLists.txt +++ b/benchmarks/cpu/CMakeLists.txt @@ -1,4 +1,10 @@ if(MLLM_BUILD_ARM_BACKEND) add_executable(Mllm-Benchmark-ARM-HPC-Sgemm arm_mllm_blas_sgemm.cpp) target_link_libraries(Mllm-Benchmark-ARM-HPC-Sgemm PRIVATE benchmark::benchmark MllmRT MllmCPUBackend) -endif() \ No newline at end of file + + add_executable(Mllm-Benchmark-Lfm2-Parallel-Linear lfm2_parallel_linear.cpp) + target_link_libraries(Mllm-Benchmark-Lfm2-Parallel-Linear PRIVATE MllmRT MllmCPUBackend) + + add_executable(Mllm-Benchmark-Lfm2-Parallel-Linear-Shared-Mx lfm2_parallel_linear_shared_mx.cpp) + target_link_libraries(Mllm-Benchmark-Lfm2-Parallel-Linear-Shared-Mx PRIVATE MllmRT MllmCPUBackend) +endif() diff --git a/benchmarks/cpu/lfm2_parallel_linear.cpp b/benchmarks/cpu/lfm2_parallel_linear.cpp new file mode 100644 index 000000000..b864c9e11 --- /dev/null +++ b/benchmarks/cpu/lfm2_parallel_linear.cpp @@ -0,0 +1,297 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mllm/backends/cpu/kernels/arm/linear/kai.hpp" +#include "mllm/mllm.hpp" + +namespace { + +using KaiHelper = mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk; +using KaiTile = KaiHelper::Tiles; + +constexpr int kInputChannels = 2048; +constexpr KaiTile kDecodeTile = KaiTile::qai8dxp1x8_qsi4c32p8x8_1x8x32; +constexpr KaiTile kPrefillTile = KaiTile::qai8dxp4x8_qsi4c32p8x8_4x8x32; + +struct ShapeCase { + std::string_view name; + std::vector output_channels; +}; + +struct Buffers { + std::vector input; + std::vector weights; + std::vector> separate_packed_weights; + std::vector merged_packed_weight; + std::vector workspace; + std::vector> separate_outputs; + std::vector> shared_outputs; + std::vector merged_output; +}; + +struct Comparison { + size_t bitwise_mismatches = 0; + float max_absolute_error = 0.0F; +}; + +uint32_t nextRandom(uint32_t& state) { + state = state * 1664525U + 1013904223U; + return state; +} + +float deterministicValue(uint32_t& state) { + const int32_t centered = static_cast((nextRandom(state) >> 8U) % 2001U) - 1000; + return static_cast(centered) / 4096.0F; +} + +int parsePositiveInt(const char* value, const char* name) { + char* end = nullptr; + const int64_t parsed = std::strtoll(value, &end, 10); + if (end == value || *end != '\0' || parsed <= 0 || parsed > std::numeric_limits::max()) { + throw std::invalid_argument(std::string(name) + " must be a positive integer"); + } + return static_cast(parsed); +} + +ShapeCase parseShape(std::string_view name) { + if (name == "gate_up") { return {.name = "gate_up", .output_channels = {10752, 10752}}; } + if (name == "qkv") { return {.name = "qkv", .output_channels = {2048, 512, 512}}; } + throw std::invalid_argument("shape must be gate_up or qkv"); +} + +size_t totalOutputChannels(const ShapeCase& shape) { + return std::accumulate(shape.output_channels.begin(), shape.output_channels.end(), size_t{0}); +} + +Buffers makeBuffers(const ShapeCase& shape, int m, KaiTile tile) { + KaiHelper kai; + Buffers buffers; + const size_t total_n = totalOutputChannels(shape); + + buffers.input.resize(static_cast(m) * kInputChannels); + buffers.weights.resize(total_n * kInputChannels); + uint32_t random_state = 0x4C464D32U; + std::generate(buffers.input.begin(), buffers.input.end(), [&] { return deterministicValue(random_state); }); + std::generate(buffers.weights.begin(), buffers.weights.end(), [&] { return deterministicValue(random_state); }); + + buffers.separate_packed_weights.reserve(shape.output_channels.size()); + buffers.separate_outputs.reserve(shape.output_channels.size()); + buffers.shared_outputs.reserve(shape.output_channels.size()); + size_t row_offset = 0; + for (const int n : shape.output_channels) { + const size_t packed_size = kai.quant_pack_rhs_size(n, kInputChannels, tile); + auto& packed = buffers.separate_packed_weights.emplace_back(packed_size); + kai.quant_pack_rhs_offline(packed.data(), buffers.weights.data() + row_offset * kInputChannels, nullptr, n, kInputChannels, + tile); + buffers.separate_outputs.emplace_back(static_cast(m) * n); + buffers.shared_outputs.emplace_back(static_cast(m) * n); + row_offset += static_cast(n); + } + + buffers.merged_packed_weight.resize(kai.quant_pack_rhs_size(static_cast(total_n), kInputChannels, tile)); + kai.quant_pack_rhs_offline(buffers.merged_packed_weight.data(), buffers.weights.data(), nullptr, static_cast(total_n), + kInputChannels, tile); + buffers.workspace.resize(kai.workspace_size(m, kInputChannels, tile)); + buffers.merged_output.resize(static_cast(m) * total_n); + return buffers; +} + +void runIndependent(const ShapeCase& shape, int m, int threads, KaiTile tile, Buffers& buffers) { + KaiHelper kai; + for (size_t index = 0; index < shape.output_channels.size(); ++index) { + kai.matmul(buffers.separate_outputs[index].data(), buffers.input.data(), buffers.separate_packed_weights[index].data(), + buffers.workspace.data(), m, kInputChannels, shape.output_channels[index], tile, threads); + } +} + +void runShared(const ShapeCase& shape, int m, int threads, KaiTile tile, Buffers& buffers) { + if (m != 1) { throw std::invalid_argument("shared-input path requires M=1"); } + std::vector projections; + projections.reserve(shape.output_channels.size()); + for (size_t index = 0; index < shape.output_channels.size(); ++index) { + projections.push_back({.dst = buffers.shared_outputs[index].data(), + .packed_weight_bias = buffers.separate_packed_weights[index].data(), + .n = shape.output_channels[index]}); + } + KaiHelper kai; + if (!kai.matmul_shared_input_m1(buffers.input.data(), projections.data(), projections.size(), buffers.workspace.data(), + kInputChannels, tile, threads)) { + throw std::runtime_error("shared-input path rejected a valid LFM2 shape"); + } +} + +void runMerged(const ShapeCase& shape, int m, int threads, KaiTile tile, Buffers& buffers) { + KaiHelper kai; + kai.matmul(buffers.merged_output.data(), buffers.input.data(), buffers.merged_packed_weight.data(), buffers.workspace.data(), + m, kInputChannels, static_cast(totalOutputChannels(shape)), tile, threads); +} + +Comparison compareSeparate(const ShapeCase& shape, int m, const std::vector>& actual, + const std::vector>& expected) { + Comparison result; + for (size_t group = 0; group < shape.output_channels.size(); ++group) { + const size_t elements = static_cast(m) * shape.output_channels[group]; + for (size_t index = 0; index < elements; ++index) { + const float lhs = actual[group][index]; + const float rhs = expected[group][index]; + if (std::bit_cast(lhs) != std::bit_cast(rhs)) { ++result.bitwise_mismatches; } + result.max_absolute_error = std::max(result.max_absolute_error, std::fabs(lhs - rhs)); + } + } + return result; +} + +Comparison compareMerged(const ShapeCase& shape, int m, const Buffers& buffers) { + Comparison result; + const size_t total_n = totalOutputChannels(shape); + size_t group_offset = 0; + for (size_t group = 0; group < shape.output_channels.size(); ++group) { + const size_t group_n = static_cast(shape.output_channels[group]); + for (int row = 0; row < m; ++row) { + for (size_t column = 0; column < group_n; ++column) { + const float lhs = buffers.merged_output[static_cast(row) * total_n + group_offset + column]; + const float rhs = buffers.separate_outputs[group][static_cast(row) * group_n + column]; + if (std::bit_cast(lhs) != std::bit_cast(rhs)) { ++result.bitwise_mismatches; } + result.max_absolute_error = std::max(result.max_absolute_error, std::fabs(lhs - rhs)); + } + } + group_offset += group_n; + } + return result; +} + +uint64_t outputHash(const ShapeCase& shape, int m, const Buffers& buffers, std::string_view variant) { + constexpr uint64_t kOffset = 1469598103934665603ULL; + constexpr uint64_t kPrime = 1099511628211ULL; + uint64_t hash = kOffset; + auto mix = [&](float value) { + hash ^= std::bit_cast(value); + hash *= kPrime; + }; + const size_t total_n = totalOutputChannels(shape); + const int rows[] = {0, m / 2, m - 1}; + size_t group_offset = 0; + for (size_t group = 0; group < shape.output_channels.size(); ++group) { + const size_t group_n = static_cast(shape.output_channels[group]); + const size_t columns[] = {0, group_n / 2, group_n - 1}; + for (const int row : rows) { + for (const size_t column : columns) { + if (variant == "merged") { + mix(buffers.merged_output[static_cast(row) * total_n + group_offset + column]); + } else { + const auto& groups = variant == "shared" ? buffers.shared_outputs : buffers.separate_outputs; + mix(groups[group][static_cast(row) * group_n + column]); + } + } + } + group_offset += group_n; + } + return hash; +} + +template +double timeMicros(Function&& function) { + const auto start = std::chrono::steady_clock::now(); + function(); + const auto end = std::chrono::steady_clock::now(); + return std::chrono::duration(end - start).count(); +} + +void runVariant(std::string_view variant, const ShapeCase& shape, int m, int threads, KaiTile tile, Buffers& buffers) { + if (variant == "independent") { + runIndependent(shape, m, threads, tile, buffers); + } else if (variant == "shared") { + runShared(shape, m, threads, tile, buffers); + } else if (variant == "merged") { + runMerged(shape, m, threads, tile, buffers); + } else { + throw std::invalid_argument("unknown benchmark variant"); + } +} + +void runPair(std::string_view pair_name, std::string_view baseline, std::string_view candidate, const ShapeCase& shape, int m, + int threads, int repeats, KaiTile tile, Buffers& buffers) { + constexpr std::string_view kSchedule = "ABBA-BAAB"; + for (int warmup = 0; warmup < 2; ++warmup) { + runVariant(baseline, shape, m, threads, tile, buffers); + runVariant(candidate, shape, m, threads, tile, buffers); + } + int sample = 0; + for (int repeat = 0; repeat < repeats; ++repeat) { + int position = 0; + for (const char selector : kSchedule) { + if (selector == '-') { continue; } + const std::string_view variant = selector == 'A' ? baseline : candidate; + const double latency_us = timeMicros([&] { runVariant(variant, shape, m, threads, tile, buffers); }); + const uint64_t hash = outputHash(shape, m, buffers, variant); + std::printf("SAMPLE pair=%.*s repeat=%d position=%d sample=%d variant=%.*s latency_us=%.3f sentinel_hash=%016llx\n", + static_cast(pair_name.size()), pair_name.data(), repeat, position, sample, + static_cast(variant.size()), variant.data(), latency_us, static_cast(hash)); + ++position; + ++sample; + } + } +} + +} // namespace + +int main(int argc, char** argv) { + try { + mllm::initializeContext(); + if (argc < 4 || argc > 5) { + std::fprintf(stderr, "usage: %s [schedule_repeats]\n", argv[0]); + return 2; + } + const ShapeCase shape = parseShape(argv[1]); + const int m = parsePositiveInt(argv[2], "M"); + const int threads = parsePositiveInt(argv[3], "threads"); + const int repeats = argc == 5 ? parsePositiveInt(argv[4], "schedule_repeats") : 2; + const KaiTile tile = m == 1 ? kDecodeTile : kPrefillTile; + + std::printf("LFM2_PARALLEL_LINEAR_SCREEN_CONFIG shape=%.*s m=%d k=%d groups=%zu total_n=%zu threads=%d " + "schedule=ABBA-BAAB repeats=%d tile=%s\n", + static_cast(shape.name.size()), shape.name.data(), m, kInputChannels, shape.output_channels.size(), + totalOutputChannels(shape), threads, repeats, m == 1 ? "dotprod_1x8" : "i8mm_4x8"); + std::printf("LFM2_PARALLEL_LINEAR_SCREEN_PROVENANCE=replica\n"); + + Buffers buffers = makeBuffers(shape, m, tile); + runIndependent(shape, m, threads, tile, buffers); + runMerged(shape, m, threads, tile, buffers); + const Comparison merged_comparison = compareMerged(shape, m, buffers); + std::printf("CORRECTNESS variant=merged bitwise_mismatches=%zu max_abs_error=%.9g\n", merged_comparison.bitwise_mismatches, + merged_comparison.max_absolute_error); + if (merged_comparison.bitwise_mismatches != 0) { return 3; } + + if (m == 1) { + runShared(shape, m, threads, tile, buffers); + const Comparison shared_comparison = compareSeparate(shape, m, buffers.shared_outputs, buffers.separate_outputs); + std::printf("CORRECTNESS variant=shared bitwise_mismatches=%zu max_abs_error=%.9g\n", + shared_comparison.bitwise_mismatches, shared_comparison.max_absolute_error); + if (shared_comparison.bitwise_mismatches != 0) { return 4; } + runPair("shared_vs_merged", "shared", "merged", shape, m, threads, repeats, tile, buffers); + runPair("independent_vs_shared", "independent", "shared", shape, m, threads, repeats, tile, buffers); + } else { + runPair("independent_vs_merged", "independent", "merged", shape, m, threads, repeats, tile, buffers); + } + + std::printf("LFM2_PARALLEL_LINEAR_SCREEN_OK\n"); + return 0; + } catch (const std::exception& error) { + std::fprintf(stderr, "LFM2_PARALLEL_LINEAR_SCREEN_ERROR %s\n", error.what()); + return 1; + } +} diff --git a/benchmarks/cpu/lfm2_parallel_linear_shared_mx.cpp b/benchmarks/cpu/lfm2_parallel_linear_shared_mx.cpp new file mode 100644 index 000000000..8044aec7b --- /dev/null +++ b/benchmarks/cpu/lfm2_parallel_linear_shared_mx.cpp @@ -0,0 +1,265 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "mllm/backends/cpu/kernels/arm/linear/kai.hpp" +#include "mllm/backends/cpu/ops/ParallelLinearOp.hpp" +#include "mllm/mllm.hpp" + +namespace { + +using KaiHelper = mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk; +using KaiTile = KaiHelper::Tiles; + +constexpr int kInputChannels = 2048; +constexpr KaiTile kPrefillTile = KaiTile::qai8dxp4x8_qsi4c32p8x8_4x8x32; + +struct ShapeCase { + std::string_view name; + std::vector output_channels; + std::vector projection_names; +}; + +struct Buffers { + mllm::Tensor input; + std::vector weights; + std::vector packed_weights; + std::vector workspace; + std::vector> independent_outputs; + std::unique_ptr parallel_op; + std::vector shared_outputs; +}; + +struct Comparison { + size_t bitwise_mismatches = 0; + float max_absolute_error = 0.0F; +}; + +uint32_t nextRandom(uint32_t& state) { + state = state * 1664525U + 1013904223U; + return state; +} + +float deterministicValue(uint32_t& state) { + const int32_t centered = static_cast((nextRandom(state) >> 8U) % 2001U) - 1000; + return static_cast(centered) / 4096.0F; +} + +int parsePositiveInt(const char* value, const char* name) { + char* end = nullptr; + const int64_t parsed = std::strtoll(value, &end, 10); + if (end == value || *end != '\0' || parsed <= 0 || parsed > std::numeric_limits::max()) { + throw std::invalid_argument(std::string(name) + " must be a positive integer"); + } + return static_cast(parsed); +} + +ShapeCase parseShape(std::string_view name) { + if (name == "gate_up") { return {.name = "gate_up", .output_channels = {10752, 10752}, .projection_names = {"w1", "w3"}}; } + if (name == "qkv") { + return {.name = "qkv", .output_channels = {2048, 512, 512}, .projection_names = {"q_proj", "k_proj", "v_proj"}}; + } + throw std::invalid_argument("shape must be gate_up or qkv"); +} + +size_t totalOutputChannels(const ShapeCase& shape) { + size_t total = 0; + for (const int n : shape.output_channels) { total += static_cast(n); } + return total; +} + +Buffers makeBuffers(const ShapeCase& shape, int m, int threads) { + KaiHelper kai; + Buffers buffers; + const size_t total_n = totalOutputChannels(shape); + + buffers.input = mllm::Tensor::empty({1, m, kInputChannels}, mllm::kFloat32, mllm::kCPU).alloc(); + buffers.weights.resize(total_n * kInputChannels); + uint32_t random_state = 0x4C464D32U; + std::generate(buffers.input.ptr(), buffers.input.ptr() + buffers.input.numel(), + [&] { return deterministicValue(random_state); }); + std::generate(buffers.weights.begin(), buffers.weights.end(), [&] { return deterministicValue(random_state); }); + + buffers.packed_weights.reserve(shape.output_channels.size()); + buffers.independent_outputs.reserve(shape.output_channels.size()); + auto parameters = mllm::ParameterFile::create(); + size_t row_offset = 0; + for (size_t index = 0; index < shape.output_channels.size(); ++index) { + const int n = shape.output_channels[index]; + const size_t packed_size_value = kai.quant_pack_rhs_size(n, kInputChannels, kPrefillTile); + if (packed_size_value > static_cast(std::numeric_limits::max())) { + throw std::overflow_error("packed weight exceeds Tensor dimension range"); + } + const int packed_size = static_cast(packed_size_value); + auto packed = mllm::Tensor::empty({packed_size}, mllm::kInt8, mllm::kCPU) + .setMemType(mllm::kParamsNormal) + .setName("screen." + shape.projection_names[index] + ".weight") + .alloc(); + kai.quant_pack_rhs_offline(packed.ptr(), buffers.weights.data() + row_offset * kInputChannels, nullptr, n, + kInputChannels, kPrefillTile); + parameters->push(packed.name(), packed); + buffers.packed_weights.push_back(std::move(packed)); + buffers.independent_outputs.emplace_back(static_cast(m) * n); + row_offset += static_cast(n); + } + buffers.workspace.resize(kai.workspace_size(m, kInputChannels, kPrefillTile)); + std::vector().swap(buffers.weights); + + mllm::aops::ParallelLinearOpOptions options{ + .in_channels = kInputChannels, + .out_channels = shape.output_channels, + .projection_names = shape.projection_names, + .bias = false, + .impl_type = mllm::aops::LinearImplTypes::kKaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp1x8_qsi4c32p8x8_1x8x32, + .kai_w4a32_decode_thread_cap = 4, + .kai_w4a32_prefill_thread_cap = 6}; + options.setThreads(threads); + buffers.parallel_op = std::make_unique(options); + buffers.parallel_op->setName("screen.parallel"); + buffers.parallel_op->load(parameters); + buffers.parallel_op->reshape({buffers.input}, buffers.shared_outputs); + buffers.parallel_op->setup({buffers.input}, buffers.shared_outputs); + return buffers; +} + +void runIndependent(const ShapeCase& shape, int m, int threads, Buffers& buffers) { + KaiHelper kai; + for (size_t index = 0; index < shape.output_channels.size(); ++index) { + kai.matmul(buffers.independent_outputs[index].data(), buffers.input.ptr(), + buffers.packed_weights[index].ptr(), buffers.workspace.data(), m, kInputChannels, + shape.output_channels[index], kPrefillTile, threads); + } +} + +void runSharedMx(const ShapeCase&, int, int, Buffers& buffers) { + buffers.parallel_op->forward({buffers.input}, buffers.shared_outputs); +} + +Comparison compareOutputs(const ShapeCase& shape, int m, const Buffers& buffers) { + Comparison result; + for (size_t group = 0; group < shape.output_channels.size(); ++group) { + const size_t elements = static_cast(m) * shape.output_channels[group]; + for (size_t index = 0; index < elements; ++index) { + const float actual = buffers.shared_outputs[group].ptr()[index]; + const float expected = buffers.independent_outputs[group][index]; + if (std::bit_cast(actual) != std::bit_cast(expected)) { ++result.bitwise_mismatches; } + result.max_absolute_error = std::max(result.max_absolute_error, std::fabs(actual - expected)); + } + } + return result; +} + +uint64_t outputHash(const ShapeCase& shape, int m, const Buffers& buffers, std::string_view variant) { + constexpr uint64_t kOffset = 1469598103934665603ULL; + constexpr uint64_t kPrime = 1099511628211ULL; + uint64_t hash = kOffset; + const int rows[] = {0, m / 2, m - 1}; + for (size_t group = 0; group < shape.output_channels.size(); ++group) { + const size_t group_n = static_cast(shape.output_channels[group]); + const size_t columns[] = {0, group_n / 2, group_n - 1}; + for (const int row : rows) { + for (const size_t column : columns) { + const size_t index = static_cast(row) * group_n + column; + const float value = variant == "shared_mx" ? buffers.shared_outputs[group].ptr()[index] + : buffers.independent_outputs[group][index]; + hash ^= std::bit_cast(value); + hash *= kPrime; + } + } + } + return hash; +} + +void runVariant(std::string_view variant, const ShapeCase& shape, int m, int threads, Buffers& buffers) { + if (variant == "independent") { + runIndependent(shape, m, threads, buffers); + } else if (variant == "shared_mx") { + runSharedMx(shape, m, threads, buffers); + } else { + throw std::invalid_argument("unknown benchmark variant"); + } +} + +template +double timeMicros(Function&& function) { + const auto start = std::chrono::steady_clock::now(); + function(); + const auto end = std::chrono::steady_clock::now(); + return std::chrono::duration(end - start).count(); +} + +void runPair(const ShapeCase& shape, int m, int threads, int repeats, Buffers& buffers) { + constexpr std::string_view kSchedule = "ABBA-BAAB"; + for (int warmup = 0; warmup < 2; ++warmup) { + runIndependent(shape, m, threads, buffers); + runSharedMx(shape, m, threads, buffers); + } + + int sample = 0; + for (int repeat = 0; repeat < repeats; ++repeat) { + int position = 0; + for (const char selector : kSchedule) { + if (selector == '-') { continue; } + const std::string_view variant = selector == 'A' ? "independent" : "shared_mx"; + const double latency_us = timeMicros([&] { runVariant(variant, shape, m, threads, buffers); }); + const uint64_t hash = outputHash(shape, m, buffers, variant); + std::printf("SAMPLE pair=independent_vs_shared_mx repeat=%d position=%d sample=%d variant=%.*s " + "latency_us=%.3f sentinel_hash=%016llx\n", + repeat, position, sample, static_cast(variant.size()), variant.data(), latency_us, + static_cast(hash)); + ++position; + ++sample; + } + } +} + +} // namespace + +int main(int argc, char** argv) { + try { + mllm::initializeContext(); + if (argc < 4 || argc > 5) { + std::fprintf(stderr, "usage: %s [schedule_repeats]\n", argv[0]); + return 2; + } + const ShapeCase shape = parseShape(argv[1]); + const int m = parsePositiveInt(argv[2], "M"); + const int threads = parsePositiveInt(argv[3], "threads"); + const int repeats = argc == 5 ? parsePositiveInt(argv[4], "schedule_repeats") : 2; + if (m < 4) { throw std::invalid_argument("shared Mx first-class CPU op screen requires M >= 4"); } + + std::printf("LFM2_PARALLEL_LINEAR_SHARED_MX_SCREEN_CONFIG shape=%.*s m=%d k=%d groups=%zu total_n=%zu " + "threads=%d schedule=ABBA-BAAB repeats=%d tile=i8mm_4x8\n", + static_cast(shape.name.size()), shape.name.data(), m, kInputChannels, shape.output_channels.size(), + totalOutputChannels(shape), threads, repeats); + std::printf("LFM2_PARALLEL_LINEAR_SHARED_MX_SCREEN_PROVENANCE=first_class_cpu_op\n"); + + Buffers buffers = makeBuffers(shape, m, threads); + runIndependent(shape, m, threads, buffers); + runSharedMx(shape, m, threads, buffers); + const Comparison comparison = compareOutputs(shape, m, buffers); + std::printf("CORRECTNESS variant=shared_mx bitwise_mismatches=%zu max_abs_error=%.9g\n", comparison.bitwise_mismatches, + comparison.max_absolute_error); + if (comparison.bitwise_mismatches != 0) { return 3; } + + runPair(shape, m, threads, repeats, buffers); + std::printf("LFM2_PARALLEL_LINEAR_SHARED_MX_SCREEN_OK\n"); + return 0; + } catch (const std::exception& error) { + std::fprintf(stderr, "LFM2_PARALLEL_LINEAR_SHARED_MX_SCREEN_ERROR %s\n", error.what()); + return 1; + } +} diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index e114c1d57..ed1cbeb82 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -8,6 +8,7 @@ add_subdirectory(minicpm4) add_subdirectory(minicpm5) add_subdirectory(qwen3) add_subdirectory(qwen3_5) +add_subdirectory(lfm2) add_subdirectory(qwen3_service) add_subdirectory(qwen3_moe) add_subdirectory(deepseek_ocr) diff --git a/examples/lfm2/CMakeLists.txt b/examples/lfm2/CMakeLists.txt new file mode 100644 index 000000000..5e7f232da --- /dev/null +++ b/examples/lfm2/CMakeLists.txt @@ -0,0 +1,3 @@ +add_executable(mllm-lfm2-runner main.cpp) +target_link_libraries(mllm-lfm2-runner PRIVATE MllmRT MllmCPUBackend) +target_include_directories(mllm-lfm2-runner PRIVATE ${MLLM_INCLUDE_DIR}) diff --git a/examples/lfm2/README.md b/examples/lfm2/README.md new file mode 100644 index 000000000..f1fa3cc0c --- /dev/null +++ b/examples/lfm2/README.md @@ -0,0 +1,123 @@ +# LFM2.5-2.6B on ARM CPU + +This example supports the text-only +[`LiquidAI/LFM2.5-2.6B`](https://huggingface.co/LiquidAI/LFM2.5-2.6B) +checkpoint. The runtime binds the official 30-layer physical schedule: 22 +stateful short-convolution layers and 8 full-attention layers. + +The attention path uses 32 query heads and 8 native KV heads. Only the eight +full-attention layers allocate KV-cache slots; the cache never expands KV +history to query-head count. Each convolution layer independently retains the +two historical FP32 samples required by its three-tap causal kernel across +prefill and decode. Batch size is 1 and this +product configuration limits the runtime cache to 2048 tokens. + +## Validate and convert + +Run from the repository root. The first audit rejects architecture, physical +layer schedule, tensor-shape, or quantization-recipe drift before conversion. + +```bash +python examples/lfm2/validate_checkpoint.py /path/to/LFM2.5-2.6B + +python -m pymllm.mobile.utils.mllm_convertor \ + --input_path /path/to/LFM2.5-2.6B \ + --output_path /path/to/lfm2.5-2.6b-w4a32-kai.mllm \ + --model_name LFM2.5-2.6B \ + --cfg_path examples/lfm2/quant_cfg_2.6B_w4a32_kai.json \ + --pipeline w4a32_kai_pipeline \ + --format v2 \ + --verbose + +python examples/lfm2/validate_converted_model.py \ + /path/to/lfm2.5-2.6b-w4a32-kai.mllm \ + /path/to/LFM2.5-2.6B +``` + +Linear weights use the existing KleidiAI dynamic-INT8-activation / INT4-weight +packing path. The depthwise convolution, norms, and lookup embedding remain +FP32. Because the checkpoint ties its output head to the embedding, conversion +retains `model.embed_tokens.weight` and creates the packed +`lm_head_out.weight` alias used by the runtime output projection. + +## Build and run + +```bash +cmake -S . -B build -DMLLM_ENABLE_EXAMPLE=ON +cmake --build build --target mllm-lfm2-runner -j + +build/bin/mllm-lfm2-runner \ + --model_path /path/to/lfm2.5-2.6b-w4a32-kai.mllm \ + --model_version v2 \ + --tokenizer_path /path/to/LFM2.5-2.6B/tokenizer.json \ + --config_path examples/lfm2/config_2.6B_w4a32_kai.json \ + --prompt "Explain how to make mobile language-model inference reliable, covering correctness, state reset, artifact verification, and performance measurement." \ + --max_new_tokens 128 \ + --min_new_tokens 128 \ + --print_token_ids +``` + +For the Android ARM build, select the repository's OpenMP thread vendor. The +CPU backend applies OpenMP only to translation units that own a parallel region; +the LFM2.5 runner does not require a model-specific OpenMP build option. + +```bash +cmake -S . -B build-android \ + -DCMAKE_TOOLCHAIN_FILE="$ANDROID_NDK_HOME/build/cmake/android.toolchain.cmake" \ + -DANDROID_ABI=arm64-v8a \ + -DANDROID_PLATFORM=android-28 \ + -DMLLM_CROSS_COMPILE=ON \ + -DMLLM_BUILD_ARM_BACKEND=ON \ + -DMLLM_ENABLE_EXAMPLE=ON \ + -DMLLM_KERNEL_USE_THREADS=ON \ + -DMLLM_KERNEL_THREADS_VENDOR_OPENMP=ON \ + -DMLLM_KERNEL_USE_THREADS_VENDOR_MLLM=OFF + +cmake --build build-android --target mllm-lfm2-runner -j +``` + +Omit `--prompt` for the interactive loop. Each prompt begins with +`resetState()`, clearing all eight logical attention slots and all 22 +short-convolution histories. The tokenizer applies the checkpoint's byte-level +BPE contract and ends the generation prompt exactly at ``. For a demo +receipt, use `demo_prompt.txt` unchanged on the host and Android, request a long +deterministic continuation with equal `max_new_tokens` and `min_new_tokens`, +and retain the printed prompt, response, token IDs, and generated-token count. +The runner deliberately streams the checkpoint output verbatim: reasoning text +and a generated `` marker are displayed rather than separated or +suppressed. This is a transparent demo behavior, not a product chat-surface +contract. + +## Product benchmark records + +Benchmark mode requires a file-backed prompt, its frozen token count, exact +model/source identities, and a fresh JSONL destination. It resets all model +state between requests, forces an exact generated-token count, and records +prefill, TTFT, decode, wall time, affinity, CPU-frequency, governor, and thermal +telemetry for every sample. + +```bash +build/bin/mllm-lfm2-runner \ + --model_path /path/to/lfm2.5-2.6b-w4a32-kai.mllm \ + --model_version v2 \ + --tokenizer_path /path/to/LFM2.5-2.6B/tokenizer.json \ + --config_path examples/lfm2/config_2.6B_w4a32_kai.json \ + --prompt_file /path/to/prompt.txt \ + --expected_prompt_tokens PROMPT_TOKENS \ + --max_new_tokens 32 \ + --benchmark_warmup 1 \ + --benchmark_samples 5 \ + --benchmark_jsonl /path/to/fresh-results.jsonl \ + --benchmark_variant MODEL_SHA256 \ + --benchmark_source_manifest SOURCE_MANIFEST_SHA256 +``` + +Use `--system_prompt` for the optional system role. `--tools_json` accepts +either one JSON tool-schema object or an array of schemas. Objects are rendered +with the pinned template's Python `json.dumps` spacing; string elements are +inserted byte-for-byte into `List of tools: [...]`. The runner returns the +model's tool-call text unchanged and does not parse or execute calls. + +This implementation and its host tests establish local source/build +correctness. They do not by themselves claim output quality, Android +execution, or device performance. diff --git a/examples/lfm2/benchmark_harness.hpp b/examples/lfm2/benchmark_harness.hpp new file mode 100644 index 000000000..e7e4db2c8 --- /dev/null +++ b/examples/lfm2/benchmark_harness.hpp @@ -0,0 +1,118 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#if defined(__ANDROID__) || defined(__linux__) +#include +#endif + +namespace mllm::examples::lfm2::benchmark { + +inline std::optional readTextFile(const std::filesystem::path& path) { + std::ifstream stream(path); + if (!stream) return std::nullopt; + std::ostringstream contents; + contents << stream.rdbuf(); + auto value = contents.str(); + while (!value.empty() && (value.back() == '\n' || value.back() == '\r')) value.pop_back(); + return value; +} + +inline std::optional readIntegerFile(const std::filesystem::path& path) { + const auto text = readTextFile(path); + if (!text.has_value()) return std::nullopt; + try { + size_t consumed = 0; + const auto value = std::stoll(*text, &consumed); + if (consumed != text->size()) return std::nullopt; + return value; + } catch (...) { return std::nullopt; } +} + +inline std::vector currentAffinityCpus() { + std::vector cpus; +#if defined(__ANDROID__) || defined(__linux__) + cpu_set_t mask; + CPU_ZERO(&mask); + if (sched_getaffinity(0, sizeof(mask), &mask) != 0) return cpus; + for (int cpu = 0; cpu < CPU_SETSIZE; ++cpu) { + if (CPU_ISSET(cpu, &mask)) cpus.push_back(cpu); + } +#endif + return cpus; +} + +inline nlohmann::json captureTelemetry() { + using Json = nlohmann::json; + Json snapshot = { +#if defined(__ANDROID__) + {"platform", "android"}, +#elif defined(__linux__) + {"platform", "linux"}, +#elif defined(__APPLE__) + {"platform", "macos"}, +#else + {"platform", "other"}, +#endif + {"captured_epoch_us", + std::chrono::duration_cast(std::chrono::system_clock::now().time_since_epoch()).count()}, + }; + + const auto affinity = currentAffinityCpus(); + snapshot["affinity_cpus"] = affinity; + snapshot["cpus"] = Json::array(); + snapshot["ceiling_vector"] = Json::array(); + snapshot["thermal_zones"] = Json::array(); + + for (const int cpu : affinity) { + const auto base = std::filesystem::path("/sys/devices/system/cpu") / ("cpu" + std::to_string(cpu)); + const auto cpufreq = base / "cpufreq"; + const auto online = cpu == 0 ? std::optional(1) : readIntegerFile(base / "online"); + const auto cpuinfo_max = readIntegerFile(cpufreq / "cpuinfo_max_freq"); + const auto scaling_max = readIntegerFile(cpufreq / "scaling_max_freq"); + const auto scaling_cur = readIntegerFile(cpufreq / "scaling_cur_freq"); + const auto governor = readTextFile(cpufreq / "scaling_governor"); + snapshot["cpus"].push_back({ + {"cpu", cpu}, + {"online", online.has_value() ? Json(*online) : Json(nullptr)}, + {"cpuinfo_max_freq", cpuinfo_max.has_value() ? Json(*cpuinfo_max) : Json(nullptr)}, + {"scaling_max_freq", scaling_max.has_value() ? Json(*scaling_max) : Json(nullptr)}, + {"scaling_cur_freq", scaling_cur.has_value() ? Json(*scaling_cur) : Json(nullptr)}, + {"scaling_governor", governor.has_value() ? Json(*governor) : Json(nullptr)}, + }); + snapshot["ceiling_vector"].push_back(scaling_max.has_value() ? Json(*scaling_max) : Json(nullptr)); + } + + const std::filesystem::path thermal_root("/sys/class/thermal"); + std::error_code error; + if (std::filesystem::exists(thermal_root, error)) { + std::vector zones; + for (const auto& entry : std::filesystem::directory_iterator(thermal_root, error)) { + if (entry.path().filename().string().starts_with("thermal_zone")) zones.push_back(entry.path()); + } + std::sort(zones.begin(), zones.end()); + for (const auto& zone : zones) { + const auto type = readTextFile(zone / "type"); + const auto temp = readIntegerFile(zone / "temp"); + snapshot["thermal_zones"].push_back({ + {"zone", zone.filename().string()}, + {"type", type.has_value() ? Json(*type) : Json(nullptr)}, + {"temp_milli_c", temp.has_value() ? Json(*temp) : Json(nullptr)}, + }); + } + } + return snapshot; +} + +} // namespace mllm::examples::lfm2::benchmark diff --git a/examples/lfm2/config_2.6B_w4a32_kai.json b/examples/lfm2/config_2.6B_w4a32_kai.json new file mode 100644 index 000000000..689e0f3a1 --- /dev/null +++ b/examples/lfm2/config_2.6B_w4a32_kai.json @@ -0,0 +1,30 @@ +{ + "architectures": ["Lfm2ForCausalLM"], + "model_type": "lfm2", + "hidden_size": 2048, + "intermediate_size": 10752, + "num_hidden_layers": 30, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "head_dim": 64, + "conv_L_cache": 3, + "conv_bias": false, + "block_auto_adjust_ff_dim": false, + "block_ffn_dim_multiplier": 1.0, + "block_multiple_of": 256, + "norm_eps": 0.00001, + "max_position_embeddings": 131072, + "vocab_size": 128000, + "tie_word_embeddings": true, + "bos_token_id": 124894, + "eos_token_id": 124900, + "pad_token_id": 124893, + "rope_parameters": {"rope_theta": 10000000.0, "rope_type": "default"}, + "layer_types": [ + "conv", "conv", "full_attention", "conv", "conv", "full_attention", "conv", "conv", "conv", "full_attention", + "conv", "conv", "conv", "full_attention", "conv", "conv", "conv", "full_attention", "conv", "conv", "conv", + "full_attention", "conv", "conv", "full_attention", "conv", "conv", "full_attention", "conv", "conv" + ], + "max_cache_length": 2048, + "linear_impl_type": "KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp1x8_qsi4c32p8x8_1x8x32" +} diff --git a/examples/lfm2/demo_prompt.txt b/examples/lfm2/demo_prompt.txt new file mode 100644 index 000000000..3f572ec11 --- /dev/null +++ b/examples/lfm2/demo_prompt.txt @@ -0,0 +1 @@ +Explain how to make mobile language-model inference reliable, covering correctness, state reset, artifact verification, and performance measurement. diff --git a/examples/lfm2/main.cpp b/examples/lfm2/main.cpp new file mode 100644 index 000000000..9b2e708e4 --- /dev/null +++ b/examples/lfm2/main.cpp @@ -0,0 +1,274 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include + +#include "benchmark_harness.hpp" + +using mllm::Argparse; + +namespace { + +auto pythonJson(const nlohmann::ordered_json& value) -> std::string { + if (value.is_array()) { + std::string output = "["; + for (size_t index = 0; index < value.size(); ++index) { + if (index != 0) output += ", "; + output += pythonJson(value[index]); + } + return output + "]"; + } + if (value.is_object()) { + std::string output = "{"; + size_t index = 0; + for (const auto& [key, item] : value.items()) { + if (index++ != 0) output += ", "; + output += nlohmann::ordered_json(key).dump() + ": " + pythonJson(item); + } + return output + "}"; + } + return value.dump(-1, ' ', false, nlohmann::ordered_json::error_handler_t::strict); +} + +} // namespace + +MLLM_MAIN({ + auto engine_args = mllm::engineArgAttach(); + auto& help = Argparse::add("-h|--help").help("Show help message"); + auto& model_path = Argparse::add("-m|--model_path").help("Converted model path").required(true); + auto& model_version = Argparse::add("-mv|--model_version").help("Model version: v1 or v2").required(true); + auto& tokenizer_path = Argparse::add("-t|--tokenizer_path").help("Tokenizer JSON path").required(true); + auto& config_path = Argparse::add("-c|--config_path").help("Runtime config path").required(true); + auto& prompt = Argparse::add("-p|--prompt").help("Run one prompt non-interactively").required(false); + auto& prompt_file = Argparse::add("--prompt_file").help("Read one benchmark prompt from a file").required(false); + auto& system_prompt = Argparse::add("--system_prompt").help("Optional system prompt").required(false); + auto& tools_json = Argparse::add("--tools_json") + .help("JSON file containing one tool schema or an array of schemas") + .required(false); + auto& max_new_tokens = Argparse::add("-g|--max_new_tokens").help("Maximum generated tokens").required(false); + auto& min_new_tokens = + Argparse::add("--min_new_tokens").help("Minimum generated tokens before EOS can stop generation").required(false); + auto& print_token_ids = Argparse::add("--print_token_ids").help("Print generated token IDs").required(false); + auto& benchmark_warmup = + Argparse::add("--benchmark_warmup").help("Unrecorded benchmark warmup requests").required(false); + auto& benchmark_samples = Argparse::add("--benchmark_samples").help("Measured benchmark requests").required(false); + auto& benchmark_jsonl = Argparse::add("--benchmark_jsonl").help("Fresh JSONL output path").required(false); + auto& benchmark_variant = + Argparse::add("--benchmark_variant").help("Bound model artifact identity").required(false); + auto& benchmark_source_manifest = + Argparse::add("--benchmark_source_manifest").help("Bound dirty-source manifest SHA-256").required(false); + auto& expected_prompt_tokens = + Argparse::add("--expected_prompt_tokens").help("Fail if tokenized prompt length differs").required(false); + + for (int index = 1; index < argc; ++index) { + if (std::string(argv[index]) == "-h" || std::string(argv[index]) == "--help") { + Argparse::printHelp(); + return 0; + } + } + Argparse::parse(argc, argv); + mllm::configEngineWithArgs(engine_args); + (void)help; + + mllm::ModelFileVersion file_version; + if (model_version.get() == "v1") { + file_version = mllm::ModelFileVersion::kV1; + } else if (model_version.get() == "v2") { + file_version = mllm::ModelFileVersion::kV2; + } else { + throw std::invalid_argument("model_version must be v1 or v2"); + } + + auto cfg = mllm::models::lfm2::Lfm2Config(config_path.get()); + int generation_limit = max_new_tokens.isSet() ? max_new_tokens.get() : 64; + if (generation_limit <= 0 || generation_limit > cfg.max_cache_length) { + throw std::invalid_argument("max_new_tokens must be between 1 and max_cache_length"); + } + int minimum_generation = min_new_tokens.isSet() ? min_new_tokens.get() : 0; + if (minimum_generation < 0 || minimum_generation > generation_limit) { + throw std::invalid_argument("min_new_tokens must be between 0 and max_new_tokens"); + } + const bool benchmark_mode = prompt_file.isSet() || benchmark_warmup.isSet() || benchmark_samples.isSet() + || benchmark_jsonl.isSet() || benchmark_variant.isSet() || benchmark_source_manifest.isSet() + || expected_prompt_tokens.isSet(); + if (prompt.isSet() && prompt_file.isSet()) throw std::invalid_argument("prompt and prompt_file are mutually exclusive"); + if (benchmark_mode) { + if (!prompt_file.isSet() || !benchmark_samples.isSet() || !benchmark_jsonl.isSet() || !benchmark_variant.isSet() + || !benchmark_source_manifest.isSet() || !expected_prompt_tokens.isSet()) { + throw std::invalid_argument("benchmark mode requires prompt_file, benchmark_samples, benchmark_jsonl, " + "benchmark_variant, benchmark_source_manifest, and expected_prompt_tokens"); + } + if (benchmark_samples.get() <= 0 || (benchmark_warmup.isSet() && benchmark_warmup.get() < 0)) { + throw std::invalid_argument("benchmark sample counts must be non-negative and measured samples must be positive"); + } + if (generation_limit < 2) throw std::invalid_argument("benchmark max_new_tokens must be at least 2"); + if (tools_json.isSet() || system_prompt.isSet()) { + throw std::invalid_argument("benchmark mode does not accept system_prompt or tools_json"); + } + if (min_new_tokens.isSet()) throw std::invalid_argument("benchmark mode forces an exact generated-token count"); + const std::filesystem::path output_path(benchmark_jsonl.get()); + std::error_code error; + if (std::filesystem::exists(output_path) && std::filesystem::file_size(output_path, error) != 0) { + throw std::invalid_argument("benchmark_jsonl must be new or empty"); + } + } + const auto model_load_start = std::chrono::steady_clock::now(); + auto parameters = mllm::load(model_path.get(), file_version); + mllm::models::lfm2::validateModelConfigMatch(cfg, parameters); + auto tokenizer = mllm::models::lfm2::Lfm2Tokenizer(tokenizer_path.get()); + auto model = mllm::models::lfm2::Lfm2ForCausalLM(cfg); + model.load(parameters); + const auto model_load_end = std::chrono::steady_clock::now(); + const auto model_load_duration_us = + std::chrono::duration_cast(model_load_end - model_load_start).count(); + fmt::print("LFM2.5-2.6B: {} layers ({} attention + {} short convolution)\n", cfg.num_hidden_layers, cfg.numAttentionLayers(), + cfg.numConvLayers()); + + int exit_code = 0; + std::vector raw_tools; + if (tools_json.isSet()) { + std::ifstream stream(tools_json.get(), std::ios::binary); + if (!stream) throw std::invalid_argument("unable to read tools_json"); + nlohmann::ordered_json tools; + stream >> tools; + if (tools.is_object()) { + raw_tools.push_back(pythonJson(tools)); + } else if (tools.is_array()) { + for (const auto& tool : tools) { raw_tools.push_back(tool.is_string() ? tool.get() : pythonJson(tool)); } + } else { + throw std::invalid_argument("tools_json must contain an object or an array of strings"); + } + } + if (benchmark_mode) { + std::ifstream prompt_stream(prompt_file.get(), std::ios::binary); + if (!prompt_stream) throw std::invalid_argument("unable to read prompt_file"); + std::string benchmark_prompt(std::istreambuf_iterator(prompt_stream), {}); + while (!benchmark_prompt.empty() && (benchmark_prompt.back() == '\n' || benchmark_prompt.back() == '\r')) { + benchmark_prompt.pop_back(); + } + if (benchmark_prompt.empty()) throw std::invalid_argument("prompt_file must not be empty"); + const auto inputs = tokenizer.convertMessage({.prompt = benchmark_prompt}); + const auto prompt_tokens = inputs.at("sequence").shape()[1]; + if (prompt_tokens != expected_prompt_tokens.get()) { + throw std::invalid_argument( + fmt::format("prompt token count {} differs from expected {}", prompt_tokens, expected_prompt_tokens.get())); + } + if (prompt_tokens + generation_limit - 1 > cfg.max_cache_length) { + throw std::invalid_argument("benchmark prompt plus generation exceeds max_cache_length"); + } + + std::ofstream jsonl(benchmark_jsonl.get(), std::ios::out | std::ios::trunc); + if (!jsonl) throw std::invalid_argument("unable to open benchmark_jsonl"); + const int warmups = benchmark_warmup.isSet() ? benchmark_warmup.get() : 0; + for (int request = 0; request < warmups + benchmark_samples.get(); ++request) { + const bool warmup = request < warmups; + model.resetState(); + const auto telemetry_before = mllm::examples::lfm2::benchmark::captureTelemetry(); + std::vector generated_token_ids; + const auto request_start = std::chrono::steady_clock::now(); + model.streamGenerate(inputs, + {{"max_length", mllm::AnyValue(generation_limit)}, + {"min_new_tokens", mllm::AnyValue(generation_limit)}, + {"do_sample", mllm::AnyValue(false)}}, + [&](int64_t token_id) { generated_token_ids.push_back(token_id); }); + const auto request_end = std::chrono::steady_clock::now(); + const auto stats = model.perfStats(); + std::vector invalid_reasons; + if (!stats.valid) invalid_reasons.push_back("invalid_performance_stats"); + if (!stats.completed) invalid_reasons.push_back("incomplete_generation"); + if (stats.prefill_tokens != prompt_tokens) invalid_reasons.push_back("prefill_token_count_mismatch"); + if (stats.generated_tokens != generation_limit || stats.decode_steps != generation_limit - 1 + || generated_token_ids.size() != static_cast(generation_limit)) { + invalid_reasons.push_back("generation_length_mismatch"); + } + nlohmann::json record = { + {"schema", "mllm.lfm25.product_benchmark.v1"}, + {"variant", benchmark_variant.get()}, + {"source_manifest_sha256", benchmark_source_manifest.get()}, + {"request_index", request}, + {"warmup", warmup}, + {"prompt_tokens", prompt_tokens}, + {"max_new_tokens", generation_limit}, + {"cpu_op_threads", mllm::Context::instance().getCpuOpThreads()}, + {"model_load_duration_us", model_load_duration_us}, + {"request_wall_duration_us", + std::chrono::duration_cast(request_end - request_start).count()}, + {"generated_token_ids", generated_token_ids}, + {"telemetry_before", telemetry_before}, + {"telemetry_after", mllm::examples::lfm2::benchmark::captureTelemetry()}, + {"stats", + {{"valid", stats.valid}, + {"completed", stats.completed}, + {"total_duration_us", stats.total_duration_us}, + {"prefill_duration_us", stats.prefill_duration_us}, + {"decode_duration_us", stats.decode_duration_us}, + {"ttft_duration_us", stats.ttft_duration_us}, + {"prefill_tokens", stats.prefill_tokens}, + {"generated_tokens", stats.generated_tokens}, + {"decode_steps", stats.decode_steps}}}, + {"invalid_reasons", invalid_reasons}, + {"status", invalid_reasons.empty() ? "ok" : "invalid"}, + }; + jsonl << record.dump() << '\n'; + jsonl.flush(); + if (!jsonl) throw std::runtime_error("failed to write benchmark_jsonl"); + if (!invalid_reasons.empty()) { + exit_code = 2; + break; + } + } + if (exit_code == 0) fmt::print("Benchmark records: {}\n", benchmark_jsonl.get()); + } else + while (true) { + std::string prompt_text = prompt.isSet() ? prompt.get() : ""; + if (!prompt.isSet()) { + fmt::print("Prompt (exit/quit to stop): "); + if (!std::getline(std::cin, prompt_text) || prompt_text == "exit" || prompt_text == "quit") break; + } + if (prompt_text.empty()) { + if (prompt.isSet()) exit_code = 1; + if (prompt.isSet()) break; + continue; + } + try { + model.resetState(); + auto inputs = tokenizer.convertMessage( + {.prompt = prompt_text, .system_prompt = system_prompt.isSet() ? system_prompt.get() : "", .tools = raw_tools}); + const auto prompt_tokens = inputs.at("sequence").shape()[1]; + if (prompt_tokens + generation_limit - 1 > cfg.max_cache_length) { + throw std::invalid_argument("prompt plus generation exceeds max_cache_length"); + } + if (prompt.isSet()) fmt::print("Prompt: {}\n", prompt_text); + fmt::print("Response: "); + mllm::models::lfm2::StreamingUtf8Decoder decoder; + int generated_tokens = 0; + for (const auto& step : model.chat(inputs, {{"max_length", mllm::AnyValue(generation_limit)}, + {"min_new_tokens", mllm::AnyValue(minimum_generation)}, + {"do_sample", mllm::AnyValue(false)}})) { + ++generated_tokens; + if (print_token_ids.isSet() && print_token_ids.get()) fmt::print(stderr, "TOKEN_ID:{}\n", step.cur_token_id); + fmt::print("{}", decoder.append(tokenizer.detokenizeBytes(step.cur_token_id))); + std::fflush(stdout); + } + fmt::print("{}\n", decoder.finish()); + fmt::print(stderr, "GENERATED_TOKEN_COUNT:{}\n", generated_tokens); + } catch (const std::exception& error) { + fmt::print(stderr, "LFM2 generation failed: {}\n", error.what()); + exit_code = 1; + } + if (prompt.isSet()) break; + } + return exit_code; +}) diff --git a/examples/lfm2/quant_cfg_2.6B_w4a32_kai.json b/examples/lfm2/quant_cfg_2.6B_w4a32_kai.json new file mode 100644 index 000000000..bc23d45eb --- /dev/null +++ b/examples/lfm2/quant_cfg_2.6B_w4a32_kai.json @@ -0,0 +1,26 @@ +{ + "^model\\.layers\\.\\d+\\.self_attn\\.q_proj\\.weight$": { + "hints": {"quant_method": "kai", "kai_matmul_triplet": "f32_qai8dxp_qsi4c32p", "kai_matmul_layout": "mxk_nxk", "kai_matmul_tile_cfg": "qai8dxp1x8_qsi4c32p8x8_1x8x32", "shape": [2048, 2048], "replace": true} + }, + "^model\\.layers\\.\\d+\\.self_attn\\.[kv]_proj\\.weight$": { + "hints": {"quant_method": "kai", "kai_matmul_triplet": "f32_qai8dxp_qsi4c32p", "kai_matmul_layout": "mxk_nxk", "kai_matmul_tile_cfg": "qai8dxp1x8_qsi4c32p8x8_1x8x32", "shape": [512, 2048], "replace": true} + }, + "^model\\.layers\\.\\d+\\.self_attn\\.out_proj\\.weight$": { + "hints": {"quant_method": "kai", "kai_matmul_triplet": "f32_qai8dxp_qsi4c32p", "kai_matmul_layout": "mxk_nxk", "kai_matmul_tile_cfg": "qai8dxp1x8_qsi4c32p8x8_1x8x32", "shape": [2048, 2048], "replace": true} + }, + "^model\\.layers\\.\\d+\\.conv\\.in_proj\\.weight$": { + "hints": {"quant_method": "kai", "kai_matmul_triplet": "f32_qai8dxp_qsi4c32p", "kai_matmul_layout": "mxk_nxk", "kai_matmul_tile_cfg": "qai8dxp1x8_qsi4c32p8x8_1x8x32", "shape": [6144, 2048], "replace": true} + }, + "^model\\.layers\\.\\d+\\.conv\\.out_proj\\.weight$": { + "hints": {"quant_method": "kai", "kai_matmul_triplet": "f32_qai8dxp_qsi4c32p", "kai_matmul_layout": "mxk_nxk", "kai_matmul_tile_cfg": "qai8dxp1x8_qsi4c32p8x8_1x8x32", "shape": [2048, 2048], "replace": true} + }, + "^model\\.layers\\.\\d+\\.feed_forward\\.(w1|w3)\\.weight$": { + "hints": {"quant_method": "kai", "kai_matmul_triplet": "f32_qai8dxp_qsi4c32p", "kai_matmul_layout": "mxk_nxk", "kai_matmul_tile_cfg": "qai8dxp1x8_qsi4c32p8x8_1x8x32", "shape": [10752, 2048], "replace": true} + }, + "^model\\.layers\\.\\d+\\.feed_forward\\.w2\\.weight$": { + "hints": {"quant_method": "kai", "kai_matmul_triplet": "f32_qai8dxp_qsi4c32p", "kai_matmul_layout": "mxk_nxk", "kai_matmul_tile_cfg": "qai8dxp1x8_qsi4c32p8x8_1x8x32", "shape": [2048, 10752], "replace": true} + }, + "^model\\.embed_tokens\\.weight$": { + "hints": {"quant_method": "kai", "kai_matmul_triplet": "f32_qai8dxp_qsi4c32p", "kai_matmul_layout": "mxk_nxk", "kai_matmul_tile_cfg": "qai8dxp1x8_qsi4c32p8x8_1x8x32", "shape": [128000, 2048], "replace": false, "rename": "lm_head_out.weight"} + } +} diff --git a/examples/lfm2/test_validators.py b/examples/lfm2/test_validators.py new file mode 100644 index 000000000..815256172 --- /dev/null +++ b/examples/lfm2/test_validators.py @@ -0,0 +1,49 @@ +# Copyright (c) MLLM Team. +# Licensed under the MIT License. +import json +import unittest +from pathlib import Path + +import validate_checkpoint +import validate_converted_model + + +class Lfm2ValidatorTest(unittest.TestCase): + def setUp(self) -> None: + self.directory = Path(__file__).parent + self.runtime = json.loads((self.directory / "config_2.6B_w4a32_kai.json").read_text()) + self.recipe = json.loads((self.directory / "quant_cfg_2.6B_w4a32_kai.json").read_text()) + + def test_official_contract_has_266_tensors(self) -> None: + validate_checkpoint.validate_config(self.runtime) + shapes = validate_checkpoint.expected_shapes() + self.assertEqual(len(shapes), 266) + self.assertEqual(sum(kind == "full_attention" for kind in validate_checkpoint.LAYER_TYPES), 8) + + def test_generation_contract_has_one_pinned_eos(self) -> None: + validate_checkpoint.validate_generation_config(validate_checkpoint.OFFICIAL_GENERATION) + drifted = dict(validate_checkpoint.OFFICIAL_GENERATION) + drifted["eos_token_id"] = [124900, 124901] + with self.assertRaisesRegex(AssertionError, "eos_token_id"): + validate_checkpoint.validate_generation_config(drifted) + + def test_quant_recipe_has_exact_tied_output_alias(self) -> None: + validate_checkpoint.validate_recipe(self.recipe, validate_checkpoint.expected_shapes()) + descriptors = validate_converted_model.expected_descriptors(self.recipe) + self.assertEqual(len(descriptors), 267) + self.assertEqual( + descriptors["model.embed_tokens.weight"][:2], + (validate_converted_model.FLOAT32, [128000, 2048]), + ) + self.assertEqual(descriptors["lm_head_out.weight"][0], validate_converted_model.BYTE) + + def test_schedule_drift_fails_closed(self) -> None: + drifted = dict(self.runtime) + drifted["layer_types"] = list(self.runtime["layer_types"]) + drifted["layer_types"][0] = "full_attention" + with self.assertRaisesRegex(AssertionError, "layer_types"): + validate_checkpoint.validate_config(drifted) + + +if __name__ == "__main__": + unittest.main() diff --git a/examples/lfm2/validate_checkpoint.py b/examples/lfm2/validate_checkpoint.py new file mode 100644 index 000000000..11aad17e1 --- /dev/null +++ b/examples/lfm2/validate_checkpoint.py @@ -0,0 +1,216 @@ +# Copyright (c) MLLM Team. +# Licensed under the MIT License. +"""Fail-closed audit for the official LFM2.5-2.6B checkpoint and W4A32 recipe.""" + +from __future__ import annotations + +import argparse +import json +import re +from contextlib import ExitStack +from pathlib import Path + +from safetensors import safe_open + + +OFFICIAL = { + "architectures": ["Lfm2ForCausalLM"], + "model_type": "lfm2", + "hidden_size": 2048, + "intermediate_size": 10752, + "num_hidden_layers": 30, + "num_attention_heads": 32, + "num_key_value_heads": 8, + "conv_L_cache": 3, + "conv_bias": False, + "block_auto_adjust_ff_dim": False, + "norm_eps": 1e-5, + "max_position_embeddings": 131072, + "vocab_size": 128000, + "tie_word_embeddings": True, + "bos_token_id": 124894, + "eos_token_id": 124900, + "pad_token_id": 124893, + "rope_parameters": {"rope_theta": 10000000.0, "rope_type": "default"}, +} +OFFICIAL_GENERATION = { + "bos_token_id": 124894, + "eos_token_id": [124900], + "pad_token_id": 124893, +} +LAYER_TYPES = [ + "conv", "conv", "full_attention", "conv", "conv", "full_attention", "conv", "conv", "conv", "full_attention", + "conv", "conv", "conv", "full_attention", "conv", "conv", "conv", "full_attention", "conv", "conv", "conv", + "full_attention", "conv", "conv", "full_attention", "conv", "conv", "full_attention", "conv", "conv", +] +KAI_FIELDS = { + "quant_method": "kai", + "kai_matmul_triplet": "f32_qai8dxp_qsi4c32p", + "kai_matmul_layout": "mxk_nxk", + "kai_matmul_tile_cfg": "qai8dxp1x8_qsi4c32p8x8_1x8x32", +} +KAI_IMPL = "KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp1x8_qsi4c32p8x8_1x8x32" + + +def expected_shapes() -> dict[str, list[int]]: + expected = { + "model.embed_tokens.weight": [128000, 2048], + "model.embedding_norm.weight": [2048], + } + for layer, layer_type in enumerate(LAYER_TYPES): + prefix = f"model.layers.{layer}" + expected.update( + { + f"{prefix}.operator_norm.weight": [2048], + f"{prefix}.ffn_norm.weight": [2048], + f"{prefix}.feed_forward.w1.weight": [10752, 2048], + f"{prefix}.feed_forward.w3.weight": [10752, 2048], + f"{prefix}.feed_forward.w2.weight": [2048, 10752], + } + ) + if layer_type == "conv": + expected.update( + { + f"{prefix}.conv.in_proj.weight": [6144, 2048], + f"{prefix}.conv.conv.weight": [2048, 1, 3], + f"{prefix}.conv.out_proj.weight": [2048, 2048], + } + ) + else: + expected.update( + { + f"{prefix}.self_attn.q_proj.weight": [2048, 2048], + f"{prefix}.self_attn.k_proj.weight": [512, 2048], + f"{prefix}.self_attn.v_proj.weight": [512, 2048], + f"{prefix}.self_attn.out_proj.weight": [2048, 2048], + f"{prefix}.self_attn.q_layernorm.weight": [64], + f"{prefix}.self_attn.k_layernorm.weight": [64], + } + ) + assert len(expected) == 266 + return expected + + +def validate_config(config: dict) -> None: + mismatches = [ + f"{name}={config.get(name)!r}, expected {value!r}" + for name, value in OFFICIAL.items() + if config.get(name) != value + ] + if config.get("layer_types") != LAYER_TYPES: + mismatches.append("layer_types differs from the official 30-layer physical schedule") + if mismatches: + raise AssertionError("Checkpoint contract mismatch: " + "; ".join(mismatches)) + + +def validate_generation_config(config: dict) -> None: + mismatches = [ + f"{name}={config.get(name)!r}, expected {value!r}" + for name, value in OFFICIAL_GENERATION.items() + if config.get(name) != value + ] + if mismatches: + raise AssertionError("Generation contract mismatch: " + "; ".join(mismatches)) + + +def validate_runtime_config(checkpoint: dict, runtime: dict) -> None: + validate_config(runtime) + for name in OFFICIAL: + if runtime.get(name) != checkpoint.get(name): + raise AssertionError(f"Runtime/checkpoint mismatch for {name}") + if runtime.get("layer_types") != checkpoint.get("layer_types"): + raise AssertionError("Runtime/checkpoint layer_types mismatch") + if runtime.get("head_dim") != 64 or runtime.get("max_cache_length") != 2048: + raise AssertionError("Runtime must bind head_dim=64 and max_cache_length=2048") + if runtime.get("linear_impl_type") != KAI_IMPL: + raise AssertionError("Runtime does not select the pinned KleidiAI W4A32 implementation") + + +def validate_recipe(recipe: dict, shapes: dict[str, list[int]]) -> None: + patterns = [(re.compile(pattern), entry["hints"]) for pattern, entry in recipe.items()] + matched_patterns = {pattern.pattern: 0 for pattern, _ in patterns} + aliases: list[str] = [] + quantized: set[str] = set() + for name, shape in shapes.items(): + matches = [(pattern, hints) for pattern, hints in patterns if pattern.fullmatch(name)] + if len(matches) > 1: + raise AssertionError(f"Multiple quantization rules match {name}") + if not matches: + continue + pattern, hints = matches[0] + matched_patterns[pattern.pattern] += 1 + quantized.add(name) + if hints.get("shape") != shape: + raise AssertionError(f"Recipe shape mismatch for {name}: {hints.get('shape')} != {shape}") + for field, value in KAI_FIELDS.items(): + if hints.get(field) != value: + raise AssertionError(f"Recipe {pattern.pattern} has invalid {field}") + if hints.get("replace") is False: + aliases.append(hints.get("rename")) + unused = [pattern for pattern, count in matched_patterns.items() if count == 0] + if unused: + raise AssertionError(f"Quantization rules match no checkpoint tensor: {unused}") + if aliases != ["lm_head_out.weight"]: + raise AssertionError("Tied embedding must create exactly one lm_head_out.weight packed alias") + intended_linears = { + name + for name, shape in shapes.items() + if len(shape) == 2 and name != "model.embed_tokens.weight" + } + if not intended_linears.issubset(quantized): + raise AssertionError(f"Unquantized intended Linear tensors: {sorted(intended_linears - quantized)}") + allowed_non_linear = {"model.embed_tokens.weight"} + unexpected = quantized - intended_linears - allowed_non_linear + if unexpected: + raise AssertionError(f"Recipe unexpectedly covers non-Linear tensors: {sorted(unexpected)}") + + +def checkpoint_shapes(checkpoint: Path) -> dict[str, list[int]]: + index_path = checkpoint / "model.safetensors.index.json" + if index_path.exists(): + index = json.loads(index_path.read_text()) + shard_names = sorted(set(index["weight_map"].values())) + elif (checkpoint / "model.safetensors").exists(): + shard_names = ["model.safetensors"] + else: + raise AssertionError("Checkpoint has no safetensors weights") + actual: dict[str, list[int]] = {} + with ExitStack() as stack: + for shard_name in shard_names: + handle = stack.enter_context(safe_open(checkpoint / shard_name, framework="pt", device="cpu")) + for name in handle.keys(): + if name in actual: + raise AssertionError(f"Duplicate tensor across shards: {name}") + actual[name] = list(handle.get_slice(name).get_shape()) + return actual + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("checkpoint", type=Path) + parser.add_argument("--quant-config", type=Path, default=Path(__file__).with_name("quant_cfg_2.6B_w4a32_kai.json")) + parser.add_argument("--runtime-config", type=Path, default=Path(__file__).with_name("config_2.6B_w4a32_kai.json")) + args = parser.parse_args() + + checkpoint_config = json.loads((args.checkpoint / "config.json").read_text()) + generation_config = json.loads((args.checkpoint / "generation_config.json").read_text()) + runtime_config = json.loads(args.runtime_config.read_text()) + recipe = json.loads(args.quant_config.read_text()) + expected = expected_shapes() + validate_config(checkpoint_config) + validate_generation_config(generation_config) + validate_runtime_config(checkpoint_config, runtime_config) + validate_recipe(recipe, expected) + actual = checkpoint_shapes(args.checkpoint) + missing = sorted(set(expected) - set(actual)) + extra = sorted(set(actual) - set(expected)) + wrong = sorted(name for name in set(expected) & set(actual) if expected[name] != actual[name]) + if missing or extra or wrong: + raise AssertionError(f"Tensor contract mismatch: missing={missing}, extra={extra}, wrong_shapes={wrong}") + print(json.dumps({"checkpoint": str(args.checkpoint), "parameters": len(actual), "attention_layers": 8, + "conv_layers": 22, "logical_cache_slots": 8}, indent=2, sort_keys=True)) + print("LFM2_CHECKPOINT_AUDIT_OK") + + +if __name__ == "__main__": + main() diff --git a/examples/lfm2/validate_converted_model.py b/examples/lfm2/validate_converted_model.py new file mode 100644 index 000000000..0a417c9f3 --- /dev/null +++ b/examples/lfm2/validate_converted_model.py @@ -0,0 +1,125 @@ +# Copyright (c) MLLM Team. +# Licensed under the MIT License. +"""Audit an LFM2.5-2.6B MLLM V2 descriptor table without loading tensor data.""" + +from __future__ import annotations + +import argparse +import json +import math +import re +import struct +from pathlib import Path + +from validate_checkpoint import expected_shapes, validate_config, validate_recipe, validate_runtime_config + + +MODEL_HEADER = struct.Struct(" int: + if in_channels <= 0 or in_channels % 32 or out_channels <= 0: + raise AssertionError("Invalid KAI W4A32 matrix dimensions") + blocks_per_row = in_channels // 32 + bytes_per_eight_rows = 8 * (blocks_per_row * 18 + 8) + return math.ceil(out_channels / 8) * bytes_per_eight_rows + + +def expected_descriptors(recipe: dict) -> dict[str, tuple[int, list[int], int]]: + source = expected_shapes() + rules = [(re.compile(pattern), entry["hints"]) for pattern, entry in recipe.items()] + expected: dict[str, tuple[int, list[int], int]] = {} + for name, shape in source.items(): + matches = [hints for pattern, hints in rules if pattern.fullmatch(name)] + if not matches: + expected[name] = (FLOAT32, shape, math.prod(shape) * 4) + continue + if len(matches) != 1: + raise AssertionError(f"Ambiguous recipe match for {name}") + hints = matches[0] + packed = packed_size(shape[0], shape[1]) + descriptor = (BYTE, [packed], packed) + if hints["replace"]: + expected[name] = descriptor + else: + expected[name] = (FLOAT32, shape, math.prod(shape) * 4) + expected[hints["rename"]] = descriptor + return expected + + +def c_string(raw: bytes) -> str: + return raw.split(b"\0", 1)[0].decode("utf-8") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("model", type=Path) + parser.add_argument("checkpoint", type=Path) + parser.add_argument("--model-name", default="LFM2.5-2.6B") + parser.add_argument("--quant-config", type=Path, default=Path(__file__).with_name("quant_cfg_2.6B_w4a32_kai.json")) + parser.add_argument("--runtime-config", type=Path, default=Path(__file__).with_name("config_2.6B_w4a32_kai.json")) + args = parser.parse_args() + + checkpoint = json.loads((args.checkpoint / "config.json").read_text()) + runtime = json.loads(args.runtime_config.read_text()) + recipe = json.loads(args.quant_config.read_text()) + validate_config(checkpoint) + validate_runtime_config(checkpoint, runtime) + validate_recipe(recipe, expected_shapes()) + expected = expected_descriptors(recipe) + + file_size = args.model.stat().st_size + actual: dict[str, tuple[int, list[int], int, int]] = {} + with args.model.open("rb") as stream: + raw_header = stream.read(MODEL_HEADER.size) + if len(raw_header) != MODEL_HEADER.size: + raise AssertionError("Truncated MLLM V2 header") + magic, version, raw_name, count, descriptor_offset = MODEL_HEADER.unpack(raw_header) + if (magic, version, c_string(raw_name), descriptor_offset) != ( + MODEL_MAGIC, + MODEL_VERSION, + args.model_name, + MODEL_HEADER.size, + ): + raise AssertionError("Invalid MLLM V2 model header") + if count != len(expected): + raise AssertionError(f"Expected {len(expected)} descriptors, file declares {count}") + for parameter_id in range(count): + raw = stream.read(PARAMETER_DESCRIPTOR.size) + if len(raw) != PARAMETER_DESCRIPTOR.size: + raise AssertionError(f"Truncated descriptor {parameter_id}") + fields = PARAMETER_DESCRIPTOR.unpack(raw) + actual_id, dtype, size, offset, rank = fields[:5] + if actual_id != parameter_id or rank > 16: + raise AssertionError(f"Invalid descriptor id/rank at {parameter_id}") + name = c_string(fields[21]) + if name in actual: + raise AssertionError(f"Duplicate descriptor: {name}") + actual[name] = (dtype, list(fields[5:21])[:rank], size, offset) + + missing = sorted(set(expected) - set(actual)) + extra = sorted(set(actual) - set(expected)) + wrong = sorted(name for name in set(expected) & set(actual) if actual[name][:3] != expected[name]) + if missing or extra or wrong: + raise AssertionError(f"Converted tensor mismatch: missing={missing}, extra={extra}, wrong={wrong}") + data_start = MODEL_HEADER.size + len(actual) * PARAMETER_DESCRIPTOR.size + next_offset = data_start + for name, (_, _, size, offset) in sorted(actual.items(), key=lambda item: item[1][3]): + if offset != next_offset: + raise AssertionError(f"Non-contiguous data before {name}: {offset} != {next_offset}") + next_offset += size + if next_offset != file_size: + raise AssertionError(f"Tensor data ends at {next_offset}, file size is {file_size}") + print(json.dumps({"model": str(args.model), "parameters": len(actual), "model_bytes": file_size, + "has_lm_head_out": "lm_head_out.weight" in actual, + "logical_cache_slots": 8}, indent=2, sort_keys=True)) + print("LFM2_CONVERTED_MODEL_AUDIT_OK") + + +if __name__ == "__main__": + main() diff --git a/mllm/backends/cpu/CMakeLists.txt b/mllm/backends/cpu/CMakeLists.txt index 1c8b32bac..86623e2b3 100644 --- a/mllm/backends/cpu/CMakeLists.txt +++ b/mllm/backends/cpu/CMakeLists.txt @@ -147,12 +147,45 @@ if(MLLM_KERNEL_USE_THREADS AND MLLM_KERNEL_THREADS_VENDOR_OPENMP) target_compile_options(MllmRT PRIVATE ${OpenMP_CXX_FLAGS}) target_include_directories(MllmRT PUBLIC ${OpenMP_CXX_INCLUDE_DIR}) - # ARM CPU kernels contain the OpenMP regions used by the mobile product - # path. Keep x86 CPU kernels on the existing MllmRT-only configuration so - # unrelated kernels are not recompiled with OpenMP as part of ARM support. if(MLLM_BUILD_ARM_BACKEND) + # OpenMP is a translation-unit concern: compile only sources that own a + # parallel region, including operators that instantiate one from a + # header-only kernel. Keep this list grouped by ownership so adding a + # parallel region does not silently broaden the whole backend target. + set(MLLM_CPU_BACKEND_OPENMP_SOURCES + # Backend operators with an in-file parallel region. + ${CMAKE_CURRENT_LIST_DIR}/ops/AvgPool1dOp.cpp + ${CMAKE_CURRENT_LIST_DIR}/ops/Conv1DOp.cpp + ${CMAKE_CURRENT_LIST_DIR}/ops/EmbeddingOp.cpp + ${CMAKE_CURRENT_LIST_DIR}/ops/GroupedQueryAttentionOp.cpp + ${CMAKE_CURRENT_LIST_DIR}/ops/LayerNormOp.cpp + ${CMAKE_CURRENT_LIST_DIR}/ops/RMSNormOp.cpp + ${CMAKE_CURRENT_LIST_DIR}/ops/SoftmaxOp.cpp + + # Operators that instantiate parallel header-only attention kernels. + ${CMAKE_CURRENT_LIST_DIR}/ops/FlashAttention2Op.cpp + ${CMAKE_CURRENT_LIST_DIR}/ops/FlashAttn2WithSinkAndSwaOp.cpp + ${CMAKE_CURRENT_LIST_DIR}/ops/PagedAttnOp.cpp + ${CMAKE_CURRENT_LIST_DIR}/ops/RadixAttnDiffDimOp.cpp + ${CMAKE_CURRENT_LIST_DIR}/ops/RadixAttnOp.cpp + ${CMAKE_CURRENT_LIST_DIR}/ops/RadixAttnWithSinkAndSwaDiffDimOp.cpp + + # Kernel translation units with an in-file parallel region. + ${CMAKE_CURRENT_LIST_DIR}/kernels/arm/cast_types.cpp + ${CMAKE_CURRENT_LIST_DIR}/kernels/arm/elementwise.cpp + ${CMAKE_CURRENT_LIST_DIR}/kernels/arm/gelu.cpp + ${CMAKE_CURRENT_LIST_DIR}/kernels/arm/linear/kai.cpp + ${CMAKE_CURRENT_LIST_DIR}/kernels/arm/mllm_blas/mllm_blas_sgemm.cpp + ${CMAKE_CURRENT_LIST_DIR}/kernels/arm/relu.cpp + ${CMAKE_CURRENT_LIST_DIR}/kernels/arm/sigmoid.cpp + ${CMAKE_CURRENT_LIST_DIR}/kernels/arm/silu.cpp + ${CMAKE_CURRENT_LIST_DIR}/kernels/common/gdn/gated_delta_net.cpp + ${CMAKE_CURRENT_LIST_DIR}/kernels/common/ggml/matmul.cpp) + + set_property( + SOURCE ${MLLM_CPU_BACKEND_OPENMP_SOURCES} + APPEND PROPERTY COMPILE_OPTIONS ${OpenMP_CXX_FLAGS}) target_link_libraries(MllmCPUBackend PUBLIC ${OpenMP_CXX_FLAGS}) - target_compile_options(MllmCPUBackend PRIVATE ${OpenMP_CXX_FLAGS}) target_include_directories(MllmCPUBackend PUBLIC ${OpenMP_CXX_INCLUDE_DIR}) endif() endif() diff --git a/mllm/backends/cpu/CPUBackend.cpp b/mllm/backends/cpu/CPUBackend.cpp index 40a5e6a08..5b59e7c91 100644 --- a/mllm/backends/cpu/CPUBackend.cpp +++ b/mllm/backends/cpu/CPUBackend.cpp @@ -9,6 +9,7 @@ #include "mllm/backends/cpu/ops/AvgPool1dOp.hpp" #include "mllm/backends/cpu/ops/CastTypeOp.hpp" #include "mllm/backends/cpu/ops/CausalMaskOp.hpp" +#include "mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.hpp" #include "mllm/backends/cpu/ops/CloneOp.hpp" #include "mllm/backends/cpu/ops/CmpOp.hpp" #include "mllm/backends/cpu/ops/ConcatOp.hpp" @@ -24,11 +25,12 @@ #include "mllm/backends/cpu/ops/FlashAttn2WithSinkAndSwaOp.hpp" #include "mllm/backends/cpu/ops/GELUOp.hpp" #include "mllm/backends/cpu/ops/GatherOp.hpp" -#include "mllm/backends/cpu/ops/GroupedQueryAttentionDecodeOp.hpp" +#include "mllm/backends/cpu/ops/GroupedQueryAttentionOp.hpp" #include "mllm/backends/cpu/ops/InterpolateOp.hpp" #include "mllm/backends/cpu/ops/LayerNorm2DOp.hpp" #include "mllm/backends/cpu/ops/MaskedScatterOp.hpp" #include "mllm/backends/cpu/ops/PadOp.hpp" +#include "mllm/backends/cpu/ops/ParallelLinearOp.hpp" #include "mllm/backends/cpu/ops/RadixAttnDiffDimOp.hpp" #include "mllm/backends/cpu/ops/RadixAttnOp.hpp" #include "mllm/backends/cpu/ops/RadixAttnWithSinkAndSwaDiffDimOp.hpp" @@ -84,7 +86,8 @@ CPUBackend::CPUBackend() : Backend(kCPU, createCPUAllocator()) { CPUConv2DOpFactory, CPULayerNorm2DOpFactory, CPUInterpolateOpFactory, CPUPadOpFactory, CPUMaskedScatterOpFactory, CPUArgsortOpFactory, CPUCloneOpFactory, CPUAvgPool1dOpFactory, CPUFlashAttention2SwaSinkOpFactory, CPURadixAttnRelaxOpFactory, CPURadixAttnSwaSinkOpFactory, CPUEqualOpFactory, CPUWhereOpFactory, - CPUGatherOpFactory, CPUGroupedQueryAttentionDecodeOpFactory>(); + CPUGatherOpFactory, CPUCausalDepthwiseConv1DOpFactory, + CPUGroupedQueryAttentionOpFactory, CPUParallelLinearOpFactory>(); } CPUBackend::~CPUBackend() { diff --git a/mllm/backends/cpu/kernels/arm/linear/kai.cpp b/mllm/backends/cpu/kernels/arm/linear/kai.cpp index 1a7d40baf..37aab1e1d 100644 --- a/mllm/backends/cpu/kernels/arm/linear/kai.cpp +++ b/mllm/backends/cpu/kernels/arm/linear/kai.cpp @@ -445,6 +445,78 @@ void KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::matmul(float* __restrict__ dst, con } } +bool KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::matmul_shared_input(const float* __restrict__ lhs_fp32, + const SharedInputProjection* projections, + size_t projection_count, void* workspace, int M, int K, + KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles tile_cfg, + int thread_count) { + if (lhs_fp32 == nullptr || projections == nullptr || projection_count < 2 || workspace == nullptr || M <= 0 || K <= 0 + || thread_count <= 0) { + return false; + } + for (size_t projection_index = 0; projection_index < projection_count; ++projection_index) { + if (projections[projection_index].dst == nullptr || projections[projection_index].packed_weight_bias == nullptr + || projections[projection_index].n <= 0) { + return false; + } + } + + const auto& ukernel = ukernels_.at(tile_cfg); + kai_run_lhs_quant_pack_qai8dxp_f32(M, K, ukernel.get_mr(), ukernel.get_kr(), ukernel.get_sr(), 0, lhs_fp32, K * sizeof(float), + workspace); + + const size_t m_step = static_cast(ukernel.get_m_step()); + const size_t n_step = static_cast(ukernel.get_n_step()); + const size_t m_tiles = (static_cast(M) + m_step - 1) / m_step; + size_t total_tiles = 0; + for (size_t projection_index = 0; projection_index < projection_count; ++projection_index) { + const size_t n_tiles = (static_cast(projections[projection_index].n) + n_step - 1) / n_step; + if (n_tiles > std::numeric_limits::max() / m_tiles) { return false; } + const size_t projection_tiles = m_tiles * n_tiles; + if (projection_tiles > std::numeric_limits::max() - total_tiles) { return false; } + total_tiles += projection_tiles; + } + + MLLM_CONDITIONAL_PARALLEL_FOR(thread_count > 1, thread_count, global_tile, 0, total_tiles, 1, { + size_t local_tile = static_cast(global_tile); + size_t projection_index = 0; + for (; projection_index < projection_count; ++projection_index) { + const size_t projection_n_tiles = (static_cast(projections[projection_index].n) + n_step - 1) / n_step; + const size_t projection_tiles = m_tiles * projection_n_tiles; + if (local_tile < projection_tiles) { break; } + local_tile -= projection_tiles; + } + + if (projection_index < projection_count) { + const auto& projection = projections[projection_index]; + const size_t projection_n_tiles = (static_cast(projection.n) + n_step - 1) / n_step; + const int m_index = static_cast((local_tile / projection_n_tiles) * m_step); + const int n_index = static_cast((local_tile % projection_n_tiles) * n_step); + const int actual_m = std::min(M - m_index, static_cast(m_step)); + const int actual_n = std::min(projection.n - n_index, static_cast(n_step)); + const size_t dst_stride = static_cast(projection.n) * sizeof(float); + const void* lhs_ptr = + static_cast(static_cast(workspace) + ukernel.get_lhs_packed_offset(m_index, K)); + const void* rhs_ptr = static_cast(reinterpret_cast(projection.packed_weight_bias) + + ukernel.get_rhs_packed_offset(n_index, K, 32)); + float* dst_ptr = reinterpret_cast(reinterpret_cast(projection.dst) + + ukernel.get_dst_offset(m_index, n_index, dst_stride)); + + ukernel.run_matmul(actual_m, actual_n, K, 32, lhs_ptr, rhs_ptr, dst_ptr, dst_stride, sizeof(float), + -std::numeric_limits::max(), std::numeric_limits::max()); + } + }); + return true; +} + +bool KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::matmul_shared_input_m1(const float* __restrict__ lhs_fp32, + const SharedInputProjection* projections, + size_t projection_count, void* workspace, int K, + KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles tile_cfg, + int thread_count) { + return matmul_shared_input(lhs_fp32, projections, projection_count, workspace, 1, K, tile_cfg, thread_count); +} + std::unordered_map KaiLinear_f32_qai8dxp_qsi4c32p_mxk_kxn::ukernels_ = { {KaiLinear_f32_qai8dxp_qsi4c32p_mxk_kxn::Tiles::qai8dxp1x8_qsi4c32p4x8_1x4x32, diff --git a/mllm/backends/cpu/kernels/arm/linear/kai.hpp b/mllm/backends/cpu/kernels/arm/linear/kai.hpp index e8f38222b..bce54b9a4 100644 --- a/mllm/backends/cpu/kernels/arm/linear/kai.hpp +++ b/mllm/backends/cpu/kernels/arm/linear/kai.hpp @@ -30,6 +30,7 @@ // fp16 - Floating-point 16-bit data type #include +#include #include #include @@ -102,6 +103,12 @@ struct KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk { qai8dxp1x4_qsi4c32p4x4_1x4, }; + struct SharedInputProjection { + float* dst; + const uint8_t* packed_weight_bias; + int n; + }; + inline bool need_pack_lhs() { return true; } inline bool need_pack_rhs() { return true; } @@ -117,6 +124,14 @@ struct KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk { void matmul(float* __restrict__ dst, const float* __restrict__ lhs_fp32, const uint8_t* packed_weight_bias, void* workspace, int M, int K, int N, KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles tile_cfg, int thread_count); + bool matmul_shared_input(const float* __restrict__ lhs_fp32, const SharedInputProjection* projections, + size_t projection_count, void* workspace, int M, int K, + KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles tile_cfg, int thread_count); + + bool matmul_shared_input_m1(const float* __restrict__ lhs_fp32, const SharedInputProjection* projections, + size_t projection_count, void* workspace, int K, + KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles tile_cfg, int thread_count); + private: void quant_nxk_qs4c32_f32(size_t n, size_t k, size_t bl, const float* rhs_f32, uint8_t* rhs_qs4c32, uint16_t* rhs_scales_bf16); diff --git a/mllm/backends/cpu/kernels/arm/linear/parallel_linear.cpp b/mllm/backends/cpu/kernels/arm/linear/parallel_linear.cpp new file mode 100644 index 000000000..16ed0ace3 --- /dev/null +++ b/mllm/backends/cpu/kernels/arm/linear/parallel_linear.cpp @@ -0,0 +1,88 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/backends/cpu/kernels/common/parallel_linear/shared_input.hpp" + +#include +#include +#include +#include + +#include "mllm/backends/cpu/kernels/arm/linear/kai.hpp" +#include "mllm/backends/cpu/kernels/common/linear/kai_w4a32_dispatch.hpp" + +namespace mllm::cpu::parallel_linear { + +namespace { + +constexpr size_t kMaximumSharedProjections = 3; +using KaiHelper = arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk; +constexpr auto kDecodeTile = KaiHelper::Tiles::qai8dxp1x8_qsi4c32p8x8_1x8x32; +constexpr auto kPrefillTile = KaiHelper::Tiles::qai8dxp4x8_qsi4c32p8x8_4x8x32; + +bool traceActivationEnabled() { + static const bool enabled = [] { + const char* value = std::getenv("MLLM_KAI_SHARED_INPUT_TRACE"); + return value != nullptr && value[0] == '1' && value[1] == '\0'; + }(); + return enabled; +} + +void traceActivation(const SharedInputPlan& plan, size_t projection_count, int32_t m, int32_t k) { + if (!traceActivationEnabled()) { return; } + + const uint32_t projection_group = static_cast(projection_count - 2); + const uint32_t activation_bit = 1U << (2U * projection_group + static_cast(m > 1)); + static std::atomic activated_groups{0}; + const uint32_t previous = activated_groups.fetch_or(activation_bit, std::memory_order_relaxed); + if ((previous & activation_bit) == 0) { + std::fprintf(stderr, "MLLM_KAI_SHARED_INPUT_ACTIVATED rhs=%zu m=%d k=%d threads=%d tile=%s\n", projection_count, m, k, + plan.thread_count, plan.kernel == SharedInputKernel::kKaiDotprod ? "dotprod_1x8" : "i8mm_4x8"); + } +} + +} // namespace + +SharedInputPlan planKaiW4A32SharedInput(int32_t m, int32_t k, int32_t requested_threads, int32_t decode_thread_cap, + int32_t prefill_thread_cap) { + if (m <= 0 || k <= 0 || requested_threads <= 0) { return {}; } + if (m > 1 && !kai_w4a32::shouldUseI8mmPrefill(m)) { return {}; } + + const auto tile = m == 1 ? kDecodeTile : kPrefillTile; + KaiHelper helper; + const size_t workspace_size = helper.workspace_size(m, k, tile); + const int32_t thread_count = kai_w4a32::threadCount(m, requested_threads, decode_thread_cap, prefill_thread_cap); + if (workspace_size == 0 || thread_count <= 0) { return {}; } + + return {.kernel = m == 1 ? SharedInputKernel::kKaiDotprod : SharedInputKernel::kKaiI8mm, + .workspace_size = workspace_size, + .thread_count = thread_count}; +} + +bool runKaiW4A32SharedInput(const SharedInputPlan& plan, const float* input, const SharedInputProjection* projections, + size_t projection_count, void* workspace, int32_t m, int32_t k) { + if (!plan.supported() || input == nullptr || projections == nullptr || projection_count < 2 + || projection_count > kMaximumSharedProjections || workspace == nullptr || m <= 0 || k <= 0 || plan.thread_count <= 0 + || (m == 1 && plan.kernel != SharedInputKernel::kKaiDotprod) || (m > 1 && plan.kernel != SharedInputKernel::kKaiI8mm)) { + return false; + } + + std::array kai_projections{}; + for (size_t index = 0; index < projection_count; ++index) { + if (projections[index].dst == nullptr || projections[index].packed_weight_bias == nullptr || projections[index].n <= 0) { + return false; + } + kai_projections[index] = { + .dst = projections[index].dst, .packed_weight_bias = projections[index].packed_weight_bias, .n = projections[index].n}; + } + + const auto tile = plan.kernel == SharedInputKernel::kKaiDotprod ? kDecodeTile : kPrefillTile; + KaiHelper helper; + if (!helper.matmul_shared_input(input, kai_projections.data(), projection_count, workspace, m, k, tile, plan.thread_count)) { + return false; + } + traceActivation(plan, projection_count, m, k); + return true; +} + +} // namespace mllm::cpu::parallel_linear diff --git a/mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.cpp b/mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.cpp new file mode 100644 index 000000000..c7600c852 --- /dev/null +++ b/mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.cpp @@ -0,0 +1,71 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.hpp" + +#include +#include + +#if defined(__aarch64__) +#include +#endif + +namespace mllm::cpu::causal_conv { + +void depthwiseCausalConvHistoryFirstF32(const float* input, const float* weight, float* state, float* output, + int batch_size, int sequence_length, int channels, int kernel_size) { + if (input == nullptr || weight == nullptr || state == nullptr || output == nullptr) { + throw std::invalid_argument("History-first depthwise causal convolution received a null pointer"); + } + if (batch_size <= 0 || sequence_length <= 0 || channels <= 0 || kernel_size <= 1) { + throw std::invalid_argument("History-first depthwise causal convolution received an invalid shape"); + } + + const int state_width = kernel_size - 1; + for (int batch = 0; batch < batch_size; ++batch) { + for (int token = 0; token < sequence_length; ++token) { + int channel = 0; +#if defined(__aarch64__) + // The K=3 fast path deinterleaves four adjacent channels into the two + // history taps and three weights while input/output stay contiguous. + // The three FMA steps intentionally match CPUConv1D's k=0,1,2 order. + if (kernel_size == 3) { + const std::size_t token_base = (static_cast(batch) * sequence_length + token) * channels; + const std::size_t batch_state_base = static_cast(batch) * channels * state_width; + for (; channel + 4 <= channels; channel += 4) { + float* state_block = state + batch_state_base + static_cast(channel) * state_width; + const float32x4x2_t history = vld2q_f32(state_block); + const float32x4x3_t taps = vld3q_f32(weight + static_cast(channel) * kernel_size); + const float32x4_t current = vld1q_f32(input + token_base + channel); + + float32x4_t value = vdupq_n_f32(0.0F); + value = vfmaq_f32(value, history.val[0], taps.val[0]); + value = vfmaq_f32(value, history.val[1], taps.val[1]); + value = vfmaq_f32(value, current, taps.val[2]); + vst1q_f32(output + token_base + channel, value); + + float32x4x2_t shifted; + shifted.val[0] = history.val[1]; + shifted.val[1] = current; + vst2q_f32(state_block, shifted); + } + } +#endif + for (; channel < channels; ++channel) { + const std::size_t state_base = (static_cast(batch) * channels + channel) * state_width; + const std::size_t input_index = (static_cast(batch) * sequence_length + token) * channels + channel; + const std::size_t weight_base = static_cast(channel) * kernel_size; + + float value = 0.0F; + for (int tap = 0; tap < state_width; ++tap) { value += state[state_base + tap] * weight[weight_base + tap]; } + value += input[input_index] * weight[weight_base + state_width]; + output[input_index] = value; + + for (int tap = 0; tap + 1 < state_width; ++tap) { state[state_base + tap] = state[state_base + tap + 1]; } + state[state_base + state_width - 1] = input[input_index]; + } + } + } +} + +} // namespace mllm::cpu::causal_conv diff --git a/mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.hpp b/mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.hpp new file mode 100644 index 000000000..b44f1cc0a --- /dev/null +++ b/mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.hpp @@ -0,0 +1,18 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +namespace mllm::cpu::causal_conv { + +// Stateful depthwise causal convolution for [B, S, C] input with a +// [B, C, K - 1] history state that is updated in place. +// +// Accumulation order is zero, historical taps in ascending order, then the +// current sample, matching CPUConv1D. Callers whose generation contract is +// bitwise sensitive depend on this order, so it is part of the contract +// rather than an implementation detail. +void depthwiseCausalConvHistoryFirstF32(const float* input, const float* weight, float* state, float* output, + int batch_size, int sequence_length, int channels, int kernel_size); + +} // namespace mllm::cpu::causal_conv diff --git a/mllm/backends/cpu/kernels/common/linear/kai_w4a32_dispatch.cpp b/mllm/backends/cpu/kernels/common/linear/kai_w4a32_dispatch.cpp new file mode 100644 index 000000000..5efe4f1bb --- /dev/null +++ b/mllm/backends/cpu/kernels/common/linear/kai_w4a32_dispatch.cpp @@ -0,0 +1,40 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/backends/cpu/kernels/common/linear/kai_w4a32_dispatch.hpp" + +#include + +#if defined(__linux__) +#include +#endif + +namespace mllm::cpu::kai_w4a32 { + +namespace { + +bool environmentFlagEnabled(const char* name) { + const char* value = std::getenv(name); + return value != nullptr && value[0] == '1' && value[1] == '\0'; +} + +} // namespace + +bool i8mmPrefillDisabled() { + static const bool disabled = environmentFlagEnabled("MLLM_KAI_PREFILL_I8MM_DISABLE"); + return disabled; +} + +bool cpuSupportsI8mm() { +#if defined(__linux__) && defined(__aarch64__) + constexpr unsigned long kHwcap2I8mm = 1UL << 13; + static const bool supported = (getauxval(AT_HWCAP2) & kHwcap2I8mm) != 0; + return supported; +#else + return false; +#endif +} + +bool shouldUseI8mmPrefill(int m) { return shouldUseI8mmPrefill(m, i8mmPrefillDisabled(), cpuSupportsI8mm()); } + +} // namespace mllm::cpu::kai_w4a32 diff --git a/mllm/backends/cpu/kernels/common/linear/kai_w4a32_dispatch.hpp b/mllm/backends/cpu/kernels/common/linear/kai_w4a32_dispatch.hpp new file mode 100644 index 000000000..0ce7a3334 --- /dev/null +++ b/mllm/backends/cpu/kernels/common/linear/kai_w4a32_dispatch.hpp @@ -0,0 +1,23 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +namespace mllm::cpu::kai_w4a32 { + +constexpr bool shouldUseI8mmPrefill(int m, bool disabled, bool cpu_supports_i8mm) { + return m >= 4 && !disabled && cpu_supports_i8mm; +} + +[[nodiscard]] bool i8mmPrefillDisabled(); + +[[nodiscard]] bool cpuSupportsI8mm(); + +[[nodiscard]] bool shouldUseI8mmPrefill(int m); + +constexpr int threadCount(int m, int requested_threads, int decode_thread_cap, int prefill_thread_cap) { + const int cap = m == 1 ? decode_thread_cap : prefill_thread_cap; + return cap > 0 && cap < requested_threads ? cap : requested_threads; +} + +} // namespace mllm::cpu::kai_w4a32 diff --git a/mllm/backends/cpu/kernels/common/parallel_linear/shared_input.hpp b/mllm/backends/cpu/kernels/common/parallel_linear/shared_input.hpp new file mode 100644 index 000000000..3e30436e8 --- /dev/null +++ b/mllm/backends/cpu/kernels/common/parallel_linear/shared_input.hpp @@ -0,0 +1,37 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +namespace mllm::cpu::parallel_linear { + +enum class SharedInputKernel : uint8_t { + kUnsupported, + kKaiDotprod, + kKaiI8mm, +}; + +struct SharedInputProjection { + float* dst = nullptr; + const uint8_t* packed_weight_bias = nullptr; + int32_t n = 0; +}; + +struct SharedInputPlan { + SharedInputKernel kernel = SharedInputKernel::kUnsupported; + size_t workspace_size = 0; + int32_t thread_count = 0; + + [[nodiscard]] bool supported() const { return kernel != SharedInputKernel::kUnsupported; } +}; + +[[nodiscard]] SharedInputPlan planKaiW4A32SharedInput(int32_t m, int32_t k, int32_t requested_threads, + int32_t decode_thread_cap, int32_t prefill_thread_cap); + +bool runKaiW4A32SharedInput(const SharedInputPlan& plan, const float* input, const SharedInputProjection* projections, + size_t projection_count, void* workspace, int32_t m, int32_t k); + +} // namespace mllm::cpu::parallel_linear diff --git a/mllm/backends/cpu/kernels/x86/parallel_linear.cpp b/mllm/backends/cpu/kernels/x86/parallel_linear.cpp new file mode 100644 index 000000000..354cf9193 --- /dev/null +++ b/mllm/backends/cpu/kernels/x86/parallel_linear.cpp @@ -0,0 +1,30 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/backends/cpu/kernels/common/parallel_linear/shared_input.hpp" + +namespace mllm::cpu::parallel_linear { + +SharedInputPlan planKaiW4A32SharedInput(int32_t m, int32_t k, int32_t requested_threads, int32_t decode_thread_cap, + int32_t prefill_thread_cap) { + (void)m; + (void)k; + (void)requested_threads; + (void)decode_thread_cap; + (void)prefill_thread_cap; + return {}; +} + +bool runKaiW4A32SharedInput(const SharedInputPlan& plan, const float* input, const SharedInputProjection* projections, + size_t projection_count, void* workspace, int32_t m, int32_t k) { + (void)plan; + (void)input; + (void)projections; + (void)projection_count; + (void)workspace; + (void)m; + (void)k; + return false; +} + +} // namespace mllm::cpu::parallel_linear diff --git a/mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.cpp b/mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.cpp new file mode 100644 index 000000000..03de5bb4a --- /dev/null +++ b/mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.cpp @@ -0,0 +1,76 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.hpp" + +#include +#include +#include +#include +#include +#include + +#include "mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.hpp" +#include "mllm/backends/cpu/kernels/common/gdn/gated_delta_net.hpp" + +namespace mllm::cpu { + +CPUCausalDepthwiseConv1DOp::CPUCausalDepthwiseConv1DOp(const aops::CausalDepthwiseConv1DOpOptions& options) + : aops::CausalDepthwiseConv1DOp(options) {} + +void CPUCausalDepthwiseConv1DOp::forward(const std::vector& inputs, std::vector& outputs) { + for (const auto& input : inputs) { + if (!input.isContiguous()) { throw std::invalid_argument("CausalDepthwiseConv1D CPU inputs must be contiguous"); } + } + if (weight_.isNil() || !weight_.isContiguous() || weight_.dtype() != kFloat32 || weight_.device() != kCPU + || weight_.shape() != Tensor::shape_t{options_.channels, 1, options_.kernel_size}) { + throw std::invalid_argument("CausalDepthwiseConv1D requires contiguous float32 [C, 1, K] weights"); + } + + const auto& input = inputs[0]; + const auto& state = inputs[1]; + auto& output = outputs[0]; + auto& updated_state = outputs[1]; + if (!options_.state_inplace) { std::memcpy(updated_state.ptr(), state.ptr(), state.bytes()); } + + const bool history_first = options_.accumulation_order == aops::CausalDepthwiseConv1DAccumulationOrder::kHistoryFirst; + if (history_first) { + causal_conv::depthwiseCausalConvHistoryFirstF32(input.ptr(), weight_.ptr(), updated_state.ptr(), + output.ptr(), input.shape()[0], input.shape()[1], + input.shape()[2], options_.kernel_size); + } else { + gdn::depthwiseCausalConvF32(input.ptr(), weight_.ptr(), updated_state.ptr(), output.ptr(), + input.shape()[0], input.shape()[1], input.shape()[2], options_.kernel_size); + } + + static const bool trace_activation = [] { + const char* value = std::getenv("MLLM_CAUSAL_CONV1D_TRACE"); + return value != nullptr && value[0] == '1' && value[1] == '\0'; + }(); + if (trace_activation) { + // One marker per accumulation order, so a device receipt shows which + // kernel the operation actually reached. + static std::atomic activated_orders{0}; + const uint32_t order_bit = 1U << static_cast(history_first); + if ((activated_orders.fetch_or(order_bit, std::memory_order_relaxed) & order_bit) == 0) { + std::fprintf(stderr, "MLLM_CAUSAL_CONV1D_ACTIVATED order=%s k=%d channels=%d\n", + aops::causalDepthwiseConv1DAccumulationOrder2Str(options_.accumulation_order), options_.kernel_size, + options_.channels); + } + } + + if (options_.bias) { + if (bias_.isNil() || !bias_.isContiguous() || bias_.dtype() != kFloat32 + || bias_.shape() != Tensor::shape_t{options_.channels}) { + throw std::invalid_argument("CausalDepthwiseConv1D bias must be contiguous float32 [C]"); + } + for (int32_t batch = 0; batch < input.shape()[0]; ++batch) { + for (int32_t token = 0; token < input.shape()[1]; ++token) { + auto* row = output.offsettedPtr({batch, token, 0}); + for (int32_t channel = 0; channel < options_.channels; ++channel) { row[channel] += bias_.ptr()[channel]; } + } + } + } +} + +} // namespace mllm::cpu diff --git a/mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.hpp b/mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.hpp new file mode 100644 index 000000000..aa073cc0d --- /dev/null +++ b/mllm/backends/cpu/ops/CausalDepthwiseConv1DOp.hpp @@ -0,0 +1,24 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include "mllm/core/aops/CausalDepthwiseConv1DOp.hpp" + +namespace mllm::cpu { + +class CPUCausalDepthwiseConv1DOp final : public aops::CausalDepthwiseConv1DOp { + public: + explicit CPUCausalDepthwiseConv1DOp(const aops::CausalDepthwiseConv1DOpOptions& options); + void forward(const std::vector& inputs, std::vector& outputs) override; +}; + +class CPUCausalDepthwiseConv1DOpFactory + : public TypedOpFactory { + public: + std::shared_ptr createOpImpl(const aops::CausalDepthwiseConv1DOpOptions& options) override { + return std::make_shared(options); + } +}; + +} // namespace mllm::cpu diff --git a/mllm/backends/cpu/ops/GroupedQueryAttentionDecodeOp.cpp b/mllm/backends/cpu/ops/GroupedQueryAttentionDecodeOp.cpp deleted file mode 100644 index 8f448301e..000000000 --- a/mllm/backends/cpu/ops/GroupedQueryAttentionDecodeOp.cpp +++ /dev/null @@ -1,104 +0,0 @@ -// Copyright (c) MLLM Team. -// Licensed under the MIT License. - -#include "mllm/backends/cpu/ops/GroupedQueryAttentionDecodeOp.hpp" - -#include -#include -#include -#include - -#include "mllm/backends/cpu/kernels/common/gqa_decode/fwd_bhsd.hpp" - -namespace mllm::cpu { -namespace { - -void groupedQueryAttentionDecodeFloat32Reference(const Tensor& query, const Tensor& key, const Tensor& value, Tensor& output) { - const auto q_shape = query.shape(); - const auto k_shape = key.shape(); - const auto v_shape = value.shape(); - const auto q_stride = query.stride(); - const auto k_stride = key.stride(); - const auto v_stride = value.stride(); - const int32_t groups = q_shape[1] / k_shape[1]; - const float scale = 1.0F / std::sqrt(static_cast(q_shape[3])); - - static thread_local std::vector probabilities; - probabilities.resize(static_cast(k_shape[2])); - - for (int32_t batch = 0; batch < q_shape[0]; ++batch) { - for (int32_t query_head = 0; query_head < q_shape[1]; ++query_head) { - const int32_t kv_head = query_head / groups; - const auto* q_head = query.coffsettedPtr({batch, query_head, 0, 0}); - const auto* k_head = key.coffsettedPtr({batch, kv_head, 0, 0}); - const auto* v_head = value.coffsettedPtr({batch, kv_head, 0, 0}); - auto* output_head = output.offsettedPtr({batch, query_head, 0, 0}); - - // Android release builds use -ffast-math; a finite sentinel keeps stable - // softmax valid under finite-math assumptions. - float maximum = std::numeric_limits::lowest(); - for (int32_t key_index = 0; key_index < k_shape[2]; ++key_index) { - const auto* key_token = k_head + static_cast(key_index) * k_stride[2]; - float score = 0.0F; - for (int32_t dim = 0; dim < q_shape[3]; ++dim) { - score += q_head[static_cast(dim) * q_stride[3]] * key_token[static_cast(dim) * k_stride[3]]; - } - probabilities[static_cast(key_index)] = score * scale; - maximum = std::max(maximum, probabilities[static_cast(key_index)]); - } - - float denominator = 0.0F; - for (int32_t key_index = 0; key_index < k_shape[2]; ++key_index) { - auto& probability = probabilities[static_cast(key_index)]; - probability = std::exp(probability - maximum); - denominator += probability; - } - const float inverse_denominator = 1.0F / denominator; - for (int32_t key_index = 0; key_index < k_shape[2]; ++key_index) { - probabilities[static_cast(key_index)] *= inverse_denominator; - } - - for (int32_t value_dim = 0; value_dim < v_shape[3]; ++value_dim) { - float result = 0.0F; - for (int32_t key_index = 0; key_index < k_shape[2]; ++key_index) { - const auto* value_token = v_head + static_cast(key_index) * v_stride[2]; - result += probabilities[static_cast(key_index)] * value_token[static_cast(value_dim) * v_stride[3]]; - } - output_head[value_dim] = result; - } - } - } -} - -} // namespace - -CPUGroupedQueryAttentionDecodeOp::CPUGroupedQueryAttentionDecodeOp(const aops::GroupedQueryAttentionDecodeOpOptions& options) - : aops::GroupedQueryAttentionDecodeOp(options) {} - -void CPUGroupedQueryAttentionDecodeOp::forward(const std::vector& inputs, std::vector& outputs) { - const auto& query = inputs[0]; - const auto& key = inputs[1]; - const auto& value = inputs[2]; - auto& output = outputs[0]; - const auto q_shape = query.shape(); - const auto k_shape = key.shape(); - const auto v_shape = value.shape(); - const auto q_stride = query.stride(); - const auto k_stride = key.stride(); - const auto v_stride = value.stride(); - const auto output_stride = output.stride(); - - static thread_local std::vector probabilities; - const int32_t group_size = q_shape[1] / k_shape[1]; - probabilities.resize(static_cast(group_size) * k_shape[2]); - - const bool completed = cpu::gqa_decode::fwdBhsdFp32( - q_shape[0], q_shape[1], k_shape[1], k_shape[2], q_shape[3], v_shape[3], query.ptr(), - {q_stride[0], q_stride[1], q_stride[2], q_stride[3]}, key.ptr(), - {k_stride[0], k_stride[1], k_stride[2], k_stride[3]}, value.ptr(), - {v_stride[0], v_stride[1], v_stride[2], v_stride[3]}, output.ptr(), - {output_stride[0], output_stride[1], output_stride[2], output_stride[3]}, probabilities.data()); - if (!completed) { groupedQueryAttentionDecodeFloat32Reference(query, key, value, output); } -} - -} // namespace mllm::cpu diff --git a/mllm/backends/cpu/ops/GroupedQueryAttentionDecodeOp.hpp b/mllm/backends/cpu/ops/GroupedQueryAttentionDecodeOp.hpp deleted file mode 100644 index 9f0598be9..000000000 --- a/mllm/backends/cpu/ops/GroupedQueryAttentionDecodeOp.hpp +++ /dev/null @@ -1,26 +0,0 @@ -// Copyright (c) MLLM Team. -// Licensed under the MIT License. - -#pragma once - -#include "mllm/core/BaseOp.hpp" -#include "mllm/core/aops/GroupedQueryAttentionDecodeOp.hpp" - -namespace mllm::cpu { - -class CPUGroupedQueryAttentionDecodeOp final : public aops::GroupedQueryAttentionDecodeOp { - public: - explicit CPUGroupedQueryAttentionDecodeOp(const aops::GroupedQueryAttentionDecodeOpOptions& options); - - void forward(const std::vector& inputs, std::vector& outputs) override; -}; - -class CPUGroupedQueryAttentionDecodeOpFactory - : public TypedOpFactory { - protected: - std::shared_ptr createOpImpl(const aops::GroupedQueryAttentionDecodeOpOptions& options) override { - return std::make_shared(options); - } -}; - -} // namespace mllm::cpu diff --git a/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp b/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp new file mode 100644 index 000000000..a9cc82b0e --- /dev/null +++ b/mllm/backends/cpu/ops/GroupedQueryAttentionOp.cpp @@ -0,0 +1,227 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/backends/cpu/ops/GroupedQueryAttentionOp.hpp" + +#include +#include +#include +#include +#include + +#include "mllm/backends/cpu/kernels/common/gqa_decode/fwd_bhsd.hpp" +#include "mllm/core/Parallel.hpp" + +namespace mllm::cpu { +namespace { + +template +void groupedQueryAttentionDirectStrided(const Tensor& query, const Tensor& key, const Tensor& value, Tensor& output) { + const auto q_shape = query.shape(); + const auto k_shape = key.shape(); + const auto v_shape = value.shape(); + const auto q_stride = query.stride(); + const auto k_stride = key.stride(); + const auto v_stride = value.stride(); + const auto o_stride = output.stride(); + const int32_t groups = q_shape[1] / k_shape[1]; + const float scale = 1.0F / std::sqrt(static_cast(q_shape[3])); + const int32_t context_offset = k_shape[2] - q_shape[2]; + const int32_t jobs = q_shape[0] * q_shape[1]; + // These tables describe arbitrary sequence strides without rebuilding the + // address expression in the value reduction. Keep their storage per caller + // thread so steady-state forwards only resize within the retained capacity. + static thread_local std::vector query_rows; + static thread_local std::vector key_rows; + static thread_local std::vector value_rows; + static thread_local std::vector output_rows; + query_rows.resize(static_cast(jobs) * q_shape[2]); + key_rows.resize(static_cast(q_shape[0]) * k_shape[1] * k_shape[2]); + value_rows.resize(static_cast(q_shape[0]) * v_shape[1] * v_shape[2]); + output_rows.resize(static_cast(jobs) * q_shape[2]); + for (int32_t batch = 0; batch < q_shape[0]; ++batch) { + for (int32_t head = 0; head < q_shape[1]; ++head) { + for (int32_t sequence = 0; sequence < q_shape[2]; ++sequence) { + const size_t row = (static_cast(batch) * q_shape[1] + head) * q_shape[2] + sequence; + query_rows[row] = query.coffsettedPtr({batch, head, sequence, 0}); + output_rows[row] = output.offsettedPtr({batch, head, sequence, 0}); + } + } + for (int32_t head = 0; head < k_shape[1]; ++head) { + for (int32_t sequence = 0; sequence < k_shape[2]; ++sequence) { + const size_t row = (static_cast(batch) * k_shape[1] + head) * k_shape[2] + sequence; + key_rows[row] = key.coffsettedPtr({batch, head, sequence, 0}); + value_rows[row] = value.coffsettedPtr({batch, head, sequence, 0}); + } + } + } + const Scalar* const* query_row_data = query_rows.data(); + const Scalar* const* key_row_data = key_rows.data(); + const Scalar* const* value_row_data = value_rows.data(); + Scalar* const* output_row_data = output_rows.data(); + + MLLM_AUTO_PARALLEL_FOR_BEGIN(job, 0, jobs, 1) { + const int32_t batch = job / q_shape[1]; + const int32_t query_head = job % q_shape[1]; + const int32_t kv_head = query_head / groups; + const size_t query_row_base = (static_cast(batch) * q_shape[1] + query_head) * q_shape[2]; + const size_t key_row_base = (static_cast(batch) * k_shape[1] + kv_head) * k_shape[2]; + const size_t value_row_base = (static_cast(batch) * v_shape[1] + kv_head) * v_shape[2]; + std::vector scores(static_cast(k_shape[2])); + for (int32_t query_index = 0; query_index < q_shape[2]; ++query_index) { + const int32_t visible_keys = context_offset + query_index + 1; + const size_t query_row_index = query_row_base + query_index; + const Scalar* query_row = query_row_data[query_row_index]; + Scalar* output_row = output_row_data[query_row_index]; + float maximum = std::numeric_limits::lowest(); + for (int32_t key_index = 0; key_index < visible_keys; ++key_index) { + float dot = 0.0F; + const Scalar* key_row = key_row_data[key_row_base + key_index]; + // Keep the contiguous dot expression separate: folding it into the + // runtime-stride induction loop changes contraction/codegen, which + // breaks callers bound to an exact generation-token oracle. + if (q_stride[3] == 1 && k_stride[3] == 1) { + for (int32_t dim = 0; dim < q_shape[3]; ++dim) { + dot += static_cast(query_row[dim]) * static_cast(key_row[dim]); + } + } else { + const Scalar* query_element = query_row; + const Scalar* key_element = key_row; + for (int32_t dim = 0; dim < q_shape[3]; ++dim) { + dot += static_cast(*query_element) * static_cast(*key_element); + query_element += q_stride[3]; + key_element += k_stride[3]; + } + } + scores[key_index] = dot * scale; + maximum = std::max(maximum, scores[key_index]); + } + + float denominator = 0.0F; + for (int32_t key_index = 0; key_index < visible_keys; ++key_index) { + scores[key_index] = std::exp(scores[key_index] - maximum); + denominator += scores[key_index]; + } + const float inverse_denominator = 1.0F / denominator; + for (int32_t value_dim = 0; value_dim < v_shape[3]; ++value_dim) { + float accumulated = 0.0F; + for (int32_t key_index = 0; key_index < visible_keys; ++key_index) { + accumulated += + (scores[key_index] * static_cast(value_row_data[value_row_base + key_index][value_dim * v_stride[3]])) + * inverse_denominator; + } + output_row[value_dim * o_stride[3]] = static_cast(accumulated); + } + } + } + MLLM_AUTO_PARALLEL_FOR_END() +} + +// Scalar fallback for the decode variant, used when the vectorized decode +// kernel declines the given geometry. +void groupedQueryAttentionDecodeFloat32Reference(const Tensor& query, const Tensor& key, const Tensor& value, + Tensor& output) { + const auto q_shape = query.shape(); + const auto k_shape = key.shape(); + const auto v_shape = value.shape(); + const auto q_stride = query.stride(); + const auto k_stride = key.stride(); + const auto v_stride = value.stride(); + // The decode kernel declines any tensor whose last-dimension stride is not 1, + // so this fallback is reached precisely when that may hold. Honour the output + // stride rather than assuming a contiguous value dimension. + const auto o_stride = output.stride(); + const int32_t groups = q_shape[1] / k_shape[1]; + const float scale = 1.0F / std::sqrt(static_cast(q_shape[3])); + + static thread_local std::vector probabilities; + probabilities.resize(static_cast(k_shape[2])); + + for (int32_t batch = 0; batch < q_shape[0]; ++batch) { + for (int32_t query_head = 0; query_head < q_shape[1]; ++query_head) { + const int32_t kv_head = query_head / groups; + const auto* q_head = query.coffsettedPtr({batch, query_head, 0, 0}); + const auto* k_head = key.coffsettedPtr({batch, kv_head, 0, 0}); + const auto* v_head = value.coffsettedPtr({batch, kv_head, 0, 0}); + auto* output_head = output.offsettedPtr({batch, query_head, 0, 0}); + + // Android release builds use -ffast-math; a finite sentinel keeps stable + // softmax valid under finite-math assumptions. + float maximum = std::numeric_limits::lowest(); + for (int32_t key_index = 0; key_index < k_shape[2]; ++key_index) { + const auto* key_token = k_head + static_cast(key_index) * k_stride[2]; + float score = 0.0F; + for (int32_t dim = 0; dim < q_shape[3]; ++dim) { + score += q_head[static_cast(dim) * q_stride[3]] * key_token[static_cast(dim) * k_stride[3]]; + } + probabilities[static_cast(key_index)] = score * scale; + maximum = std::max(maximum, probabilities[static_cast(key_index)]); + } + + float denominator = 0.0F; + for (int32_t key_index = 0; key_index < k_shape[2]; ++key_index) { + auto& probability = probabilities[static_cast(key_index)]; + probability = std::exp(probability - maximum); + denominator += probability; + } + const float inverse_denominator = 1.0F / denominator; + for (int32_t key_index = 0; key_index < k_shape[2]; ++key_index) { + probabilities[static_cast(key_index)] *= inverse_denominator; + } + + for (int32_t value_dim = 0; value_dim < v_shape[3]; ++value_dim) { + float result = 0.0F; + for (int32_t key_index = 0; key_index < k_shape[2]; ++key_index) { + const auto* value_token = v_head + static_cast(key_index) * v_stride[2]; + result += probabilities[static_cast(key_index)] * value_token[static_cast(value_dim) * v_stride[3]]; + } + output_head[static_cast(value_dim) * o_stride[3]] = result; + } + } + } +} + +void groupedQueryAttentionDecodeNativeKV(const Tensor& query, const Tensor& key, const Tensor& value, Tensor& output) { + const auto q_shape = query.shape(); + const auto k_shape = key.shape(); + const auto v_shape = value.shape(); + const auto q_stride = query.stride(); + const auto k_stride = key.stride(); + const auto v_stride = value.stride(); + const auto output_stride = output.stride(); + + static thread_local std::vector probabilities; + const int32_t group_size = q_shape[1] / k_shape[1]; + probabilities.resize(static_cast(group_size) * k_shape[2]); + + const bool completed = cpu::gqa_decode::fwdBhsdFp32( + q_shape[0], q_shape[1], k_shape[1], k_shape[2], q_shape[3], v_shape[3], query.ptr(), + {q_stride[0], q_stride[1], q_stride[2], q_stride[3]}, key.ptr(), + {k_stride[0], k_stride[1], k_stride[2], k_stride[3]}, value.ptr(), + {v_stride[0], v_stride[1], v_stride[2], v_stride[3]}, output.ptr(), + {output_stride[0], output_stride[1], output_stride[2], output_stride[3]}, probabilities.data()); + if (!completed) { groupedQueryAttentionDecodeFloat32Reference(query, key, value, output); } +} + +} // namespace + +CPUGroupedQueryAttentionOp::CPUGroupedQueryAttentionOp(const aops::GroupedQueryAttentionOpOptions& options) + : aops::GroupedQueryAttentionOp(options) {} + +void CPUGroupedQueryAttentionOp::forward(const std::vector& inputs, std::vector& outputs) { + switch (options_.implementation) { + case aops::GroupedQueryAttentionImplementation::kDirectStrided: + if (inputs[0].dtype() == kFloat32) { + groupedQueryAttentionDirectStrided(inputs[0], inputs[1], inputs[2], outputs[0]); + } else { + groupedQueryAttentionDirectStrided(inputs[0], inputs[1], inputs[2], outputs[0]); + } + return; + case aops::GroupedQueryAttentionImplementation::kDecodeNativeKV: + groupedQueryAttentionDecodeNativeKV(inputs[0], inputs[1], inputs[2], outputs[0]); + return; + } + throw std::invalid_argument("Unsupported CPU GroupedQueryAttention implementation"); +} + +} // namespace mllm::cpu diff --git a/mllm/backends/cpu/ops/GroupedQueryAttentionOp.hpp b/mllm/backends/cpu/ops/GroupedQueryAttentionOp.hpp new file mode 100644 index 000000000..e1ea785c5 --- /dev/null +++ b/mllm/backends/cpu/ops/GroupedQueryAttentionOp.hpp @@ -0,0 +1,24 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include "mllm/core/aops/GroupedQueryAttentionOp.hpp" + +namespace mllm::cpu { + +class CPUGroupedQueryAttentionOp final : public aops::GroupedQueryAttentionOp { + public: + explicit CPUGroupedQueryAttentionOp(const aops::GroupedQueryAttentionOpOptions& options); + void forward(const std::vector& inputs, std::vector& outputs) override; +}; + +class CPUGroupedQueryAttentionOpFactory + : public TypedOpFactory { + public: + std::shared_ptr createOpImpl(const aops::GroupedQueryAttentionOpOptions& options) override { + return std::make_shared(options); + } +}; + +} // namespace mllm::cpu diff --git a/mllm/backends/cpu/ops/LinearOp.cpp b/mllm/backends/cpu/ops/LinearOp.cpp index 7d93754a1..f28fb53c3 100644 --- a/mllm/backends/cpu/ops/LinearOp.cpp +++ b/mllm/backends/cpu/ops/LinearOp.cpp @@ -6,12 +6,9 @@ #include #include -#if defined(__linux__) -#include -#endif - #include "mllm/backends/cpu/ops/LinearOp.hpp" #include "mllm/backends/cpu/kernels/Kernels.hpp" +#include "mllm/backends/cpu/kernels/common/linear/kai_w4a32_dispatch.hpp" #include "mllm/core/DataTypes.hpp" #include "mllm/core/aops/LinearOp.hpp" @@ -32,19 +29,8 @@ bool environmentFlagEnabled(const char* name) { return value != nullptr && value[0] == '1' && value[1] == '\0'; } -bool cpuSupportsI8mm() { -#if defined(__linux__) && defined(__aarch64__) - constexpr unsigned long kHwcap2I8mm = 1UL << 13; - static const bool supported = (getauxval(AT_HWCAP2) & kHwcap2I8mm) != 0; - return supported; -#else - return false; -#endif -} - KaiW4A32Tile selectKaiW4A32PrefillTile(int m) { - static const bool disabled = environmentFlagEnabled("MLLM_KAI_PREFILL_I8MM_DISABLE"); - if (detail::shouldUseKaiW4A32I8mmPrefill(m, disabled, cpuSupportsI8mm())) { return kKaiW4A32I8mmTile; } + if (kai_w4a32::shouldUseI8mmPrefill(m)) { return kKaiW4A32I8mmTile; } return kKaiW4A32DotProdTile; } @@ -61,7 +47,7 @@ void traceKaiW4A32PrefillTile(KaiW4A32Tile tile, int m, int k, int n, int thread if (tile == kKaiW4A32I8mmTile) { std::fprintf(stderr, "MLLM_KAI_PREFILL_I8MM_ACTIVATED m=%d k=%d n=%d threads=%d\n", m, k, n, threads); } else { - const char* reason = environmentFlagEnabled("MLLM_KAI_PREFILL_I8MM_DISABLE") ? "disabled" : "unsupported"; + const char* reason = kai_w4a32::i8mmPrefillDisabled() ? "disabled" : "unsupported"; std::fprintf(stderr, "MLLM_KAI_PREFILL_DOTPROD_FALLBACK reason=%s m=%d k=%d n=%d threads=%d\n", reason, m, k, n, threads); } } @@ -72,6 +58,11 @@ void traceKaiW4A32PrefillTile(KaiW4A32Tile tile, int m, int k, int n, int thread CPULinearOp::CPULinearOp(const aops::LinearOpOptions& options) : LinearOp(options) {} +int CPULinearOp::kaiW4A32ThreadCount(int m) const { + return kai_w4a32::threadCount(m, options_.getThreads(), options_.kai_w4a32_decode_thread_cap, + options_.kai_w4a32_prefill_thread_cap); +} + Tensor CPULinearOp::acquireKaiWorkspace(int32_t workspace_size, int m) { if (m != 1) { return Tensor::empty({workspace_size}, kInt8, kCPU).alloc(); } @@ -248,10 +239,11 @@ void CPULinearOp::forward(const std::vector& inputs, std::vector int32_t work_space_size = kai_helper.workspace_size( M, K, ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles::qai8dxp1x8_qsi4c32p4x8_1x4x32); auto workspace = acquireKaiWorkspace(work_space_size, M); + const int thread_count = kaiW4A32ThreadCount(M); kai_helper.matmul(o.ptr(), input.ptr(), weight_.ptr(), workspace.ptr(), M, K, N, ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles::qai8dxp1x8_qsi4c32p4x8_1x4x32, - options_.getThreads()); + thread_count); return; } case aops::LinearImplTypes::kKaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp1x8_qsi4c32p8x8_1x8x32: { @@ -262,7 +254,8 @@ void CPULinearOp::forward(const std::vector& inputs, std::vector KaiW4A32Helper kai_helper; const auto tile = selectKaiW4A32PrefillTile(M); - traceKaiW4A32PrefillTile(tile, M, K, N, options_.getThreads()); + const int thread_count = kaiW4A32ThreadCount(M); + traceKaiW4A32PrefillTile(tile, M, K, N, thread_count); // FIXME: // Can be optimized for better performance. @@ -270,7 +263,7 @@ void CPULinearOp::forward(const std::vector& inputs, std::vector auto workspace = acquireKaiWorkspace(work_space_size, M); kai_helper.matmul(o.ptr(), input.ptr(), weight_.ptr(), workspace.ptr(), M, K, - N, tile, options_.getThreads()); + N, tile, thread_count); return; } case aops::LinearImplTypes::kKaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp4x8_qsi4c32p4x8_8x4x32: { @@ -287,10 +280,11 @@ void CPULinearOp::forward(const std::vector& inputs, std::vector int32_t work_space_size = kai_helper.workspace_size( M, K, ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles::qai8dxp4x8_qsi4c32p4x8_8x4x32); auto workspace = acquireKaiWorkspace(work_space_size, M); + const int thread_count = kaiW4A32ThreadCount(M); kai_helper.matmul(o.ptr(), input.ptr(), weight_.ptr(), workspace.ptr(), M, K, N, ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles::qai8dxp4x8_qsi4c32p4x8_8x4x32, - options_.getThreads()); + thread_count); return; } case aops::LinearImplTypes::kKaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp4x8_qsi4c32p4x8_16x4x32: { @@ -305,10 +299,11 @@ void CPULinearOp::forward(const std::vector& inputs, std::vector int32_t work_space_size = kai_helper.workspace_size( M, K, ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles::qai8dxp4x8_qsi4c32p4x8_16x4x32); auto workspace = acquireKaiWorkspace(work_space_size, M); + const int thread_count = kaiW4A32ThreadCount(M); kai_helper.matmul(o.ptr(), input.ptr(), weight_.ptr(), workspace.ptr(), M, K, N, ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles::qai8dxp4x8_qsi4c32p4x8_16x4x32, - options_.getThreads()); + thread_count); return; } case aops::LinearImplTypes::kKaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp4x8_qsi4c32p8x8_4x8x32: { @@ -325,10 +320,11 @@ void CPULinearOp::forward(const std::vector& inputs, std::vector int32_t work_space_size = kai_helper.workspace_size( M, K, ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles::qai8dxp4x8_qsi4c32p8x8_4x8x32); auto workspace = acquireKaiWorkspace(work_space_size, M); + const int thread_count = kaiW4A32ThreadCount(M); kai_helper.matmul(o.ptr(), input.ptr(), weight_.ptr(), workspace.ptr(), M, K, N, ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles::qai8dxp4x8_qsi4c32p8x8_4x8x32, - options_.getThreads()); + thread_count); return; } case aops::LinearImplTypes::kKaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp1x4_qsi4c32p4x4_1x4: { @@ -343,10 +339,11 @@ void CPULinearOp::forward(const std::vector& inputs, std::vector int32_t work_space_size = kai_helper.workspace_size( M, K, ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles::qai8dxp1x4_qsi4c32p4x4_1x4); auto workspace = acquireKaiWorkspace(work_space_size, M); + const int thread_count = kaiW4A32ThreadCount(M); kai_helper.matmul(o.ptr(), input.ptr(), weight_.ptr(), workspace.ptr(), M, K, N, ::mllm::cpu::arm::KaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk::Tiles::qai8dxp1x4_qsi4c32p4x4_1x4, - options_.getThreads()); + thread_count); return; } #endif diff --git a/mllm/backends/cpu/ops/LinearOp.hpp b/mllm/backends/cpu/ops/LinearOp.hpp index e189530ee..1519772c8 100644 --- a/mllm/backends/cpu/ops/LinearOp.hpp +++ b/mllm/backends/cpu/ops/LinearOp.hpp @@ -7,14 +7,6 @@ namespace mllm::cpu { -namespace detail { - -constexpr bool shouldUseKaiW4A32I8mmPrefill(int m, bool disabled, bool cpu_supports_i8mm) { - return m >= 4 && !disabled && cpu_supports_i8mm; -} - -} // namespace detail - class CPULinearOp final : public aops::LinearOp { public: explicit CPULinearOp(const aops::LinearOpOptions& options); @@ -28,6 +20,8 @@ class CPULinearOp final : public aops::LinearOp { private: Tensor acquireKaiWorkspace(int32_t workspace_size, int m); + [[nodiscard]] int kaiW4A32ThreadCount(int m) const; + Tensor kai_decode_workspace_; }; diff --git a/mllm/backends/cpu/ops/ParallelLinearOp.cpp b/mllm/backends/cpu/ops/ParallelLinearOp.cpp new file mode 100644 index 000000000..2330a443a --- /dev/null +++ b/mllm/backends/cpu/ops/ParallelLinearOp.cpp @@ -0,0 +1,102 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/backends/cpu/ops/ParallelLinearOp.hpp" + +#include +#include +#include + +#include "mllm/backends/cpu/kernels/common/parallel_linear/shared_input.hpp" + +namespace mllm::cpu { + +CPUParallelLinearOp::CPUParallelLinearOp(const aops::ParallelLinearOpOptions& options) : aops::ParallelLinearOp(options) { + fallback_ops_.reserve(options_.out_channels.size()); + for (const int32_t out_channels : options_.out_channels) { + aops::LinearOpOptions child_options{.in_channels = options_.in_channels, + .out_channels = out_channels, + .bias = options_.bias, + .impl_type = options_.impl_type, + .kai_w4a32_decode_thread_cap = options_.kai_w4a32_decode_thread_cap, + .kai_w4a32_prefill_thread_cap = options_.kai_w4a32_prefill_thread_cap}; + child_options.setThreads(options_.getThreads()); + fallback_ops_.push_back(std::make_unique(child_options)); + } +} + +void CPUParallelLinearOp::load(const ParameterFile::ptr_t& ploader) { + aops::ParallelLinearOp::load(ploader); + for (size_t index = 0; index < fallback_ops_.size(); ++index) { + fallback_ops_[index]->weight() = weights_[index]; + if (options_.bias) { fallback_ops_[index]->bias() = biases_[index]; } + } +} + +Tensor CPUParallelLinearOp::acquireKaiWorkspace(int32_t workspace_size, int m) { + if (m != 1) { return Tensor::empty({workspace_size}, kInt8, kCPU).alloc(); } + + if (kai_decode_workspace_.isNil() || kai_decode_workspace_.numel() < static_cast(workspace_size)) { + kai_decode_workspace_ = Tensor::empty({workspace_size}, kInt8, kCPU).alloc(); + } + return kai_decode_workspace_; +} + +bool CPUParallelLinearOp::tryForwardSharedInputKai(const Tensor& input, std::vector& outputs) { + constexpr size_t kMaximumSharedProjections = 3; + constexpr auto kRequiredImpl = aops::LinearImplTypes::kKaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp1x8_qsi4c32p8x8_1x8x32; + + if (input.isNil() || input.device() != kCPU || input.dtype() != kFloat32 || !input.isContiguous() || input.rank() < 2 + || input.size(-1) != options_.in_channels || options_.bias || options_.impl_type != kRequiredImpl + || outputs.size() != weights_.size() || weights_.size() < 2 || weights_.size() > kMaximumSharedProjections) { + return false; + } + for (size_t index = 0; index + 2 < input.shape().size(); ++index) { + if (input.shape()[index] != 1) { return false; } + } + const int32_t m = input.size(-2); + if (m <= 0) { return false; } + for (size_t index = 0; index < weights_.size(); ++index) { + if (weights_[index].isNil() || weights_[index].device() != kCPU || outputs[index].isNil() + || outputs[index].dtype() != kFloat32 || outputs[index].device() != kCPU || !outputs[index].isContiguous() + || outputs[index].rank() != input.rank() || outputs[index].size(-2) != m + || outputs[index].size(-1) != options_.out_channels[index]) { + return false; + } + for (size_t dimension = 0; dimension + 2 < input.shape().size(); ++dimension) { + if (outputs[index].shape()[dimension] != input.shape()[dimension]) { return false; } + } + } + + const auto plan = + parallel_linear::planKaiW4A32SharedInput(m, options_.in_channels, options_.getThreads(), + options_.kai_w4a32_decode_thread_cap, options_.kai_w4a32_prefill_thread_cap); + if (!plan.supported() || plan.workspace_size > static_cast(std::numeric_limits::max())) { return false; } + + std::array projections{}; + for (size_t index = 0; index < weights_.size(); ++index) { + projections[index] = { + .dst = outputs[index].ptr(), + .packed_weight_bias = reinterpret_cast(weights_[index].ptr()), + .n = options_.out_channels[index], + }; + } + + auto workspace = acquireKaiWorkspace(static_cast(plan.workspace_size), m); + return parallel_linear::runKaiW4A32SharedInput(plan, input.ptr(), projections.data(), weights_.size(), + workspace.ptr(), m, options_.in_channels); +} + +void CPUParallelLinearOp::forward(const std::vector& inputs, std::vector& outputs) { + const auto& input = inputs[0]; + if (tryForwardSharedInputKai(input, outputs)) { return; } + for (size_t index = 0; index < fallback_ops_.size(); ++index) { + // Keep fallback Linear execution aligned if the parent op's thread policy + // is adjusted after construction. + fallback_ops_[index]->options().setThreads(options_.getThreads()); + std::vector child_outputs = {outputs[index]}; + fallback_ops_[index]->forward({input}, child_outputs); + } +} + +} // namespace mllm::cpu diff --git a/mllm/backends/cpu/ops/ParallelLinearOp.hpp b/mllm/backends/cpu/ops/ParallelLinearOp.hpp new file mode 100644 index 000000000..67a5754ae --- /dev/null +++ b/mllm/backends/cpu/ops/ParallelLinearOp.hpp @@ -0,0 +1,39 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include +#include + +#include "mllm/backends/cpu/ops/LinearOp.hpp" +#include "mllm/core/aops/ParallelLinearOp.hpp" + +namespace mllm::cpu { + +class CPUParallelLinearOp final : public aops::ParallelLinearOp { + public: + explicit CPUParallelLinearOp(const aops::ParallelLinearOpOptions& options); + + void load(const ParameterFile::ptr_t& ploader) override; + void forward(const std::vector& inputs, std::vector& outputs) override; + + private: + bool tryForwardSharedInputKai(const Tensor& input, std::vector& outputs); + Tensor acquireKaiWorkspace(int32_t workspace_size, int m); + + std::vector> fallback_ops_; + // Only the decode-sized (m == 1) buffer is retained; a prefill workspace is + // two orders of magnitude larger and must not stay resident for the rest of + // the process, which CPULinearOp already avoids the same way. + Tensor kai_decode_workspace_; +}; + +class CPUParallelLinearOpFactory : public TypedOpFactory { + public: + std::shared_ptr createOpImpl(const aops::ParallelLinearOpOptions& options) override { + return std::make_shared(options); + } +}; + +} // namespace mllm::cpu diff --git a/mllm/compile/ir/GeneratedRTTIKind.hpp b/mllm/compile/ir/GeneratedRTTIKind.hpp index b08d6c103..726ec8657 100644 --- a/mllm/compile/ir/GeneratedRTTIKind.hpp +++ b/mllm/compile/ir/GeneratedRTTIKind.hpp @@ -1,4 +1,4 @@ -// Auto generated: 2026-08-13 00:19:31 +// Auto generated: 2026-08-21 15:20:12 // do not modify this file #pragma once @@ -41,7 +41,9 @@ enum NodeKind : uint32_t { RK_Op_LinalgIROp_ViewOp, RK_Op_LinalgIROp_SplitOp, RK_Op_LinalgIROp_FlashAttention2Op, - RK_Op_LinalgIROp_GroupedQueryAttentionDecodeOp, + RK_Op_LinalgIROp_CausalDepthwiseConv1DOp, + RK_Op_LinalgIROp_GroupedQueryAttentionOp, + RK_Op_LinalgIROp_ParallelLinearOp, RK_Op_LinalgIROp_RepeatOp, RK_Op_LinalgIROp_PermuteOp, RK_Op_LinalgIROp_Conv1DOp, diff --git a/mllm/compile/ir/NodeRTTIClassOfImpl.hpp b/mllm/compile/ir/NodeRTTIClassOfImpl.hpp index c107bed4b..0a08e255b 100644 --- a/mllm/compile/ir/NodeRTTIClassOfImpl.hpp +++ b/mllm/compile/ir/NodeRTTIClassOfImpl.hpp @@ -1,4 +1,4 @@ -// Auto generated: 2026-08-13 00:19:31 +// Auto generated: 2026-08-21 15:20:12 // do not modify this file #pragma once namespace mllm::ir { @@ -93,9 +93,16 @@ struct NodeRTTIClassOfImpl { #define RTTI_RK_OP_LINALGIROP_FLASHATTENTION2OP_IMPL(v) \ return (v)->getKind() >= RK_Op_LinalgIROp_FlashAttention2Op && (v)->getKind() <= RK_Op_LinalgIROp_FlashAttention2Op -#define RTTI_RK_OP_LINALGIROP_GROUPEDQUERYATTENTIONDECODEOP_IMPL(v) \ - return (v)->getKind() >= RK_Op_LinalgIROp_GroupedQueryAttentionDecodeOp \ - && (v)->getKind() <= RK_Op_LinalgIROp_GroupedQueryAttentionDecodeOp +#define RTTI_RK_OP_LINALGIROP_CAUSALDEPTHWISECONV1DOP_IMPL(v) \ + return (v)->getKind() >= RK_Op_LinalgIROp_CausalDepthwiseConv1DOp \ + && (v)->getKind() <= RK_Op_LinalgIROp_CausalDepthwiseConv1DOp + +#define RTTI_RK_OP_LINALGIROP_GROUPEDQUERYATTENTIONOP_IMPL(v) \ + return (v)->getKind() >= RK_Op_LinalgIROp_GroupedQueryAttentionOp \ + && (v)->getKind() <= RK_Op_LinalgIROp_GroupedQueryAttentionOp + +#define RTTI_RK_OP_LINALGIROP_PARALLELLINEAROP_IMPL(v) \ + return (v)->getKind() >= RK_Op_LinalgIROp_ParallelLinearOp && (v)->getKind() <= RK_Op_LinalgIROp_ParallelLinearOp #define RTTI_RK_OP_LINALGIROP_REPEATOP_IMPL(v) \ return (v)->getKind() >= RK_Op_LinalgIROp_RepeatOp && (v)->getKind() <= RK_Op_LinalgIROp_RepeatOp diff --git a/mllm/compile/ir/linalg/Op.cpp b/mllm/compile/ir/linalg/Op.cpp index c1e28c887..a33cd581b 100644 --- a/mllm/compile/ir/linalg/Op.cpp +++ b/mllm/compile/ir/linalg/Op.cpp @@ -66,7 +66,9 @@ LINALG_AOPS_DECL(OpTypes::kSTFT, STFTOp); LINALG_AOPS_DECL(OpTypes::kISTFT, ISTFTOp); LINALG_AOPS_DECL(OpTypes::kFlashAttention2, FlashAttention2Op); -LINALG_AOPS_DECL(OpTypes::kGroupedQueryAttentionDecode, GroupedQueryAttentionDecodeOp); +LINALG_AOPS_DECL(OpTypes::kCausalDepthwiseConv1D, CausalDepthwiseConv1DOp); +LINALG_AOPS_DECL(OpTypes::kGroupedQueryAttention, GroupedQueryAttentionOp); +LINALG_AOPS_DECL(OpTypes::kParallelLinear, ParallelLinearOp); LINALG_AOPS_DECL(OpTypes::kRepeat, RepeatOp); LINALG_AOPS_DECL(OpTypes::kPermute, PermuteOp); diff --git a/mllm/compile/ir/linalg/Op.hpp b/mllm/compile/ir/linalg/Op.hpp index 81cba50a1..e79765cc7 100644 --- a/mllm/compile/ir/linalg/Op.hpp +++ b/mllm/compile/ir/linalg/Op.hpp @@ -35,7 +35,9 @@ class X2XOp; class ViewOp; class SplitOp; class FlashAttention2Op; -class GroupedQueryAttentionDecodeOp; +class CausalDepthwiseConv1DOp; +class GroupedQueryAttentionOp; +class ParallelLinearOp; class RepeatOp; class PermuteOp; class Conv1DOp; @@ -198,7 +200,9 @@ LINALG_AOPS_DEFINE(ViewOp, VIEWOP); LINALG_AOPS_DEFINE(SplitOp, SPLITOP); LINALG_AOPS_DEFINE(FlashAttention2Op, FLASHATTENTION2OP); -LINALG_AOPS_DEFINE(GroupedQueryAttentionDecodeOp, GROUPEDQUERYATTENTIONDECODEOP); +LINALG_AOPS_DEFINE(CausalDepthwiseConv1DOp, CAUSALDEPTHWISECONV1DOP); +LINALG_AOPS_DEFINE(GroupedQueryAttentionOp, GROUPEDQUERYATTENTIONOP); +LINALG_AOPS_DEFINE(ParallelLinearOp, PARALLELLINEAROP); LINALG_AOPS_DEFINE(RepeatOp, REPEATOP); LINALG_AOPS_DEFINE(PermuteOp, PERMUTEOP); diff --git a/mllm/compile/ir/rtti_kind_gen.py b/mllm/compile/ir/rtti_kind_gen.py index 6d9c7531c..03029b23b 100644 --- a/mllm/compile/ir/rtti_kind_gen.py +++ b/mllm/compile/ir/rtti_kind_gen.py @@ -247,7 +247,9 @@ def define_lianlg_ir(ir: dict): op.derive(Cls("ViewOp")) op.derive(Cls("SplitOp")) op.derive(Cls("FlashAttention2Op")) - op.derive(Cls("GroupedQueryAttentionDecodeOp")) + op.derive(Cls("CausalDepthwiseConv1DOp")) + op.derive(Cls("GroupedQueryAttentionOp")) + op.derive(Cls("ParallelLinearOp")) op.derive(Cls("RepeatOp")) op.derive(Cls("PermuteOp")) op.derive(Cls("Conv1DOp")) diff --git a/mllm/compile/jit/binary/LinalgIRSerialization.cpp b/mllm/compile/jit/binary/LinalgIRSerialization.cpp index 1f8f7c213..307556400 100644 --- a/mllm/compile/jit/binary/LinalgIRSerialization.cpp +++ b/mllm/compile/jit/binary/LinalgIRSerialization.cpp @@ -11,7 +11,9 @@ #include "mllm/core/aops/TransposeOp.hpp" #include "mllm/core/aops/CastTypeOp.hpp" #include "mllm/core/aops/FlashAttention2Op.hpp" -#include "mllm/core/aops/GroupedQueryAttentionDecodeOp.hpp" +#include "mllm/core/aops/CausalDepthwiseConv1DOp.hpp" +#include "mllm/core/aops/GroupedQueryAttentionOp.hpp" +#include "mllm/core/aops/ParallelLinearOp.hpp" #include "mllm/core/aops/KVCacheOp.hpp" #include "mllm/core/aops/MultimodalRoPEOp.hpp" #include "mllm/core/aops/VisionRoPEOp.hpp" @@ -70,7 +72,9 @@ nlohmann::json dumpLinalgIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { CASE(Split) CASE(STFT) CASE(FlashAttention2) - CASE(GroupedQueryAttentionDecode) + CASE(CausalDepthwiseConv1D) + CASE(GroupedQueryAttention) + CASE(ParallelLinear) CASE(Repeat) CASE(Permute) CASE(Conv1D) @@ -137,7 +141,34 @@ nlohmann::json dumpLinearOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { return {{"in_channels", options.in_channels}, {"out_channels", options.out_channels}, {"bias", options.bias}, - {"impl_type", LinearImplTypes2Str(options.impl_type)}}; + {"impl_type", LinearImplTypes2Str(options.impl_type)}, + {"kai_w4a32_decode_thread_cap", options.kai_w4a32_decode_thread_cap}, + {"kai_w4a32_prefill_thread_cap", options.kai_w4a32_prefill_thread_cap}}; +} + +nlohmann::json dumpCausalDepthwiseConv1DOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { + const auto options = static_cast(op->getAOp())->options(); + return {{"channels", options.channels}, + {"kernel_size", options.kernel_size}, + {"bias", options.bias}, + {"state_inplace", options.state_inplace}, + {"accumulation_order", aops::causalDepthwiseConv1DAccumulationOrder2Str(options.accumulation_order)}}; +} + +nlohmann::json dumpGroupedQueryAttentionOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { + const auto options = static_cast(op->getAOp())->options(); + return {{"implementation", aops::groupedQueryAttentionImplementation2Str(options.implementation)}}; +} + +nlohmann::json dumpParallelLinearOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { + const auto options = static_cast(op->getAOp())->options(); + return {{"in_channels", options.in_channels}, + {"out_channels", options.out_channels}, + {"projection_names", options.projection_names}, + {"bias", options.bias}, + {"impl_type", aops::LinearImplTypes2Str(options.impl_type)}, + {"kai_w4a32_decode_thread_cap", options.kai_w4a32_decode_thread_cap}, + {"kai_w4a32_prefill_thread_cap", options.kai_w4a32_prefill_thread_cap}}; } nlohmann::json dumpRoPEOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { return {}; } @@ -216,7 +247,6 @@ nlohmann::json dumpFlashAttention2OpIROptions(const ir::linalg::LinalgIROp::ptr_ {"D", options.D}, {"hp_exp", options.hp_exp}, {"causal_mask", options.causal_mask}}; } -nlohmann::json dumpGroupedQueryAttentionDecodeOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { return {}; } nlohmann::json dumpRepeatOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op) { auto options = ((aops::RepeatOp*)op->getAOp())->options(); diff --git a/mllm/compile/jit/binary/LinalgIRSerialization.hpp b/mllm/compile/jit/binary/LinalgIRSerialization.hpp index 7c6a314ab..a3c2901d3 100644 --- a/mllm/compile/jit/binary/LinalgIRSerialization.hpp +++ b/mllm/compile/jit/binary/LinalgIRSerialization.hpp @@ -37,7 +37,9 @@ nlohmann::json dumpViewOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); nlohmann::json dumpSplitOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); nlohmann::json dumpSTFTOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); nlohmann::json dumpFlashAttention2OpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); -nlohmann::json dumpGroupedQueryAttentionDecodeOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); +nlohmann::json dumpCausalDepthwiseConv1DOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); +nlohmann::json dumpGroupedQueryAttentionOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); +nlohmann::json dumpParallelLinearOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); nlohmann::json dumpRepeatOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); nlohmann::json dumpPermuteOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); nlohmann::json dumpConv1DOpIROptions(const ir::linalg::LinalgIROp::ptr_t& op); diff --git a/mllm/compile/jit/interpreter/AopsFromJson.cpp b/mllm/compile/jit/interpreter/AopsFromJson.cpp index 3ca8eb710..f6c5c7a78 100644 --- a/mllm/compile/jit/interpreter/AopsFromJson.cpp +++ b/mllm/compile/jit/interpreter/AopsFromJson.cpp @@ -22,7 +22,9 @@ #include "mllm/core/aops/SplitOp.hpp" #include "mllm/core/aops/STFTOp.hpp" #include "mllm/core/aops/FlashAttention2Op.hpp" -#include "mllm/core/aops/GroupedQueryAttentionDecodeOp.hpp" +#include "mllm/core/aops/CausalDepthwiseConv1DOp.hpp" +#include "mllm/core/aops/GroupedQueryAttentionOp.hpp" +#include "mllm/core/aops/ParallelLinearOp.hpp" #include "mllm/core/aops/RepeatOp.hpp" #include "mllm/core/aops/PermuteOp.hpp" #include "mllm/core/aops/GELUOp.hpp" @@ -106,6 +108,12 @@ BaseOp::ptr_t aopsFromJson(const nlohmann::json& json) { return __flashAttention2FromJson(json); } else if (op_type == "GroupedQueryAttentionDecode") { return __groupedQueryAttentionDecodeFromJson(json); + } else if (op_type == "CausalDepthwiseConv1D") { + return __causalDepthwiseConv1DFromJson(json); + } else if (op_type == "GroupedQueryAttention") { + return __groupedQueryAttentionFromJson(json); + } else if (op_type == "ParallelLinear") { + return __parallelLinearFromJson(json); } else if (op_type == "Repeat") { return __repeatFromJson(json); } else if (op_type == "Permute") { @@ -232,6 +240,12 @@ BaseOp::ptr_t __linearFromJson(const nlohmann::json& json) { std::string impl_type_str = opts["impl_type"]; options.impl_type = aops::str2LinearImplTypes(impl_type_str); } + if (opts.contains("kai_w4a32_decode_thread_cap")) { + options.kai_w4a32_decode_thread_cap = opts["kai_w4a32_decode_thread_cap"]; + } + if (opts.contains("kai_w4a32_prefill_thread_cap")) { + options.kai_w4a32_prefill_thread_cap = opts["kai_w4a32_prefill_thread_cap"]; + } } DeviceTypes backend = DeviceTypes::kCPU; @@ -598,13 +612,69 @@ BaseOp::ptr_t __flashAttention2FromJson(const nlohmann::json& json) { return op; } +// Compatibility entry for graphs serialized while decode-only grouped-query +// attention was a separate operation. It now reconstructs as the +// DecodeNativeKV implementation, which reaches the same kernel. BaseOp::ptr_t __groupedQueryAttentionDecodeFromJson(const nlohmann::json& json) { - aops::GroupedQueryAttentionDecodeOpOptions options; + aops::GroupedQueryAttentionOpOptions options; + options.implementation = aops::GroupedQueryAttentionImplementation::kDecodeNativeKV; DeviceTypes backend = DeviceTypes::kCPU; if (json.contains("backend")) { backend = str2DeviceType(json["backend"]); } - return Context::instance().getBackend(backend)->createOp(OpTypes::kGroupedQueryAttentionDecode, options); + return Context::instance().getBackend(backend)->createOp(OpTypes::kGroupedQueryAttention, options); +} + +BaseOp::ptr_t __causalDepthwiseConv1DFromJson(const nlohmann::json& json) { + aops::CausalDepthwiseConv1DOpOptions options; + if (json.contains("op_options")) { + const auto& opts = json["op_options"]; + if (opts.contains("channels")) options.channels = opts["channels"]; + if (opts.contains("kernel_size")) options.kernel_size = opts["kernel_size"]; + if (opts.contains("bias")) options.bias = opts["bias"]; + if (opts.contains("state_inplace")) options.state_inplace = opts["state_inplace"]; + if (opts.contains("accumulation_order")) { + options.accumulation_order = + aops::str2CausalDepthwiseConv1DAccumulationOrder(opts["accumulation_order"].get()); + } + } + DeviceTypes backend = DeviceTypes::kCPU; + if (json.contains("backend")) { backend = str2DeviceType(json["backend"]); } + return Context::instance().getBackend(backend)->createOp(OpTypes::kCausalDepthwiseConv1D, options); +} + +BaseOp::ptr_t __groupedQueryAttentionFromJson(const nlohmann::json& json) { + aops::GroupedQueryAttentionOpOptions options; + if (json.contains("op_options") && json["op_options"].contains("implementation")) { + options.implementation = + aops::str2GroupedQueryAttentionImplementation(json["op_options"]["implementation"].get()); + } + DeviceTypes backend = DeviceTypes::kCPU; + if (json.contains("backend")) { backend = str2DeviceType(json["backend"]); } + return Context::instance().getBackend(backend)->createOp(OpTypes::kGroupedQueryAttention, options); +} + +BaseOp::ptr_t __parallelLinearFromJson(const nlohmann::json& json) { + aops::ParallelLinearOpOptions options; + if (json.contains("op_options")) { + const auto& opts = json["op_options"]; + if (opts.contains("in_channels")) options.in_channels = opts["in_channels"]; + if (opts.contains("out_channels")) options.out_channels = opts["out_channels"].get>(); + if (opts.contains("projection_names")) { + options.projection_names = opts["projection_names"].get>(); + } + if (opts.contains("bias")) options.bias = opts["bias"]; + if (opts.contains("impl_type")) options.impl_type = aops::str2LinearImplTypes(opts["impl_type"]); + if (opts.contains("kai_w4a32_decode_thread_cap")) { + options.kai_w4a32_decode_thread_cap = opts["kai_w4a32_decode_thread_cap"]; + } + if (opts.contains("kai_w4a32_prefill_thread_cap")) { + options.kai_w4a32_prefill_thread_cap = opts["kai_w4a32_prefill_thread_cap"]; + } + } + DeviceTypes backend = DeviceTypes::kCPU; + if (json.contains("backend")) { backend = str2DeviceType(json["backend"]); } + return Context::instance().getBackend(backend)->createOp(OpTypes::kParallelLinear, options); } BaseOp::ptr_t __repeatFromJson(const nlohmann::json& json) { diff --git a/mllm/compile/jit/interpreter/AopsFromJson.hpp b/mllm/compile/jit/interpreter/AopsFromJson.hpp index 90eccbed8..d93e1a478 100644 --- a/mllm/compile/jit/interpreter/AopsFromJson.hpp +++ b/mllm/compile/jit/interpreter/AopsFromJson.hpp @@ -35,6 +35,9 @@ BaseOp::ptr_t __splitFromJson(const nlohmann::json& json); BaseOp::ptr_t __stftFromJson(const nlohmann::json& json); BaseOp::ptr_t __flashAttention2FromJson(const nlohmann::json& json); BaseOp::ptr_t __groupedQueryAttentionDecodeFromJson(const nlohmann::json& json); +BaseOp::ptr_t __causalDepthwiseConv1DFromJson(const nlohmann::json& json); +BaseOp::ptr_t __groupedQueryAttentionFromJson(const nlohmann::json& json); +BaseOp::ptr_t __parallelLinearFromJson(const nlohmann::json& json); BaseOp::ptr_t __repeatFromJson(const nlohmann::json& json); BaseOp::ptr_t __permuteFromJson(const nlohmann::json& json); BaseOp::ptr_t __conv2dFromJson(const nlohmann::json& json); diff --git a/mllm/core/OpTypes.hpp b/mllm/core/OpTypes.hpp index a770532c9..0a3263003 100644 --- a/mllm/core/OpTypes.hpp +++ b/mllm/core/OpTypes.hpp @@ -97,8 +97,14 @@ enum class OpTypes : int32_t { kSigmoid = 75, - // Phase-specific native KV-head attention. - kGroupedQueryAttentionDecode = 76, + // 76 is retired: the decode-only grouped-query attention operation became + // the DecodeNativeKV implementation of kGroupedQueryAttention. Do not reuse + // the value, so an older serialized graph can never alias a new operation. + + // Reusable model primitives introduced by hybrid mobile decoders. + kCausalDepthwiseConv1D = 77, + kGroupedQueryAttention = 78, + kParallelLinear = 79, // Dynamic Op Start for user to register there own ops. kDynamicOp_Start = 4096, @@ -184,7 +190,9 @@ inline std::string optype2Str(OpTypes type) { case OpTypes::kEqual: return "Equal"; case OpTypes::kWhere: return "Where"; case OpTypes::kSigmoid: return "Sigmoid"; - case OpTypes::kGroupedQueryAttentionDecode: return "GroupedQueryAttentionDecode"; + case OpTypes::kCausalDepthwiseConv1D: return "CausalDepthwiseConv1D"; + case OpTypes::kGroupedQueryAttention: return "GroupedQueryAttention"; + case OpTypes::kParallelLinear: return "ParallelLinear"; case OpTypes::kDynamicOp_Start: return "DynamicOp_Start"; case OpTypes::kOpType_End: return "OpType_End"; default: return "Unknown"; diff --git a/mllm/core/aops/CausalDepthwiseConv1DOp.cpp b/mllm/core/aops/CausalDepthwiseConv1DOp.cpp new file mode 100644 index 000000000..554bb8f7a --- /dev/null +++ b/mllm/core/aops/CausalDepthwiseConv1DOp.cpp @@ -0,0 +1,82 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/core/aops/CausalDepthwiseConv1DOp.hpp" + +#include + +#include "mllm/compile/ir/graph/Op.hpp" +#include "mllm/compile/ir/linalg/Op.hpp" +#include "mllm/compile/ir/tensor/Op.hpp" +#include "mllm/core/Tensor.hpp" +#include "mllm/utils/Common.hpp" + +namespace mllm::aops { + +CausalDepthwiseConv1DOp::CausalDepthwiseConv1DOp(const CausalDepthwiseConv1DOpOptions& options) + : BaseOp(OpTypes::kCausalDepthwiseConv1D), options_(options) {} + +void CausalDepthwiseConv1DOp::load(const ParameterFile::ptr_t& ploader) { + weight_ = ploader->pull(getName() + ".weight"); + if (options_.bias) { bias_ = ploader->pull(getName() + ".bias"); } + if (ploader->version() == ModelFileVersion::kV1) { + weight_ = weight_.view({options_.channels, 1, options_.kernel_size}); + if (options_.bias) { bias_ = bias_.view({options_.channels}); } + } +} + +void CausalDepthwiseConv1DOp::trace(void* trace_context, const std::vector& inputs, std::vector& outputs) { + auto* ir_ctx = static_cast(trace_context); + if (weight_ && !ir_ctx->lookupSymbolTable(getName() + ".weight")) { + ir::IRWriterGuard guard(ir_ctx, ir_ctx->lookupSymbolTable("init")->cast_()->getTopRegion()); + ir_ctx->create(ir_ctx->create(weight_)); + if (options_.bias) { ir_ctx->create(ir_ctx->create(bias_)); } + } + auto i_irs = ir::tensor::wrapTensors2TensorIR(ir_ctx, inputs); + auto o_irs = ir::tensor::wrapTensors2TensorIR(ir_ctx, outputs); + ir_ctx->create(shared_from_this(), i_irs, o_irs); +} + +void CausalDepthwiseConv1DOp::forward(const std::vector& inputs, std::vector& outputs) { + NYI("CausalDepthwiseConv1DOp::forward not implemented in aops base."); +} + +void CausalDepthwiseConv1DOp::reshape(const std::vector& inputs, std::vector& outputs) { + if (inputs.size() != 2) { throw std::invalid_argument("CausalDepthwiseConv1D expects input and state"); } + const auto& input = inputs[0]; + const auto& state = inputs[1]; + if (options_.channels <= 0 || options_.kernel_size <= 1) { + throw std::invalid_argument("CausalDepthwiseConv1D options require positive channels and kernel_size > 1"); + } + if (input.rank() != 3 || input.shape()[2] != options_.channels) { + throw std::invalid_argument("CausalDepthwiseConv1D input must have [B, S, C] shape"); + } + const Tensor::shape_t expected_state = {input.shape()[0], options_.channels, options_.kernel_size - 1}; + if (state.shape() != expected_state) { + throw std::invalid_argument("CausalDepthwiseConv1D state must have [B, C, K - 1] shape"); + } + for (const auto& tensor : inputs) { + if (tensor.dtype() != kFloat32 || tensor.device() != input.device()) { + throw std::invalid_argument("CausalDepthwiseConv1D requires float32 inputs on one device"); + } + } + outputs.emplace_back(Tensor::empty(input.shape(), input.dtype(), input.device())); + outputs.emplace_back(options_.state_inplace ? state : Tensor::empty(state.shape(), state.dtype(), state.device())); +} + +void CausalDepthwiseConv1DOp::setup(const std::vector& inputs, std::vector& outputs) { + if (options_.state_inplace) { + outputs[0].alloc(); + } else { + BaseOp::setup(inputs, outputs); + } +} + +ParameterFile::ptr_t CausalDepthwiseConv1DOp::getParams() { + auto params = ParameterFile::create(); + params->push(getName() + ".weight", weight_); + if (options_.bias) { params->push(getName() + ".bias", bias_); } + return params; +} + +} // namespace mllm::aops diff --git a/mllm/core/aops/CausalDepthwiseConv1DOp.hpp b/mllm/core/aops/CausalDepthwiseConv1DOp.hpp new file mode 100644 index 000000000..fe693ad62 --- /dev/null +++ b/mllm/core/aops/CausalDepthwiseConv1DOp.hpp @@ -0,0 +1,63 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +#include "mllm/core/BaseOp.hpp" +#include "mllm/core/ParameterFile.hpp" + +namespace mllm::aops { + +enum class CausalDepthwiseConv1DAccumulationOrder : int32_t { + kCurrentFirst = 0, + kHistoryFirst = 1, +}; + +inline const char* causalDepthwiseConv1DAccumulationOrder2Str(CausalDepthwiseConv1DAccumulationOrder order) { + switch (order) { + case CausalDepthwiseConv1DAccumulationOrder::kCurrentFirst: return "CurrentFirst"; + case CausalDepthwiseConv1DAccumulationOrder::kHistoryFirst: return "HistoryFirst"; + } + return "Unknown"; +} + +inline CausalDepthwiseConv1DAccumulationOrder str2CausalDepthwiseConv1DAccumulationOrder(const std::string& value) { + if (value == "CurrentFirst") return CausalDepthwiseConv1DAccumulationOrder::kCurrentFirst; + if (value == "HistoryFirst") return CausalDepthwiseConv1DAccumulationOrder::kHistoryFirst; + throw std::invalid_argument("Unknown CausalDepthwiseConv1D accumulation order: " + value); +} + +struct CausalDepthwiseConv1DOpOptions : public BaseOpOptions { + int32_t channels = 0; + int32_t kernel_size = 0; + bool bias = false; + bool state_inplace = false; + CausalDepthwiseConv1DAccumulationOrder accumulation_order = CausalDepthwiseConv1DAccumulationOrder::kCurrentFirst; +}; + +class CausalDepthwiseConv1DOp : public BaseOp { + public: + explicit CausalDepthwiseConv1DOp(const CausalDepthwiseConv1DOpOptions& options); + + void load(const ParameterFile::ptr_t& ploader) override; + void trace(void* trace_context, const std::vector& inputs, std::vector& outputs) override; + void forward(const std::vector& inputs, std::vector& outputs) override; + void reshape(const std::vector& inputs, std::vector& outputs) override; + void setup(const std::vector& inputs, std::vector& outputs) override; + ParameterFile::ptr_t getParams() override; + + inline Tensor& weight() { return weight_; } + inline Tensor& bias() { return bias_; } + inline const CausalDepthwiseConv1DOpOptions& options() const { return options_; } + + protected: + Tensor weight_; + Tensor bias_; + CausalDepthwiseConv1DOpOptions options_; +}; + +} // namespace mllm::aops diff --git a/mllm/core/aops/GroupedQueryAttentionDecodeOp.cpp b/mllm/core/aops/GroupedQueryAttentionDecodeOp.cpp deleted file mode 100644 index 0643fa360..000000000 --- a/mllm/core/aops/GroupedQueryAttentionDecodeOp.cpp +++ /dev/null @@ -1,63 +0,0 @@ -// Copyright (c) MLLM Team. -// Licensed under the MIT License. - -#include "mllm/core/aops/GroupedQueryAttentionDecodeOp.hpp" - -#include "mllm/compile/ir/linalg/Op.hpp" -#include "mllm/core/BaseOp.hpp" -#include "mllm/core/Tensor.hpp" -#include "mllm/utils/Common.hpp" - -namespace mllm::aops { - -GroupedQueryAttentionDecodeOp::GroupedQueryAttentionDecodeOp(const GroupedQueryAttentionDecodeOpOptions& options) - : BaseOp(OpTypes::kGroupedQueryAttentionDecode), options_(options) {} - -void GroupedQueryAttentionDecodeOp::load(const ParameterFile::ptr_t& ploader) { MLLM_EMPTY_SCOPE; } - -void GroupedQueryAttentionDecodeOp::trace(void* trace_context, const std::vector& inputs, - std::vector& outputs) { - auto* ir_ctx = static_cast(trace_context); - auto i_irs = ir::tensor::wrapTensors2TensorIR(ir_ctx, inputs); - auto o_irs = ir::tensor::wrapTensors2TensorIR(ir_ctx, outputs); - ir_ctx->create(shared_from_this(), i_irs, o_irs); -} - -void GroupedQueryAttentionDecodeOp::forward(const std::vector& inputs, std::vector& outputs) { - NYI("GroupedQueryAttentionDecodeOp::forward not implemented in aops base."); -} - -void GroupedQueryAttentionDecodeOp::reshape(const std::vector& inputs, std::vector& outputs) { - MLLM_RT_ASSERT_EQ(inputs.size(), 3); - const auto& query = inputs[0]; - const auto& key = inputs[1]; - const auto& value = inputs[2]; - const auto& q_shape = query.shape(); - const auto& k_shape = key.shape(); - const auto& v_shape = value.shape(); - - MLLM_RT_ASSERT_EQ(q_shape.size(), 4); - MLLM_RT_ASSERT_EQ(k_shape.size(), 4); - MLLM_RT_ASSERT_EQ(v_shape.size(), 4); - MLLM_RT_ASSERT_EQ(q_shape[0], k_shape[0]); - MLLM_RT_ASSERT_EQ(q_shape[0], v_shape[0]); - MLLM_RT_ASSERT_EQ(q_shape[2], 1); - MLLM_RT_ASSERT_EQ(k_shape[1], v_shape[1]); - MLLM_RT_ASSERT_EQ(k_shape[2], v_shape[2]); - MLLM_RT_ASSERT_EQ(q_shape[3], k_shape[3]); - MLLM_RT_ASSERT(q_shape[1] > 0 && k_shape[1] > 0 && q_shape[1] % k_shape[1] == 0); - MLLM_RT_ASSERT(k_shape[2] > 0 && v_shape[3] > 0); - MLLM_RT_ASSERT_EQ(query.dtype(), kFloat32); - MLLM_RT_ASSERT_EQ(query.dtype(), key.dtype()); - MLLM_RT_ASSERT_EQ(query.dtype(), value.dtype()); - MLLM_RT_ASSERT_EQ(query.device(), key.device()); - MLLM_RT_ASSERT_EQ(query.device(), value.device()); - - outputs.emplace_back(Tensor::empty({q_shape[0], q_shape[1], 1, v_shape[3]}, query.dtype(), query.device())); -} - -void GroupedQueryAttentionDecodeOp::setup(const std::vector& inputs, std::vector& outputs) { - BaseOp::setup(inputs, outputs); -} - -} // namespace mllm::aops diff --git a/mllm/core/aops/GroupedQueryAttentionDecodeOp.hpp b/mllm/core/aops/GroupedQueryAttentionDecodeOp.hpp deleted file mode 100644 index 73efe9be0..000000000 --- a/mllm/core/aops/GroupedQueryAttentionDecodeOp.hpp +++ /dev/null @@ -1,36 +0,0 @@ -// Copyright (c) MLLM Team. -// Licensed under the MIT License. - -#pragma once - -#include "mllm/core/BaseOp.hpp" -#include "mllm/core/ParameterFile.hpp" - -namespace mllm::aops { - -struct GroupedQueryAttentionDecodeOpOptions : public BaseOpOptions {}; - -// Single-token grouped-query attention over native KV-head BHSD cache views. -// Inputs: query [B, Hq, 1, Dqk], key [B, Hkv, S, Dqk], -// value [B, Hkv, S, Dv]. Output: [B, Hq, 1, Dv]. -class GroupedQueryAttentionDecodeOp : public BaseOp { - public: - explicit GroupedQueryAttentionDecodeOp(const GroupedQueryAttentionDecodeOpOptions& options); - - void load(const ParameterFile::ptr_t& ploader) override; - - void trace(void* trace_context, const std::vector& inputs, std::vector& outputs) override; - - void forward(const std::vector& inputs, std::vector& outputs) override; - - void reshape(const std::vector& inputs, std::vector& outputs) override; - - void setup(const std::vector& inputs, std::vector& outputs) override; - - inline const GroupedQueryAttentionDecodeOpOptions& options() const { return options_; } - - protected: - GroupedQueryAttentionDecodeOpOptions options_; -}; - -} // namespace mllm::aops diff --git a/mllm/core/aops/GroupedQueryAttentionOp.cpp b/mllm/core/aops/GroupedQueryAttentionOp.cpp new file mode 100644 index 000000000..a41b19411 --- /dev/null +++ b/mllm/core/aops/GroupedQueryAttentionOp.cpp @@ -0,0 +1,66 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/core/aops/GroupedQueryAttentionOp.hpp" + +#include + +#include "mllm/compile/ir/linalg/Op.hpp" +#include "mllm/core/Tensor.hpp" +#include "mllm/utils/Common.hpp" + +namespace mllm::aops { + +GroupedQueryAttentionOp::GroupedQueryAttentionOp(const GroupedQueryAttentionOpOptions& options) + : BaseOp(OpTypes::kGroupedQueryAttention), options_(options) {} + +void GroupedQueryAttentionOp::load(const ParameterFile::ptr_t& ploader) { MLLM_EMPTY_SCOPE; } + +void GroupedQueryAttentionOp::trace(void* trace_context, const std::vector& inputs, std::vector& outputs) { + auto* ir_ctx = static_cast(trace_context); + auto i_irs = ir::tensor::wrapTensors2TensorIR(ir_ctx, inputs); + auto o_irs = ir::tensor::wrapTensors2TensorIR(ir_ctx, outputs); + ir_ctx->create(shared_from_this(), i_irs, o_irs); +} + +void GroupedQueryAttentionOp::forward(const std::vector& inputs, std::vector& outputs) { + NYI("GroupedQueryAttentionOp::forward not implemented in aops base."); +} + +void GroupedQueryAttentionOp::reshape(const std::vector& inputs, std::vector& outputs) { + if (inputs.size() != 3) { throw std::invalid_argument("GroupedQueryAttention expects query, key, and value"); } + const auto& query = inputs[0]; + const auto& key = inputs[1]; + const auto& value = inputs[2]; + const auto q_shape = query.shape(); + const auto k_shape = key.shape(); + const auto v_shape = value.shape(); + if (q_shape.size() != 4 || k_shape.size() != 4 || v_shape.size() != 4 || q_shape[0] <= 0 || q_shape[0] != k_shape[0] + || q_shape[0] != v_shape[0] || q_shape[1] <= 0 || k_shape[1] <= 0 || k_shape[1] != v_shape[1] + || q_shape[1] % k_shape[1] != 0 || q_shape[2] <= 0 || k_shape[2] <= 0 || k_shape[2] != v_shape[2] + || k_shape[2] < q_shape[2] || q_shape[3] <= 0 || q_shape[3] != k_shape[3] || v_shape[3] <= 0) { + throw std::invalid_argument("GroupedQueryAttention expects compatible [B, H, S, D] tensors"); + } + if (query.dtype() != key.dtype() || query.dtype() != value.dtype() + || (query.dtype() != kFloat32 && query.dtype() != kFloat16)) { + throw std::invalid_argument("GroupedQueryAttention requires matching float32 or float16 inputs"); + } + if (query.device() != key.device() || query.device() != value.device()) { + throw std::invalid_argument("GroupedQueryAttention inputs must be on the same device"); + } + if (options_.implementation == GroupedQueryAttentionImplementation::kDecodeNativeKV) { + if (q_shape[2] != 1) { + throw std::invalid_argument("GroupedQueryAttention DecodeNativeKV accepts a single query position only"); + } + if (query.dtype() != kFloat32) { + throw std::invalid_argument("GroupedQueryAttention DecodeNativeKV supports float32 only"); + } + } + outputs.emplace_back(Tensor::empty({q_shape[0], q_shape[1], q_shape[2], v_shape[3]}, value.dtype(), value.device())); +} + +void GroupedQueryAttentionOp::setup(const std::vector& inputs, std::vector& outputs) { + BaseOp::setup(inputs, outputs); +} + +} // namespace mllm::aops diff --git a/mllm/core/aops/GroupedQueryAttentionOp.hpp b/mllm/core/aops/GroupedQueryAttentionOp.hpp new file mode 100644 index 000000000..884818775 --- /dev/null +++ b/mllm/core/aops/GroupedQueryAttentionOp.hpp @@ -0,0 +1,61 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +#include "mllm/core/BaseOp.hpp" + +namespace mllm::aops { + +// Grouped-query attention variants over native KV-head history. They differ in +// reduction order and in the shapes they accept, so the variant is part of the +// operation contract: callers bound to an exact generation-token oracle cannot +// be migrated between them silently. +enum class GroupedQueryAttentionImplementation : int32_t { + // Any query length. Masks per query position and accumulates in the order + // established by the eager reference. + kDirectStrided = 0, + // Single query position only. Uses the dedicated decode kernel, which is + // faster but reduces in a different order. + kDecodeNativeKV = 1, +}; + +inline const char* groupedQueryAttentionImplementation2Str(GroupedQueryAttentionImplementation implementation) { + switch (implementation) { + case GroupedQueryAttentionImplementation::kDirectStrided: return "DirectStrided"; + case GroupedQueryAttentionImplementation::kDecodeNativeKV: return "DecodeNativeKV"; + } + return "Unknown"; +} + +inline GroupedQueryAttentionImplementation str2GroupedQueryAttentionImplementation(const std::string& value) { + if (value == "DirectStrided") return GroupedQueryAttentionImplementation::kDirectStrided; + if (value == "DecodeNativeKV") return GroupedQueryAttentionImplementation::kDecodeNativeKV; + throw std::invalid_argument("Unknown GroupedQueryAttention implementation: " + value); +} + +struct GroupedQueryAttentionOpOptions : public BaseOpOptions { + GroupedQueryAttentionImplementation implementation = GroupedQueryAttentionImplementation::kDirectStrided; +}; + +class GroupedQueryAttentionOp : public BaseOp { + public: + explicit GroupedQueryAttentionOp(const GroupedQueryAttentionOpOptions& options); + + void load(const ParameterFile::ptr_t& ploader) override; + void trace(void* trace_context, const std::vector& inputs, std::vector& outputs) override; + void forward(const std::vector& inputs, std::vector& outputs) override; + void reshape(const std::vector& inputs, std::vector& outputs) override; + void setup(const std::vector& inputs, std::vector& outputs) override; + + inline const GroupedQueryAttentionOpOptions& options() const { return options_; } + + protected: + GroupedQueryAttentionOpOptions options_; +}; + +} // namespace mllm::aops diff --git a/mllm/core/aops/LinearOp.hpp b/mllm/core/aops/LinearOp.hpp index 4b0535066..dd149bb58 100644 --- a/mllm/core/aops/LinearOp.hpp +++ b/mllm/core/aops/LinearOp.hpp @@ -168,6 +168,8 @@ struct LinearOpOptions : public BaseOpOptions { int32_t out_channels; bool bias; LinearImplTypes impl_type; + int32_t kai_w4a32_decode_thread_cap = 0; + int32_t kai_w4a32_prefill_thread_cap = 0; LinearImplTypes qnn_impl_type = LinearImplTypes::kQNN_tensor_symm_w8a16; // specify Linear type when using QNN }; diff --git a/mllm/core/aops/ParallelLinearOp.cpp b/mllm/core/aops/ParallelLinearOp.cpp new file mode 100644 index 000000000..48e3f3824 --- /dev/null +++ b/mllm/core/aops/ParallelLinearOp.cpp @@ -0,0 +1,122 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/core/aops/ParallelLinearOp.hpp" + +#include + +#include "mllm/compile/ir/graph/Op.hpp" +#include "mllm/compile/ir/linalg/Op.hpp" +#include "mllm/compile/ir/tensor/Op.hpp" +#include "mllm/core/Tensor.hpp" +#include "mllm/utils/Common.hpp" + +namespace mllm::aops { + +ParallelLinearOp::ParallelLinearOp(const ParallelLinearOpOptions& options) + : BaseOp(OpTypes::kParallelLinear), options_(options) {} + +// Parameters keep the original per-projection checkpoint names, so the fused +// operation resolves them in its own parent scope rather than under its own +// layer name. Two fused operations registered in one parent must therefore not +// declare the same projection names, or they would resolve to the same +// tensors; validateProjectionNames rejects the locally detectable violations. +std::string ParallelLinearOp::projectionParameterName(size_t index, const char* suffix) const { + const auto separator = getName().rfind('.'); + const std::string parent = separator == std::string::npos ? std::string{} : getName().substr(0, separator + 1); + return parent + options_.projection_names.at(index) + suffix; +} + +void ParallelLinearOp::validateProjectionNames() const { + if (options_.projection_names.size() != options_.out_channels.size() || options_.projection_names.size() < 2) { + throw std::invalid_argument("ParallelLinear requires matching projection names and at least two outputs"); + } + for (size_t index = 0; index < options_.projection_names.size(); ++index) { + const auto& name = options_.projection_names[index]; + if (name.empty() || name.find('.') != std::string::npos) { + throw std::invalid_argument("ParallelLinear projection names must be non-empty and scope-local: " + name); + } + for (size_t other = 0; other < index; ++other) { + if (options_.projection_names[other] == name) { + throw std::invalid_argument("ParallelLinear projection names must be unique: " + name); + } + } + } +} + +void ParallelLinearOp::load(const ParameterFile::ptr_t& ploader) { + validateProjectionNames(); + weights_.clear(); + biases_.clear(); + weights_.reserve(options_.projection_names.size()); + if (options_.bias) { biases_.reserve(options_.projection_names.size()); } + for (size_t index = 0; index < options_.projection_names.size(); ++index) { + auto weight = ploader->pull(projectionParameterName(index, ".weight")); + if (ploader->version() == ModelFileVersion::kV1 + && (options_.impl_type == LinearImplTypes::kDefault || options_.impl_type == LinearImplTypes::kBLAS + || options_.impl_type == LinearImplTypes::kGGUF || options_.impl_type == LinearImplTypes::kMllmBlas)) { + weight = weight.view({options_.out_channels[index], options_.in_channels}); + } + weights_.push_back(std::move(weight)); + if (options_.bias) { + auto bias = ploader->pull(projectionParameterName(index, ".bias")); + if (ploader->version() == ModelFileVersion::kV1) { bias = bias.view({options_.out_channels[index]}); } + biases_.push_back(std::move(bias)); + } + } +} + +void ParallelLinearOp::trace(void* trace_context, const std::vector& inputs, std::vector& outputs) { + auto* ir_ctx = static_cast(trace_context); + if (!weights_.empty()) { + ir::IRWriterGuard guard(ir_ctx, ir_ctx->lookupSymbolTable("init")->cast_()->getTopRegion()); + for (size_t index = 0; index < weights_.size(); ++index) { + const auto weight_name = projectionParameterName(index, ".weight"); + if (!ir_ctx->lookupSymbolTable(weight_name)) { + ir_ctx->create(ir_ctx->create(weights_[index])); + } + if (options_.bias) { + const auto bias_name = projectionParameterName(index, ".bias"); + if (!ir_ctx->lookupSymbolTable(bias_name)) { + ir_ctx->create(ir_ctx->create(biases_[index])); + } + } + } + } + auto i_irs = ir::tensor::wrapTensors2TensorIR(ir_ctx, inputs); + auto o_irs = ir::tensor::wrapTensors2TensorIR(ir_ctx, outputs); + ir_ctx->create(shared_from_this(), i_irs, o_irs); +} + +void ParallelLinearOp::forward(const std::vector& inputs, std::vector& outputs) { + NYI("ParallelLinearOp::forward not implemented in aops base."); +} + +void ParallelLinearOp::reshape(const std::vector& inputs, std::vector& outputs) { + if (inputs.size() != 1 || inputs[0].rank() < 2 || inputs[0].size(-1) != options_.in_channels) { + throw std::invalid_argument("ParallelLinear expects one [..., M, K] input with the configured K"); + } + if (options_.in_channels <= 0) { throw std::invalid_argument("ParallelLinear options are incomplete"); } + validateProjectionNames(); + for (const int32_t channels : options_.out_channels) { + if (channels <= 0) { throw std::invalid_argument("ParallelLinear output channels must be positive"); } + auto shape = inputs[0].shape(); + shape.back() = channels; + outputs.emplace_back(Tensor::empty(shape, inputs[0].dtype(), inputs[0].device())); + } +} + +void ParallelLinearOp::setup(const std::vector& inputs, std::vector& outputs) { + BaseOp::setup(inputs, outputs); +} + +ParameterFile::ptr_t ParallelLinearOp::getParams() { + auto params = ParameterFile::create(); + for (size_t index = 0; index < weights_.size(); ++index) { + params->push(projectionParameterName(index, ".weight"), weights_[index]); + if (options_.bias) { params->push(projectionParameterName(index, ".bias"), biases_[index]); } + } + return params; +} + +} // namespace mllm::aops diff --git a/mllm/core/aops/ParallelLinearOp.hpp b/mllm/core/aops/ParallelLinearOp.hpp new file mode 100644 index 000000000..7d7243743 --- /dev/null +++ b/mllm/core/aops/ParallelLinearOp.hpp @@ -0,0 +1,49 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +#include "mllm/core/BaseOp.hpp" +#include "mllm/core/ParameterFile.hpp" +#include "mllm/core/aops/LinearOp.hpp" + +namespace mllm::aops { + +struct ParallelLinearOpOptions : public BaseOpOptions { + int32_t in_channels = 0; + std::vector out_channels; + std::vector projection_names; + bool bias = false; + LinearImplTypes impl_type = LinearImplTypes::kDefault; + int32_t kai_w4a32_decode_thread_cap = 0; + int32_t kai_w4a32_prefill_thread_cap = 0; +}; + +class ParallelLinearOp : public BaseOp { + public: + explicit ParallelLinearOp(const ParallelLinearOpOptions& options); + + void load(const ParameterFile::ptr_t& ploader) override; + void trace(void* trace_context, const std::vector& inputs, std::vector& outputs) override; + void forward(const std::vector& inputs, std::vector& outputs) override; + void reshape(const std::vector& inputs, std::vector& outputs) override; + void setup(const std::vector& inputs, std::vector& outputs) override; + ParameterFile::ptr_t getParams() override; + + [[nodiscard]] inline const ParallelLinearOpOptions& options() const { return options_; } + + protected: + [[nodiscard]] std::string projectionParameterName(size_t index, const char* suffix) const; + + void validateProjectionNames() const; + + std::vector weights_; + std::vector biases_; + ParallelLinearOpOptions options_; +}; + +} // namespace mllm::aops diff --git a/mllm/models/common/rope_tables.hpp b/mllm/models/common/rope_tables.hpp new file mode 100644 index 000000000..00a4d66d4 --- /dev/null +++ b/mllm/models/common/rope_tables.hpp @@ -0,0 +1,79 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +#include "mllm/core/Tensor.hpp" + +namespace mllm::models::common { + +// Analytical RoPE tables for models whose rotation is performed by the +// registered nn::RoPE operation and whose only model-side responsibility is +// materializing the immutable frequency table and the per-request sin/cos +// buffers. +// +// This is request orchestration that produces constant operation inputs, not a +// tensor operation, so it stays on the model side rather than under +// nn/llm_components. +// +// This is the plain default-RoPE contract with no attention scaling. Models +// that need a scaled or otherwise reparameterized table should not reuse these +// helpers, because the scaling factor changes the frozen numerical result. + +inline auto makeRoPEInvFreq(int32_t head_dim, float rope_theta) -> Tensor { + if (head_dim <= 1 || head_dim % 2 != 0) { + throw std::invalid_argument("makeRoPEInvFreq requires an even head_dim greater than one"); + } + if (!std::isfinite(rope_theta) || rope_theta <= 0.0F) { + throw std::invalid_argument("makeRoPEInvFreq requires a finite positive rope_theta"); + } + auto inv_freq = Tensor::empty({head_dim / 2}, kFloat32, kCPU).alloc(); + auto* data = inv_freq.ptr(); + for (int32_t index = 0; index < head_dim / 2; ++index) { + data[index] = 1.0F / std::pow(rope_theta, 2.0F * static_cast(index) / static_cast(head_dim)); + } + return inv_freq; +} + +// position_ids is [B, S] int64. The returned sin/cos are [B, S, head_dim] with +// each half-dimension angle duplicated into both halves, matching the +// rotate-half layout consumed by nn::RoPE. +inline auto makeRotaryPosEmbedding(const Tensor& position_ids, const Tensor& inv_freq) -> std::pair { + if (position_ids.isNil() || inv_freq.isNil()) { + throw std::invalid_argument("makeRotaryPosEmbedding inputs must not be nil"); + } + if (position_ids.shape().size() != 2 || position_ids.dtype() != kInt64 || position_ids.device() != kCPU) { + throw std::invalid_argument("makeRotaryPosEmbedding expects a rank-2 int64 CPU position_ids tensor"); + } + if (inv_freq.shape().size() != 1 || inv_freq.dtype() != kFloat32 || inv_freq.device() != kCPU) { + throw std::invalid_argument("makeRotaryPosEmbedding expects a rank-1 float32 CPU inv_freq tensor"); + } + + const auto batch = position_ids.shape()[0]; + const auto sequence = position_ids.shape()[1]; + const auto half_dim = inv_freq.shape()[0]; + const auto head_dim = half_dim * 2; + auto sin = Tensor::empty({batch, sequence, head_dim}, kFloat32, kCPU).alloc(); + auto cos = Tensor::empty({batch, sequence, head_dim}, kFloat32, kCPU).alloc(); + const auto* positions = position_ids.ptr(); + const auto* frequencies = inv_freq.ptr(); + auto* sin_data = sin.ptr(); + auto* cos_data = cos.ptr(); + for (int32_t b = 0; b < batch; ++b) { + for (int32_t s = 0; s < sequence; ++s) { + for (int32_t d = 0; d < half_dim; ++d) { + const auto angle = static_cast(positions[b * sequence + s]) * frequencies[d]; + const auto offset = (b * sequence + s) * head_dim; + sin_data[offset + d] = sin_data[offset + d + half_dim] = std::sin(angle); + cos_data[offset + d] = cos_data[offset + d + half_dim] = std::cos(angle); + } + } + } + return {sin, cos}; +} + +} // namespace mllm::models::common diff --git a/mllm/models/lfm2/configuration_lfm2.hpp b/mllm/models/lfm2/configuration_lfm2.hpp new file mode 100644 index 000000000..067c2feca --- /dev/null +++ b/mllm/models/lfm2/configuration_lfm2.hpp @@ -0,0 +1,170 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include +#include +#include + +#include "mllm/core/ParameterFile.hpp" +#include "mllm/core/aops/LinearOp.hpp" +#include "mllm/engine/ConfigFile.hpp" + +namespace mllm::models::lfm2 { + +struct Lfm2Config : protected ConfigFile { + Lfm2Config() = default; + + explicit Lfm2Config(const std::string& file_path) : ConfigFile(file_path) { + const auto& cfg = data(); + hidden_size = cfg.at("hidden_size"); + intermediate_size = cfg.at("intermediate_size"); + num_hidden_layers = cfg.at("num_hidden_layers"); + num_attention_heads = cfg.at("num_attention_heads"); + // head_dim's default divides by this below, which runs before validate(). + if (num_attention_heads <= 0) { throw std::invalid_argument("LFM2 num_attention_heads must be positive"); } + num_key_value_heads = cfg.at("num_key_value_heads"); + max_position_embeddings = cfg.at("max_position_embeddings"); + vocab_size = cfg.at("vocab_size"); + conv_L_cache = cfg.at("conv_L_cache"); + conv_bias = cfg.at("conv_bias"); + norm_eps = cfg.at("norm_eps"); + tie_word_embeddings = cfg.at("tie_word_embeddings"); + bos_token_id = cfg.at("bos_token_id"); + eos_token_id = cfg.at("eos_token_id"); + pad_token_id = cfg.at("pad_token_id"); + block_auto_adjust_ff_dim = cfg.value("block_auto_adjust_ff_dim", false); + block_ffn_dim_multiplier = cfg.value("block_ffn_dim_multiplier", 1.0F); + block_multiple_of = cfg.value("block_multiple_of", 256); + head_dim = cfg.value("head_dim", hidden_size / num_attention_heads); + max_cache_length = cfg.value("max_cache_length", 2048); + + const auto& rope = cfg.at("rope_parameters"); + rope_theta = rope.at("rope_theta"); + rope_type = rope.value("rope_type", std::string("default")); + for (const auto& layer_type : cfg.at("layer_types")) { layer_types.push_back(layer_type.get()); } + + const auto linear_impl_name = cfg.value("linear_impl_type", std::string("Default")); + linear_impl_type = aops::str2LinearImplTypes(linear_impl_name); + if (linear_impl_type == aops::LinearImplTypes::kDefault && linear_impl_name != "Default") { + throw std::invalid_argument("LFM2 contains an unsupported linear_impl_type: " + linear_impl_name); + } + validate(); + } + + int32_t hidden_size = 2048; + int32_t intermediate_size = 10752; + int32_t num_hidden_layers = 30; + int32_t num_attention_heads = 32; + int32_t num_key_value_heads = 8; + int32_t head_dim = 64; + int32_t max_position_embeddings = 131072; + int32_t vocab_size = 128000; + int32_t conv_L_cache = 3; + bool conv_bias = false; + float norm_eps = 1.0e-5F; + float rope_theta = 10000000.0F; + std::string rope_type = "default"; + bool block_auto_adjust_ff_dim = false; + float block_ffn_dim_multiplier = 1.0F; + int32_t block_multiple_of = 256; + bool tie_word_embeddings = true; + int64_t bos_token_id = 124894; + int64_t eos_token_id = 124900; + int64_t pad_token_id = 124893; + int32_t max_cache_length = 2048; + std::vector layer_types; + aops::LinearImplTypes linear_impl_type = aops::LinearImplTypes::kDefault; + + [[nodiscard]] bool isAttentionLayer(int32_t layer_idx) const { + return layer_types.at(static_cast(layer_idx)) == "full_attention"; + } + [[nodiscard]] int32_t numAttentionLayers() const { + int32_t count = 0; + for (const auto& type : layer_types) count += type == "full_attention"; + return count; + } + [[nodiscard]] int32_t numConvLayers() const { return num_hidden_layers - numAttentionLayers(); } + [[nodiscard]] int32_t attentionSlotForPhysicalLayer(int32_t physical_layer) const { + if (!isAttentionLayer(physical_layer)) { throw std::invalid_argument("LFM2 physical layer is not attention"); } + int32_t slot = 0; + for (int32_t layer = 0; layer < physical_layer; ++layer) slot += isAttentionLayer(layer); + return slot; + } + + private: + void validate() const { + if (hidden_size <= 0 || intermediate_size <= 0 || num_hidden_layers <= 0 || num_attention_heads <= 0 + || num_key_value_heads <= 0 || head_dim <= 0 || max_position_embeddings <= 0 || vocab_size <= 0 || max_cache_length <= 0 + || conv_L_cache <= 1 || num_attention_heads % num_key_value_heads != 0 + || hidden_size != num_attention_heads * head_dim) { + throw std::invalid_argument("LFM2 contains invalid model dimensions"); + } + if (layer_types.size() != static_cast(num_hidden_layers)) { + throw std::invalid_argument("LFM2 layer_types size must equal num_hidden_layers"); + } + for (const auto& type : layer_types) { + if (type != "conv" && type != "full_attention") { + throw std::invalid_argument("LFM2 layer_types contains unsupported value: " + type); + } + } + if (!std::isfinite(norm_eps) || norm_eps <= 0.0F || !std::isfinite(rope_theta) || rope_theta <= 0.0F) { + throw std::invalid_argument("LFM2 norm_eps and rope_theta must be finite and positive"); + } + if (rope_type != "default") { throw std::invalid_argument("LFM2 CPU currently supports default RoPE only"); } + if (conv_bias) { throw std::invalid_argument("LFM2 CPU currently supports conv_bias=false only"); } + if (!tie_word_embeddings) { throw std::invalid_argument("LFM2 CPU requires tied embeddings"); } + if (block_auto_adjust_ff_dim) { throw std::invalid_argument("LFM2 CPU does not support adjusted FFN dimensions"); } + } +}; + +inline auto officialLayerTypes() -> const std::vector& { + static const std::vector schedule = { + "conv", "conv", "full_attention", "conv", "conv", "full_attention", + "conv", "conv", "conv", "full_attention", "conv", "conv", + "conv", "full_attention", "conv", "conv", "conv", "full_attention", + "conv", "conv", "conv", "full_attention", "conv", "conv", + "full_attention", "conv", "conv", "full_attention", "conv", "conv"}; + return schedule; +} + +inline auto matchesOfficialRuntimeContract(const Lfm2Config& cfg) -> bool { + constexpr auto kKai = aops::LinearImplTypes::kKaiLinear_f32_qai8dxp_qsi4c32p_mxk_nxk_qai8dxp1x8_qsi4c32p8x8_1x8x32; + return cfg.hidden_size == 2048 && cfg.intermediate_size == 10752 && cfg.num_hidden_layers == 30 + && cfg.num_attention_heads == 32 && cfg.num_key_value_heads == 8 && cfg.head_dim == 64 + && cfg.max_position_embeddings == 131072 && cfg.vocab_size == 128000 && cfg.conv_L_cache == 3 && !cfg.conv_bias + && cfg.norm_eps == 1.0e-5F && cfg.rope_theta == 10000000.0F && cfg.rope_type == "default" + && !cfg.block_auto_adjust_ff_dim && cfg.tie_word_embeddings && cfg.bos_token_id == 124894 && cfg.eos_token_id == 124900 + && cfg.pad_token_id == 124893 && cfg.max_cache_length == 2048 && cfg.layer_types == officialLayerTypes() + && cfg.numAttentionLayers() == 8 && cfg.linear_impl_type == kKai; +} + +inline void validateModelConfigMatch(const Lfm2Config& cfg, const ParameterFile::ptr_t& parameter_file) { + constexpr auto kEmbedding = "model.embed_tokens.weight"; + if (!matchesOfficialRuntimeContract(cfg)) { + throw std::invalid_argument("LFM2 model/config mismatch: CPU runner supports only official LFM2.5-2.6B W4A32"); + } + if (parameter_file == nullptr || !parameter_file->has(kEmbedding)) { + throw std::invalid_argument(std::string("LFM2 model/config mismatch: missing ") + kEmbedding); + } + const auto embedding = parameter_file->pull(kEmbedding); + if (embedding.dtype() != kFloat32) { + throw std::invalid_argument("LFM2 model/config mismatch: embedding must remain float32"); + } + const auto expected_numel = static_cast(cfg.vocab_size) * static_cast(cfg.hidden_size); + if (parameter_file->version() == ModelFileVersion::kV1) { + if (embedding.numel() != expected_numel) { throw std::invalid_argument("LFM2 model/config embedding element mismatch"); } + } else { + const auto shape = embedding.shape(); + if (shape.size() != 2 || shape[0] != cfg.vocab_size || shape[1] != cfg.hidden_size) { + throw std::invalid_argument("LFM2 model/config embedding shape mismatch"); + } + } + if (!parameter_file->has("lm_head_out.weight")) { + throw std::invalid_argument("LFM2 model/config mismatch: converted W4A32 model is missing tied lm_head_out.weight alias"); + } +} + +} // namespace mllm::models::lfm2 diff --git a/mllm/models/lfm2/modeling_lfm2.hpp b/mllm/models/lfm2/modeling_lfm2.hpp new file mode 100644 index 000000000..4c0f650e3 --- /dev/null +++ b/mllm/models/lfm2/modeling_lfm2.hpp @@ -0,0 +1,330 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include +#include +#include + +#include "mllm/core/Tensor.hpp" +#include "mllm/models/ARGeneration.hpp" +#include "mllm/models/common/rope_tables.hpp" +#include "mllm/models/lfm2/configuration_lfm2.hpp" +#include "mllm/nn/Functional.hpp" +#include "mllm/nn/Module.hpp" +#include "mllm/nn/Nn.hpp" +#include "mllm/nn/lmcache/KVHeadStaticCache.hpp" + +namespace mllm::models::lfm2 { + +// Rotation itself remains the registered nn::RoPE operation; only the +// immutable analytical tables are materialized here. +using common::makeRoPEInvFreq; +using common::makeRotaryPosEmbedding; + +inline constexpr int32_t kLfm2KaiDecodeThreadCap = 4; +inline constexpr int32_t kLfm2KaiPrefillThreadCap = 6; + +inline auto makeLfm2LinearOptions(int32_t in_channels, int32_t out_channels, bool bias, aops::LinearImplTypes impl_type) + -> aops::LinearOpOptions { + return {.in_channels = in_channels, + .out_channels = out_channels, + .bias = bias, + .impl_type = impl_type, + // Source-bound OnePlus 13T screening selected six workers for I8MM + // prefill and four workers to avoid decode GEMV oversubscription. + .kai_w4a32_decode_thread_cap = kLfm2KaiDecodeThreadCap, + .kai_w4a32_prefill_thread_cap = kLfm2KaiPrefillThreadCap}; +} + +class Lfm2MLP final : public nn::Module { + public: + Lfm2MLP() = default; + Lfm2MLP(const std::string& name, const Lfm2Config& cfg) : nn::Module(name) { + gate_up_proj_ = reg( + "gate_up_proj", cfg.hidden_size, std::vector{cfg.intermediate_size, cfg.intermediate_size}, + std::vector{"w1", "w3"}, false, cfg.linear_impl_type, kLfm2KaiDecodeThreadCap, kLfm2KaiPrefillThreadCap); + w2_ = reg("w2", makeLfm2LinearOptions(cfg.intermediate_size, cfg.hidden_size, false, cfg.linear_impl_type)); + silu_ = reg("silu"); + } + std::vector forward(const std::vector& inputs, const std::vector&) override { + auto gate_up = gate_up_proj_(inputs[0]); + return {w2_(silu_(gate_up[0]) * gate_up[1])}; + } + + private: + nn::ParallelLinear gate_up_proj_; + nn::Linear w2_; + nn::SiLU silu_; +}; + +class Lfm2Attention final : public nn::Module { + public: + Lfm2Attention() = default; + Lfm2Attention(const std::string& name, const Lfm2Config& cfg) : nn::Module(name) { + hidden_size_ = cfg.hidden_size; + head_dim_ = cfg.head_dim; + query_heads_ = cfg.num_attention_heads; + kv_heads_ = cfg.num_key_value_heads; + qkv_proj_ = reg( + "qkv_proj", hidden_size_, std::vector{query_heads_ * head_dim_, kv_heads_ * head_dim_, kv_heads_ * head_dim_}, + std::vector{"q_proj", "k_proj", "v_proj"}, false, cfg.linear_impl_type, kLfm2KaiDecodeThreadCap, + kLfm2KaiPrefillThreadCap); + out_proj_ = + reg("out_proj", makeLfm2LinearOptions(query_heads_ * head_dim_, hidden_size_, false, cfg.linear_impl_type)); + q_layernorm_ = reg("q_layernorm", cfg.norm_eps, false); + k_layernorm_ = reg("k_layernorm", cfg.norm_eps, false); + q_rope_ = reg("q_rope", cfg.rope_theta, cfg.max_position_embeddings, head_dim_); + k_rope_ = reg("k_rope", cfg.rope_theta, cfg.max_position_embeddings, head_dim_); + gqa_ = reg("gqa", aops::GroupedQueryAttentionImplementation::kDirectStrided); + } + + std::vector forward(const std::vector& inputs, const std::vector& args) override { + const auto& x = inputs[0]; + const int32_t batch = x.shape()[0]; + const int32_t sequence = x.shape()[1]; + auto qkv = qkv_proj_(x); + auto query = qkv[0].view({batch, sequence, query_heads_, head_dim_}); + auto key = qkv[1].view({batch, sequence, kv_heads_, head_dim_}); + auto value = qkv[2].view({batch, sequence, kv_heads_, head_dim_}); + query = q_layernorm_(query).transpose(1, 2); + key = k_layernorm_(key).transpose(1, 2); + value = value.transpose(1, 2); + query = q_rope_(query, inputs[1], inputs[2]); + key = k_rope_(key, inputs[1], inputs[2]); + + auto* cache = args.at(0).get(); + auto updated = cache->updateKVCache(logical_cache_slot_, key, value); + // LFM2 keeps the established direct accumulation order for both prefill + // and decode. MiniCPM5's dedicated decode kernel is faster, but the + // OnePlus exact-token gate showed that its different reduction order + // changes LFM2 generation beginning at token 12. + auto output = gqa_(query, updated[0], updated[1]); + output = output.transpose(1, 2).contiguous().view({batch, sequence, query_heads_ * head_dim_}); + return {out_proj_(output)}; + } + + int32_t logical_cache_slot_ = 0; + + private: + int32_t hidden_size_ = 0; + int32_t head_dim_ = 0; + int32_t query_heads_ = 0; + int32_t kv_heads_ = 0; + nn::ParallelLinear qkv_proj_; + nn::Linear out_proj_; + nn::RMSNorm q_layernorm_; + nn::RMSNorm k_layernorm_; + nn::RoPE q_rope_; + nn::RoPE k_rope_; + nn::GroupedQueryAttention gqa_; +}; + +class Lfm2ShortConv final : public nn::Module { + public: + Lfm2ShortConv() = default; + Lfm2ShortConv(const std::string& name, const Lfm2Config& cfg) : nn::Module(name) { + hidden_size_ = cfg.hidden_size; + kernel_size_ = cfg.conv_L_cache; + in_proj_ = + reg("in_proj", makeLfm2LinearOptions(hidden_size_, 3 * hidden_size_, cfg.conv_bias, cfg.linear_impl_type)); + conv_ = reg("conv", hidden_size_, kernel_size_, cfg.conv_bias, true, + aops::CausalDepthwiseConv1DAccumulationOrder::kHistoryFirst); + out_proj_ = + reg("out_proj", makeLfm2LinearOptions(hidden_size_, hidden_size_, cfg.conv_bias, cfg.linear_impl_type)); + } + + // The causal kernel consumes only K - 1 historical samples. The former + // generic concat/Conv1D path stored K samples but never read the oldest one. + void resetState(int32_t batch_size) { state_ = Tensor::zeros({batch_size, hidden_size_, kernel_size_ - 1}, kFloat32, kCPU); } + + [[nodiscard]] const Tensor& state() const { return state_; } + + std::vector forward(const std::vector& inputs, const std::vector&) override { + const auto& input = inputs[0]; + const int32_t batch = input.shape()[0]; + const int32_t sequence = input.shape()[1]; + if (input.dtype() != kFloat32 || input.device() != kCPU) { + throw std::invalid_argument("LFM2 short convolution currently requires float32 CPU activations"); + } + if (state_.isNil() || state_.shape()[0] != batch) resetState(batch); + + auto projected = in_proj_(input); + auto b = projected[{kAll, kAll, {0, hidden_size_}}].contiguous(); + auto c = projected[{kAll, kAll, {hidden_size_, 2 * hidden_size_}}].contiguous(); + auto x = projected[{kAll, kAll, {2 * hidden_size_, 3 * hidden_size_}}].contiguous(); + auto bx = b * x; + auto [convolved, updated_state] = conv_(bx, state_); + state_ = std::move(updated_state); + return {out_proj_(c * convolved)}; + } + + private: + int32_t hidden_size_ = 0; + int32_t kernel_size_ = 0; + nn::Linear in_proj_; + nn::CausalDepthwiseConv1D conv_; + nn::Linear out_proj_; + Tensor state_; +}; + +class Lfm2AttentionDecoder final : public nn::Module { + public: + Lfm2AttentionDecoder() = default; + Lfm2AttentionDecoder(const std::string& name, const Lfm2Config& cfg) : nn::Module(name) { + self_attn_ = reg("self_attn", cfg); + feed_forward_ = reg("feed_forward", cfg); + operator_norm_ = reg("operator_norm", cfg.norm_eps, false); + ffn_norm_ = reg("ffn_norm", cfg.norm_eps, false); + } + std::vector forward(const std::vector& inputs, const std::vector& args) override { + auto residual = inputs[0]; + auto hidden = residual + self_attn_(operator_norm_(residual), inputs[1], inputs[2], args.at(0))[0]; + return {hidden + feed_forward_(ffn_norm_(hidden))[0]}; + } + Lfm2Attention self_attn_; + + private: + Lfm2MLP feed_forward_; + nn::RMSNorm operator_norm_; + nn::RMSNorm ffn_norm_; +}; + +class Lfm2ConvDecoder final : public nn::Module { + public: + Lfm2ConvDecoder() = default; + Lfm2ConvDecoder(const std::string& name, const Lfm2Config& cfg) : nn::Module(name) { + conv_ = reg("conv", cfg); + feed_forward_ = reg("feed_forward", cfg); + operator_norm_ = reg("operator_norm", cfg.norm_eps, false); + ffn_norm_ = reg("ffn_norm", cfg.norm_eps, false); + } + std::vector forward(const std::vector& inputs, const std::vector&) override { + auto residual = inputs[0]; + auto hidden = residual + conv_(operator_norm_(residual))[0]; + return {hidden + feed_forward_(ffn_norm_(hidden))[0]}; + } + Lfm2ShortConv conv_; + + private: + Lfm2MLP feed_forward_; + nn::RMSNorm operator_norm_; + nn::RMSNorm ffn_norm_; +}; + +class Lfm2Model final : public nn::Module { + public: + Lfm2Model() = default; + Lfm2Model(const std::string& name, const Lfm2Config& cfg) : nn::Module(name) { + embed_tokens_ = reg("embed_tokens", cfg.vocab_size, cfg.hidden_size); + int32_t attention_index = 0; + int32_t conv_index = 0; + for (int32_t physical_layer = 0; physical_layer < cfg.num_hidden_layers; ++physical_layer) { + const auto layer_name = "layers." + std::to_string(physical_layer); + if (cfg.isAttentionLayer(physical_layer)) { + auto layer = reg(layer_name, cfg); + layer.self_attn_.logical_cache_slot_ = attention_index; + attention_layers_.push_back(std::move(layer)); + layer_kind_.push_back(0); + layer_dispatch_.push_back(attention_index++); + } else { + conv_layers_.push_back(reg(layer_name, cfg)); + layer_kind_.push_back(1); + layer_dispatch_.push_back(conv_index++); + } + } + embedding_norm_ = reg("embedding_norm", cfg.norm_eps, false); + } + + void resetConvStates(int32_t batch_size) { + for (auto& layer : conv_layers_) layer.conv_.resetState(batch_size); + } + + std::vector forward(const std::vector& inputs, const std::vector& args) override { + auto hidden = embed_tokens_(inputs[0]); + for (size_t layer = 0; layer < layer_kind_.size(); ++layer) { + if (layer_kind_[layer] == 0) { + hidden = attention_layers_[layer_dispatch_[layer]](hidden, inputs[1], inputs[2], args.at(0))[0]; + } else { + hidden = conv_layers_[layer_dispatch_[layer]](hidden)[0]; + } + } + return {embedding_norm_(hidden)}; + } + + private: + nn::Embedding embed_tokens_; + nn::RMSNorm embedding_norm_; + std::vector attention_layers_; + std::vector conv_layers_; + std::vector layer_kind_; + std::vector layer_dispatch_; +}; + +class Lfm2ForCausalLM final : public ARGeneration, public nn::Module { + public: + explicit Lfm2ForCausalLM(const Lfm2Config& cfg) + : kv_cache_(cfg.max_cache_length, cfg.numAttentionLayers(), cfg.num_key_value_heads, cfg.head_dim) { + eos_token_id_ = cfg.eos_token_id; + max_length_ = cfg.max_cache_length; + model_ = reg("model", cfg); + // The converter creates this packed alias from the tied embedding. Keeping + // it explicit lets the mobile W4A32 output projection use the standard Linear path. + lm_head_ = + reg("lm_head_out", makeLfm2LinearOptions(cfg.hidden_size, cfg.vocab_size, false, cfg.linear_impl_type)); + registerBuffer("inv_freq", makeRoPEInvFreq(cfg.head_dim, cfg.rope_theta)); + } + + ARGenerationOutputPast forward(const ARGenerationOutputPast& input, const ARGenerationArgs&) override { + auto sequence = input.at("sequence"); + const auto shape = sequence.shape(); + if (shape.size() != 2 || shape[0] != 1 || shape[1] <= 0 || sequence.dtype() != kInt64 || sequence.device() != kCPU) { + throw std::invalid_argument("LFM2 CPU expects a non-empty rank-2 int64 CPU sequence with batch size 1"); + } + const int32_t sequence_length = shape[1]; + const int32_t cached_tokens = kv_cache_.getCurrentSeqCnt(0); + if (sequence_length > max_length_ || cached_tokens > max_length_ - sequence_length) { + throw std::invalid_argument("LFM2 sequence exceeds the configured cache capacity"); + } + + // Generation orchestration: position IDs and analytical sin/cos buffers + // are inputs to the registered nn::RoPE operations, not backend kernels. + Tensor position_ids; + if (input.count("position_ids")) { + position_ids = input.at("position_ids"); + if (sequence_length == 1) { + const auto previous = *position_ids.offsettedPtr({0, position_ids.shape()[1] - 1}); + position_ids = Tensor::empty({1, 1}, kInt64, kCPU).alloc(); + *position_ids.ptr() = previous + 1; + } + } else { + if (cached_tokens != 0) { + throw std::invalid_argument( + "LFM2 continuation with a non-empty KV cache requires the position_ids returned by the previous step"); + } + position_ids = Tensor::empty({1, sequence_length}, kInt64, kCPU).alloc(); + for (int32_t index = 0; index < sequence_length; ++index) position_ids.ptr()[index] = index; + } + auto [sin, cos] = makeRotaryPosEmbedding(position_ids, getBuffer("inv_freq")); + sequence = model_(sequence, sin, cos, AnyValue(&kv_cache_))[0]; + sequence = sequence[{kAll, {sequence.shape()[1] - 1}, kAll}]; + sequence = lm_head_(sequence); + return {{"sequence", sequence}, {"position_ids", position_ids}}; + } + + void resetState(int32_t batch_size = 1) { + if (batch_size != 1) { throw std::invalid_argument("LFM2 CPU currently supports batch size 1 only"); } + kv_cache_.clearCache(); + model_.resetConvStates(batch_size); + } + [[nodiscard]] nn::KVHeadStaticCache& kvCache() { return kv_cache_; } + + private: + Lfm2Model model_; + nn::Linear lm_head_; + nn::KVHeadStaticCache kv_cache_; +}; + +} // namespace mllm::models::lfm2 diff --git a/mllm/models/lfm2/tokenization_lfm2.hpp b/mllm/models/lfm2/tokenization_lfm2.hpp new file mode 100644 index 000000000..2a8f52090 --- /dev/null +++ b/mllm/models/lfm2/tokenization_lfm2.hpp @@ -0,0 +1,235 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. +#pragma once + +#include +#include +#include +#include +#include +#include + +#include "mllm/models/ARGeneration.hpp" +#include "mllm/preprocessor/StreamingUtf8Decoder.hpp" +#include "mllm/preprocessor/tokenizers/AutoTokenizer.hpp" +#include "mllm/preprocessor/tokenizers/BPE.hpp" +#include "mllm/preprocessor/tokenizers/Unicode.hpp" + +namespace mllm::models::lfm2 { + +using StreamingUtf8Decoder = preprocessor::StreamingUtf8Decoder; + +// Implements the checkpoint regex. The material difference from the older +// Qwen pattern is the numeric branch: LFM2 groups one to three digits. +inline bool tokenizerMatch(const std::wstring& input, size_t& pos, std::wstring& matched) { + if (pos >= input.size()) return false; + static const std::wstring contractions[] = {L"'s", L"'t", L"'re", L"'ve", L"'m", L"'ll", L"'d"}; + for (const auto& contraction : contractions) { + bool match = pos + contraction.size() <= input.size(); + for (size_t index = 0; match && index < contraction.size(); ++index) { + match = std::towlower(input[pos + index]) == contraction[index]; + } + if (match) { + matched = input.substr(pos, contraction.size()); + pos += contraction.size(); + return true; + } + } + + { + const auto original = pos; + const auto start = pos; + if (!preprocessor::isLetter(input[pos]) && !preprocessor::isDigit(input[pos]) && input[pos] != L'\r' + && input[pos] != L'\n') { + ++pos; + } + if (pos < input.size() && preprocessor::isLetter(input[pos])) { + while (pos < input.size() && preprocessor::isLetter(input[pos])) ++pos; + matched = input.substr(start, pos - start); + return true; + } + pos = original; + } + + if (preprocessor::isDigit(input[pos])) { + const auto start = pos; + while (pos < input.size() && pos - start < 3 && preprocessor::isDigit(input[pos])) ++pos; + matched = input.substr(start, pos - start); + return true; + } + + { + const auto original = pos; + const auto start = pos; + if (input[pos] == L' ') ++pos; + if (pos < input.size() && !std::iswspace(input[pos]) && !preprocessor::isLetter(input[pos]) + && !preprocessor::isDigit(input[pos])) { + while (pos < input.size() && !std::iswspace(input[pos]) && !preprocessor::isLetter(input[pos]) + && !preprocessor::isDigit(input[pos])) { + ++pos; + } + while (pos < input.size() && (input[pos] == L'\r' || input[pos] == L'\n')) ++pos; + matched = input.substr(start, pos - start); + return true; + } + pos = original; + } + + { + const auto start = pos; + auto scan = pos; + size_t last_break = std::wstring::npos; + while (scan < input.size() && std::iswspace(input[scan])) { + if (input[scan] == L'\r' || input[scan] == L'\n') last_break = scan + 1; + ++scan; + } + if (last_break != std::wstring::npos) { + pos = last_break; + matched = input.substr(start, pos - start); + return true; + } + } + + if (std::iswspace(input[pos])) { + const auto start = pos; + while (pos < input.size() && std::iswspace(input[pos])) ++pos; + if (pos >= input.size()) { + matched = input.substr(start, pos - start); + return true; + } + if (pos - start > 1) { + --pos; + matched = input.substr(start, pos - start); + return true; + } + pos = start; + } + if (std::iswspace(input[pos])) { + const auto start = pos; + while (pos < input.size() && std::iswspace(input[pos])) ++pos; + matched = input.substr(start, pos - start); + return true; + } + return false; +} + +inline bool tokenizerRegex(const std::string& input, std::vector& pieces) { + const auto wide = preprocessor::utf8string2WideString(input); + size_t pos = 0; + while (pos < wide.size()) { + std::wstring matched; + if (tokenizerMatch(wide, pos, matched)) { + pieces.push_back(matched); + } else { + pieces.push_back(wide.substr(pos++, 1)); + } + } + return true; +} + +struct Lfm2Message { + std::string prompt; + std::string system_prompt; + // Each item is a raw JSON tool schema. The pinned template accepts string + // tools verbatim, which keeps key order and JSON bytes caller-controlled. + std::vector tools; + + static auto render(const Lfm2Message& message) -> std::string { + std::string rendered = "<|startoftext|>"; + if (!message.system_prompt.empty() || !message.tools.empty()) { + rendered += "<|im_start|>system\n"; + rendered += message.system_prompt; + if (!message.tools.empty()) { + if (!message.system_prompt.empty()) rendered += '\n'; + rendered += "List of tools: ["; + for (size_t index = 0; index < message.tools.size(); ++index) { + if (index != 0) rendered += ", "; + rendered += message.tools[index]; + } + rendered += ']'; + } + rendered += "<|im_end|>\n"; + } + rendered += "<|im_start|>user\n" + message.prompt + "<|im_end|>\n<|im_start|>assistant\n"; + return rendered; + } +}; + +class Lfm2Tokenizer final : public preprocessor::AutoTokenizer { + public: + explicit Lfm2Tokenizer(const std::string& tokenizer_json) { + preprocessor::initLocal(); + preprocessor::makeBytes2UnicodeMap(bytes_to_unicode_); + for (const auto& [byte, codepoint] : bytes_to_unicode_) unicode_to_bytes_.insert({codepoint, byte}); + bpe_.initFromSentencePieceJson(tokenizer_json); + for (const auto* token : {L"<|pad|>", L"<|startoftext|>", L"<|endoftext|>", L"<|fim_pre|>", L"<|fim_mid|>", L"<|fim_suf|>", + L"<|im_start|>", L"<|im_end|>", L"", L"", L"<|tool_list_start|>", + L"<|tool_list_end|>", L"<|tool_call_start|>", L"<|tool_call_end|>"}) { + special_tokens_trie_.add(token); + } + } + + std::vector _tokenize(const std::string& input) override { + std::vector pieces; + tokenizerRegex(input, pieces); + std::vector tokens; + for (const auto& piece : pieces) { + std::wstring mapped; + for (const auto byte : preprocessor::wideString2Utf8String(piece)) { + mapped.push_back(bytes_to_unicode_.at(static_cast(byte))); + } + const auto bpe_tokens = bpe_._bpe(mapped); + tokens.insert(tokens.end(), bpe_tokens.begin(), bpe_tokens.end()); + } + return tokens; + } + + std::vector tokenize(const std::string& input) override { + const auto pieces = special_tokens_trie_.split(preprocessor::utf8string2WideString(input)); + std::vector tokens; + for (const auto& piece : pieces) { + if (special_tokens_trie_.isSpecialToken(piece)) { + tokens.push_back(piece); + } else { + const auto normal = _tokenize(preprocessor::wideString2Utf8String(piece)); + tokens.insert(tokens.end(), normal.begin(), normal.end()); + } + } + return tokens; + } + + std::wstring _detokenize(int64_t id) override { return bpe_._lookup_inverse_vocab(id); } + std::string detokenizeBytes(int64_t id) { + const auto token = _detokenize(id); + std::string bytes; + for (const auto codepoint : token) { + const auto found = unicode_to_bytes_.find(codepoint); + if (found == unicode_to_bytes_.end()) throw std::runtime_error("LFM2 tokenizer encountered unknown byte symbol"); + bytes.push_back(static_cast(found->second)); + } + return bytes; + } + std::wstring detokenize(int64_t id) override { return preprocessor::utf8string2WideString(detokenizeBytes(id)); } + + Tensor convert2Ids(const std::vector& tokens) override { return idsTensor(tokens, kExtraInput); } + ARGenerationOutputPast convertMessage(const Lfm2Message& message) { + const auto rendered = Lfm2Message::render(message); + return {{"sequence", idsTensor(tokenize(rendered), kNormal)}}; + } + + private: + Tensor idsTensor(const std::vector& tokens, TensorMemTypes mem_type) { + auto result = Tensor::empty({1, static_cast(tokens.size())}, kInt64, kCPU) + .setMemType(mem_type) + .setName("lfm2-tokenizer-i0") + .alloc(); + auto* ids = result.ptr(); + for (size_t index = 0; index < tokens.size(); ++index) ids[index] = bpe_._lookup_vocab(tokens[index]); + return result; + } + preprocessor::BPE bpe_; + std::unordered_map bytes_to_unicode_; + std::unordered_map unicode_to_bytes_; +}; + +} // namespace mllm::models::lfm2 diff --git a/mllm/models/minicpm5/modeling_minicpm5.hpp b/mllm/models/minicpm5/modeling_minicpm5.hpp index a405ebb32..bf976c4be 100644 --- a/mllm/models/minicpm5/modeling_minicpm5.hpp +++ b/mllm/models/minicpm5/modeling_minicpm5.hpp @@ -10,6 +10,7 @@ #include "mllm/mllm.hpp" #include "mllm/models/ARGeneration.hpp" +#include "mllm/models/common/rope_tables.hpp" #include "mllm/models/minicpm5/configuration_minicpm5.hpp" #include "mllm/nn/Functional.hpp" #include "mllm/nn/Module.hpp" @@ -20,39 +21,10 @@ namespace mllm::models::minicpm5 { -inline auto makeMiniCPM5RoPEInvFreq(int32_t output_dim, float rope_theta) -> Tensor { - auto inv_freq = Tensor::empty({output_dim / 2}, kFloat32, kCPU).alloc(); - for (int32_t dim = 0; dim < output_dim / 2; ++dim) { - inv_freq.ptr()[dim] = 1.0F / std::pow(rope_theta, 2.0F * static_cast(dim) / output_dim); - } - return inv_freq; -} - -inline auto makeMiniCPM5RotaryPosEmbedding(const Tensor& position_ids, const Tensor& inv_freq) -> std::pair { - const int32_t batch = position_ids.shape()[0]; - const int32_t sequence = position_ids.shape()[1]; - const int32_t half_dim = inv_freq.shape()[0]; - const int32_t dim = half_dim * 2; - auto sin_embedding = Tensor::empty({batch, sequence, dim}, kFloat32, kCPU).alloc(); - auto cos_embedding = Tensor::empty({batch, sequence, dim}, kFloat32, kCPU).alloc(); - - for (int32_t batch_index = 0; batch_index < batch; ++batch_index) { - for (int32_t sequence_index = 0; sequence_index < sequence; ++sequence_index) { - const auto position = position_ids.ptr()[batch_index * sequence + sequence_index]; - for (int32_t index = 0; index < half_dim; ++index) { - const float frequency = static_cast(position) * inv_freq.ptr()[index]; - const float sine = std::sin(frequency); - const float cosine = std::cos(frequency); - const auto offset = (batch_index * sequence + sequence_index) * dim + index; - sin_embedding.ptr()[offset] = sine; - sin_embedding.ptr()[offset + half_dim] = sine; - cos_embedding.ptr()[offset] = cosine; - cos_embedding.ptr()[offset + half_dim] = cosine; - } - } - } - return {sin_embedding, cos_embedding}; -} +// Rotation itself remains the registered nn::RoPE operation; only the +// immutable analytical tables are materialized here. +using common::makeRoPEInvFreq; +using common::makeRotaryPosEmbedding; class MiniCPM5MLP final : public nn::Module { public: @@ -91,7 +63,7 @@ class MiniCPM5Attention final : public nn::Module { o_proj_ = reg("o_proj", query_heads_ * head_dim_, hidden_size_, config.attention_bias, config.linear_impl_type); q_rope_ = reg("q_rope", config.rope_theta, config.max_position_embeddings, config.head_dim); k_rope_ = reg("k_rope", config.rope_theta, config.max_position_embeddings, config.head_dim); - gqa_decode_ = reg("gqa_decode"); + gqa_decode_ = reg("gqa_decode", aops::GroupedQueryAttentionImplementation::kDecodeNativeKV); } std::vector forward(const std::vector& inputs, const std::vector& args) override { @@ -123,7 +95,7 @@ class MiniCPM5Attention final : public nn::Module { nn::Linear o_proj_; nn::RoPE q_rope_; nn::RoPE k_rope_; - nn::GroupedQueryAttentionDecode gqa_decode_; + nn::GroupedQueryAttention gqa_decode_; int32_t hidden_size_ = 0; int32_t head_dim_ = 0; int32_t query_heads_ = 0; @@ -188,7 +160,7 @@ class MiniCPM5ForCausalLM final : public ARGeneration, public nn::Module { max_length_ = config.max_cache_length; model_ = reg("model", config); lm_head_ = reg("lm_head", config.hidden_size, config.vocab_size, false, config.linear_impl_type); - registerBuffer("inv_freq", makeMiniCPM5RoPEInvFreq(config.head_dim, config.rope_theta)); + registerBuffer("inv_freq", makeRoPEInvFreq(config.head_dim, config.rope_theta)); } ARGenerationOutputPast forward(const ARGenerationOutputPast& input, const ARGenerationArgs& args) override { @@ -211,7 +183,7 @@ class MiniCPM5ForCausalLM final : public ARGeneration, public nn::Module { auto position_ids = Tensor::empty({1, shape[1]}, kInt64, kCPU).alloc(); for (int32_t index = 0; index < shape[1]; ++index) { position_ids.ptr()[index] = cached_tokens + index; } - auto [sine, cosine] = makeMiniCPM5RotaryPosEmbedding(position_ids, getBuffer("inv_freq")); + auto [sine, cosine] = makeRotaryPosEmbedding(position_ids, getBuffer("inv_freq")); auto hidden = model_(sequence, sine, cosine, AnyValue(&kv_cache_))[0]; hidden = hidden[{kAll, {hidden.shape()[1] - 1}, kAll}]; diff --git a/mllm/models/minicpm5/tokenization_minicpm5.hpp b/mllm/models/minicpm5/tokenization_minicpm5.hpp index 89f189556..4be331fd5 100644 --- a/mllm/models/minicpm5/tokenization_minicpm5.hpp +++ b/mllm/models/minicpm5/tokenization_minicpm5.hpp @@ -14,86 +14,14 @@ #include #include "mllm/models/ARGeneration.hpp" +#include "mllm/preprocessor/StreamingUtf8Decoder.hpp" #include "mllm/preprocessor/tokenizers/AutoTokenizer.hpp" #include "mllm/preprocessor/tokenizers/BPE.hpp" #include "mllm/preprocessor/tokenizers/Unicode.hpp" namespace mllm::models::minicpm5 { -class MiniCPM5StreamingUtf8Decoder { - public: - std::string append(std::string_view bytes) { - pending_.append(bytes.data(), bytes.size()); - return drain(false); - } - std::string finish() { return drain(true); } - void reset() { pending_.clear(); } - - private: - static constexpr std::string_view kReplacementCharacter = "\xEF\xBF\xBD"; - - static bool isContinuationByte(unsigned char byte) { return byte >= 0x80 && byte <= 0xBF; } - static size_t sequenceLength(unsigned char lead) { - if (lead <= 0x7F) return 1; - if (lead >= 0xC2 && lead <= 0xDF) return 2; - if (lead >= 0xE0 && lead <= 0xEF) return 3; - if (lead >= 0xF0 && lead <= 0xF4) return 4; - return 0; - } - static bool isValidSecondByte(unsigned char lead, unsigned char second) { - if (!isContinuationByte(second)) return false; - if (lead == 0xE0) return second >= 0xA0; - if (lead == 0xED) return second <= 0x9F; - if (lead == 0xF0) return second >= 0x90; - if (lead == 0xF4) return second <= 0x8F; - return true; - } - - std::string drain(bool flush) { - std::string output; - size_t offset = 0; - while (offset < pending_.size()) { - const auto lead = static_cast(pending_[offset]); - const size_t length = sequenceLength(lead); - if (length == 1) { - output.push_back(pending_[offset++]); - continue; - } - if (length == 0) { - output.append(kReplacementCharacter); - ++offset; - continue; - } - const size_t available = pending_.size() - offset; - bool valid_prefix = true; - for (size_t index = 1; index < std::min(available, length); ++index) { - const auto byte = static_cast(pending_[offset + index]); - if ((index == 1 && !isValidSecondByte(lead, byte)) || (index > 1 && !isContinuationByte(byte))) { - valid_prefix = false; - break; - } - } - if (!valid_prefix) { - output.append(kReplacementCharacter); - ++offset; - continue; - } - if (available < length) { - if (flush) { - output.append(kReplacementCharacter); - offset = pending_.size(); - } - break; - } - output.append(pending_, offset, length); - offset += length; - } - pending_.erase(0, offset); - return output; - } - - std::string pending_; -}; +using MiniCPM5StreamingUtf8Decoder = preprocessor::StreamingUtf8Decoder; // Equivalent to MiniCPM5's two Split pre-tokenizers: numeric runs are first // isolated into chunks of at most three digits, then the Qwen-style pattern is diff --git a/mllm/models/qwen3_5/tokenization_qwen3_5.hpp b/mllm/models/qwen3_5/tokenization_qwen3_5.hpp index e041fbbb9..2fd7d532a 100644 --- a/mllm/models/qwen3_5/tokenization_qwen3_5.hpp +++ b/mllm/models/qwen3_5/tokenization_qwen3_5.hpp @@ -11,6 +11,7 @@ #include #include "mllm/preprocessor/tokenizers/BPE.hpp" +#include "mllm/preprocessor/StreamingUtf8Decoder.hpp" #include "mllm/models/ARGeneration.hpp" #include "mllm/models/qwen3_5/image_preprocessor_qwen3_5.hpp" #include "mllm/models/qwen3_5/multimodal_qwen3_5.hpp" @@ -20,89 +21,7 @@ namespace mllm::models::qwen3_5 { -class Qwen3_5StreamingUtf8Decoder { - public: - std::string append(std::string_view bytes) { - if (!bytes.empty()) { pending_.append(bytes.data(), bytes.size()); } - return drain(false); - } - - std::string finish() { return drain(true); } - - void reset() { pending_.clear(); } - - private: - static constexpr std::string_view kReplacementCharacter = "\xEF\xBF\xBD"; - - static bool isContinuationByte(unsigned char byte) { return byte >= 0x80 && byte <= 0xBF; } - - static size_t sequenceLength(unsigned char lead) { - if (lead <= 0x7F) return 1; - if (lead >= 0xC2 && lead <= 0xDF) return 2; - if (lead >= 0xE0 && lead <= 0xEF) return 3; - if (lead >= 0xF0 && lead <= 0xF4) return 4; - return 0; - } - - static bool isValidSecondByte(unsigned char lead, unsigned char second) { - if (!isContinuationByte(second)) return false; - if (lead == 0xE0) return second >= 0xA0; - if (lead == 0xED) return second <= 0x9F; - if (lead == 0xF0) return second >= 0x90; - if (lead == 0xF4) return second <= 0x8F; - return true; - } - - std::string drain(bool flush) { - std::string output; - size_t offset = 0; - - while (offset < pending_.size()) { - const auto lead = static_cast(pending_[offset]); - const size_t sequence_length = sequenceLength(lead); - if (sequence_length == 1) { - output.push_back(pending_[offset++]); - continue; - } - if (sequence_length == 0) { - output.append(kReplacementCharacter); - ++offset; - continue; - } - - const size_t available = pending_.size() - offset; - const size_t prefix_length = std::min(available, sequence_length); - bool valid_prefix = true; - for (size_t index = 1; index < prefix_length; ++index) { - const auto byte = static_cast(pending_[offset + index]); - if ((index == 1 && !isValidSecondByte(lead, byte)) || (index > 1 && !isContinuationByte(byte))) { - valid_prefix = false; - break; - } - } - if (!valid_prefix) { - output.append(kReplacementCharacter); - ++offset; - continue; - } - if (available < sequence_length) { - if (flush) { - output.append(kReplacementCharacter); - offset = pending_.size(); - } - break; - } - - output.append(pending_, offset, sequence_length); - offset += sequence_length; - } - - pending_.erase(0, offset); - return output; - } - - std::string pending_; -}; +using Qwen3_5StreamingUtf8Decoder = preprocessor::StreamingUtf8Decoder; // Reuse the Qwen3 regex pattern — same BPE tokenization scheme. // (?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| diff --git a/mllm/nn/Functional.cpp b/mllm/nn/Functional.cpp index 856b4ec38..7ad399c67 100644 --- a/mllm/nn/Functional.cpp +++ b/mllm/nn/Functional.cpp @@ -6,7 +6,7 @@ #include "mllm/core/aops/ElewiseOps.hpp" #include "mllm/core/aops/FlashAttention2Op.hpp" #include "mllm/core/aops/GatherOp.hpp" -#include "mllm/core/aops/GroupedQueryAttentionDecodeOp.hpp" +#include "mllm/core/aops/GroupedQueryAttentionOp.hpp" #include "mllm/core/aops/MatMulOp.hpp" #include "mllm/core/aops/LinearOp.hpp" #include "mllm/core/aops/ReduceOps.hpp" @@ -89,8 +89,14 @@ Tensor flashAttention2(const Tensor& Q, const Tensor& K, const Tensor& V) { } Tensor groupedQueryAttentionDecode(const Tensor& query, const Tensor& key, const Tensor& value) { - return Context::instance().buildOpAndSubmitTask(OpTypes::kGroupedQueryAttentionDecode, - aops::GroupedQueryAttentionDecodeOpOptions{}, {query, key, value})[0]; + return groupedQueryAttention(query, key, value, aops::GroupedQueryAttentionImplementation::kDecodeNativeKV); +} + +Tensor groupedQueryAttention(const Tensor& query, const Tensor& key, const Tensor& value, + aops::GroupedQueryAttentionImplementation implementation) { + return Context::instance().buildOpAndSubmitTask(OpTypes::kGroupedQueryAttention, + aops::GroupedQueryAttentionOpOptions{.implementation = implementation}, + {query, key, value})[0]; } Tensor softmax(const Tensor& x, int32_t dim) { diff --git a/mllm/nn/Functional.hpp b/mllm/nn/Functional.hpp index 691bddd26..c3bccfad1 100644 --- a/mllm/nn/Functional.hpp +++ b/mllm/nn/Functional.hpp @@ -12,6 +12,7 @@ #include "mllm/core/aops/SplitOp.hpp" #include "mllm/core/aops/PadOp.hpp" #include "mllm/core/aops/InterpolateOp.hpp" +#include "mllm/core/aops/GroupedQueryAttentionOp.hpp" #include "mllm/core/aops/RadixAttnWithSinkAndSwaDiffDimOp.hpp" #include "mllm/engine/Context.hpp" @@ -112,6 +113,10 @@ Tensor flashAttention2(const Tensor& Q, const Tensor& K, const Tensor& V); Tensor groupedQueryAttentionDecode(const Tensor& query, const Tensor& key, const Tensor& value); +Tensor groupedQueryAttention( + const Tensor& query, const Tensor& key, const Tensor& value, + aops::GroupedQueryAttentionImplementation implementation = aops::GroupedQueryAttentionImplementation::kDirectStrided); + Tensor softmax(const Tensor& x, int32_t dim); Tensor log(const Tensor& x); diff --git a/mllm/nn/Layer.hpp b/mllm/nn/Layer.hpp index 6459a4933..bfccf4360 100644 --- a/mllm/nn/Layer.hpp +++ b/mllm/nn/Layer.hpp @@ -93,6 +93,13 @@ class Layer { return {outs[0], outs[1], outs[2]}; \ } +#define MLLM_LAYER_ANY_INPUTS_ANY_OUTPUTS_FORWARD \ + template \ + std::vector operator()(Args&&... args) { \ + auto inputs = std::vector{std::forward(args)...}; \ + return __main(inputs); \ + } + #define MLLM_LAYER_ENABLE_INPLACE_ATTRIBUTE(__CXX_CLASS_NAME__) \ inline __CXX_CLASS_NAME__& inplace() { \ auto& opts = const_cast<::mllm::aops::__CXX_CLASS_NAME__##OpOptions&>( \ diff --git a/mllm/nn/Nn.hpp b/mllm/nn/Nn.hpp index f47481c01..a1492d9ae 100644 --- a/mllm/nn/Nn.hpp +++ b/mllm/nn/Nn.hpp @@ -13,7 +13,8 @@ #include "mllm/nn/layers/Sigmoid.hpp" // IWYU pragma: export #include "mllm/nn/layers/Embedding.hpp" // IWYU pragma: export #include "mllm/nn/layers/GELU.hpp" // IWYU pragma: export -#include "mllm/nn/layers/GroupedQueryAttentionDecode.hpp" // IWYU pragma: export +#include "mllm/nn/layers/CausalDepthwiseConv1D.hpp" // IWYU pragma: export +#include "mllm/nn/layers/GroupedQueryAttention.hpp" // IWYU pragma: export #include "mllm/nn/layers/QuickGELU.hpp" // IWYU pragma: export #include "mllm/nn/layers/ReLU.hpp" // IWYU pragma: export #include "mllm/nn/layers/LayerNorm.hpp" // IWYU pragma: export @@ -25,6 +26,7 @@ #include "mllm/nn/layers/RoPE.hpp" // IWYU pragma: export #include "mllm/nn/layers/MultimodalRoPE.hpp" // IWYU pragma: export #include "mllm/nn/layers/Param.hpp" // IWYU pragma: export +#include "mllm/nn/layers/ParallelLinear.hpp" // IWYU pragma: export #include "mllm/nn/layers/KVCache.hpp" // IWYU pragma: export #include "mllm/nn/layers/Conv1D.hpp" // IWYU pragma: export #include "mllm/nn/layers/AvgPool1d.hpp" // IWYU pragma: export diff --git a/mllm/nn/layers/CausalDepthwiseConv1D.cpp b/mllm/nn/layers/CausalDepthwiseConv1D.cpp new file mode 100644 index 000000000..461dd00ac --- /dev/null +++ b/mllm/nn/layers/CausalDepthwiseConv1D.cpp @@ -0,0 +1,27 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/nn/layers/CausalDepthwiseConv1D.hpp" + +namespace mllm::nn { + +CausalDepthwiseConv1D::CausalDepthwiseConv1D() + : Layer(OpTypes::kCausalDepthwiseConv1D, aops::CausalDepthwiseConv1DOpOptions{}) {} + +CausalDepthwiseConv1D::CausalDepthwiseConv1D(int32_t channels, int32_t kernel_size, bool bias, bool state_inplace, + aops::CausalDepthwiseConv1DAccumulationOrder accumulation_order) + : Layer(OpTypes::kCausalDepthwiseConv1D, aops::CausalDepthwiseConv1DOpOptions{.channels = channels, + .kernel_size = kernel_size, + .bias = bias, + .state_inplace = state_inplace, + .accumulation_order = accumulation_order}) {} + +Tensor CausalDepthwiseConv1D::weight() const { + return std::static_pointer_cast(impl()->getInstancedOp())->weight(); +} + +Tensor CausalDepthwiseConv1D::bias() const { + return std::static_pointer_cast(impl()->getInstancedOp())->bias(); +} + +} // namespace mllm::nn diff --git a/mllm/nn/layers/CausalDepthwiseConv1D.hpp b/mllm/nn/layers/CausalDepthwiseConv1D.hpp new file mode 100644 index 000000000..ddd76c143 --- /dev/null +++ b/mllm/nn/layers/CausalDepthwiseConv1D.hpp @@ -0,0 +1,23 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include "mllm/core/aops/CausalDepthwiseConv1DOp.hpp" +#include "mllm/nn/Layer.hpp" + +namespace mllm::nn { + +class CausalDepthwiseConv1D : public Layer { + public: + CausalDepthwiseConv1D(); + CausalDepthwiseConv1D(int32_t channels, int32_t kernel_size, bool bias, bool state_inplace, + aops::CausalDepthwiseConv1DAccumulationOrder accumulation_order); + + [[nodiscard]] Tensor weight() const; + [[nodiscard]] Tensor bias() const; + + MLLM_LAYER_ANY_INPUTS_2_OUTPUTS_FORWARD +}; + +} // namespace mllm::nn diff --git a/mllm/nn/layers/GroupedQueryAttention.cpp b/mllm/nn/layers/GroupedQueryAttention.cpp new file mode 100644 index 000000000..0e629c091 --- /dev/null +++ b/mllm/nn/layers/GroupedQueryAttention.cpp @@ -0,0 +1,14 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/nn/layers/GroupedQueryAttention.hpp" + +namespace mllm::nn { + +GroupedQueryAttention::GroupedQueryAttention() + : Layer(OpTypes::kGroupedQueryAttention, aops::GroupedQueryAttentionOpOptions{}) {} + +GroupedQueryAttention::GroupedQueryAttention(aops::GroupedQueryAttentionImplementation implementation) + : Layer(OpTypes::kGroupedQueryAttention, aops::GroupedQueryAttentionOpOptions{.implementation = implementation}) {} + +} // namespace mllm::nn diff --git a/mllm/nn/layers/GroupedQueryAttention.hpp b/mllm/nn/layers/GroupedQueryAttention.hpp new file mode 100644 index 000000000..99eca8769 --- /dev/null +++ b/mllm/nn/layers/GroupedQueryAttention.hpp @@ -0,0 +1,19 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include "mllm/core/aops/GroupedQueryAttentionOp.hpp" +#include "mllm/nn/Layer.hpp" + +namespace mllm::nn { + +class GroupedQueryAttention : public Layer { + public: + GroupedQueryAttention(); + explicit GroupedQueryAttention(aops::GroupedQueryAttentionImplementation implementation); + + MLLM_LAYER_ANY_INPUTS_1_OUTPUTS_FORWARD +}; + +} // namespace mllm::nn diff --git a/mllm/nn/layers/GroupedQueryAttentionDecode.cpp b/mllm/nn/layers/GroupedQueryAttentionDecode.cpp deleted file mode 100644 index 2d8e2e1e4..000000000 --- a/mllm/nn/layers/GroupedQueryAttentionDecode.cpp +++ /dev/null @@ -1,14 +0,0 @@ -// Copyright (c) MLLM Team. -// Licensed under the MIT License. - -#include "mllm/nn/layers/GroupedQueryAttentionDecode.hpp" - -namespace mllm::nn { - -GroupedQueryAttentionDecode::GroupedQueryAttentionDecode() - : GroupedQueryAttentionDecode(aops::GroupedQueryAttentionDecodeOpOptions{}) {} - -GroupedQueryAttentionDecode::GroupedQueryAttentionDecode(const aops::GroupedQueryAttentionDecodeOpOptions& options) - : Layer(OpTypes::kGroupedQueryAttentionDecode, options) {} - -} // namespace mllm::nn diff --git a/mllm/nn/layers/GroupedQueryAttentionDecode.hpp b/mllm/nn/layers/GroupedQueryAttentionDecode.hpp deleted file mode 100644 index 99121bc29..000000000 --- a/mllm/nn/layers/GroupedQueryAttentionDecode.hpp +++ /dev/null @@ -1,20 +0,0 @@ -// Copyright (c) MLLM Team. -// Licensed under the MIT License. - -#pragma once - -#include "mllm/core/aops/GroupedQueryAttentionDecodeOp.hpp" -#include "mllm/nn/Layer.hpp" - -namespace mllm::nn { - -class GroupedQueryAttentionDecode : public Layer { - public: - GroupedQueryAttentionDecode(); - - explicit GroupedQueryAttentionDecode(const aops::GroupedQueryAttentionDecodeOpOptions& options); - - MLLM_LAYER_ANY_INPUTS_1_OUTPUTS_FORWARD -}; - -} // namespace mllm::nn diff --git a/mllm/nn/layers/ParallelLinear.cpp b/mllm/nn/layers/ParallelLinear.cpp new file mode 100644 index 000000000..2e6c015b9 --- /dev/null +++ b/mllm/nn/layers/ParallelLinear.cpp @@ -0,0 +1,25 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include "mllm/nn/layers/ParallelLinear.hpp" + +#include + +#include "mllm/core/aops/ParallelLinearOp.hpp" + +namespace mllm::nn { + +ParallelLinear::ParallelLinear() : Layer(OpTypes::kParallelLinear, aops::ParallelLinearOpOptions{}) {} + +ParallelLinear::ParallelLinear(int32_t in_channels, std::vector out_channels, + std::vector projection_names, bool bias, aops::LinearImplTypes impl_type, + int32_t decode_thread_cap, int32_t prefill_thread_cap) + : Layer(OpTypes::kParallelLinear, aops::ParallelLinearOpOptions{.in_channels = in_channels, + .out_channels = std::move(out_channels), + .projection_names = std::move(projection_names), + .bias = bias, + .impl_type = impl_type, + .kai_w4a32_decode_thread_cap = decode_thread_cap, + .kai_w4a32_prefill_thread_cap = prefill_thread_cap}) {} + +} // namespace mllm::nn diff --git a/mllm/nn/layers/ParallelLinear.hpp b/mllm/nn/layers/ParallelLinear.hpp new file mode 100644 index 000000000..a8e56bae7 --- /dev/null +++ b/mllm/nn/layers/ParallelLinear.hpp @@ -0,0 +1,26 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +#include "mllm/core/aops/LinearOp.hpp" +#include "mllm/nn/Layer.hpp" + +namespace mllm::nn { + +class ParallelLinear : public Layer { + public: + ParallelLinear(); + + ParallelLinear(int32_t in_channels, std::vector out_channels, std::vector projection_names, + bool bias = true, aops::LinearImplTypes impl_type = aops::LinearImplTypes::kDefault, + int32_t decode_thread_cap = 0, int32_t prefill_thread_cap = 0); + + MLLM_LAYER_ANY_INPUTS_ANY_OUTPUTS_FORWARD +}; + +} // namespace mllm::nn diff --git a/mllm/nn/llm_components/GroupedQueryAttention.hpp b/mllm/nn/llm_components/GroupedQueryAttention.hpp index 91a6bbbe7..35327d029 100644 --- a/mllm/nn/llm_components/GroupedQueryAttention.hpp +++ b/mllm/nn/llm_components/GroupedQueryAttention.hpp @@ -12,6 +12,34 @@ namespace mllm::nn::llm_components { +inline void validateGroupedQueryAttention(const Tensor& query, const Tensor& key, const Tensor& value) { + if (query.isNil() || key.isNil() || value.isNil()) { + throw std::invalid_argument("groupedQueryAttention inputs must not be nil"); + } + const auto q_shape = query.shape(); + const auto k_shape = key.shape(); + const auto v_shape = value.shape(); + if (q_shape.size() != 4 || k_shape.size() != 4 || v_shape.size() != 4 || q_shape[0] <= 0 || q_shape[0] != k_shape[0] + || q_shape[0] != v_shape[0] || q_shape[1] <= 0 || k_shape[1] <= 0 || k_shape[1] != v_shape[1] + || q_shape[1] % k_shape[1] != 0 || q_shape[2] <= 0 || k_shape[2] <= 0 || k_shape[2] != v_shape[2] || q_shape[3] <= 0 + || q_shape[3] != k_shape[3] || v_shape[3] <= 0) { + throw std::invalid_argument("groupedQueryAttention expects compatible [B, query_heads/KV_heads, sequence, dim] tensors"); + } + if (query.dtype() != key.dtype() || query.dtype() != value.dtype()) { + throw std::invalid_argument("groupedQueryAttention inputs must have the same dtype"); + } + if (query.device() != key.device() || query.device() != value.device()) { + throw std::invalid_argument("groupedQueryAttention inputs must be on the same device"); + } + if (query.device() != kCPU) { throw std::invalid_argument("groupedQueryAttention currently supports CPU only"); } + if (query.dtype() != kFloat32 && query.dtype() != kFloat16) { + throw std::invalid_argument("groupedQueryAttention supports float32 and float16 only"); + } + if (k_shape[2] < q_shape[2]) { + throw std::invalid_argument("groupedQueryAttention key sequence cannot be shorter than query"); + } +} + inline Tensor groupedQueryAttentionEager(const Tensor& query, const Tensor& key, const Tensor& value) { const auto q_shape = query.shape(); const auto k_shape = key.shape(); @@ -64,31 +92,9 @@ inline Tensor groupedQueryAttentionEager(const Tensor& query, const Tensor& key, // its shared KV head directly, so the persistent or temporary full KV history // is never expanded to query-head count. inline Tensor groupedQueryAttention(const Tensor& query, const Tensor& key, const Tensor& value) { - if (query.isNil() || key.isNil() || value.isNil()) { - throw std::invalid_argument("groupedQueryAttention inputs must not be nil"); - } + validateGroupedQueryAttention(query, key, value); const auto q_shape = query.shape(); const auto k_shape = key.shape(); - const auto v_shape = value.shape(); - if (q_shape.size() != 4 || k_shape.size() != 4 || v_shape.size() != 4 || q_shape[0] <= 0 || q_shape[0] != k_shape[0] - || q_shape[0] != v_shape[0] || q_shape[1] <= 0 || k_shape[1] <= 0 || k_shape[1] != v_shape[1] - || q_shape[1] % k_shape[1] != 0 || q_shape[2] <= 0 || k_shape[2] <= 0 || k_shape[2] != v_shape[2] || q_shape[3] <= 0 - || q_shape[3] != k_shape[3] || v_shape[3] <= 0) { - throw std::invalid_argument("groupedQueryAttention expects compatible [B, query_heads/KV_heads, sequence, dim] tensors"); - } - if (query.dtype() != key.dtype() || query.dtype() != value.dtype()) { - throw std::invalid_argument("groupedQueryAttention inputs must have the same dtype"); - } - if (query.device() != key.device() || query.device() != value.device()) { - throw std::invalid_argument("groupedQueryAttention inputs must be on the same device"); - } - if (query.device() != kCPU) { throw std::invalid_argument("groupedQueryAttention currently supports CPU only"); } - if (query.dtype() != kFloat32 && query.dtype() != kFloat16) { - throw std::invalid_argument("groupedQueryAttention supports float32 and float16 only"); - } - - const int32_t context_offset = k_shape[2] - q_shape[2]; - if (context_offset < 0) { throw std::invalid_argument("groupedQueryAttention key sequence cannot be shorter than query"); } if (query.dtype() == kFloat32 && q_shape[2] == 1) { return functional::groupedQueryAttentionDecode(query, key, value); } diff --git a/mllm/preprocessor/StreamingUtf8Decoder.hpp b/mllm/preprocessor/StreamingUtf8Decoder.hpp new file mode 100644 index 000000000..8ce8cbd64 --- /dev/null +++ b/mllm/preprocessor/StreamingUtf8Decoder.hpp @@ -0,0 +1,97 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include +#include +#include + +namespace mllm::preprocessor { + +// Incrementally validates byte-level tokenizer output and emits only complete +// UTF-8 sequences. Invalid sequences and incomplete final sequences are +// replaced with U+FFFD. +class StreamingUtf8Decoder { + public: + std::string append(std::string_view bytes) { + pending_.append(bytes.data(), bytes.size()); + return drain(false); + } + + std::string finish() { return drain(true); } + + void reset() { pending_.clear(); } + + private: + static constexpr std::string_view kReplacementCharacter = "\xEF\xBF\xBD"; + + static bool isContinuationByte(unsigned char byte) { return byte >= 0x80 && byte <= 0xBF; } + + static size_t sequenceLength(unsigned char lead) { + if (lead <= 0x7F) return 1; + if (lead >= 0xC2 && lead <= 0xDF) return 2; + if (lead >= 0xE0 && lead <= 0xEF) return 3; + if (lead >= 0xF0 && lead <= 0xF4) return 4; + return 0; + } + + static bool isValidSecondByte(unsigned char lead, unsigned char second) { + if (!isContinuationByte(second)) return false; + if (lead == 0xE0) return second >= 0xA0; + if (lead == 0xED) return second <= 0x9F; + if (lead == 0xF0) return second >= 0x90; + if (lead == 0xF4) return second <= 0x8F; + return true; + } + + std::string drain(bool flush) { + std::string output; + size_t offset = 0; + while (offset < pending_.size()) { + const auto lead = static_cast(pending_[offset]); + const size_t sequence_length = sequenceLength(lead); + if (sequence_length == 1) { + output.push_back(pending_[offset++]); + continue; + } + if (sequence_length == 0) { + output.append(kReplacementCharacter); + ++offset; + continue; + } + + const size_t available = pending_.size() - offset; + const size_t prefix_length = std::min(available, sequence_length); + bool valid_prefix = true; + for (size_t index = 1; index < prefix_length; ++index) { + const auto byte = static_cast(pending_[offset + index]); + if ((index == 1 && !isValidSecondByte(lead, byte)) || (index > 1 && !isContinuationByte(byte))) { + valid_prefix = false; + break; + } + } + if (!valid_prefix) { + output.append(kReplacementCharacter); + ++offset; + continue; + } + if (available < sequence_length) { + if (flush) { + output.append(kReplacementCharacter); + offset = pending_.size(); + } + break; + } + + output.append(pending_, offset, sequence_length); + offset += sequence_length; + } + pending_.erase(0, offset); + return output; + } + + std::string pending_; +}; + +} // namespace mllm::preprocessor diff --git a/mllm/preprocessor/tokenizers/BPE.cpp b/mllm/preprocessor/tokenizers/BPE.cpp index b71e50e47..1a33e5fd8 100644 --- a/mllm/preprocessor/tokenizers/BPE.cpp +++ b/mllm/preprocessor/tokenizers/BPE.cpp @@ -27,6 +27,8 @@ bool BPE::initFromSentencePieceJson(const std::string& file_path) { return false; } + ignore_merges_ = json_data["model"].value("ignore_merges", false); + for (const auto& [key, value] : json_data["model"]["vocab"].items()) { auto str = utf8string2WideString(key); vocab_.insert({ @@ -74,6 +76,11 @@ bool BPE::initFromSentencePieceJson(const std::string& file_path) { std::vector BPE::_bpe(const std::wstring& token) { // TODO check cache + // Checkpoints with ignore_merges keep whole vocabulary entries intact; the + // merge table cannot reconstruct all of them, so re-deriving would change + // the ids the checkpoint's own tokenizer produces. + if (ignore_merges_ && vocab_.count(token)) return {token}; + std::vector word; for (const auto& w : token) word.emplace_back(std::wstring{w}); diff --git a/mllm/preprocessor/tokenizers/BPE.hpp b/mllm/preprocessor/tokenizers/BPE.hpp index 765de0394..c45394ea4 100644 --- a/mllm/preprocessor/tokenizers/BPE.hpp +++ b/mllm/preprocessor/tokenizers/BPE.hpp @@ -39,6 +39,11 @@ class BPE { std::unordered_map vocab_; std::unordered_map vocab_inverse_; std::unordered_map, int64_t, BPEPairHash> bpe_ranks_; + // HuggingFace `model.ignore_merges`. When the checkpoint sets it, a token + // that is already a vocabulary entry is emitted whole instead of being + // rebuilt from the merge table. Roughly 2% of such entries are unreachable + // by merges alone, so ignoring the flag silently changes the token ids. + bool ignore_merges_ = false; }; } // namespace mllm::preprocessor \ No newline at end of file diff --git a/tests/cpu/CausalDepthwiseConvKernelTest.hpp b/tests/cpu/CausalDepthwiseConvKernelTest.hpp new file mode 100644 index 000000000..2c81b52e9 --- /dev/null +++ b/tests/cpu/CausalDepthwiseConvKernelTest.hpp @@ -0,0 +1,101 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. +#pragma once + +#include + +#include +#include +#include +#include +#include + +#include "mllm/backends/cpu/kernels/common/causal_conv/depthwise_causal_conv.hpp" +#include "mllm/mllm.hpp" + +#include "KernelTestHelper.hpp" + +class CausalDepthwiseConvKernelTest : public KernelTest { + public: + CausalDepthwiseConvKernelTest() = default; + ~CausalDepthwiseConvKernelTest() override = default; + + bool testHistoryFirstK3Once(const std::unordered_map& cfg) { + constexpr int kKernel = 3; + const int batch = cfg.at("B"); + const int sequence = cfg.at("S"); + const int channels = cfg.at("C"); + const bool non_zero_history = cfg.at("non_zero_history") != 0; + + const auto element_count = static_cast(batch) * sequence * channels; + const auto state_count = static_cast(batch) * channels * (kKernel - 1); + const std::vector input = makeBuffer(element_count, channels + sequence); + const std::vector weight = makeBuffer(static_cast(channels) * kKernel, kKernel); + const std::vector initial_state = + non_zero_history ? makeBuffer(state_count, 19) : std::vector(state_count, 0.0F); + + auto kernel_state = initial_state; + std::vector kernel_output(element_count, 0.0F); + mllm::cpu::causal_conv::depthwiseCausalConvHistoryFirstF32(input.data(), weight.data(), kernel_state.data(), + kernel_output.data(), batch, sequence, channels, kKernel); + + auto reference_state = initial_state; + std::vector reference_output(element_count, 0.0F); + referenceDepthwiseCausalConvHistoryFirst(input, weight, reference_state, reference_output, batch, sequence, channels, + kKernel); + + if (kernel_output != reference_output) { + mllm::print("history-first output mismatch for B=", batch, "S=", sequence, "C=", channels, + "history=", non_zero_history ? "non-zero" : "zero"); + return false; + } + if (kernel_state != reference_state) { + mllm::print("history-first state mismatch for B=", batch, "S=", sequence, "C=", channels, + "history=", non_zero_history ? "non-zero" : "zero"); + return false; + } + return true; + } + + bool testHistoryFirstK3(const std::vector>& cfgs) { + for (const auto& cfg : cfgs) { + if (!testHistoryFirstK3Once(cfg)) { return false; } + } + return true; + } + + private: + static 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; + } + + static 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; + } + + static void referenceDepthwiseCausalConvHistoryFirst(const std::vector& input, const std::vector& weight, + std::vector& state, std::vector& output, int batch_size, + int sequence_length, int channels, int kernel_size) { + const int state_width = kernel_size - 1; + for (int batch = 0; batch < batch_size; ++batch) { + for (int token = 0; token < sequence_length; ++token) { + for (int channel = 0; channel < channels; ++channel) { + const auto state_base = (static_cast(batch) * channels + channel) * state_width; + const auto element = (static_cast(batch) * sequence_length + token) * channels + channel; + const auto weight_base = static_cast(channel) * kernel_size; + + float value = 0.0F; + for (int tap = 0; tap < state_width; ++tap) { value += state[state_base + tap] * weight[weight_base + tap]; } + value += input[element] * weight[weight_base + state_width]; + output[element] = value; + + for (int tap = 0; tap + 1 < state_width; ++tap) { state[state_base + tap] = state[state_base + tap + 1]; } + state[state_base + state_width - 1] = input[element]; + } + } + } + } +}; diff --git a/tests/cpu/KernelTest.cpp b/tests/cpu/KernelTest.cpp index 9f8d613ee..277bb51d7 100644 --- a/tests/cpu/KernelTest.cpp +++ b/tests/cpu/KernelTest.cpp @@ -533,6 +533,41 @@ TEST_F(ElementwiseKernelTest, DivScalarInt32) { true); } +//===----------------------------------------------------------------------===// +// Causal depthwise convolution +//===----------------------------------------------------------------------===// +#include "CausalDepthwiseConvKernelTest.hpp" +TEST_F(CausalDepthwiseConvKernelTest, HistoryFirstK3MatchesScalarReferenceBitwise) { + EXPECT_EQ(testHistoryFirstK3({ + // Scalar path. + {{"B", 1}, {"S", 1}, {"C", 1}, {"non_zero_history", 0}}, + {{"B", 1}, {"S", 1}, {"C", 1}, {"non_zero_history", 1}}, + + // One exact NEON vector block. + {{"B", 1}, {"S", 28}, {"C", 4}, {"non_zero_history", 0}}, + {{"B", 1}, {"S", 28}, {"C", 4}, {"non_zero_history", 1}}, + + // NEON vector loop with scalar tails. + {{"B", 2}, {"S", 2}, {"C", 3}, {"non_zero_history", 0}}, + {{"B", 2}, {"S", 2}, {"C", 3}, {"non_zero_history", 1}}, + {{"B", 2}, {"S", 2}, {"C", 5}, {"non_zero_history", 0}}, + {{"B", 2}, {"S", 2}, {"C", 5}, {"non_zero_history", 1}}, + + // LFM production-width exact block and tail coverage. + {{"B", 1}, {"S", 225}, {"C", 2045}, {"non_zero_history", 0}}, + {{"B", 1}, {"S", 225}, {"C", 2045}, {"non_zero_history", 1}}, + {{"B", 2}, {"S", 28}, {"C", 2048}, {"non_zero_history", 0}}, + {{"B", 2}, {"S", 28}, {"C", 2048}, {"non_zero_history", 1}}, + }), + true); +} + +//===----------------------------------------------------------------------===// +// Parallel linear +//===----------------------------------------------------------------------===// +#include "ParallelLinearKernelTest.hpp" +TEST_F(ParallelLinearKernelTest, KaiW4A32DispatchPolicy) { EXPECT_EQ(testKaiW4A32DispatchPolicy(), true); } + //===----------------------------------------------------------------------===// // CausalMaskOp //===----------------------------------------------------------------------===// diff --git a/tests/cpu/ParallelLinearKernelTest.hpp b/tests/cpu/ParallelLinearKernelTest.hpp new file mode 100644 index 000000000..336cba91f --- /dev/null +++ b/tests/cpu/ParallelLinearKernelTest.hpp @@ -0,0 +1,37 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#pragma once + +#include "mllm/backends/cpu/kernels/common/linear/kai_w4a32_dispatch.hpp" +#include "mllm/backends/cpu/kernels/common/parallel_linear/shared_input.hpp" + +#include "KernelTestHelper.hpp" + +class ParallelLinearKernelTest : public KernelTest { + public: + ParallelLinearKernelTest() = default; + ~ParallelLinearKernelTest() override = default; + + bool testKaiW4A32DispatchPolicy() { + using mllm::cpu::kai_w4a32::shouldUseI8mmPrefill; + using mllm::cpu::kai_w4a32::threadCount; + + if (shouldUseI8mmPrefill(3, false, true) || shouldUseI8mmPrefill(4, true, true) || shouldUseI8mmPrefill(4, false, false) + || !shouldUseI8mmPrefill(4, false, true)) { + return false; + } + if (threadCount(1, 8, 4, 6) != 4 || threadCount(28, 8, 4, 6) != 6 || threadCount(1, 2, 4, 6) != 2 + || threadCount(28, 8, 0, 0) != 8) { + return false; + } + + const auto decode_plan = mllm::cpu::parallel_linear::planKaiW4A32SharedInput(1, 64, 8, 4, 6); +#if defined(MLLM_HOST_ARCH_ARM64) || defined(MLLM_HOST_ARCH_ARM) + return decode_plan.supported() && decode_plan.kernel == mllm::cpu::parallel_linear::SharedInputKernel::kKaiDotprod + && decode_plan.workspace_size > 0 && decode_plan.thread_count == 4; +#else + return !decode_plan.supported(); +#endif + } +}; diff --git a/tests/nn/CMakeLists.txt b/tests/nn/CMakeLists.txt index fbd04e07e..f90afe27a 100644 --- a/tests/nn/CMakeLists.txt +++ b/tests/nn/CMakeLists.txt @@ -14,4 +14,12 @@ add_executable(Mllm-Test-Nn-GroupedQueryAttention GroupedQueryAttentionTest.cpp) target_link_libraries(Mllm-Test-Nn-GroupedQueryAttention PRIVATE gtest_main MllmRT MllmCPUBackend) target_include_directories(Mllm-Test-Nn-GroupedQueryAttention PRIVATE ${MLLM_INCLUDE_DIR}) +add_executable(Mllm-Test-Nn-CausalDepthwiseConv1D CausalDepthwiseConv1DTest.cpp) +target_link_libraries(Mllm-Test-Nn-CausalDepthwiseConv1D PRIVATE gtest_main MllmRT MllmCPUBackend) +target_include_directories(Mllm-Test-Nn-CausalDepthwiseConv1D PRIVATE ${MLLM_INCLUDE_DIR}) + +add_executable(Mllm-Test-Nn-ParallelLinear ParallelLinearTest.cpp) +target_link_libraries(Mllm-Test-Nn-ParallelLinear PRIVATE gtest_main MllmRT MllmCPUBackend) +target_include_directories(Mllm-Test-Nn-ParallelLinear PRIVATE ${MLLM_INCLUDE_DIR}) + include(GoogleTest) diff --git a/tests/nn/CausalDepthwiseConv1DTest.cpp b/tests/nn/CausalDepthwiseConv1DTest.cpp new file mode 100644 index 000000000..a6bcde795 --- /dev/null +++ b/tests/nn/CausalDepthwiseConv1DTest.cpp @@ -0,0 +1,211 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "mllm/compile/ir/Trace.hpp" +#include "mllm/compile/ir/linalg/Op.hpp" +#include "mllm/compile/jit/binary/LinalgIRSerialization.hpp" +#include "mllm/compile/jit/interpreter/AopsFromJson.hpp" +#include "mllm/mllm.hpp" +#include "mllm/nn/Nn.hpp" + +namespace { + +using mllm::Tensor; + +class CausalDepthwiseConv1DTest : public testing::Test { + protected: + static void SetUpTestSuite() { mllm::initializeContext(); } +}; + +class CausalDepthwiseConv1DModule final : public mllm::nn::Module { + public: + CausalDepthwiseConv1DModule(std::string name, bool bias, bool state_inplace) : Module(std::move(name)), bias_(bias) { + conv_ = reg("conv", 2, 3, bias, state_inplace, + mllm::aops::CausalDepthwiseConv1DAccumulationOrder::kHistoryFirst); + } + + std::vector forward(const std::vector& inputs, const std::vector&) override { + auto [output, state] = conv_(inputs[0], inputs[1]); + return {output, state}; + } + + [[nodiscard]] bool hasBias() const { return bias_; } + + private: + mllm::nn::CausalDepthwiseConv1D conv_; + bool bias_; +}; + +template +auto findOp(const mllm::ir::node_ptr_t& node) -> typename OpType::ptr_t { + if (node->isa_()) { return node->cast_(); } + if (!node->isa_()) { return nullptr; } + for (const auto& region : node->cast_()->regions()) { + for (const auto& op : region->ops()) { + if (auto found = findOp(op)) { return found; } + } + } + return nullptr; +} + +Tensor tensor(const Tensor::shape_t& shape, const std::vector& values) { + return Tensor::fromVector(values, shape, mllm::kFloat32, mllm::kCPU); +} + +Tensor parameter(const std::string& name, const Tensor::shape_t& shape, const std::vector& values) { + auto result = Tensor::empty(shape, mllm::kFloat32, mllm::kCPU).setMemType(mllm::kParamsNormal).setName(name).alloc(); + std::copy(values.begin(), values.end(), result.ptr()); + return result; +} + +void loadParameters(CausalDepthwiseConv1DModule& module, const std::string& module_name, const std::vector& weights, + const std::vector& bias = {}) { + auto parameters = mllm::ParameterFile::create(); + parameters->push(module_name + ".conv.weight", parameter(module_name + ".conv.weight", {2, 1, 3}, weights)); + if (module.hasBias()) { parameters->push(module_name + ".conv.bias", parameter(module_name + ".conv.bias", {2}, bias)); } + module.load(parameters); +} + +std::pair, std::vector> referenceHistoryFirstK3(const std::vector& input, + const std::vector& weights, + std::vector state, + const std::vector& bias, int batch, + int sequence, int channels) { + constexpr int kKernelSize = 3; + constexpr int kHistorySize = kKernelSize - 1; + std::vector output(static_cast(batch) * sequence * channels); + for (int batch_index = 0; batch_index < batch; ++batch_index) { + for (int token = 0; token < sequence; ++token) { + for (int channel = 0; channel < channels; ++channel) { + const auto state_base = (static_cast(batch_index) * channels + channel) * kHistorySize; + const auto element = (static_cast(batch_index) * sequence + token) * channels + channel; + const auto weight_base = static_cast(channel) * kKernelSize; + float value = 0.0F; + for (int tap = 0; tap < kHistorySize; ++tap) { value += state[state_base + tap] * weights[weight_base + tap]; } + value += input[element] * weights[weight_base + kHistorySize]; + output[element] = value + (bias.empty() ? 0.0F : bias[channel]); + state[state_base] = state[state_base + 1]; + state[state_base + 1] = input[element]; + } + } + } + return {output, state}; +} + +void expectExact(const Tensor& actual, const std::vector& expected) { + ASSERT_EQ(actual.numel(), expected.size()); + EXPECT_EQ(actual.toVector(), expected); +} + +TEST_F(CausalDepthwiseConv1DTest, EagerMatchesReferenceAndMutatesStateInPlace) { + constexpr char kModuleName[] = "causal_depthwise_conv_eager"; + const std::vector weights = {0.5F, 0.25F, 2.0F, -1.0F, 0.5F, -0.25F}; + const std::vector bias = {0.125F, -0.25F}; + const std::vector input_values = {1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F}; + const std::vector initial_state = {0.5F, -0.5F, 1.0F, -1.0F}; + auto module = CausalDepthwiseConv1DModule(kModuleName, true, true); + loadParameters(module, kModuleName, weights, bias); + auto input = tensor({1, 3, 2}, input_values); + auto state = tensor({1, 2, 2}, initial_state); + const auto* state_storage = state.ptr(); + + const auto outputs = module(input, state); + const auto [expected_output, expected_state] = referenceHistoryFirstK3(input_values, weights, initial_state, bias, 1, 3, 2); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[0].shape(), input.shape()); + EXPECT_EQ(outputs[1].shape(), state.shape()); + EXPECT_EQ(outputs[1].ptr(), state_storage); + expectExact(outputs[0], expected_output); + expectExact(outputs[1], expected_state); +} + +TEST_F(CausalDepthwiseConv1DTest, NonInplaceOutputPreservesInputState) { + constexpr char kModuleName[] = "causal_depthwise_conv_copy_state"; + const std::vector weights = {0.5F, 0.25F, 2.0F, -1.0F, 0.5F, -0.25F}; + const std::vector input_values = {1.0F, 2.0F, 3.0F, 4.0F}; + const std::vector initial_state = {0.5F, -0.5F, 1.0F, -1.0F}; + auto module = CausalDepthwiseConv1DModule(kModuleName, false, false); + loadParameters(module, kModuleName, weights); + auto input = tensor({1, 2, 2}, input_values); + auto state = tensor({1, 2, 2}, initial_state); + const auto* state_storage = state.ptr(); + + const auto outputs = module(input, state); + const auto [expected_output, expected_state] = referenceHistoryFirstK3(input_values, weights, initial_state, {}, 1, 2, 2); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_NE(outputs[1].ptr(), state_storage); + expectExact(state, initial_state); + expectExact(outputs[0], expected_output); + expectExact(outputs[1], expected_state); +} + +TEST_F(CausalDepthwiseConv1DTest, ChunkedExecutionMatchesOneShot) { + const std::vector weights = {0.5F, 0.25F, 2.0F, -1.0F, 0.5F, -0.25F}; + const std::vector full_input = {1.0F, 2.0F, 3.0F, 4.0F, 5.0F, 6.0F, 7.0F, 8.0F}; + const std::vector initial_state = {0.5F, -0.5F, 1.0F, -1.0F}; + + auto one_shot = CausalDepthwiseConv1DModule("causal_depthwise_conv_one_shot", false, true); + loadParameters(one_shot, "causal_depthwise_conv_one_shot", weights); + const auto one_shot_outputs = one_shot(tensor({1, 4, 2}, full_input), tensor({1, 2, 2}, initial_state)); + + auto chunked = CausalDepthwiseConv1DModule("causal_depthwise_conv_chunked", false, true); + loadParameters(chunked, "causal_depthwise_conv_chunked", weights); + auto chunk_state = tensor({1, 2, 2}, initial_state); + const auto first = chunked(tensor({1, 2, 2}, {1.0F, 2.0F, 3.0F, 4.0F}), chunk_state); + const auto second = chunked(tensor({1, 2, 2}, {5.0F, 6.0F, 7.0F, 8.0F}), first[1]); + + const auto expected_output = one_shot_outputs[0].toVector(); + const auto first_output = first[0].toVector(); + const auto second_output = second[0].toVector(); + EXPECT_TRUE(std::equal(first_output.begin(), first_output.end(), expected_output.begin())); + EXPECT_TRUE(std::equal(second_output.begin(), second_output.end(), expected_output.begin() + first_output.size())); + EXPECT_EQ(second[1].toVector(), one_shot_outputs[1].toVector()); +} + +TEST_F(CausalDepthwiseConv1DTest, RejectsInvalidStateGeometry) { + auto module = CausalDepthwiseConv1DModule("causal_depthwise_conv_invalid", false, true); + EXPECT_THROW((void)module(tensor({1, 2, 2}, {1.0F, 2.0F, 3.0F, 4.0F}), tensor({1, 2, 1}, {0.0F, 0.0F})), + std::invalid_argument); +} + +TEST_F(CausalDepthwiseConv1DTest, TraceAndSerializationPreserveStateSemantics) { + CausalDepthwiseConv1DModule module("causal_depthwise_conv_trace", true, true); + auto ir_context = mllm::ir::trace(module, Tensor::empty({1, 2, 2}, mllm::kFloat32, mllm::kCPU), + Tensor::empty({1, 2, 2}, mllm::kFloat32, mllm::kCPU)); + auto op = findOp(ir_context->topLevelOp()); + ASSERT_NE(op, nullptr); + EXPECT_EQ(op->getAOp()->getOpType(), mllm::OpTypes::kCausalDepthwiseConv1D); + + const auto options = mllm::jit::binary::dumpLinalgIROptions(op); + EXPECT_EQ(options.at("channels"), 2); + EXPECT_EQ(options.at("kernel_size"), 3); + EXPECT_EQ(options.at("bias"), true); + EXPECT_EQ(options.at("state_inplace"), true); + EXPECT_EQ(options.at("accumulation_order"), "HistoryFirst"); + + const auto restored = mllm::jit::interpreter::aopsFromJson( + nlohmann::json{{"op_type", "CausalDepthwiseConv1D"}, {"backend", "CPU"}, {"op_options", options}}); + ASSERT_NE(restored, nullptr); + ASSERT_EQ(restored->getOpType(), mllm::OpTypes::kCausalDepthwiseConv1D); + const auto restored_options = std::static_pointer_cast(restored)->options(); + EXPECT_EQ(restored_options.channels, 2); + EXPECT_EQ(restored_options.kernel_size, 3); + EXPECT_TRUE(restored_options.bias); + EXPECT_TRUE(restored_options.state_inplace); + EXPECT_EQ(restored_options.accumulation_order, mllm::aops::CausalDepthwiseConv1DAccumulationOrder::kHistoryFirst); +} + +} // namespace diff --git a/tests/nn/GroupedQueryAttentionTest.cpp b/tests/nn/GroupedQueryAttentionTest.cpp index d57b084bc..af66bb718 100644 --- a/tests/nn/GroupedQueryAttentionTest.cpp +++ b/tests/nn/GroupedQueryAttentionTest.cpp @@ -37,14 +37,24 @@ class GroupedQueryAttentionDecodeTraceModule final : public mllm::nn::Module { } }; -mllm::ir::linalg::GroupedQueryAttentionDecodeOp::ptr_t findGroupedQueryAttentionDecodeOp(const mllm::ir::node_ptr_t& node) { - if (node->isa_()) { - return node->cast_(); +class GroupedQueryAttentionTraceModule final : public mllm::nn::Module { + public: + GroupedQueryAttentionTraceModule() : Module("gqa_trace") {} + + std::vector forward(const std::vector& inputs, const std::vector& args) override { + return {mllm::nn::functional::groupedQueryAttention(inputs[0], inputs[1], inputs[2])}; + } +}; + + +mllm::ir::linalg::GroupedQueryAttentionOp::ptr_t findGroupedQueryAttentionOp(const mllm::ir::node_ptr_t& node) { + if (node->isa_()) { + return node->cast_(); } if (!node->isa_()) { return nullptr; } for (const auto& region : node->cast_()->regions()) { for (const auto& op : region->ops()) { - if (auto found = findGroupedQueryAttentionDecodeOp(op)) { return found; } + if (auto found = findGroupedQueryAttentionOp(op)) { return found; } } } return nullptr; @@ -102,6 +112,73 @@ Tensor gqaReference(const Tensor& query, const Tensor& key, const Tensor& value) return output; } +Tensor gqaLegacyDirectStridedReference(const Tensor& query, const Tensor& key, const Tensor& value) { + const auto q_shape = query.shape(); + const auto k_shape = key.shape(); + const auto v_shape = value.shape(); + const int32_t groups = q_shape[1] / k_shape[1]; + const float scale = 1.0F / std::sqrt(static_cast(q_shape[3])); + const int32_t context_offset = k_shape[2] - q_shape[2]; + const int32_t jobs = q_shape[0] * q_shape[1]; + auto output = Tensor::zeros({q_shape[0], q_shape[1], q_shape[2], v_shape[3]}, mllm::kFloat32, mllm::kCPU); + + std::vector query_rows(static_cast(jobs) * q_shape[2]); + std::vector key_rows(static_cast(q_shape[0]) * k_shape[1] * k_shape[2]); + std::vector value_rows(static_cast(q_shape[0]) * v_shape[1] * v_shape[2]); + std::vector output_rows(static_cast(jobs) * q_shape[2]); + for (int32_t batch = 0; batch < q_shape[0]; ++batch) { + for (int32_t head = 0; head < q_shape[1]; ++head) { + for (int32_t sequence = 0; sequence < q_shape[2]; ++sequence) { + const size_t row = (static_cast(batch) * q_shape[1] + head) * q_shape[2] + sequence; + query_rows[row] = query.coffsettedPtr({batch, head, sequence, 0}); + output_rows[row] = output.offsettedPtr({batch, head, sequence, 0}); + } + } + for (int32_t head = 0; head < k_shape[1]; ++head) { + for (int32_t sequence = 0; sequence < k_shape[2]; ++sequence) { + const size_t row = (static_cast(batch) * k_shape[1] + head) * k_shape[2] + sequence; + key_rows[row] = key.coffsettedPtr({batch, head, sequence, 0}); + value_rows[row] = value.coffsettedPtr({batch, head, sequence, 0}); + } + } + } + + for (int32_t job = 0; job < jobs; ++job) { + const int32_t batch = job / q_shape[1]; + const int32_t query_head = job % q_shape[1]; + const int32_t kv_head = query_head / groups; + std::vector scores(static_cast(k_shape[2])); + for (int32_t query_index = 0; query_index < q_shape[2]; ++query_index) { + const int32_t visible_keys = context_offset + query_index + 1; + float maximum = std::numeric_limits::lowest(); + for (int32_t key_index = 0; key_index < visible_keys; ++key_index) { + float dot = 0.0F; + const size_t query_row = (static_cast(batch) * q_shape[1] + query_head) * q_shape[2] + query_index; + const size_t key_row = (static_cast(batch) * k_shape[1] + kv_head) * k_shape[2] + key_index; + for (int32_t dim = 0; dim < q_shape[3]; ++dim) { dot += query_rows[query_row][dim] * key_rows[key_row][dim]; } + scores[key_index] = dot * scale; + maximum = std::max(maximum, scores[key_index]); + } + + float denominator = 0.0F; + for (int32_t key_index = 0; key_index < visible_keys; ++key_index) { + scores[key_index] = std::exp(scores[key_index] - maximum); + denominator += scores[key_index]; + } + for (int32_t value_dim = 0; value_dim < v_shape[3]; ++value_dim) { + float accumulated = 0.0F; + for (int32_t key_index = 0; key_index < visible_keys; ++key_index) { + const size_t value_row = (static_cast(batch) * v_shape[1] + kv_head) * v_shape[2] + key_index; + accumulated += (scores[key_index] / denominator) * value_rows[value_row][value_dim]; + } + const size_t output_row = (static_cast(batch) * q_shape[1] + query_head) * q_shape[2] + query_index; + output_rows[output_row][value_dim] = accumulated; + } + } + } + return output; +} + void expectNear(Tensor actual, Tensor expected, float tolerance = 1e-5F) { ASSERT_EQ(actual.shape(), expected.shape()); const auto actual_cpu = actual.to(mllm::kCPU).contiguous(); @@ -121,6 +198,56 @@ TEST_F(GroupedQueryAttentionTest, MatchesRepeatedKVReference) { expectNear(actual, expected); } +TEST_F(GroupedQueryAttentionTest, RegisteredDirectStridedMatchesReference) { + auto query = sequential({1, 4, 3, 5}, 0.07F); + auto key = sequential({1, 2, 3, 5}, 0.11F); + auto value = sequential({1, 2, 3, 5}, 0.13F); + const auto actual = mllm::nn::functional::groupedQueryAttention(query, key, value); + const auto expected = gqaReference(query, key, value); + + expectNear(actual, expected, 1e-6F); +} + +TEST_F(GroupedQueryAttentionTest, DirectStridedMatchesLegacyReductionAtLfm25Geometry) { + struct Case { + int32_t query_length; + int32_t key_length; + }; + + // Cover both LFM2.5 prefill and decode geometry. Android production builds + // compile the backend with -ffast-math while this reference remains in the + // test translation unit, so the portable contract is tight numerical + // agreement; the full-model device gate separately freezes generated token + // IDs against the exact incumbent artifact. + for (const auto test_case : {Case{28, 28}, Case{1, 225}}) { + auto query = sequential({1, 32, test_case.query_length, 64}, 0.007F); + auto key = sequential({1, 8, test_case.key_length, 64}, 0.011F); + auto value = sequential({1, 8, test_case.key_length, 64}, 0.013F); + const auto actual = mllm::nn::functional::groupedQueryAttention(query, key, value); + const auto expected = gqaLegacyDirectStridedReference(query, key, value); + + ASSERT_NO_FATAL_FAILURE(expectNear(actual, expected, 1e-6F)) + << "query_length=" << test_case.query_length << " key_length=" << test_case.key_length; + } +} + +TEST_F(GroupedQueryAttentionTest, DirectStridedOpTraceAndSerializationRoundTrip) { + GroupedQueryAttentionTraceModule module; + auto ir_ctx = mllm::ir::trace(module, Tensor::empty({1, 4, 2, 5}, mllm::kFloat32, mllm::kCPU), + Tensor::empty({1, 2, 4, 5}, mllm::kFloat32, mllm::kCPU), + Tensor::empty({1, 2, 4, 3}, mllm::kFloat32, mllm::kCPU)); + auto ir_op = findGroupedQueryAttentionOp(ir_ctx->topLevelOp()); + ASSERT_NE(ir_op, nullptr); + EXPECT_EQ(ir_op->getAOp()->getOpType(), mllm::OpTypes::kGroupedQueryAttention); + + const auto options = mllm::jit::binary::dumpLinalgIROptions(ir_op); + EXPECT_EQ(options.at("implementation"), "DirectStrided"); + const auto restored = mllm::jit::interpreter::aopsFromJson( + nlohmann::json{{"op_type", "GroupedQueryAttention"}, {"backend", "CPU"}, {"op_options", options}}); + ASSERT_NE(restored, nullptr); + EXPECT_EQ(restored->getOpType(), mllm::OpTypes::kGroupedQueryAttention); +} + TEST_F(GroupedQueryAttentionTest, SupportsOneKVHeadAndRejectsIllegalGeometry) { auto query = sequential({1, 4, 1, 4}, 0.09F); auto key = sequential({1, 1, 2, 4}, 0.12F); @@ -178,22 +305,59 @@ TEST_F(GroupedQueryAttentionTest, DecodeStaysFiniteAtMiniCPM5ProductGeometry) { expectNear(actual, expected, 2e-5F); } +TEST_F(GroupedQueryAttentionTest, DecodeStaysFiniteAtLfm25ProductGeometry) { + auto query = sequential({1, 32, 1, 64}, 0.007F); + auto key_buffer = sequential({1, 8, 2048, 64}, 0.011F); + auto value_buffer = sequential({1, 8, 2048, 64}, 0.013F); + auto key = key_buffer[{mllm::kAll, mllm::kAll, {mllm::kAll, 201}, mllm::kAll}]; + auto value = value_buffer[{mllm::kAll, mllm::kAll, {mllm::kAll, 201}, mllm::kAll}]; + + const auto actual = mllm::nn::llm_components::groupedQueryAttention(query, key, value); + const auto expected = gqaReference(query, key, value); + + EXPECT_EQ(actual.shape(), (Tensor::shape_t{1, 32, 1, 64})); + for (int index = 0; index < actual.numel(); ++index) { EXPECT_TRUE(std::isfinite(actual.ptr()[index])); } + expectNear(actual, expected, 2e-5F); +} + TEST_F(GroupedQueryAttentionTest, DecodeOpTraceAndSerializationRoundTrip) { GroupedQueryAttentionDecodeTraceModule module; auto ir_ctx = mllm::ir::trace(module, Tensor::empty({1, 4, 1, 5}, mllm::kFloat32, mllm::kCPU), Tensor::empty({1, 2, 4, 5}, mllm::kFloat32, mllm::kCPU), Tensor::empty({1, 2, 4, 3}, mllm::kFloat32, mllm::kCPU)); - auto ir_op = findGroupedQueryAttentionDecodeOp(ir_ctx->topLevelOp()); + auto ir_op = findGroupedQueryAttentionOp(ir_ctx->topLevelOp()); ASSERT_NE(ir_op, nullptr); ASSERT_NE(ir_op->getAOp(), nullptr); - EXPECT_EQ(ir_op->getAOp()->getOpType(), mllm::OpTypes::kGroupedQueryAttentionDecode); + EXPECT_EQ(ir_op->getAOp()->getOpType(), mllm::OpTypes::kGroupedQueryAttention); const auto options = mllm::jit::binary::dumpLinalgIROptions(ir_op); - EXPECT_TRUE(options.empty()); - const nlohmann::json encoded = {{"op_type", "GroupedQueryAttentionDecode"}, {"backend", "CPU"}, {"op_options", options}}; + EXPECT_EQ(options.at("implementation"), "DecodeNativeKV"); + const nlohmann::json encoded = {{"op_type", "GroupedQueryAttention"}, {"backend", "CPU"}, {"op_options", options}}; + const auto restored = mllm::jit::interpreter::aopsFromJson(encoded); + ASSERT_NE(restored, nullptr); + EXPECT_EQ(restored->getOpType(), mllm::OpTypes::kGroupedQueryAttention); +} + +// Graphs serialized before the decode-only operation was folded into +// GroupedQueryAttention must still reconstruct. +TEST_F(GroupedQueryAttentionTest, LegacyDecodeOpTypeStringStillReconstructs) { + const nlohmann::json encoded = { + {"op_type", "GroupedQueryAttentionDecode"}, {"backend", "CPU"}, {"op_options", nlohmann::json::object()}}; const auto restored = mllm::jit::interpreter::aopsFromJson(encoded); ASSERT_NE(restored, nullptr); - EXPECT_EQ(restored->getOpType(), mllm::OpTypes::kGroupedQueryAttentionDecode); + EXPECT_EQ(restored->getOpType(), mllm::OpTypes::kGroupedQueryAttention); } +TEST_F(GroupedQueryAttentionTest, SupportsTransposedNonContiguousHeadViews) { + auto query_bshd = sequential({1, 3, 4, 5}, 0.07F); + auto key_bshd = sequential({1, 3, 2, 5}, 0.11F); + auto value_bshd = sequential({1, 3, 2, 5}, 0.13F); + auto query = query_bshd.transpose(1, 2); + auto key = key_bshd.transpose(1, 2); + auto value = value_bshd.transpose(1, 2); + + const auto actual = mllm::nn::functional::groupedQueryAttention(query, key, value); + const auto expected = gqaReference(query.contiguous(), key.contiguous(), value.contiguous()); + expectNear(actual, expected); +} } // namespace diff --git a/tests/nn/ParallelLinearTest.cpp b/tests/nn/ParallelLinearTest.cpp new file mode 100644 index 000000000..5a477b594 --- /dev/null +++ b/tests/nn/ParallelLinearTest.cpp @@ -0,0 +1,139 @@ +// Copyright (c) MLLM Team. +// Licensed under the MIT License. + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include "mllm/compile/ir/Trace.hpp" +#include "mllm/compile/ir/linalg/Op.hpp" +#include "mllm/compile/jit/binary/LinalgIRSerialization.hpp" +#include "mllm/compile/jit/interpreter/AopsFromJson.hpp" +#include "mllm/core/aops/ParallelLinearOp.hpp" +#include "mllm/mllm.hpp" +#include "mllm/nn/Nn.hpp" + +namespace { + +using mllm::Tensor; + +class ParallelLinearTest : public testing::Test { + protected: + static void SetUpTestSuite() { mllm::initializeContext(); } +}; + +class ParallelLinearModule final : public mllm::nn::Module { + public: + ParallelLinearModule(std::string name, bool bias) : Module(std::move(name)) { + projections_ = + reg("pair", 2, std::vector{2, 1}, std::vector{"left", "right"}, bias, + mllm::aops::LinearImplTypes::kGGUF, 4, 6); + } + + std::vector forward(const std::vector& inputs, const std::vector&) override { + return projections_(inputs[0]); + } + + private: + mllm::nn::ParallelLinear projections_; +}; + +template +auto findOp(const mllm::ir::node_ptr_t& node) -> typename OpType::ptr_t { + if (node->isa_()) { return node->cast_(); } + if (!node->isa_()) { return nullptr; } + for (const auto& region : node->cast_()->regions()) { + for (const auto& op : region->ops()) { + if (auto found = findOp(op)) { return found; } + } + } + return nullptr; +} + +Tensor parameter(const std::string& name, const Tensor::shape_t& shape, const std::vector& values) { + auto result = Tensor::empty(shape, mllm::kFloat32, mllm::kCPU).setMemType(mllm::kParamsNormal).setName(name).alloc(); + std::copy(values.begin(), values.end(), result.ptr()); + return result; +} + +TEST_F(ParallelLinearTest, EagerOwnsSiblingParametersAndMatchesIndependentProjections) { + ParallelLinearModule module("parallel_linear_eager", true); + auto parameters = mllm::ParameterFile::create(); + parameters->push("parallel_linear_eager.left.weight", + parameter("parallel_linear_eager.left.weight", {2, 2}, {1.0F, 2.0F, 3.0F, 4.0F})); + parameters->push("parallel_linear_eager.left.bias", parameter("parallel_linear_eager.left.bias", {2}, {0.5F, -0.5F})); + parameters->push("parallel_linear_eager.right.weight", parameter("parallel_linear_eager.right.weight", {1, 2}, {5.0F, 6.0F})); + parameters->push("parallel_linear_eager.right.bias", parameter("parallel_linear_eager.right.bias", {1}, {1.0F})); + module.load(parameters); + + auto input = Tensor::fromVector({2.0F, 3.0F}, {1, 1, 2}, mllm::kFloat32, mllm::kCPU); + const auto outputs = module(input); + + ASSERT_EQ(outputs.size(), 2); + EXPECT_EQ(outputs[0].shape(), (Tensor::shape_t{1, 1, 2})); + EXPECT_EQ(outputs[1].shape(), (Tensor::shape_t{1, 1, 1})); + EXPECT_EQ(outputs[0].toVector(), (std::vector{8.5F, 17.5F})); + EXPECT_EQ(outputs[1].toVector(), (std::vector{29.0F})); +} + +TEST_F(ParallelLinearTest, RejectsInvalidProjectionAndInputContracts) { + auto reshapeWith = [](std::vector out_channels, std::vector projection_names, + const Tensor::shape_t& input_shape = {1, 1, 2}) { + auto op = std::make_shared( + mllm::aops::ParallelLinearOpOptions{.in_channels = 2, + .out_channels = std::move(out_channels), + .projection_names = std::move(projection_names), + .bias = false, + .impl_type = mllm::aops::LinearImplTypes::kGGUF}); + std::vector inputs = {Tensor::empty(input_shape, mllm::kFloat32, mllm::kCPU)}; + std::vector outputs; + op->reshape(inputs, outputs); + }; + + EXPECT_THROW(reshapeWith({2, 1}, {"same", "same"}), std::invalid_argument); + EXPECT_THROW(reshapeWith({2, 1}, {"left", "nested.right"}), std::invalid_argument); + EXPECT_THROW(reshapeWith({2, 1}, {"left", ""}), std::invalid_argument); + EXPECT_THROW(reshapeWith({2}, {"left"}), std::invalid_argument); + EXPECT_THROW(reshapeWith({2, 0}, {"left", "right"}), std::invalid_argument); + EXPECT_THROW(reshapeWith({2, 1}, {"left", "right"}, {1, 1, 3}), std::invalid_argument); + EXPECT_NO_THROW(reshapeWith({2, 1}, {"left", "right"})); +} + +TEST_F(ParallelLinearTest, TraceAndSerializationPreserveProjectionContract) { + ParallelLinearModule module("parallel_linear_trace", false); + auto ir_context = mllm::ir::trace(module, Tensor::empty({1, 1, 2}, mllm::kFloat32, mllm::kCPU)); + auto op = findOp(ir_context->topLevelOp()); + ASSERT_NE(op, nullptr); + EXPECT_EQ(op->getAOp()->getOpType(), mllm::OpTypes::kParallelLinear); + + const auto serialized = mllm::jit::binary::dumpLinalgIROptions(op); + EXPECT_EQ(serialized.at("in_channels"), 2); + EXPECT_EQ(serialized.at("out_channels"), (std::vector{2, 1})); + EXPECT_EQ(serialized.at("projection_names"), (std::vector{"left", "right"})); + EXPECT_EQ(serialized.at("bias"), false); + EXPECT_EQ(serialized.at("impl_type"), "GGUF"); + EXPECT_EQ(serialized.at("kai_w4a32_decode_thread_cap"), 4); + EXPECT_EQ(serialized.at("kai_w4a32_prefill_thread_cap"), 6); + + const auto restored = mllm::jit::interpreter::aopsFromJson( + nlohmann::json{{"op_type", "ParallelLinear"}, {"backend", "CPU"}, {"op_options", serialized}}); + ASSERT_NE(restored, nullptr); + ASSERT_EQ(restored->getOpType(), mllm::OpTypes::kParallelLinear); + const auto restored_options = std::static_pointer_cast(restored)->options(); + EXPECT_EQ(restored_options.in_channels, 2); + EXPECT_EQ(restored_options.out_channels, (std::vector{2, 1})); + EXPECT_EQ(restored_options.projection_names, (std::vector{"left", "right"})); + EXPECT_FALSE(restored_options.bias); + EXPECT_EQ(restored_options.impl_type, mllm::aops::LinearImplTypes::kGGUF); + EXPECT_EQ(restored_options.kai_w4a32_decode_thread_cap, 4); + EXPECT_EQ(restored_options.kai_w4a32_prefill_thread_cap, 6); +} + +} // namespace